When the loss will not go down

A training run that is not learning tells you almost nothing about why, and the temptation is to start changing the architecture. This section is the checklist to work through instead, in order: is the loss at step zero what uniform guessing would give, can the model overfit a single batch, does the data look right when you decode it back, is the learning rate within a factor of ten, and do gradients reach the parameters at all. Each step rules out a whole class of cause, and the answer is almost never the architecture.

In this order#

1. 1.Initial loss: close to log of the number of classes for classification? If not, check initialisation and labels.
2. 2.Overfit a single batch: train on the same small batch for a few hundred steps. The loss should go to nearly zero. If it cannot, the model or the loss is broken, and no amount of data will fix it.
3. 3.Look at what goes in: decode a batch back into text or images and look at it with the labels. Misaligned labels, wrong normalisation and empty inputs show up immediately.
4. 4.Learning rate: too high and loss jumps or explodes; too low and it barely moves. Try ten times higher and lower.
5. 5.Gradients: log the total gradient norm. Zero means something is detached or not trainable; enormous means instability, fixed with clipping, warmup or a lower rate.

Step 1: the loss at initialisation#

An untrained classifier should be uninformed, not wrong. Uninformed means equal probability on every class, and the cross-entropy of that is the natural log of the number of classes. This number is free, takes one forward pass, and rules out a surprising amount.

classes    expected initial loss     if you see instead
      2              0.6931              much higher: labels or initialisation
     10              2.3026              exactly 0: the model is predicting the label
     65              4.1744              nan: bad init, or a loss applied to a softmax
 50,257             10.8249              much lower: leakage, or the target is the input

Step 2: overfit one batch#

Take sixteen examples and train on only those for a few hundred steps. A correctly wired model must drive the loss to essentially zero, because memorising sixteen examples requires no generalisation at all. If it cannot, no quantity of data will help and the fault is in the model, the loss or the optimiser.

import torch, torch.nn as nn, torch.nn.functional as F

X = torch.randn(16, 20); Y = torch.randint(0, 5, (16,))   # labels are random noise

model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 5))
opt = torch.optim.AdamW(model.parameters(), lr=1e-2)
for _ in range(400):
    loss = F.cross_entropy(model(X), Y)
    opt.zero_grad(); loss.backward(); opt.step()
print(loss.item())          # 0.00009   memorised, as it must be

# Now the classic wiring bug: only some parameters handed to the optimiser.
model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 5))
opt = torch.optim.AdamW(model[0].parameters(), lr=1e-2)    # layer 2 will never update
# ... same loop ...
print(loss.item())          # 0.00569   60x higher, and still falling. Suspicious, not obvious.

Note how weak the failure signal is. The broken model still reached 0.00569, which on an unfamiliar problem reads as a success. The test is only decisive if you know what the healthy number looks like, which is why you run the healthy version first and write the number down.

Step 3: look at what actually goes in#

Decode a batch back into text or images and print it beside its label. This catches the bugs that no gradient statistic can: targets not shifted, images normalised twice, labels off by one from a category mapping, an empty string where a document should be, augmentation that has rotated the digits into different digits. It takes two minutes and it is the step people skip because it is not clever.

Step 4 and 5: learning rate and gradients#

x, y = batch(train)                      # one fixed batch
for step in range(300):
    loss = F.cross_entropy(model(x).view(-1, V), y.view(-1))
    opt.zero_grad(); loss.backward()
    norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)   # returns the norm before clipping
    opt.step()
    if step % 50 == 0:
        print(step, f"{loss.item():.3f}", f"grad norm {norm:.2f}")

The learning rate is swept by factors of ten, not by fine adjustment: try it a hundred times smaller and a hundred times larger and see which direction improves things. Anything finer is wasted before the order of magnitude is right. A useful default to start from is 3e-4 for AdamW on a transformer and 1e-3 for smaller networks.

Gradient norms distinguish causes that the loss curve cannot. Print them per parameter rather than in total, because a total norm of 0.4 is perfectly consistent with one layer receiving nothing.

for name, p in model.named_parameters():
    print(name, None if p.grad is None else round(p.grad.norm().item(), 5))

# 0.weight   0.41170        healthy
# 0.bias     0.13331
# 2.weight   0.82731
# 2.bias     0.35066

# with requires_grad = False on layer 0:
# 0.weight   None           <- not "small". None. It is frozen.

What you see                       Cause                             Fix
p.grad is None                     frozen, detached, or not in the graph  check requires_grad
p.grad is exactly 0.0              dead ReLU, or a saturated activation   lower lr, change init
early layers tiny, late layers big vanishing gradients                    residuals, normalisation
norm jumps 100x on some batches    a handful of extreme examples          clip, then inspect them
norm healthy, loss flat            the loss does not measure the task     re-read the loss

NaNs#

• The log of zero or a division by zero in a custom loss. Add a small epsilon, or use the built-in stable functions such as cross_entropy with logits.
• float16 overflow in mixed precision. Use bfloat16 where the hardware supports it, or a gradient scaler.
• An exploding learning rate. The loss rises for a few steps before becoming NaN; look back in the log.
• torch.autograd.set_detect_anomaly(True) reports the operation that produced the first NaN, at a large speed cost. Use it only while debugging.

The first of those is worth seeing, because it is why you are told to hand the loss raw logits rather than probabilities. A hand-written loss takes the log of a probability that is allowed to be exactly zero; the built-in one never forms the probability at all.

p = torch.tensor([0.0, 1.0])
print(torch.log(p)[0].item())          # -inf   then -inf * 0 in the backward pass is nan
print(torch.log(p + 1e-8)[0].item())   # -18.42  survivable

# The built-in version works in log space throughout and never forms p at all:
print(F.cross_entropy(torch.tensor([[-1e4, 1e4]]), torch.tensor([0])).item())   # 20000.0

When a NaN appears mid-run rather than at the start, it is almost always instability rather than a formula. Look back a few hundred steps in the log: an exploding run rises before it breaks, and the step where the loss first started climbing is the one to explain. If the loss was flat and then NaN in a single step, suspect a single pathological batch and print it.

if not torch.isfinite(loss):
    torch.save({"x": xb, "y": yb, "step": step}, "bad_batch.pt")
    raise RuntimeError(f"non-finite loss at step {step}")