Skip to content

v-code01/normnull

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

3 Commits
 
 
 
 
 
 
 
 
 
 

Repository files navigation

normnull

The backward pass of a normalization layer is an exact scaled orthogonal projector, and LayerNorm and RMSNorm differ by exactly one dimension: the mean.

Take a normalization layer with no affine and its Jacobian at a point x with normalized output yhat (so ||yhat||^2 = d). The backward pass multiplies the upstream gradient by this Jacobian, and it has an exact linear-algebraic form.

The two projectors

RMSNorm, y = x / rms(x):

J = (1/r) * (I - yhat yhat^T / d) .

Symmetric, rank d - 1, an orthogonal projector (times 1/r) onto the complement of span{x}. Its single null direction is x, the scale direction. Spectrum: 0 once, and 1/r with multiplicity d - 1.

LayerNorm, y = (x - mean(x)) / std(x):

J = (1/sigma) * (I - 11^T / d - yhat yhat^T / d) .

Since 1 and yhat are orthogonal and both normalize to unit length, the bracket is an exact orthogonal projector onto the (d - 2)-dimensional complement of span{1, yhat}. Rank d - 2, null directions x (scale) and 1 (the all-ones, mean, DC direction). Spectrum: 0 with multiplicity 2, and 1/sigma with multiplicity d - 2.

Verified from the autograd Jacobian at machine precision (fp64): both are symmetric to 1e-17, RMSNorm has rank d - 1 and one zero eigenvalue, LayerNorm has rank d - 2 and two zero eigenvalues, and the nonzero spectrum of each is a single value (spread < 1e-15). J @ x is zero for both; J @ 1 is zero only for LayerNorm.

The exact difference is the mean direction

Both normalizers annihilate the scale direction in their backward pass. LayerNorm annihilates one more direction, the all-ones vector, and RMSNorm does not. That is the entire difference between the two backward passes: a single dimension, the mean.

The consequence survives the affine gain

A real module applies an elementwise gain and bias after the norm. The gain multiplies the upstream gradient before it reaches the Jacobian, so it changes which vector J acts on but not J's row space. The input-gradient invariants therefore survive: the input gradient dx is always orthogonal to yhat, and for LayerNorm dx is additionally mean-free, sum(dx) = 0. So under LayerNorm every gradient the layer sends backward has zero mean across the feature dimension, a conserved quantity of the whole backward pass; under RMSNorm the mean component flows freely. Measured with the real affine gain and bias at width 2048 and 4096, on CPU and on the M4 GPU: LayerNorm's sum(dx) sits at the fp32 reduction floor (about 1e-5 to 0), while RMSNorm's is six orders larger (about 8 to 22); both dx . yhat are at the floor.

Layout

  • projector.py: the Jacobian structure of both norms from autograd, verified to machine precision (symmetry, rank, two-point spectrum, the scale and mean directions).
  • meanfree.py: the gain-invariant input-gradient consequence, LayerNorm mean-free and RMSNorm not, at production width on CPU and the M4 GPU.
  • test_normnull.py: the ranks, the null directions, the projector idempotency, the flat spectrum, and the mean-free gradient as tests.

Reproduce

python projector.py
python meanfree.py
python test_normnull.py

Takeaway

Choosing RMSNorm over LayerNorm, as Llama, Gemma, and Qwen do, keeps the DC gradient direction alive: the backward pass no longer forces every gradient to be mean-free. That single dimension is the exact optimizer-visible difference between the two normalizers, and it is a property of the Jacobian, not of any trained model.

About

A normalization layer's backward pass is an exact scaled orthogonal projector: RMSNorm rank d-1 (kills scale), LayerNorm rank d-2 (kills scale and the mean/DC direction). The difference is one dimension, the mean. Gain-invariant, measured on the M4 GPU: LayerNorm's input gradient is always mean-free, RMSNorm's is not.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages