What breaks when you stack 96 blocks
The pieces are all in hand: attention, position, several heads, feed-forward. Now they get tied into a block and stacked. Tying them takes two lines, and those two lines are the difference between twenty layers and ninety-six.
x = LN(x + Attn(x)); x = LN(x + FFN(x)) # Post-LN, the original paper
x = x + Attn(LN(x)); x = x + FFN(LN(x)) # Pre-LN, the current default
The only difference is whether the norm sits outside or inside the
residual. Do what part four did: stack them and measure the gradient at each
layer. d=128, four heads, 4x expansion, PyTorch’s default initialisation, in
float64.
At twenty layers, nothing happens
layer 1 grad layer 20 grad ratio (seed 0) ratio (median of 6)
Post-LN 1.78e-01 3.45e-01 0.516 0.595
Pre-LN 1.89e-01 1.15e-01 1.646 1.618
no residual 2.99e-01 9.86e-01 0.303 0.360
All three sit inside a factor of three, removing the residual connections
included, and changing the seed even reorders them - nothing next to the 1e9
spread part four got out of a bad initialisation.
Part six says why. Normalisation already resets the scale at every layer.
With an LN in place the forward pass does not die even without residuals. So
a twenty-layer stack cannot tell these three apart.
The depth has to go up.
Up to ninety-six
The number below is the first layer’s gradient divided by the last layer’s. At
1 the two ends match; below it the early layers starve, above it the late
ones do. Median of six seeds.
L Post-LN Pre-LN no residual
6 0.671 1.20 0.773
24 0.534 1.69 0.432
48 0.372 2.23 0.245
96 0.306 2.89 0.031
The three rows tell different stories.
Only the no-residual stack breaks. From 0.773 to 0.031: at ninety-six
layers the first block receives a thirty-second of what the last one does. And it
accelerates - it holds at 0.245 through forty-eight layers and falls away
sharply after. An LN keeps the forward pass alive but does not keep the
backward pass balanced across layers.
Post-LN and Pre-LN tilt in opposite directions. Post-LN’s 0.306 starves
the early layers by 3.3x; Pre-LN’s 2.89 starves the late ones by 2.9x. Both grow
with depth, and both are still inside a factor of three at ninety-six layers.
So what carries depth is the residual connection, and where the norm goes is the dial that tilts the gradient one way or the other on top of it. They are an order of magnitude apart.
Post-LN starving its early layers is the reason learning-rate warmup is known to be necessary for it. Note though that what is measured here is only the gradient profile at initialisation. How it changes as training proceeds is a question this experiment does not answer.
The residual stream grows
Pre-LN comes with one more property. Every block only ever adds to the stream, so the stream keeps growing.
L layer 1 sd last sd factor sqrt(L)
6 1.04 1.21 1.17 2.45
24 1.03 1.91 1.85 4.90
96 1.04 3.59 3.46 9.80
The growth is usually quoted as sqrt(L): add L independent things and the
variance multiplies by L, so the standard deviation multiplies by sqrt(L).
Measured, it is 3.46, not 9.80 - an exponent of 0.272, about half of
0.5.
The premise does not hold. sqrt(L) is the calculation for terms comparable
in size to the stream, and at default initialisation a sublayer’s output is
far smaller than the stream. Checking it is easy: scale up the initialisation of
Wo and W2.
output scale 1 factor 3.46 exponent 0.272
output scale 3 factor 9.21 exponent 0.487
output scale 10 factor 14.19 exponent 0.581
At three times the scale the exponent is 0.487, right on 0.5. sqrt(L) is
not wrong but conditional, and at real initialisation the condition fails.
The stream does grow all the same, which is why a Pre-LN network puts one more
LN at the very end. The experiment above includes it. Without it the output
leaves at whatever size the stream reached.
So
- Twenty layers distinguish nothing. Normalisation resets the scale every layer, so even a residual-free stack looks healthy
- Ninety-six layers separate them. Without residuals the first layer’s gradient
is
1/32of the last’s, and it turns sharp past forty-eight - Post-LN and Pre-LN tilt opposite ways,
0.306against2.89. Opposite in direction but both inside a factor of three - a different order of magnitude from the residual’s thirty-two - The Pre-LN stream grows, but not by
sqrt(L). Measured exponent0.272; scale the sublayer outputs by three and it becomes0.487, matching the calculation - What is measured is the gradient profile at initialisation. Training is a separate question
Next time all of it runs at once. Having looked at the pieces and the wiring, the smallest thing that actually predicts characters gets built end to end.
Comments