Exercises — Layer normalization

Across a row. LayerNorm takes one example (one row) and averages over all its $H$ features. It never looks at other examples. (BatchNorm is the one that goes down a column — same feature, many examples.) See the red arrow in the figure above.

• ${\hat{x}}_i$ = the standardized activation (mean 0, variance 1 over the row). Not learnable — computed from the data.
• ${\gamma }_i$ = the learnable scale. Learnable.
• ${\beta }_i$ = the learnable shift. Learnable. So the network stretches ($\gamma$) and slides ($\beta$) the normalized values to whatever range suits the next layer.

$ϵ$ guards the division, which blows up when ${\sigma }^2\to 0$ (a near-constant row). Placing it under the root, $\sqrt{{\sigma }^2+ϵ}$, guarantees the denominator is at least $\sqrt{ϵ}>0$ even when ${\sigma }^2=0$. If you added it after the root, $\sqrt{{\sigma }^2}+ϵ$, then when ${\sigma }^2$ is exactly $0$ the denominator is $ϵ$ (fine) — but the standard convention puts it inside so the whole quantity stays a clean "variance plus floor." Either way the point is: never divide by zero.

• Mean: $\mu =\frac{2+4+6+8}{4}=\frac{20}{4}=5$.
• Deviations: $(2-5,4-5,6-5,8-5)=(-3,-1,1,3)$.
• Variance: squares $(9,1,1,9)$, mean $=\frac{20}{4}=5$, so ${\sigma }^2=5$, $\sqrt{5}\approx \mathrm{2.2360679...}$
• Standardize: $\hat{\mathbf{x}}=\left(\frac{-3}{\sqrt{5}},\frac{-1}{\sqrt{5}},\frac{1}{\sqrt{5}},\frac{3}{\sqrt{5}}\right)\approx (-1.3416,-0.4472,0.4472,1.3416)$. Check: sum $=0$ ✓ and mean of squares $=\frac{9+1+1+9}{5\cdot 4}=\frac{20}{20}=1$ ✓.

$y_i=3{\hat{x}}_i-2$: $\mathbf{y}\approx (3(-1.3416)-2,3(-0.4472)-2,3(0.4472)-2,3(1.3416)-2)=(-6.0249,-3.3416,-0.6584,2.0249).$ Mean of $\mathbf{y}$: since each ${\hat{x}}_i$ averages to 0, $\text{mean}(\mathbf{y})=3\cdot 0-2=-2=\beta$. The affine step deliberately gives the row a new mean ($\beta =-2$) and a new standard deviation ($\gamma =3$).

• Mean: $\mu =7$. Deviations: all $0$. Variance: ${\sigma }^2=0$.
• Standardize: every numerator is $0$, so ${\hat{x}}_i=\frac{0}{\sqrt{0+10^{-5}}}=0$. $\hat{\mathbf{x}}=(0,0,0,0).$ Without $ϵ$ this would be $\frac{0}{0}=$ NaN. With $ϵ$, a constant row cleanly maps to all zeros — no crash. This is exactly the degenerate case the parent's third mistake warned about.

New mean: ${\mu }^{\prime }=\frac{1}{H}\sum (a{x}_i+b)=a\mu +b$. New deviation: $x_i^{\prime }-{\mu }^{\prime }=(a{x}_i+b)-(a\mu +b)=a(x_i-\mu )$. New variance: ${\sigma }^{\prime 2}=\frac{1}{H}\sum a^2(x_i-\mu {)}^2=a^2{\sigma }^2$, so $\sqrt{{\sigma }^{\prime 2}}=a\sqrt{{\sigma }^2}$ (since $a>0$). Then ${\hat{x}}_i^{\prime }=\frac{a(x_i-\mu )}{a\sqrt{{\sigma }^2}}=\frac{x_i-\mu }{\sqrt{{\sigma }^2}}={\hat{x}}_i.$ The $a$ and $b$ cancel exactly. Interpretation: if an earlier layer suddenly scales or offsets all its outputs, LayerNorm erases that change — the next layer sees identical inputs. This is the stability that fights Vanishing and exploding gradients.

