Skip to content

v-code01/varsplit

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

3 Commits
 
 
 
 
 
 
 
 
 
 

Repository files navigation

varsplit

A Pre-LN transformer's residual-stream variance grows with depth as a stationary MLP term plus a rising attention term, so the attention and MLP contributions cross over and a single depth-uniform residual scale cannot keep them balanced.

At initialization a Pre-LN residual stream accumulates variance: each sublayer adds its output, V_{l+1} = V_l + delta_l, and the total grows about linearly with depth (measured V_L / L near 1 out to 64 layers). The interesting part is where the increment comes from.

The split

  • The MLP increment is stationary. The MLP is a per-token map; it does not see how tokens relate, so its variance contribution has no depth trend. Measured across 24 seeds, its slope against layer index is +0.0003 with a 95% interval of [-0.0013, +0.0018], straddling zero.
  • The attention increment rises. The attention output is a mixture across tokens, so its variance depends on how correlated the tokens have become. At init, attention drives the tokens toward a shared direction (rank collapse), so the mean pairwise token correlation climbs from 0 toward 0.8 over depth, and the attention increment climbs with it. Its slope is +0.010 with a 95% interval of [+0.007, +0.013], strictly positive, and it correlates with the token correlation across layers.

The contribution ratio inverts

Because one term is flat and the other rises, the two cross over. The MLP increment dominates the first few layers (about 5x the attention increment at layer 1) and the attention increment dominates the deep layers (about 3x the MLP increment by layer 48). The per-layer variance is built mostly by MLPs early and mostly by attention late.

Why depth-uniform residual scaling is mis-calibrated

Schemes that keep the forward pass at unit scale by shrinking the residual branches (a single beta in DeepNet, a symmetric 1/sqrt(2L)) apply the same factor to every sublayer at every depth. A single scale alpha multiplies every contribution by alpha^2, so it cannot flatten a contribution that is itself non-stationary: the late attention increment stays about 2.4x the early one whatever alpha is. Keeping the per-layer variance contribution constant needs a layer-dependent scale, alpha_l proportional to 1 / sqrt(delta_attn(l)), and it has to differ between the attention branch, which rises, and the MLP branch, which does not. A uniform residual scale tuned at shallow depth silently over- or under-damps one branch as the network gets deeper.

Layout

  • varstream.py: an initialized Pre-LN causal stack, its per-sublayer variance increments, and the token correlation that drives the attention term, all seed-averaged.
  • trends.py: the depth-slope of each increment with confidence intervals (MLP stationary, attention rising), the link between the attention increment and the token correlation, and the contribution ratio a uniform scale cannot remove.
  • test_varsplit.py: the linear growth, the stationary-versus-rising split, the correlation link, and the ratio inversion as tests.

Reproduce

python varstream.py
python trends.py
python test_varsplit.py

Takeaway

The residual-stream variance of a Pre-LN transformer is not built uniformly with depth. Its MLP part is depth-flat and its attention part rises with rank collapse, their roles invert, and any residual-branch scaling that is uniform in depth is calibrated for one regime and wrong for the other.

About

A Pre-LN transformer's residual-stream variance grows as a stationary MLP term plus a rising attention term (driven by rank-collapse token correlation): the two invert, MLP-dominated early and attention-dominated late, so a depth-uniform residual scale (DeepNet, 1/sqrt(2L)) provably cannot keep the per-layer variance balanced. Slopes with seed CIs.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages