Exercises — Residual connections and layer norm placement

(a) is Post-LN — the $\text{LN}$ wraps the whole sum, so it is applied after the residual add. In figure s01 this is the bottom row: the white LN box sits directly on the cyan pipe. Read it inside-out: first $x_i+F(x_i)$ (residual), then $\text{LN}$ last. (b) is Pre-LN — the $\text{LN}$ sits inside the branch, touching only the copy fed to $F$; the raw $x_i$ passes through untouched. This is the top row of s01, where the cyan bar runs unbroken and the LN box hangs below it on the branch. The bare $x_i +$ at the front is the un-normalized highway.

Rule of thumb: whatever is added last, un-normalized, is the residual stream. In (b) $x_i$ is added last untouched → clean highway.

$\frac{dy}{dx}=1+\frac{dF}{dx}$ The ==$+1$== term is the guarantee. Even if $F$ has learned nothing (so $\frac{dF}{dx}\approx 0$), the derivative is still $1$, not $0$. That "$1$" is the pipe going straight through (the unbroken cyan bar in the top row of s01) — the gradient cannot vanish along the residual path.

Across the feature axis, per example independently: $\mu =\frac{1}{d}{\sum }_{i=1}^d{x}_i$ over the $d$ features of that one token. Preferred because Transformers see variable-length sequences with padding — batch statistics would wobble with sequence length and padding, whereas per-example feature statistics are stable regardless of what the rest of the batch looks like. See 4.1.01-Self-attention-mechanism for where these token vectors come from.

Step — WHAT: mean. $\mu =\frac{2+4+4+6}{4}=\frac{16}{4}=4$. Step — variance. Deviations: $[-2,0,0,2]$. Squares: $[4,0,0,4]$. ${\sigma }^2=\frac{4+0+0+4}{4}=2$, so $\sigma =\sqrt{2}\approx 1.4142$. Step — normalize. $\frac{x-\mu }{\sigma }=\left[\frac{-2}{\sqrt{2}},0,0,\frac{2}{\sqrt{2}}\right]=[-\sqrt{2},0,0,\sqrt{2}]\approx [-1.4142,0,0,1.4142]$. WHY it looks like this: subtracting the mean centered the cloud on $0$; dividing by $\sigma$ set its spread to $1$. The two middle entries were exactly average, so they map to $0$. (We used $ϵ=0$ here purely because $\sigma >0$ made it safe; L2.4 shows why real code never dares set $ϵ=0$.)

$x_{i+1}=[1,3]+[0.5,-1]=[1.5,2]$. If $F$ output zeros: $x_{i+1}=[1,3]+[0,0]=[1,3]$ — identical to the input. This is the whole point: an untrained Pre-LN layer is a pass-through, so stacking many of them at initialization does nothing harmful. This is why Pre-LN starts training gently (see 4.2.02-Training-stability-and-convergence).

Residual sum: $z=[1,3]+[0.5,-1]=[1.5,2]$. Mean: $\mu =\frac{1.5+2}{2}=1.75$. Deviations: $[-0.25,0.25]$. ${\sigma }^2=\frac{0.0625+0.0625}{2}=0.0625$, $\sigma =0.25$. Normalize: $\frac{[-0.25,0.25]}{0.25}=[-1,1]$. So $x_{i+1}=[-1,1]$. Notice even with a tiny branch output, Post-LN rescales the whole stream every layer — the raw magnitudes of $x_i$ are erased. That is the "squeezed pipe" behaviour (the LN box on the cyan bar in the bottom row of s01).

(i) The break. Take $x=[5,5,5,5]$ — all features identical. Then $\mu =5$, every deviation is $0$, so ${\sigma }^2=0$ and $\sigma =0$. With $ϵ=0$ the normalization divides $\frac{0}{0}$ — undefined (NaN in real code). This is the ==$\sigma \to 0$ degenerate case==. (ii) With $ϵ=10^{-5}$. Numerator is $x-\mu =[0,0,0,0]$, denominator is $\sqrt{0+10^{-5}}=\sqrt{10^{-5}}\approx 0.003162$. So output $=\frac{[0,0,0,0]}{0.003162}=[0,0,0,0]$ — clean, finite, no NaN. $ϵ$ turned a $\frac{0}{0}$ into a well-defined $\frac{0}{\text{tiny}}=0$. (iii) Backward. Because the forward pass divides by $\sqrt{{\sigma }^2+ϵ}$, every derivative of LN carries that same $\sqrt{{\sigma }^2+ϵ}$ in its denominator (the chain rule differentiates through it). So the LN gradient scale is really $\frac{\gamma }{\sqrt{{\sigma }^2+ϵ}}$, not $\frac{\gamma }{\sigma }$ — the $ϵ$ caps the gradient magnitude when $\sigma \to 0$: instead of the gradient exploding as $\frac{\gamma }{\sigma }\to \infty$, it saturates at $\frac{\gamma }{\sqrt{ϵ}}$. That bound is exactly what stops the $\sigma \to 0$ case from destroying training.