We need ${\gamma }_i{\hat{x}}_i+{\beta }_i=x_i$. Since $x_i=\mu +\sqrt{{\sigma }^2}{\hat{x}}_i$ (rearranging the standardize step), matching gives ${\gamma }_i=\sqrt{{\sigma }^2}=\sqrt{5}\approx 2.2361,{\beta }_i=\mu =5.$ Concretely with $\hat{\mathbf{x}}\approx (-1.3416,\dots )$: $\sqrt{5}\cdot (-1.3416)+5\approx 2$ ✓, recovering $x_1=2$. So the affine step is expressive enough to make LayerNorm a no-op — proving $\gamma ,\beta$ are not redundant.

BatchNorm's statistics for a feature are averaged down the column — over all examples in the batch. With batch size 1 the column has a single number, so its "variance" is $0$ and the mean is just that number: standardizing gives $0$, destroying all information (or NaN without $ϵ$). LayerNorm averages across the row of that single example's $H$ features, which is always well-defined regardless of batch size. That is precisely why Transformers and RNNs adopt LayerNorm. See Internal covariate shift for the original motivation both methods share.

RMS: $\frac{1}{H}\sum x_j^2=\frac{4+16+36+64}{4}=\frac{120}{4}=30$, so $\sqrt{30}\approx 5.4772$. $\bar{\mathbf{x}}=\left(\frac{2}{\sqrt{30}},\frac{4}{\sqrt{30}},\frac{6}{\sqrt{30}},\frac{8}{\sqrt{30}}\right)\approx (0.3651,0.7303,1.0954,1.4606).$ Compare first entries: LayerNorm gave ${\hat{x}}_1\approx -1.3416$; RMSNorm gives ${\bar{x}}_1\approx 0.3651$. They differ because RMSNorm skips centering ($-\mu$), so it keeps the sign/offset of the raw value; LayerNorm re-centers first, flipping $x_1=2$ (below the mean $5$) to a negative standardized value. RMSNorm is cheaper (no mean) and often works comparably in Transformers.

This is $(2,4,6,8)\times 50$ — a pure scaling of Ex. 2.1 by $a=50$. By the invariance proof (Ex. 3.1), LayerNorm erases the factor of $50$: $\mathbf{z}=\hat{\mathbf{x}}\approx (-1.3416,-0.4472,0.4472,1.3416),$ identical to Ex. 2.1. Lesson: even if the residual branch $F$ produces huge activations, the LayerNorm hands the next block bounded numbers — this is why LayerNorm placement tames Vanishing and exploding gradients in deep stacks.

Since $\hat{\mathbf{x}}$ has mean 0 and std 1, applying $y_i=\gamma {\hat{x}}_i+\beta$ gives mean $\beta$ and std $|\gamma |$. To hit mean $10$, std $4$: $\beta =10,\gamma =4.$ Check on Ex. 2.1's $\hat{\mathbf{x}}$: $y_1=4(-1.3416)+10\approx 4.6337$, and the four outputs $(4.6337,8.2111,11.7889,15.3663)$ have mean $10$ and std $4$. The affine step is just a linear map from the standard $(\text{mean }0,\text{std }1)$ template to any $(\text{mean }\beta ,\text{std }\gamma )$ you like.

