Guides / Matrix calculus

Batch norm backward: (dx̂ − mean(dx̂) − x̂ mean(dx̂ ⊙ x̂))/√(σ² + ε), one feature at a time down the batch

Batch normalization standardizes each feature using the mean and variance of that feature over the minibatch, then scales by γ and shifts by β. The gradients of γ and β are sums over the batch. The gradient to the input is where the work is, because the batch mean μ and variance σ² are computed from every example in the batch, so every example's output depends on every other example's input. The result is the layer-norm formula turned on its side: for each feature, the gradient arriving at x̂, minus its mean over the batch, minus a multiple of x̂, all over √(σ² + ε). The notation is that of the practice set: X is N × d with rows as examples, and dA is the gradient of the loss with the shape of A, as in the batched backprop guide. The layer norm backward guide derives the same formula along the other axis. The mistakes come from holding the batch statistics constant, from taking the mean over the wrong axis, and from differentiating a variance that divides by N − 1 when the forward pass divided by N.

Why it works

Take one feature, a column x of N numbers. Its forward pass is a small graph: μ from x, σ² from x and μ, and x̂ from all three. The gradient reaching σ² is dσ² = −Σ_n dx̂_nx̂_n / (2(σ² + ε)), and the gradient reaching μ is dμ = −Σ_n dx̂_n / √(σ² + ε). The path from μ through σ² carries nothing, because deviations from the mean sum to 0. Each x_n then collects three terms: its direct path to x̂_n, a share dμ/N of the mean, and a share dσ² · 2(x_n − μ)/N of the variance. Put over one denominator, those three terms are the short answer: the μ path becomes the mean of dx̂, and the σ² path becomes the multiple of x̂.

It is the layer norm gradient with rows and columns swapped. Layer norm takes its statistics along each row, over one example's features; batch norm takes them down each column, over one feature's examples. So X̂ = LN(X^T)^T, and the two checks carry over: shifting one feature by the same amount on every example leaves X̂ unchanged, so each column of dX sums to 0; scaling it leaves X̂ unchanged when ε = 0, so each column of dX is orthogonal to that column of X̂ up to a term of order ε.

Example 1: the forward pass

A batch of four examples with two features, X = 01656185, with γ = (2, 0.5) and β = (1, 0). Take ε = 0 for now. Find μ, σ², X̂ and Y.

1. μ = (20/4, 12/4) = (5, 3)Column means: the mean over the four examples for each feature.
2. column 1 − 5 = (−5, 1, 1, 3)The deviations of feature 1; they sum to 0.
3. σ₁² = 36/4 = 925 + 1 + 1 + 9 = 36, divided by N = 4, not by N − 1.
4. column 2 − 3 = (−2, 2, −2, 2)Feature 2's deviations.
5. σ₂² = 16/4 = 4So the two features have standard deviations 3 and 2.
6. X̂ = −5/3−11/311/3−111Each column divided by its own standard deviation. Each column of X̂ has mean 0 and squared length N = 4.
7. Y = −7/3−0.55/30.55/3−0.530.5Column 1 is 2x̂ + 1 and column 2 is 0.5x̂: one scale and one shift per feature, shared by every example.

Example 2: the gradients of γ, β and X̂

The layer above sends back dY = 110220−11. Find dγ, dβ and dX̂.

1. Y_nj = γ_jX̂_nj + β_jγ_j and β_j appear in all N rows of column j, so their gradients sum over the batch.
2. dγ₁ = −5/3 + 2/3 − 1 = −2Column 1 of dY ⊙ X̂ is (−5/3, 0, 2/3, −1); add it up.
3. dγ₂ = −1 + 2 + 0 + 1 = 2Column 2 of dY ⊙ X̂.
4. dβ = (2, 4)The column sums of dY, dY^T1.
5. dX̂ = 20.50140−20.5dY diag(γ): column 1 times 2, column 2 times 0.5. From here on each feature is handled on its own.

Example 3: the input gradient for feature 1

Continue with column 1: dx̂ = (2, 0, 4, −2), x̂ = (−5/3, 1/3, 1/3, 1) and √(σ² + ε) = 3. Find dx and check it.

1. mean(dx̂) = 4/4 = 1The μ term, a mean over the batch.
2. dx̂ ⊙ x̂ = (−10/3, 0, 4/3, −2)For the σ² term …
3. mean(dx̂ ⊙ x̂) = −4/4 = −1… which is x̂^Tdx̂/N.
4. dx = (dx̂ − 1 + x̂)/3Subtracting (−1)x̂ adds x̂.
5. dx̂ − 1 + x̂Entry by entry: 2 − 1 − 5/3, 0 − 1 + 1/3, 4 − 1 + 1/3, −2 − 1 + 1 …
6. = (−2, −2, 10, −6)/3… in thirds.
7. dx = (−2, −2, 10, −6)/9Divide by 3.
8. sum 0, dot with x̂ 0Both checks pass: −2 − 2 + 10 − 6 = 0, and (10 − 2 + 10 − 18)/27 = 0. With ε = 10⁻⁵ the dot product is about −1.5 × 10⁻⁶, the size of the ε term. Central finite differences of Σ dY ⊙ Y agree with every entry, with and without ε.

