Debugging neural network training

The big idea

A normal program with a bug usually crashes or prints something obviously wrong. A neural network with a bug usually trains anyway. The loss goes down, a little slower than it could; accuracy is decent, a few points below what it should be. Shuffled labels, a forgotten normalization step, a wrong axis in a softmax, or a learning rate ten times too high can all produce a model that looks plausible. Neural network bugs fail silently.

So debugging training is less about reading stack traces and more like a pilot's pre-flight checklist: a fixed sequence of cheap checks, each of which rules out a whole class of problems before you spend hours (or GPU-days) on a full run. The single most useful check is to make the model memorize a tiny batch. If it cannot, something is broken, and no amount of tuning will fix it.

How it works

Step 1: Look at the data first

Before any model code, look at actual examples, after every preprocessing step, exactly as the model will receive them. Print a batch: are images the right way up, are pixel values in the range you expect, do labels match inputs, are classes balanced? Many "model" bugs are data bugs: labels misaligned after a shuffle or join, images normalized twice, text truncated, padding counted as content. In the code, shuffling labels against inputs gives a model that reaches 47.2 percent on its own training labels by memorizing, and 11.1 percent on real validation data, which is chance for 10 classes.

Step 2: Check the loss at step 0

Before training, you know what the loss should be. A 10-class classifier with small random output weights predicts roughly uniform probabilities, so its cross-entropy should be about -ln(1/10) = ln(10) = 2.303. In the code it is 2.301. If it were 15, something (output scale, label encoding, a missing softmax) is off. Also check that a simple baseline, such as predicting the most common class, gives the score you expect.

Step 3: Overfit one small batch

Take 8 to 32 examples and train on only them, with no regularization, for a few hundred steps. A correct pipeline drives the loss to nearly zero: in the code, from 2.282 to 0.0006 on 8 examples. If the loss stalls, the bug is in the model, loss or optimizer wiring, not in the data volume or hyperparameters. This test takes seconds and catches most plumbing errors: detached tensors, wrong loss inputs, frozen parameters, gradients not reaching a layer.

Step 4: Read the loss curves

Always plot training and validation loss together. The shape tells you which knob to turn: a slow straight decline means a learning rate too low; a noisy or rising curve means too high; a widening gap means overfitting; a flat line at the step-0 value means nothing is learning. Log the gradient norm too: a sudden spike just before a loss spike points to exploding gradients ( Vanishing and exploding gradients, residual connections ).

Step 5: Chase NaNs and infinities

A loss of NaN is a crash in disguise. Common causes: a learning rate too high, log of zero or division by zero in a hand-written loss, exploding gradients, unscaled inputs, or overflow in low precision. The fix is to find the first step and the first layer where a non-finite value appears: check inputs and labels for NaN, assert finiteness after each layer on a failing batch, lower the learning rate, add gradient clipping, and use the framework's numerically stable fused losses (cross-entropy from logits rather than log of a softmax). In the code, the same network on raw 0 to 16 pixel values spikes to a loss of 24.9 and ends stuck at 2.3, chance level; with standardized inputs it trains to 100 percent training accuracy. A learning rate of 50 on standardized inputs ends at a loss of 24.4.

Step 6: Change one thing at a time

Start from a small, known-good setup, often a published recipe or a simple baseline, and add complexity one piece at a time, checking that each addition helps. When something breaks, you know which change did it. Fix random seeds while debugging so runs are comparable, and keep a log of every experiment's settings and results.

Going deeper

Code: four checks on a digits classifier

import numpy as np
from sklearn.datasets import load_digits

Xraw, y = load_digits(return_X_y=True)                  # pixel values 0 to 16
X = (Xraw - Xraw.mean(0)) / (Xraw.std(0) + 1e-8)       # standardized
Xtr, ytr, Xva, yva = X[:1400], y[:1400], X[1400:], y[1400:]

