Training in practice, interpretability and robustness

Knowing the algorithms is not the same as getting a network to train. This lesson collects the practical knowledge that separates a model that works from a pile of NaNs: a systematic debugging procedure, the mechanics of mixed-precision training and memory budgets, reproducibility, and the frameworks and hardware that real projects use. It then asks what a trained network has learned and how far to trust it, covering attribution methods for interpretability, adversarial examples, and distribution shift. AI-480 covers hardware and distributed training in depth, AI-420 production systems, and AI-430 interpretability and robustness from the safety side; this lesson gives the working knowledge every practitioner needs.

A debugging procedure

Neural network bugs rarely crash. A wrong sign, a missing normalization or a mislabelled class usually produces a model that trains, just badly, and the failure is hard to tell apart from bad hyperparameters. Andrej Karpathy's widely cited recipe for training neural networks starts from that observation and argues for building up slowly, verifying each step. A workable procedure is as follows.

• Look at the data first: inspect examples and labels, check class balance, duplicates and leakage between splits, and confirm that preprocessing does what you think by viewing its output.
• Set up the full pipeline with a trivial model and a simple baseline, and check that the evaluation code gives the expected numbers for the baseline.
• Check the loss at initialization. With K balanced classes and a small output layer, cross-entropy should start near log K; a much larger value means the output scale or the labels are wrong.
• Overfit one small batch. A correct model should drive the training loss on a few examples to nearly zero; if it cannot, there is a bug in the model, the loss or the gradient flow.
• Then scale up the data and model, add regularization, and tune, changing one thing at a time and keeping a log of every run.

# Two sanity checks before any long training run: the loss at initialization, and overfitting one small batch.
import numpy as np
rng = np.random.default_rng(466)
K, d, H = 10, 50, 128
X = rng.normal(size=(16, d)); y = rng.integers(0, K, 16)        # one batch of 16 examples, 10 classes
W1 = rng.normal(0, np.sqrt(2 / d), (d, H)); b1 = np.zeros(H)
W2 = rng.normal(0, 0.01, (H, K)); b2 = np.zeros(K)               # small output layer: near-uniform start
def forward():
    h = np.maximum(0, X @ W1 + b1); s = h @ W2 + b2
    s = s - s.max(1, keepdims=True); logp = s - np.log(np.exp(s).sum(1, keepdims=True))
    return h, logp
h, logp = forward()
print(f"initial loss {-logp[np.arange(16), y].mean():.4f}  vs  log K = {np.log(K):.4f}")
for step in range(1, 301):
    h, logp = forward()
    g = np.exp(logp); g[np.arange(16), y] -= 1; g /= 16
    gh = (g @ W2.T) * (h > 0)
    W2 -= 0.5 * h.T @ g; b2 -= 0.5 * g.sum(0); W1 -= 0.5 * X.T @ gh; b1 -= 0.5 * gh.sum(0)
    if step in (50, 300):
        print(f"step {step}: loss on the batch {-logp[np.arange(16), y].mean():.4f}, accuracy {np.mean(logp.argmax(1) == y):.2f}")

initial loss 2.3234  vs  log K = 2.3026
step 50: loss on the batch 0.0049, accuracy 1.00
step 300: loss on the batch 0.0007, accuracy 1.00

Other checks earn their keep. Plot the norms of activations and gradients per layer at initialization and during training; a layer whose values are thousands of times larger or smaller than its neighbours' points to an initialization or normalization problem (lesson 4). Watch the ratio of update size to weight size, which should be around \(10^{-3}\) per step for typical learning rates. Visualize predictions on validation examples, especially wrong ones. If the loss suddenly becomes NaN, check for a too-high learning rate, missing gradient clipping, a logarithm or division of zero, or float16 overflow, and find the first step and layer where a non-finite value appears.

Mixed precision

Modern accelerators perform 16-bit matrix arithmetic several times faster than 32-bit and halve the memory traffic, so large networks are trained in mixed precision (Paulius Micikevicius and colleagues, 2018). Weights are kept in a float32 master copy; the forward and backward passes run in a 16-bit format; and the update is applied to the master copy, because updates are often smaller than float16's precision can represent relative to the weight (adding \(10^{-4}\) to 1 in float16 leaves it unchanged). Two 16-bit formats are used. IEEE float16 has 10 bits of precision but a narrow range, so small gradients underflow to zero; the fix is loss scaling, multiplying the loss by a large factor before backpropagation and dividing the gradients by it afterwards, with the factor adjusted dynamically to avoid overflow. bfloat16, developed at Google, keeps float32's 8-bit exponent and so its range, with only 7 bits of precision; it needs no loss scaling and is now the usual choice on hardware that supports it. Eight-bit and even four-bit floating-point formats are increasingly used for parts of training and inference on the newest hardware, as AI-480 describes.

# Mixed precision: float16's narrow range loses small gradients unless the loss is scaled up first.
import numpy as np
f16 = np.finfo(np.float16)
print(f"float16: max {float(f16.max):.0f}, smallest normal {float(f16.tiny):.2e}, machine epsilon {float(f16.eps):.2e}")
bf = 2.0 ** -126
print(f"bfloat16 keeps float32's range (smallest normal about {bf:.1e}) with only 8 bits of precision")
rng = np.random.default_rng(466)
grads = np.abs(rng.normal(0, 1e-6, 100000)) * 10 ** rng.uniform(-2, 0, 100000)   # many tiny gradient values
for scale in (1, 1024, 65536):
    g16 = (grads * scale).astype(np.float16)
    lost = np.mean(g16 == 0)
    back = g16.astype(np.float32) / scale
    rel = np.abs(back[g16 != 0] - grads[g16 != 0]) / grads[g16 != 0]
    print(f"loss scale {scale:6}: {lost:6.1%} of gradients flushed to zero, median relative error of the rest {np.median(rel):.1e}")
print("1 + 1e-4 in float16:", np.float16(1) + np.float16(1e-4), " (why master weights are kept in float32)")

float16: max 65504, smallest normal 6.10e-05, machine epsilon 9.77e-04
bfloat16 keeps float32's range (smallest normal about 1.2e-38) with only 8 bits of precision
loss scale      1:  37.1% of gradients flushed to zero, median relative error of the rest 8.3e-02
loss scale   1024:   0.1% of gradients flushed to zero, median relative error of the rest 2.9e-04
loss scale  65536:   0.0% of gradients flushed to zero, median relative error of the rest 1.7e-04
1 + 1e-4 in float16: 1.0  (why master weights are kept in float32)

Memory and its budget

Training memory has four parts: the weights, their gradients, the optimizer state and the activations stored for the backward pass. With mixed precision and Adam, the first three cost about 16 bytes per parameter: a 16-bit copy of the weights and of the gradients, a float32 master copy, and Adam's two float32 moment estimates. A model with a billion parameters therefore needs about 16 GB before any activations, which is why the memory techniques of AI-480 exist: sharding optimizer state across devices (ZeRO and FSDP), offloading, and model parallelism. Activations scale with batch size, sequence length, width and depth, and for long sequences they dominate. Gradient checkpointing (Tianqi Chen and colleagues, 2016) stores only some activations, such as each block's input, and recomputes the rest during the backward pass, trading about one extra forward pass for a large saving.

# Where training memory goes for a 1.3-billion-parameter model with mixed precision and Adam.
N = 1.3e9
weights16, grads16 = 2 * N, 2 * N             # 16-bit working copy of weights and gradients
master32 = 4 * N                              # float32 master weights
adam32 = 8 * N                                # two float32 moment estimates
states = weights16 + grads16 + master32 + adam32
print(f"weights, gradients and optimizer state: {states / 2 ** 30:.1f} GiB  ({states / N:.0f} bytes per parameter)")
L, d, T, B = 24, 2048, 2048, 8                 # layers, width, sequence length, batch
act_per_layer = 34 * T * B * d                 # a commonly used rough count of bytes stored per layer
acts = L * act_per_layer
print(f"activations, rough estimate: {acts / 2 ** 30:.1f} GiB")
ckpt = act_per_layer * 2 + L * T * B * d * 2   # keep only each layer's input, recompute the rest one layer at a time
print(f"with gradient checkpointing at every layer: about {ckpt / 2 ** 30:.1f} GiB, for roughly one extra forward pass of compute")

weights, gradients and optimizer state: 19.4 GiB  (16 bytes per parameter)
activations, rough estimate: 25.5 GiB
with gradient checkpointing at every layer: about 3.6 GiB, for roughly one extra forward pass of compute

Reproducibility and experiment management

Deep learning results vary with the random seed, sometimes by more than the effect being claimed, so report results over several seeds with their spread, and compare methods with equal tuning budgets (Gábor Melis, Chris Dyer and Phil Blunsom, 2018, found that a well-tuned LSTM matched several newer architectures). Fix seeds for every source of randomness, but know that GPU kernels for some operations are non-deterministic unless deterministic modes are enabled, at some cost in speed. Record the code version, data version, configuration and environment for every run, with an experiment tracker or at least a disciplined log. Keep a held-out test set untouched until the end. These habits are the deep learning version of the evaluation discipline of CMPUT-367, and they matter more, not less, when each run is expensive.

Frameworks and hardware

