Guides / Matrix calculus

Layer norm backward: (g′ − mean(g′) − x̂ mean(g′ ⊙ x̂))/σ, and the two directions the loss can't see

Layer normalization centers each token's features, divides by their standard deviation, then scales by γ and shifts by β. The gradients of γ and β take one line each. The gradient to the input is where the work is, because the mean μ and the standard deviation σ are themselves functions of every feature. The result fits on one line: the gradient arriving at x̂, minus its mean, minus a multiple of x̂, all over σ. The conventions are those of the matrix calculus conventions guide, with x one token as a column. The batched version follows the batched backprop guide, with rows as tokens. The mistakes come from treating μ and σ as constants, from dividing only one term by σ, and from leaving γ's gradient unsummed over a batch.

Short answer

For x ∈ ℝⁿ: μ = mean(x), v = mean((x − μ1)²), σ = √(v + ε), x̂ = (x − μ1)/σ and y = γ ⊙ x̂ + β. Given g = ∇_yL: ∇_γL = g ⊙ x̂ and ∇_βL = g. With g′ = γ ⊙ g, which is ∇_x̂L:

∇_xL = (g′ − mean(g′)1 − x̂ mean(g′ ⊙ x̂))/σ.

It always sums to 0, and it is orthogonal to x̂ up to a term of order ε. In a batch each row uses its own μ and σ, and the gradients of γ and β sum over the rows.

Why it works

The Jacobian of x̂ with respect to x is (I − 11^T/n − x̂x̂^T/n)/σ. The I is the direct path. The 11^T/n term comes from μ: every feature moves the mean by 1/n, and the centering subtracts that move from every entry. The x̂x̂^T/n term comes from σ, whose gradient is x̂/n. The Jacobian is symmetric, so applying it to g′ gives the formula, and each outer product turns into a mean: 11^Tg′/n = mean(g′)1.

The two subtractions say what the loss can't see. Adding a constant to every feature leaves x̂ unchanged, so the gradient has no component along 1. Scaling the centered features would also leave x̂ unchanged if ε were 0, so the gradient along x̂ is only of order ε. That gives two checks on any answer: its entries sum to 0, and its dot product with x̂ is close to 0.

Example 1: the forward pass

Take one token with four features, x = (4, 0, 4, 0), and γ = (1, 0.5, 2, 1), β = (0, 0, 0, 1). Ignore ε for now. Find μ, v, σ, x̂ and y.

μ = 8/4 = 2The mean of the four features.
x − μ1 = (2, −2, 2, −2)The centered features sum to 0.
v = (4 + 4 + 4 + 4)/4 = 4, σ = 2The variance divides by n, not n − 1.
x̂ = (1, −1, 1, −1)Mean 0, and ‖x̂‖² = 4 = n.
y = γ ⊙ x̂ + β = (1, −0.5, 2, 0)Entry by entry: 1 × 1, 0.5 × (−1), 2 × 1, 1 × (−1) + 1. With ε = 10⁻⁵, σ becomes 2.0000025 and x̂ changes in the sixth digit.

Example 2: the gradients of γ and β

The layer above sends back g = (2, 0, 0.5, 1). Find ∇_γL, ∇_βL and g′.

y_i = γ_ix̂_i + β_iEach γ_i and β_i appears in one output only.
∇_γL = g ⊙ x̂ = (2, 0, 0.5, −1)The coefficient of γ_i is x̂_i, not y_i.
∇_βL = g = (2, 0, 0.5, 1)The coefficient of β_i is 1.
g′ = γ ⊙ g = (2, 0, 1, 1)The gradient at x̂: ∂y/∂x̂ = diag(γ), so it is an elementwise product. From here on only g′ is needed.

Example 3: the gradient to the input

Continue to ∇_xL, and check it against the two directions the loss can't see.

mean(g′) = 4/4 = 1The μ term.
g′ ⊙ x̂ = (2, 0, 1, −1)For the σ term …
mean(g′ ⊙ x̂) = 2/4 = 0.5… which is x̂^Tg′/n.
g′ − 1 · 1 = (1, −1, 0, 0)Subtract the mean from every entry.
− 0.5 x̂ → (0.5, −0.5, −0.5, 0.5)0.5 x̂ = (0.5, −0.5, 0.5, −0.5).
∇_xL = (0.25, −0.25, −0.25, 0.25)Divide everything by σ = 2.
sum 0, dot with x̂ 0Both checks pass. With ε = 10⁻⁵ the dot product is 2.5 × 10⁻⁶, the size the ε term predicts. Central finite differences of g^Ty agree with every entry, with and without ε.

