LayerNorm, RMSNorm, and Residual Scale
Normalization prepares a token vector for the next computation. LayerNorm centers its features and rescales them; RMSNorm rescales them without centering. In a pre-norm block, each branch reads this normalized view while the residual path keeps the original vector.
Before each branch, and after the stack → normalization. View the model diagram →
Work within one token’s vector
A transformer has one vector per token. Normalization here works across the features inside one vector. It does not average “the”, “cat”, and “sat” together, and it does not mix examples in a batch.
For a tensor shaped $[B,n,d]$, that means $Bn$ independent calculations, each over $d$ features. A batch with only one example works the same way.
LayerNorm: center, then rescale
Start with $x=(1,2,3,4)$. For this hand calculation, set aside the learned scale and shift and the small constant $\epsilon$.
- Find the mean: $(1+2+3+4)/4=2.5$.
- Subtract it: the centered vector is $(-1.5,-0.5,0.5,1.5)$.
- Divide by the standard deviation: the mean squared centered value is $1.25$, so divide each entry by $\sqrt{1.25}$.
The result is approximately $(-1.342,-0.447,0.447,1.342)$. At this point it has mean zero and variance one. LayerNorm then applies a learned scale $\gamma_j$ and shift $\beta_j$ to each feature:
$\mu$ is the feature mean and $\sigma^2$ is the mean squared distance from it. The variance divides by $d$, not $d-1$. The learned $\gamma$ and $\beta$ are shared across token positions; the mean and variance are computed separately for every token.
After that learned scale and shift, the output need not have mean zero or variance one. LayerNorm also does not keep running batch statistics. See the Layer Normalization paper.
RMSNorm: rescale without subtracting the mean
For the same vector, the mean of the original squared features is $(1+4+9+16)/4=7.5$. RMSNorm divides by $\sqrt{7.5}$, giving approximately $(0.365,0.730,1.095,1.461)$. All four features stay positive because no mean was subtracted.
The common form uses a learned scale and no learned shift. The RMSNorm paper describes this alternative. The figure compares both operations on the same input.
Both use ε = 10⁻⁵, learned scale 1, and no learned shift. Each column follows one feature.
| Vector | 1 | 2 | 3 | 4 |
|---|---|---|---|---|
| Input | 1 | 2 | 3 | 4 |
| LayerNorm | −1.342 | −0.447 | 0.447 | 1.342 |
| RMSNorm | 0.365 | 0.730 | 1.095 | 1.461 |
Input mean = 2.5 · centered variance = 1.25 · mean square = 7.5
Ignoring $\epsilon$, both methods remove a positive rescaling of the input. Only LayerNorm also removes a shared additive shift. They control different properties of the vector.
The small constant changes the edge cases
If all features equal the same constant, LayerNorm's centered vector is zero and its output equals $\beta$. If all features are zero, RMSNorm returns zero in its scale-only form. A positive $\epsilon$ prevents division by zero and affects derivatives when the variance is tiny. Its placement inside the square root matters.
In low precision, computing the mean square or variance in higher precision can prevent loss of accuracy. A cast back to the model's activation dtype may follow. “Uses BF16” does not imply every internal reduction runs in BF16. Inspect the actual primitive before reproducing numerical behavior.
Pre-norm and post-norm have different paths
Pre-norm normalizes the input to a branch, computes the update, then adds it to the untouched residual vector. Post-norm computes and adds the update first, then normalizes the sum. Write the branch as $F$ and normalization as $N$:
The difference is which path passes through normalization. In pre-norm, the residual path bypasses it. In post-norm, normalization also transforms the residual sum. The block diagram follows the pre-norm arrangement.
Optional: what changes in the gradient path?
For pre-norm, the derivative is $I+J_FJ_N$ at the appropriate inputs. The residual path contributes an identity term. For post-norm, the derivative is $J_N(I+J_F)$, so normalization transforms the gradient through the sum too.
An identity term alone does not guarantee a well-behaved total gradient; the branch term can reinforce or cancel it. Initialization, learning rate, final normalization, and depth still matter. On Layer Normalization in the Transformer Architecture analyzes these training dynamics.
A reference implementation exposes the axes
import numpy as np
def layer_norm(x, gamma, beta, eps=1e-5):
mean = x.mean(axis=-1, keepdims=True)
variance = ((x - mean) ** 2).mean(axis=-1, keepdims=True)
return gamma * (x - mean) / np.sqrt(variance + eps) + beta
def rms_norm(x, gamma, eps=1e-5):
mean_square = (x ** 2).mean(axis=-1, keepdims=True)
return gamma * x / np.sqrt(mean_square + eps)Changing axis=-1 to include the sequence axis changes the function and can introduce dependencies between tokens, including future information in a supposedly causal model. The shape of gamma is not enough to establish the normalization axes; read the reduction itself.
Try it: Adding 100 to every feature leaves which normalized vector unchanged, ignoring epsilon and affine parameters?
LayerNorm: centering subtracts the added constant. RMSNorm generally changes because both its numerator and its root-mean-square denominator change without removing the offset.
Continue through depth
The block chapter places these operations in a complete forward pass. Attention Residuals later asks whether earlier layer outputs should be mixed with learned weights rather than a fixed sum.