Real projects use a framework rather than NumPy. PyTorch, from Meta and now governed by the PyTorch Foundation, builds the graph as Python code runs and dominates research; its compile mode traces and optimizes models for speed. JAX, from Google, transforms pure functions (differentiation with grad, compilation with jit, vectorization with vmap, parallelization across devices) and is widely used for large-scale training on TPUs. TensorFlow and Keras remain in production use. All of them provide automatic differentiation (lesson 2), GPU and TPU execution, standard layers, optimizers and data loaders. GPUs, led by NVIDIA's data-centre parts, and Google's TPUs are the main training hardware; their performance depends on keeping matrix units busy, which favours large batched matrix multiplications with dimensions that are multiples of 8 or 16, and on memory bandwidth, which often limits operations such as normalization and attention more than arithmetic does. AI-480 develops the roofline model that quantifies this.

Interpretability: what has the network learned?

A network with millions of parameters does not explain itself. Interpretability methods try to identify which inputs drove a decision, what individual units or directions represent, and what algorithms the network implements. Attribution methods assign each input feature a share of responsibility for an output. The simplest is the gradient of the output score with respect to the input, a saliency map (Karen Simonyan, Andrea Vedaldi and Andrew Zisserman, 2014), which says how sensitive the score is to each feature locally. Gradient times input weights sensitivity by the feature's value. Integrated gradients (Mukund Sundararajan, Ankur Taly and Qiqi Yan, 2017) average the gradient along a straight path from a baseline, such as an all-black image, to the input, and multiply by the difference from the baseline; they satisfy completeness, meaning the attributions sum exactly to the difference between the output at the input and at the baseline.

# Attributions for a small fixed network: gradient, gradient x input, and integrated gradients.
import numpy as np
rng = np.random.default_rng(466)
W1 = rng.normal(size=(4, 6)); w2 = rng.normal(size=6)
def f(x): return np.tanh(x @ W1) @ w2                      # a scalar score
def grad(x): return W1 @ ((1 - np.tanh(x @ W1) ** 2) * w2)
x = np.array([1.0, -0.5, 2.0, 0.0]); base = np.zeros(4)
ig = np.zeros(4)
for k in range(200):                                         # midpoint rule along the straight path from the baseline
    a = (k + 0.5) / 200
    ig += grad(base + a * (x - base)) / 200
ig *= (x - base)
print("score f(x) =", round(f(x), 4), " f(baseline) =", round(f(base), 4))
print("gradient          ", np.round(grad(x), 3))
print("gradient x input  ", np.round(grad(x) * x, 3), " sum", round(float(np.sum(grad(x) * x)), 3))
print("integrated grads  ", np.round(ig, 3), " sum", round(float(ig.sum()), 3), "(completeness: equals f(x) - f(baseline))")

score f(x) = 0.7139  f(baseline) = 0.0
gradient           [ 0.823 -3.448 -1.123  0.852]
gradient x input   [ 0.823  1.724 -2.247  0.   ]  sum 0.301
integrated grads   [ 0.871  2.067 -2.225  0.   ]  sum 0.714 (completeness: equals f(x) - f(baseline))

Attribution maps must be used with care. Julius Adebayo and colleagues (2018) showed that some popular saliency methods produce nearly the same maps for a trained network and for one with randomized weights, so those maps reflect the input's structure more than the model's reasoning; sanity checks of that kind should accompany any attribution method. Results depend on the baseline chosen, and attention weights (lesson 8) are not by themselves faithful explanations either. Beyond attribution, probing classifiers test whether information such as part of speech can be read out of a layer's activations, and mechanistic interpretability tries to reverse-engineer the circuits a network uses, for example by decomposing activations into interpretable features with sparse autoencoders (Anthropic's work of 2023–2024 on dictionary learning in language models). AI-430 covers these methods.

Adversarial examples

Christian Szegedy and colleagues (2014) discovered that imperceptibly small, carefully chosen changes to an image could make a high-accuracy classifier output any label with high confidence. Ian Goodfellow, Jonathon Shlens and Christian Szegedy (2015) gave a simple explanation: the models are too linear. For a linear score \(\mathbf{w}^{\top}\mathbf{x}\), changing every coordinate by at most \(\epsilon\) in the direction of the sign of the corresponding weight changes the score by \(\epsilon\lVert\mathbf{w}\rVert_1\), which grows with the dimension even though each change is tiny. Images have hundreds of thousands of dimensions, so per-pixel changes too small to see add up to a large change in the score. Random noise of the same size, by contrast, mostly cancels out.

# Robustness of a linear classifier: small changes to every input dimension add up in high dimension.
import numpy as np
rng = np.random.default_rng(466)
for d in (10, 100, 1000):
    w = rng.normal(size=d) / np.sqrt(d)
    X = rng.normal(size=(5000, d)); y = np.sign(X @ w)
    margin = y * (X @ w)                                   # how far each point is from the boundary (score units)
    for eps in (0.01, 0.05):
        # the worst perturbation with every coordinate changed by at most eps shifts the score by eps * ||w||_1
        acc = np.mean(margin > eps * np.abs(w).sum())
        print(f"d = {d:4}, per-coordinate change {eps}: accuracy under worst-case perturbation {acc:.3f}")
