DEV Community

Devanshu Biswas
Devanshu Biswas

Posted on

A Layer Norm's Jacobian Is a Projector, So 27 Post-LN Blocks Attenuate the Gradient Once and by 0.6676, Not 27 Times

This page was written to say that Post-LN decays the gradient geometrically with depth. The measurement killed it, and the reason is a line of algebra rather than a training run. A layer norm's Jacobian is 1/σ times an orthogonal projector, and a projector is idempotent: apply it once and the mean and the radial component are gone, apply it twenty-seven more times and nothing further leaves.

The control needs no training at all. Set every branch output matrix to zero, so each block is a literal no-op, and read the gradient arriving at block 0 against the gradient at the top.

Pre-LN, branches deleted       1.0000    deviation 0.0e+0
bare residual, deleted         1.0000    deviation 0.0e+0
Post-LN, branches deleted      0.6676    27 interior blocks move it 5.0e-16
Enter fullscreen mode Exit fullscreen mode

The attenuation is real and it fires once. Put the branches back and nothing vanishes anywhere either. At L = 32, the gradient reaching block 0 relative to the top:

placement ‖g₀‖/‖g_L‖ at L = 32
Pre-LN 5.37x
Post-LN 4.66x
no normaliser 35.51x

Every one is above 1, because each block's branch adds to the stream gradient on the way down. The diagnostic everybody reaches for is hunting a vanishing gradient in a place that does not have one.

Then the claim this page came to kill. There are two ways to start a residual branch small: multiply the branch by c, or initialise its output matrix at c times its usual scale. The branch is linear in that matrix, so at initialisation these are the same function - powers of two agree bit for bit at 0.0e+0, everything else to 6.2e-16, which is the number to quote since engines round differently. They are not the same experiment. Fitted as a slope, the function-space step scales as c^2.0000000000000018 under SGD and c^1.0000004 under Adam, whose update is scale-invariant. At c = 1/8 the two parameterisations move the function by 8x different amounts under Adam and 64x under SGD from an identical start, and after 150 Adam steps the init-scaled branch is 3.35x the larger. Shrinking the initialisation does not shrink the updates at all.

What the measurement contradicted

I had "check the gradient norms at initialisation" written up as the diagnostic worth keeping. Scored as a ranker over 90 trained configurations it gets AUC 0.582 against a coin's 0.500 - beaten by one optimiser step (0.774), by ten (0.776), by the training loss you were going to watch anyway (0.886), and by α√L, a closed form needing no data, no forward pass and no training, at 0.604. "Deeper is better" scores 0.396, below the coin.

The number is not broken; the comparison was. Hold the placement and the learning rate fixed, so the ratio only arbitrates comparisons it can actually see, and the same ratio scores 0.802 over 405 pairs.

And the recommendation is undercut by its own noise floor: seed spread 0.0657 against a held-out sample spread of 0.0162, so of three placements where α = 1/√L looks like it helps, only 2 clear the spread of the runs the claim is made from.

The data is a fixed teacher network, so p*(y|x) is a forward pass rather than an estimate and the held-out score is an exact mean KL - 0 for the teacher, floor 0.7685 nats over 4 classes. Gradients checked against independently written central differences to 3.2e-8, 67 in-page assertions and 1,617 in Node. One file, no build step, fetching nothing off itself: https://dev48.infy.uk/dl/day72-residual-scaling-and-depth-stability.html

One deep-learning idea a day, computed rather than quoted: https://dev48.infy.uk/deeplearningfromzero.php

Top comments (0)