Example 4: a batch of two tokens

Add a second token, x = (3, −1, −1, 3), with upstream gradient (1, 1, 0, 2), and keep γ. Find its row of ∇_XL, then ∇_γL and ∇_βL for the batch.

μ = 1, σ = 2, x̂ = (1, −1, −1, 1)This row's own statistics: centered (2, −2, −2, 2), variance 4.
g′ = (1, 0.5, 0, 2)γ ⊙ g for this row.
mean(g′) = 0.8753.5/4.
mean(g′ ⊙ x̂) = 0.625(1 − 0.5 − 0 + 2)/4.
(g′ − 0.875 − 0.625x̂)/2This row of ∇_XL. No other row enters: each token is normalized alone.
= (−0.25, 0.125, −0.125, 0.25)It sums to 0, and its dot product with x̂ is 0.
(2, 0, 0.5, −1) + (1, −1, 0, 2)γ is shared by both rows, so their g ⊙ x̂ add …
∇_γL = (3, −1, 0.5, 1)… in matrix form (G ⊙ X̂)^T1, length n, the shape of γ.
∇_βL = (3, 1, 0.5, 3)G^T1, the column sums of G. Finite differences of the batch loss agree.

Common mistakes, and the exact line they happen on

1. Treating μ and σ as constants

The input gradient of Example 3.

g′ = (2, 0, 1, 1), σ = 2Right so far.
∇_xL = g′/σ = (1, 0, 0.5, 0.5)This is the mistake. Per feature, layer norm looks like "subtract a number, divide by a number". But μ and σ move when x moves, and those paths give the other two terms. This answer sums to 2, so it claims that shifting every feature by the same amount changes the loss, which it can't.

2. Dividing only the first term by σ

The same gradient, written out term by term.

mean(g′) = 1, mean(g′ ⊙ x̂) = 0.5Right so far.
g′/σ − 1 − 0.5x̂This is the mistake. The 1/σ multiplies the whole Jacobian, so every term. This gives (−0.5, −0.5, −1, 0), which sums to −2 instead of 0. It agrees with the right answer only when σ = 1, so it looks fine on rows that are already close to unit variance.

3. γ's gradient left as a matrix

The batch of Example 4.

row k contributes g^(k) ⊙ x̂^(k)Right so far.
∇_γL = G ⊙ X̂This is the mistake. That is 2 × 4, one row per token, but γ has length 4. The input gradient really is the per-token formula row by row, because each token has its own row of X; γ is shared, so its rows add to (3, −1, 0.5, 1). Array code may broadcast the update and quietly give each token its own γ.

Practice

Find μ, v and x̂ for x = (1, 3, 1, 3), ignoring ε.
Show answer
μ = 2, v = 1, x̂ = (−1, 1, −1, 1).
For Example 1's token, the upstream gradient is g′ = (1, 1, 1, 1) at x̂. What is ∇_xL, and why?
Show answer
0. A gradient that is the same in every entry asks for a shift, and a shift doesn't change x̂: g′ − mean(g′)1 = 0 and mean(g′ ⊙ x̂) = 0.
For the same token, g′ = x̂ = (1, −1, 1, −1). What is ∇_xL with ε = 0?
Show answer
0. mean(g′) = 0 and mean(g′ ⊙ x̂) = 1, so the bracket is x̂ − x̂. Asking for a larger x̂ asks for a rescale, which layer norm undoes.
A batch has upstream rows (1, 0, 2), (0, 1, 1) and (−1, 1, 0). Find ∇_βL.
Show answer
(0, 2, 3), the column sums of G.
RMSNorm drops the mean: x̂ = x/r with r = √(mean(x²)), and ∇_xL = (g′ − x̂ mean(g′ ⊙ x̂))/r. Find it for x = (2, −2, 2, −2) and g′ = (1, 0, 0, 0). Does it sum to 0?
Show answer
r = 2, x̂ = (1, −1, 1, −1), mean(g′ ⊙ x̂) = 0.25, so ∇_xL = (0.375, 0.125, −0.125, 0.125). It sums to 0.5: without the centering, a shift does change the output.