print("random perturbations of the same size barely matter:")
d = 1000; w = rng.normal(size=d) / np.sqrt(d); X = rng.normal(size=(5000, d)); y = np.sign(X @ w)
noise = 0.05 * np.sign(rng.normal(size=X.shape))
print(f"  d = 1000, random +-0.05 on every coordinate: accuracy {np.mean(np.sign((X + noise) @ w) == y):.3f}")

d =   10, per-coordinate change 0.01: accuracy under worst-case perturbation 0.979
d =   10, per-coordinate change 0.05: accuracy under worst-case perturbation 0.900
d =  100, per-coordinate change 0.01: accuracy under worst-case perturbation 0.936
d =  100, per-coordinate change 0.05: accuracy under worst-case perturbation 0.704
d = 1000, per-coordinate change 0.01: accuracy under worst-case perturbation 0.815
d = 1000, per-coordinate change 0.05: accuracy under worst-case perturbation 0.206
random perturbations of the same size barely matter:
  d = 1000, random +-0.05 on every coordinate: accuracy 0.983

Adversarial robustness matters for any model deployed where someone may want to fool it, such as content filters, malware detectors and autonomous vehicles, and it connects to robustness more broadly. The main defence with consistent evidence is adversarial training, which trains on worst-case perturbed examples found during training (Aleksander Madry and colleagues, 2018); it improves robustness within the perturbation budget it was trained for, at a cost in clean accuracy and training time. Certified defences such as randomized smoothing give provable guarantees for small perturbations. Many proposed defences were later shown to fail against adaptive evaluations (Nicholas Carlini and colleagues), so robustness claims should be tested rigorously, for instance with standardized benchmarks. For language models, related concerns include jailbreaks and prompt injection, covered in AI-430.

Distribution shift

A more common threat than an adversary is the world changing. Models are trained on one distribution and deployed on another: new hospitals with different scanners, new cameras, new slang, seasons, populations or products. Covariate shift changes the input distribution, label shift changes class frequencies, and concept shift changes the relationship between inputs and labels. A particularly dangerous case is shortcut learning (Robert Geirhos and colleagues, 2020), where a network latches onto a feature that happens to correlate with the label in the training data but is not the real cause: a pneumonia detector that learned which hospital's scanner produced an image, or a classifier that recognized cows by their grassy backgrounds. Such models look excellent on held-out data from the same distribution and fail when the correlation breaks.

# Distribution shift: a model that relied on a spurious feature fails when the correlation breaks.
import numpy as np
rng = np.random.default_rng(466)
def data(n, spurious_agreement):
    y = rng.integers(0, 2, n)
    core = y + rng.normal(0, 1.0, n)                       # genuinely predictive but noisy
    agree = rng.uniform(size=n) < spurious_agreement
    spur = np.where(agree, y, 1 - y) + rng.normal(0, 0.1, n)   # nearly noiseless, but only correlated
    return np.c_[core, spur, np.ones(n)], y
Xtr, ytr = data(5000, 0.95)                                # in training, the spurious feature agrees 95% of the time
w = np.zeros(3)
for _ in range(3000):                                      # logistic regression by gradient descent
    p = 1 / (1 + np.exp(-Xtr @ w)); w -= 0.5 * Xtr.T @ (p - ytr) / len(ytr)
print("weights (core, spurious, bias):", np.round(w, 2))
for agree in (0.95, 0.5, 0.05):
    X, y = data(5000, agree)
    print(f"test set where spurious feature agrees {agree:.0%} of the time: accuracy {np.mean((X @ w > 0) == y):.3f}")

weights (core, spurious, bias): [ 0.99  5.94 -3.46]
test set where spurious feature agrees 95% of the time: accuracy 0.947
test set where spurious feature agrees 50% of the time: accuracy 0.512
test set where spurious feature agrees 5% of the time: accuracy 0.071

The defences are mostly about evaluation and data. Test on data from the deployment setting, including deliberately shifted sets such as different sites or time periods, and report performance by subgroup and by condition. Monitor inputs and outputs after deployment for drift (AI-420). Collect diverse training data, and use augmentation to break spurious correlations where they are known. Methods such as group distributionally robust optimization, which minimizes the worst loss over predefined groups, help when the spurious feature can be identified. Calibrated uncertainty and the ability to abstain on unfamiliar inputs make failures safer. No method guarantees robustness to shifts no one anticipated, which is why monitoring and human oversight remain part of any serious deployment.

Worked examples