보이는 딥러닝 part 12 of 13

What breaks when you stack 96 blocks

guide / / 4 sections

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
1e0 1e-1 1e-2 1 25 49 73 96 block gradient size no residual Post-LN Pre-LN
Gradient magnitude per block in a 96-block stack. Drop the residual connections and the early blocks sink to 0.026, a thirty-third of the last block. With residuals, putting the norm before or after leaves the profile flat within a factor of three.

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/32 of the last’s, and it turns sharp past forty-eight
  • Post-LN and Pre-LN tilt opposite ways, 0.306 against 2.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 exponent 0.272; scale the sublayer outputs by three and it becomes 0.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