Write ${\hat{x}}_1=\frac{x_1-\mu }{\sqrt{{\sigma }^2}}$ with $\mu =\frac{1}{H}\sum x_j$, ${\sigma }^2=\frac{1}{H}\sum (x_j-\mu {)}^2$. Both numerator and denominator are homogeneous of degree 1 in $\mathbf{x}$: scaling $\mathbf{x}\mapsto a\mathbf{x}$ makes $x_1-\mu \mapsto a(x_1-\mu )$ and $\sqrt{{\sigma }^2}\mapsto a\sqrt{{\sigma }^2}$. So ${\hat{x}}_1$ is homogeneous of degree 0 — invariant. A degree-0 function's gradient is homogeneous of degree $-1$: $\frac{\partial {\hat{x}}_1}{\partial x_1}{|}_{a\mathbf{x}}=\frac{1}{a}\frac{\partial {\hat{x}}_1}{\partial x_1}{|}_{\mathbf{x}}$... but the derivative is taken with respect to the scaled coordinate, and by the chain rule that extra $a$ cancels — the derivative of ${\hat{x}}_1$ w.r.t. the input entry is scale-free. Explicitly, differentiating gives $\frac{\partial {\hat{x}}_1}{\partial x_1}=\frac{1}{\sqrt{{\sigma }^2}}\left[\left(1-\frac{1}{H}\right)-\frac{(x_1-\mu {)}^2}{H{\sigma }^2}\right].$ For $\mathbf{x}=(2,4,6,8)$: ${\sigma }^2=5$, $x_1-\mu =-3$, so $\frac{\partial {\hat{x}}_1}{\partial x_1}=\frac{1}{\sqrt{5}}\left[\frac{3}{4}-\frac{9}{4\cdot 5}\right]=\frac{1}{\sqrt{5}}\left[0.75-0.45\right]=\frac{0.30}{\sqrt{5}}\approx \mathrm{0.1342.}$ For $\mathbf{x}=(20,40,60,80)$ ($a=10$): ${\sigma }^2=500$, $x_1-\mu =-30$, giving $\frac{1}{\sqrt{500}}[0.75-\frac{900}{4\cdot 500}]=\frac{0.30}{\sqrt{500}}\cdot ...$ wait — recompute: $\frac{900}{2000}=0.45$, so $\frac{1}{\sqrt{500}}(0.30)=\frac{0.30}{22.360}\approx 0.01342$. The raw derivative shrinks by $a=10$ — but the gradient the optimizer sees, $\frac{\partial {\hat{x}}_1}{\partial (\text{scaled }{x}_1)}$, is exactly $0.1342$ again once you account for the coordinate scaling. Takeaway: LayerNorm makes the effective learning signal insensitive to the overall magnitude of a layer's activations — a key reason it stabilizes optimization.

We know ${\sum }_j{\hat{x}}_j^2=H$ (from the parent's variance identity, since mean-square $=1$). For any single index $i$, ${\hat{x}}_i^2\le {\sum }_j{\hat{x}}_j^2$; the other $H-1$ terms are $\ge 0$, but we can do better. The most extreme case puts all "mass" in one coordinate against the mean-zero constraint ${\sum }_j{\hat{x}}_j=0$. Maximizing ${\hat{x}}_i^2$ subject to $\sum {\hat{x}}_j=0$ and $\sum {\hat{x}}_j^2=H$ gives the other $H-1$ entries equal, and solving yields ${\hat{x}}_i^2\le H-1$, so $|{\hat{x}}_i|\le \sqrt{H-1}$. For $H=4$: bound is $\sqrt{3}\approx 1.732$. Our largest standardized value is ${\hat{x}}_4=\frac{3}{\sqrt{5}}\approx 1.3416<1.732$ ✓. So no standardized activation can explode — a built-in ceiling that helps prevent Vanishing and exploding gradients.

1. Batch independence (Ex. 3.3): generation uses batch size 1; BatchNorm's per-column statistics collapse, LayerNorm's per-row statistics are always well-defined.
2. Train = inference (parent, mistake 4): LayerNorm needs no running averages, so the model behaves identically while training and while decoding — no train/test mismatch.
3. Scale stability around residuals (Ex. 4.2): deep Transformer stacks with residual branches can produce large activations; LayerNorm's scale-invariance (Ex. 3.1) hands each block bounded inputs, controlling Vanishing and exploding gradients. (Bonus: swapping to RMSNorm keeps points 1–3 while saving the mean computation.)

• Axis: LayerNorm normalizes across the row (features of one example); BatchNorm down the column.
• Recipe: subtract $\mu$, divide by $\sqrt{{\sigma }^2+ϵ}$ (divisor $H$, not $H-1$), then $\gamma \hat{x}+\beta$.
• Invariance: immune to shifting/scaling the whole vector (Ex. 3.1) — not to per-feature scaling (L3 trap).
• $\gamma ,\beta$ can fully undo normalization ⇒ not redundant (Ex. 3.2).
• Bounds: $|{\hat{x}}_i|\le \sqrt{H-1}$ (Ex. 5.2). Constant row → all zeros thanks to $ϵ$ (Ex. 2.3).
• Why Transformers use it: batch-independent, train = inference, tames residual blow-ups.