Differentiate the sum: $\frac{\partial x_{i+1}}{\partial x_i}=I+\frac{\partial F}{\partial (\text{LN})}\cdot \frac{\partial \text{LN}}{\partial x_i}.$

• The $I$ (identity) comes from the bare $x_i$ term — the untouched highway (the unbroken cyan bar, top row of s01). It is a constant, not shrinking with depth.
• The second term is the branch: it can be small or large, but it is added to $I$, never multiplied into it. When you chain this over $n$ layers, you get $\prod (I+\text{small})$, which stays near $I$ if the branches are small — gradients neither vanish nor explode. See 5.1.03-Gradient-flow-in-deep-networks.

The exact matrix. Let $z$ be the input to LN, with mean $\mu$, variance ${\sigma }^2$, and ${\hat{z}}_a=\frac{z_a-\mu }{\sqrt{{\sigma }^2+ϵ}}$ the normalized entries. Entry $(a,b)$ of the Jacobian is $(J_{\text{LN}}{)}_{ab}=\frac{\partial \text{LN}(z{)}_a}{\partial z_b}=\frac{{\gamma }_a}{\sqrt{{\sigma }^2+ϵ}}\left(\underset{\text{diagonal}}{\underbrace{{\delta }_{ab}}}-\underset{\text{mean removal}}{\underbrace{\frac{1}{d}}}-\underset{\text{variance removal}}{\underbrace{\frac{1}{d}{\hat{z}}_a{\hat{z}}_b}}\right).$ Reading the three pieces:

• ${\delta }_{ab}$ (the Kronecker delta, $=1$ if $a=b$ else $0$) is the plain identity part — this is why people call $J_{\text{LN}}$ "diagonal-ish."
• $-\frac{1}{d}$ is a constant subtracted from every entry — this is the mean-removal coupling ($\mu$ depends on all features, so changing one $z_b$ shifts every output). It fills in the off-diagonal.
• $-\frac{1}{d}{\hat{z}}_a{\hat{z}}_b$ is the variance-removal coupling — changing $z_b$ changes $\sigma$, which rescales every output. This is the second off-diagonal contribution and it does not cancel in general; it only vanishes when the $\hat{z}$ are orthogonal to the perturbation, which is not the typical case.
• The whole matrix is scaled by $\frac{{\gamma }_a}{\sqrt{{\sigma }^2+ϵ}}$ — this is where $ϵ$ enters the gradient, capping the scale when $\sigma \to 0$ (see L2.4).

Why fragile. In Pre-LN the identity is added ($I+\dots$). In Post-LN the identity is inside a product with $J_{\text{LN}}$. Chaining $n$ layers gives ${\prod }_i{J}_{\text{LN}}^{(i)}(I+\dots )$ — every one of those $J_{\text{LN}}$ factors multiplies together. Because each carries the $\frac{\gamma }{\sqrt{{\sigma }^2+ϵ}}$ scale that depends on that layer's activation statistics, the product drifts far from $1$ and does so unevenly across depth. That uneven, multiplicative chain is why Post-LN needs learning-rate warmup.

Deep-layer scale $=\frac{1}{\sqrt{9}}=\frac{1}{3}\approx 0.333$. Shallow-layer scale $=\frac{1}{\sqrt{0.25}}=\frac{1}{0.5}=2$. Ratio (shallow / deep) $=\frac{2}{1/3}=2\times 3=6$. The gradient is amplified 6× more in the shallow layer than the deep one purely because of differing $\sigma$. That non-uniformity across depth is precisely Xiong et al.'s "gradient imbalance" — it is why a single global learning rate is dangerous for Post-LN, and why warmup (ramping the LR slowly) is used to let $\sigma$ settle first.

Pre-LN: $x_{i+1}=x_i+b$ (the branch adds $b$ every layer).

• $x_1=[0,0]+[0.1,0.1]=[0.1,0.1]$
• $x_2=[0.1,0.1]+[0.1,0.1]=[0.2,0.2]$
• $x_3=[0.2,0.2]+[0.1,0.1]=[0.3,0.3]$ The stream grows linearly: $x_n=n\cdot b$. Because the residual path is never normalized in Pre-LN (the unbroken cyan bar in s01), each layer's contribution is permanently added into a shared, growing stream — later layers can read everything earlier layers wrote. This additive-memory picture is the "residual stream" interpretation.