def train(X, y, steps, lr, seed=0):
    """64-32-10 MLP, softmax cross-entropy, full-batch gradient descent. Returns losses, params."""
    rng = np.random.default_rng(seed)
    W1, W2 = rng.normal(0, np.sqrt(2 / 64), (64, 32)), rng.normal(0, 0.01, (32, 10))
    losses = []
    for _ in range(steps):
        h = np.maximum(0, X @ W1)
        z = h @ W2; z -= z.max(1, keepdims=True)
        p = np.exp(z) / np.exp(z).sum(1, keepdims=True)
        losses.append(-np.mean(np.log(p[np.arange(len(y)), y] + 1e-12)))
        d = p.copy(); d[np.arange(len(y)), y] -= 1; d /= len(y)
        gW2, gW1 = h.T @ d, X.T @ ((d @ W2.T) * (h > 0))
        W1, W2 = W1 - lr * gW1, W2 - lr * gW2
    return losses, (W1, W2)
acc = lambda P, X, y: np.mean((np.maximum(0, X @ P[0]) @ P[1]).argmax(1) == y)

# Check 1: the loss at step 0 should be about -ln(1/10) for 10 balanced classes
L, _ = train(Xtr, ytr, 1, 0.1)
print(f"initial loss {L[0]:.3f}, expected ln(10) = {np.log(10):.3f}")

# Check 2: overfit a single small batch; a correct pipeline drives the loss towards 0
L, _ = train(Xtr[:8], ytr[:8], 300, 0.5)
print(f"overfit 8 examples: loss {L[0]:.3f} -> {L[-1]:.4f}")

# Check 3: the same settings on raw, unscaled pixels, and a learning rate far too high
with np.errstate(all="ignore"):
    L, _ = train(Xraw[:1400], ytr, 200, 0.5)
    print(f"raw pixels, lr 0.5: loss {L[0]:.3g} -> {L[-1]:.3g} after 200 steps (max {max(L):.3g})")
    L, _ = train(Xtr, ytr, 50, 50.0)
    print(f"standardized, lr 50: loss after 50 steps {L[-1]:.3g}")

# Check 4: a data bug, labels shuffled against the inputs
L, P = train(Xtr, ytr, 500, 0.5)
print(f"\ncorrect labels:  train acc {acc(P, Xtr, ytr):.3f}  val acc {acc(P, Xva, yva):.3f}")
yshuf = np.random.default_rng(1).permutation(ytr)
L, P = train(Xtr, yshuf, 500, 0.5)
print(f"shuffled labels: train acc {acc(P, Xtr, yshuf):.3f}  val acc {acc(P, Xva, yva):.3f}  (chance is 0.100)")

initial loss 2.301, expected ln(10) = 2.303
overfit 8 examples: loss 2.282 -> 0.0006
raw pixels, lr 0.5: loss 2.29 -> 2.3 after 200 steps (max 24.9)
standardized, lr 50: loss after 50 steps 24.4

correct labels:  train acc 1.000  val acc 0.927
shuffled labels: train acc 0.472  val acc 0.111  (chance is 0.100)

A checklist

1. Inspect raw and preprocessed examples; check label alignment and class balance.
2. Compute simple baselines (majority class, linear model) to know what "good" means.
3. Verify the initial loss against its expected value.
4. Overfit one small batch to near-zero loss.
5. Train the full model; plot training and validation curves and gradient norms from the start.
6. If it underfits: bigger model, longer training, higher learning rate, fewer regularizers. If it overfits: more data, augmentation, regularization, early stopping.
7. Only then tune hyperparameters, preferably with random search ( Hyperparameter tuning ).
8. Look at the model's worst errors on validation data; they often reveal label noise or data problems.

Mode and pipeline bugs

A family of bugs comes from training and evaluation behaving differently: forgetting model.eval() (dropout and batch norm stay in training mode), preprocessing validation data differently from training data, or leaking validation examples into training. Another family comes from the data pipeline: augmentations applied to labels incorrectly, shuffling that is off, or a loader that silently repeats the same batch. When validation numbers are much better or much worse than expected, suspect these first.

Trade-offs

Common mistakes

• Starting with a large model and full dataset before checking that a tiny version can learn.
• Watching only training loss, so overfitting goes unnoticed.
• Changing several hyperparameters at once and not knowing which one mattered.
• Ignoring a step-0 loss that is far from its expected value.
• Trusting the data pipeline without ever looking at a batch as the model sees it.

In the wild

Check yourself

Recap