On Layer Normalization in the Transformer Architecture

Where the original transformer puts the norm

The transformer of Vaswani et al. normalizes after the residual addition, a placement now called Post-LN: the sublayer output is added to its input and the sum is normalized before being handed to the next block.

$$x_{l+1} = \operatorname{LayerNorm}\big(x_l + \mathcal{F}_l(x_l)\big)$$

Here $\mathcal{F}_l$ is the block's attention or feed-forward function, and the LayerNorm sits on the main path that carries information from layer to layer. That last fact — the norm is on the path, not beside it — is the whole story of why Post-LN is hard to train.

Why the gradient misbehaves at initialization

At initialization the weights inside $\mathcal{F}_l$ are small random numbers, so the block barely perturbs its input and $\mathcal{F}_l(x_l) \approx 0$, which means the pre-norm sum is essentially just $x_l$ passed through.

$$x_l + \mathcal{F}_l(x_l) \approx x_l$$

The backward signal $\partial \mathcal{L}/\partial x_l$ must therefore travel through the Jacobian $\partial \operatorname{LayerNorm}(x_l + \mathcal{F}(x_l))/\partial x_l$ at every layer, and the paper shows that for Post-LN this Jacobian scales the gradient by a factor that grows with depth. Stack $L$ of these and the per-layer factors compound multiplicatively; the gradient reaching the early layers and the parameters near the output can be larger by a factor on the order of $\sqrt{L}$ than what a well-behaved network would produce, so the first optimizer steps are effectively operating on a badly miscalibrated landscape. The standard remedy is warmup — starting at a tiny learning rate and ramping it over thousands of steps — which exists precisely to survive this early window.

Moving the norm into the branch

Pre-LN makes one structural change: it normalizes the input to the sublayer and leaves the residual addition untouched, so the main path is now a clean sum.

$$x_{l+1} = x_l + \mathcal{F}_l\big(\operatorname{LayerNorm}(x_l)\big)$$

Because no normalization sits between $x_l$ and $x_{l+1}$ on the main path, differentiating the recurrence yields an identity term that nothing can suppress, plus the sublayer's own contribution.

$$\frac{\partial \mathcal{L}}{\partial x_l} = \frac{\partial \mathcal{L}}{\partial x_{l+1}} \left( I + \frac{\partial \mathcal{F}_l(\operatorname{LayerNorm}(x_l))}{\partial x_l} \right)$$

The $I$ is what matters: whatever the sublayer's Jacobian does — vanish, explode, rotate — the gradient still has an unobstructed identity route back through the residual stream, so the deep compounding that wrecks Post-LN initialization cannot happen. Unrolling the recurrence, the gradient magnitude at each layer is bounded independently of depth rather than scaling with it.

What this buys in practice

The consequence is operational, not just theoretical: Pre-LN transformers train from scratch with a constant or simply-decayed learning rate and no warmup, while Post-LN networks of the same depth diverge under the same schedule unless warmup tames those first steps. Removing warmup removes a fragile, dataset-specific hyperparameter — the warmup length that works for one model size is wrong for another — which is a large part of why Pre-LN scales gracefully as you add layers.

Who uses which

BERT is Post-LN, which is one reason it is sensitive to warmup and learning rate. Essentially every large decoder trained since is Pre-LN: GPT-2 and GPT-3, LLaMA, Mistral, DeepSeek, and Gemini all normalize inside the branch. When you read a modern architecture diagram and see the LayerNorm (or RMSNorm) feeding the attention and MLP rather than following them, you are looking at this paper's recommendation.

The tradeoff that didn't win

Pre-LN is not strictly better. When Post-LN is tuned carefully — the right warmup, the right schedule — it can reach a slightly lower final loss, because normalizing after the residual keeps the activations at a fixed scale at every depth and gives the optimizer a stronger, better-conditioned signal late in training. Pre-LN lets the residual stream's variance grow with depth, which is exactly what makes its gradients tame but also slightly blunts that late-training signal. The field adopted Pre-LN anyway, because at the scale where a single run costs millions of dollars, a training process that simply does not diverge is worth far more than a marginal improvement in final perplexity that you can only reach if every knob is set perfectly.