Layer 1: $z=[0,0]+[0.1,-0.1]=[0.1,-0.1]$. $\mu =0$. ${\sigma }^2=\frac{0.01+0.01}{2}=0.01$, $\sigma =0.1$. Normalize: $\frac{[0.1,-0.1]}{0.1}=[1,-1]$. So $x_1=[1,-1]$. Layer 2: $z=[1,-1]+[0.1,-0.1]=[1.1,-1.1]$. $\mu =0$, ${\sigma }^2=\frac{1.21+1.21}{2}=1.21$, $\sigma =1.1$. Normalize: $\frac{[1.1,-1.1]}{1.1}=[1,-1]$. So $x_2=[1,-1]$. Contrast: Post-LN re-normalizes to the same magnitude every layer — $x_1=x_2=[1,-1]$, no linear accumulation. The stream cannot "grow"; its scale is reset each layer. That's the fundamental representational difference: Pre-LN accumulates, Post-LN overwrites the scale.

Mechanism: Post-LN gradients chain multiplicatively through ${\prod }_i{J}_{\text{LN}}^{(i)}$. With 6 layers that product stays near $1$; with 48 layers the same per-layer imperfections compound, and the depth-dependent $\frac{\gamma }{\sqrt{{\sigma }^2+ϵ}}$ imbalance (L3.3) makes the deepest layers receive wildly-scaled updates. A single large LR then overshoots → divergence. Fixes:

1. Learning-rate warmup — ramp LR from ~0 over a few thousand steps so activation statistics ($\sigma$) settle before big updates (the standard Post-LN remedy).
2. Switch to Pre-LN — moves the identity out of the product into a sum ($I+\dots$), giving depth-uniform gradients. See 6.2.05-Layer-scaling-and-initialization for a third fix (scaling the branch at init).

Proof: Pre-LN is $x_{i+1}=x_i+F(\text{LN}(x_i))$. If $F(\cdot )\equiv 0$ then $x_{i+1}=x_i+0=x_i$. Identity, for any input — QED. Consequence: If you initialize the final linear projection inside each branch to zero, then $F$ outputs $0$ at step 0, so the entire deep stack is one giant identity map at initialization. The forward pass is therefore perfectly well-behaved no matter how deep, and gradients begin flowing through the clean $I$ highway before the branches have learned anything. Training then "grows" corrections from zero — the safest possible start. This is the logic behind branch/residual scaling schemes in 6.2.05-Layer-scaling-and-initialization.

Bound. By the triangle inequality, $\Vert I+B_i\Vert \le \Vert I\Vert +\Vert B_i\Vert =1+\Vert B_i\Vert \le 1+\delta$. Operator norm is submultiplicative ($\Vert MN\Vert \le \Vert M\Vert \Vert N\Vert$), so chaining all $n$ factors: $\Vert {\prod }_{i=1}^n(I+B_i)\Vert \le {\prod }_{i=1}^n\Vert I+B_i\Vert \le {\prod }_{i=1}^n(1+\delta )=(1+\delta {)}^n.$ This is finite for any finite $n$ and, crucially, grows only geometrically at rate $1+\delta$ — controllable by keeping each branch contribution ($\delta$) small (exactly what zero/scaled init in L5.1 does). Contrast Post-LN, where the analogous product carries the extra $J_{\text{LN}}^{(i)}$ factors and lacks this clean shape. Evaluate. $(1+0.02{)}^{48}=1.02^{48}\approx 2.587$. So the total backward gain across 48 layers is at most ~$2.6\times$ — no explosion.

Because Pre-LN places the identity in an added position ($I+\text{branch}$) so gradients chain as a benign $(1+\delta {)}^n$ product, whereas Post-LN's identity sits inside a multiplicative LN-Jacobian chain that drifts unevenly with depth.

Pre-LN keeps the residual highway additive (never normalized), giving depth-uniform gradient flow.

Post-LN puts LayerNorm ::: after the residual add — $\text{LN}(x+F(x))$ Pre-LN puts LayerNorm ::: inside the branch, before F — $x+F(\text{LN}(x))$ The term guaranteeing gradient survival in $y=x+F(x)$ is ::: the +1 (identity) in $dy/dx=1+dF/dx$ LayerNorm divides variance by ::: $d$ (population form — the features ARE the whole population, nothing to correct for) Why nonzero $ϵ$ in $\sqrt{{\sigma }^2+ϵ}$? ::: to avoid dividing by zero when $\sigma \to 0$ (all features equal), and to cap the gradient scale at $\gamma /\sqrt{ϵ}$ $J_{\text{LN}}$ off-diagonal terms come from ::: mean removal ($-1/d$) and variance removal ($-{\hat{z}}_a{\hat{z}}_b/d$) Pre-LN gradient chains as $(1+\delta {)}^n$ because the identity is ::: added, not multiplied Does Pre-LN eliminate warmup? ::: No — it only reduces dependence on the warmup choice