A Transformer starts training and the loss drops for a few dozen steps, then
stops. It sits flat at a number suspiciously close to the cost of guessing
uniformly — ln(vocab_size), which is 3.611 for a 37-character vocabulary and
around 10.8 for a 50k token vocabulary. Samples are pure noise, and the predicted
distribution is nearly flat at every position. The mirror version of the same
bug is louder: the loss climbs, then prints nan or a number like 2.327e+10,
and the run dies. A smaller copy of the same code, two blocks instead of six,
trains fine, which sends the team hunting for a data bug that is not there.
Changing the seed changes nothing. Adding more steps changes nothing.
A deeper stack than you have trained before: six or twelve blocks instead of two, or the same depth at a new learning rate after somebody "tuned" it. Module 2.8's ablation runs reproduce every version of this on a laptop CPU in five seconds per run. Three ingredients combine to produce it, and they are usually blamed one at a time. Learning rate: at 0.001 all variants trained to within 0.04 of each other, at 0.01 the pre-norm model read 0.256 while the post-norm model read 2.892 and the no-norm model read 1.434, and at 0.03 the no-norm run reached 2.327e+10. Norm placement: post-norm blocks — the original 2017 arrangement, where the norm sits after the residual add — need a warmup schedule to survive rates that pre-norm blocks tolerate directly. Missing pieces: with the residual connections removed at six blocks, the gradient reaching block 1 fell from 5.951e-02 to 1.785e-07 and the loss parked at 2.982, which is the flat-prediction signature with a different cause. Linked modules: 2.8, 2.4 (making training work), 2.2 (blame flow), 2.7.
A loss pinned at ln(vocab_size) means the model is predicting every option with
equal probability, so the first job is to decide whether it cannot learn or is
not being told. Run the uniform-prediction check: take one batch, print the
softmax over the output at a few positions, and print the entropy. Entropy at or
near ln(vocab_size) with a flat distribution says the output layer is producing
nothing, so the fault is upstream of it. Next, print the gradient norm per layer,
lowest block first, exactly as 2.7 printed it per time step. A first block
receiving something like 1.785e-07 while the last block receives 1e-01 is a
broken path, and the residual connections are the first thing to check, since
they are the road the gradient travels. If gradients look healthy and the loss
still will not move, the fault is the step size: halve the learning rate and run
three hundred steps. If the loss now falls, you were overshooting; if it climbs
or prints nan, you were exploding. Then vary depth as the deciding experiment.
Two blocks training fine while six blocks stall is the clearest possible evidence
that the problem is depth-related — norm placement, residual path or warmup — and
not the data, the tokeniser or the corpus. Check the initialisation last: weights
scaled far from the usual small values will reproduce all of these symptoms.
Move to a known-good recipe and change one knob at a time from there. The recipe:
pre-norm blocks, with layer normalisation inside each branch before attention and
before the feed-forward, residual connections on both branches, a final norm
before the output head, AdamW, and a learning rate near 3e-3 for a tiny model or
near 3e-4 for a large one, with a few hundred warmup steps. On module 2.8's
six-block model that recipe read 0.256 while post-norm read 2.892 at the same
rate. If you must keep post-norm to match a paper or a checkpoint, add linear
warmup over the first several hundred steps and lower the peak rate, because
post-norm without warmup is the classic version of this failure. Verify the
residual path with the gradient-per-layer print before you touch anything else —
without residuals no learning rate will save you. Add gradient clipping at a norm
near 1.0 to convert the explosive version into a survivable one, but do not
mistake clipping for a fix: it controls the blow-up and does nothing for a stall.
Keep ln(vocab_size) printed beside the loss so "flat" is visible instead of
inferred. Change one thing per run and record depth, learning rate, norm
placement, final loss and the gradient into block 1 in a small table, the way
module 2.8's ablation table is laid out.
Start every new architecture at two blocks and a gentle learning rate, get a
falling loss, then scale depth — a stall that appears when you add layers is a
different bug from one that was there all along. Print ln(vocab_size) at the
top of every log so a flat model is obvious in one glance. Keep the
gradient-norm-per-layer print in your training script permanently, guarded behind
a flag, because it costs nothing and answers the hardest question in this failure.
Prefer pre-norm for anything you are writing yourself, and if you copy post-norm
from a paper, copy its warmup schedule with it. Never tune two knobs in one run;
2.4's discipline about changing one thing at a time is worth more at depth six
than it was at depth two.
Q: "Your six-layer Transformer's loss sticks at the entropy of a uniform guess
and never moves, but the two-layer version trains. Walk me through it." Strong
answer: name the number first — a loss at ln(vocab_size) means uniform
predictions — then say the depth clue rules out the data. Give the two
measurements in order: entropy of the output distribution, and gradient norm per
layer from the lowest block up. Name the three suspects and how each is
distinguished: broken residual path, which shows a collapsed gradient into block
1; norm placement, which shows up as post-norm surviving 1e-03 and dying at
1e-02; and learning rate, which is settled by halving it. Finish with the fix
ladder and one honest caveat — gradient clipping saves an exploding run and does
nothing at all for a stalled one.