Example 4: feature 2, in training and at inference

Find dx for column 2, where dx̂ = (0.5, 1, 0, 0.5), x̂ = (−1, 1, −1, 1) and the standard deviation is 2. Then find the same column of dX at inference, when the layer uses running statistics μ̄₂ = 3 and σ̄₂² = 4.

1. mean(dx̂) = 2/4 = 0.5The μ term.
2. mean(dx̂ ⊙ x̂) = 1/4(−0.5 + 1 − 0 + 0.5)/4.
3. dx̂ − 0.5 − 0.25x̂Gives (0.25, 0.25, −0.25, −0.25) …
4. dx = (1, 1, −1, −1)/8… divided by 2. It sums to 0 and is orthogonal to x̂.
5. inference: dx = 0.5 dy/2With fixed statistics, X_n2 reaches only Y_n2, with slope γ₂/√(σ̄₂² + ε).
6. = (0.25, 0.5, 0, 0.25)The column of dY diag(γ) D̄. The statistics are the same as the batch's, yet the answers differ: here μ̄ and σ̄² are constants, so the centering terms are gone and the column sums to 1. Finite differences with the statistics frozen agree.

Common mistakes, and the exact line they happen on

1. Holding the batch mean and variance constant

Feature 1 of Example 3, with the inference gradient used in training.

1. dx̂ = (2, 0, 4, −2), σ = 3Right so far.
2. dx = dx̂/σThis is the mistake. In training, μ and σ² are computed from the current batch, so they move when x moves. This keeps only the direct path and gives (2, 0, 4, −2)/3, which sums to 4/3: it claims that shifting feature 1 on every example changes the loss, when X̂ doesn't change at all.

2. Taking the mean over the features instead of the batch

The forward pass of Example 1, written with an axis argument.

1. X is 4 × 2, rows examplesRight so far.
2. μ = (0.5, 5.5, 3.5, 6.5)This is the mistake. These are the row means, one per example, which is what layer norm uses. Batch norm needs one mean per feature, the column means (5, 3). Here the wrong μ has length 4 and fails to broadcast against the rows; when N = d it has the right length and runs without an error.

3. Differentiating an N − 1 variance

Feature 1 again, with the derivative of the variance taken from the sample-variance formula.

1. ∂σ²/∂x_n = 2(x_n − μ)/NRight, for the forward pass as defined: it divides by N.
2. … = 2(x_n − μ)/(N − 1)This is the mistake. It makes the x̂ term N/(N − 1) = 4/3 times too large and gives (−11, −5, 31, −15)/27. That still sums to 0, so the first check passes, but its dot product with x̂ is 4/9 instead of 0. The forward and backward passes have to use the same variance.

Practice

1. One feature takes the values (2, 2, 6, 6) over a batch of four. Find μ, σ² and x̂, ignoring ε.

   Show answer

   μ = 4, σ² = 4, x̂ = (−1, −1, 1, 1).

2. For feature 1 of Example 1, the gradient at x̂ is dx̂ = (3, 3, 3, 3). What is dx, and why?

   Show answer

   0. dx̂ − mean(dx̂)1 = 0 and mean(dx̂ ⊙ x̂) = 3 · mean(x̂) = 0. A gradient that is the same on every example asks for a shift of the whole feature, and the centering undoes any shift.

3. A batch of three has γ = (1, 2) and upstream rows (1, 0), (2, 1) and (0, −1). Find dβ and dX̂.

   Show answer

   dβ = (3, 0), the column sums. dX̂ has rows (1, 0), (2, 2) and (0, −2): column 2 is doubled.

4. At inference a feature has γ_j = 3 and running variance σ̄_j² = 9, with ε ignored. Its upstream column is (1, −2, 0, 4). Find its column of dX. Does it sum to 0?

   Show answer

   (1, −2, 0, 4), since γ_j/σ̄_j = 1. It sums to 3: with fixed statistics there is no centering, so nothing forces the sum to 0.

5. A network is trained with a batch of one. What does a batch-norm layer output, and what gradient reaches its input and its γ?

   Show answer

   Each feature's mean is its only value and its variance is 0, so X̂ = 0 (ε keeps it from being 0/0) and the output is β^T whatever the input. So dX = 0 and dγ = 0; only β gets a gradient, dβ = dY^T.