Pre-Norm vs Post-Norm › Worked Example: Tracing One Gradient Step

Abstract gradient flow arguments are most convincing when you can trace them through concrete numbers. Let's work through a simplified example with a two-layer network, small enough to compute by hand but structured exactly like a real transformer.

Suppose our model dimension is $d = 2$, and we have a single-token sequence. The input to layer 2 is the output of layer 1, call it $\mathbf{x} = [2.0, -1.0]^T$. The loss gradient arriving at the output of layer 2 is $\frac{\partial L}{\partial \mathbf{x}'_2} = [0.5, 0.5]^T$.

In the pre-norm case, the output of layer 2 is:

$$\mathbf{x}'_2 = \mathbf{x}_2 + \text{Sublayer}(\text{LayerNorm}(\mathbf{x}_2))$$

The gradient with respect to the input to layer 2 has two additive parts. The residual path contributes $\frac{\partial L}{\partial \mathbf{x}'_2}$ directly: the gradient $[0.5, 0.5]^T$ passes through unchanged. The sublayer path contributes whatever the sublayer's Jacobian carries. Suppose the sublayer Jacobian for this input is the matrix $J_s = 0.1 \cdot I$ (a small, well-behaved transformation). The LayerNorm Jacobian is $J_{LN}$, but notice it only appears in the sublayer branch. The total gradient at layer 2's input is:

$$\frac{\partial L}{\partial \mathbf{x}_2} = [0.5, 0.5]^T + J_s \cdot J_{LN} \cdot [0.5, 0.5]^T \approx [0.5, 0.5]^T + [0.05, 0.05]^T = [0.55, 0.55]^T$$

The residual contribution dominates, and the LayerNorm Jacobian only affects the smaller sublayer branch.

In the post-norm case, the output of layer 2 is:

$$\mathbf{x}'_2 = \text{LayerNorm}(\mathbf{x}_2 + \text{Sublayer}(\mathbf{x}_2))$$

Now the LayerNorm Jacobian $J_{LN}$ sits between the incoming gradient and everything else. The gradient at layer 2's input becomes:

$$\frac{\partial L}{\partial \mathbf{x}_2} = J_{LN} \cdot [0.5, 0.5]^T \cdot (1 + J_s)$$

The key difference is that $J_{LN}$ now multiplies the full gradient before it splits into the residual and sublayer contributions. In a well-behaved layer, $J_{LN}$ has eigenvalues close to 1 and this is fine. But in practice, LayerNorm Jacobians can have eigenvalues significantly less than 1 in certain directions, especially early in training when activations are poorly distributed. Stack 24 such layers, and those sub-unit eigenvalues multiply together, creating the gradient attenuation we visualized earlier.

The key takeaway from this numerical walk-through: pre-norm's $+1$ term ensures the incoming gradient always reaches the input intact. Even if the sublayer branch contributes nothing useful, the residual branch delivers the full gradient signal. Post-norm has no such guarantee.