Layer Normalization: Pre-LN vs Post-LN

Positional encoding showed how attention knows where tokens sit. Layer normalization controls a different part of the transformer: the scale of the residual stream that carries token features through dozens of blocks.

Imagine a warehouse routing model that decides which shelf a package should visit next. At each processing station, the model adds a small update to the package's feature vector: current location, priority, weight, and destination zip. After fifty stations, poorly scaled updates can compound. Some features may become far larger than others; a downstream softmax may become sharply saturated and gradients may weaken. Deep networks face this scale-control risk. Layer Normalization helps manage it, and its placement inside a transformer block changes the optimization behavior.

This article assumes you've seen residual connections before. If you need a quick refresher: a residual block computes output = input + sublayer(input). The input travels along a skip path, and the sublayer (attention or a feed-forward network) adds a small change. Layer Normalization decides whether you normalize the input before the sublayer, or normalize the sum after it.

---

Why scale drift becomes a training problem

Every residual block adds a change to a hidden state. Residual updates aren't literal multipliers, but a toy multiplicative drift makes the risk easy to see: a scale factor of 1.1 repeated across fifty blocks gives 1.1^50, roughly 117. Backward gradients face a related issue because they multiply block Jacobians; their product can become too large or too small.

Layer Normalization helps by re-centering and rescaling each token's hidden state. Think of it as a calibration step. Before learned scale and shift are applied, the normalized vector sits near zero with a spread close to one. The model then learns scale parameters to stretch those values back into the range it needs.

---

LayerNorm by hand

Before the formula, here is a concrete example. Take a single token's hidden state, a four-dimensional vector:

Step 1: compute the mean. (3 + 1 + (-1) + 5) / 4 = 2.0

Step 2: subtract the mean from every entry. [1.0, -1.0, -3.0, 3.0]

Step 3: compute the variance. (1^2 + (-1)^2 + (-3)^2 + 3^2) / 4 = (1 + 1 + 9 + 9) / 4 = 5.0

Step 4: divide by the standard deviation. sqrt(5.0) ≈ 2.236 [0.45, -0.45, -1.34, 1.34]

Step 5: scale and shift with learned parameters. If γ = [1.0, 1.0, 1.0, 1.0] and β = [0.0, 0.0, 0.0, 0.0], the output is the same normalized vector. In practice, the model learns the γ and β that help it predict the next token.

That's the entire operation. It happens independently for every token in the sequence. Other tokens in the batch have no role in its statistics, so the same rule applies during training and during autoregressive generation with batch size one.

The general formula

For a vector x ∈ ℝ^d (a single token's hidden state), LayerNorm computes:^[1]

$\text{LayerNorm}(x) = \gamma \odot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta$

where:

• μ = (1/d) Σ x_i: mean over the hidden dimension
• σ² = (1/d) Σ (x_i - μ)²: variance over the hidden dimension
• γ, β ∈ ℝ^d: learnable scale and shift parameters
• ε ≈ 10^-5: tiny constant to prevent division by zero

The subtraction centers the values around zero. The division squeezes them to a similar scale. The learned γ (scale) and β (shift) let the model decide whether the normalized values should be stretched, compressed, or offset for the next layer.

Key distinction from BatchNorm

LayerNorm normalizes across the feature dimension for each token independently, while Batch Normalization (BatchNorm) normalizes across the batch dimension. This makes LayerNorm independent of batch size, which is critical for variable-length sequences and autoregressive generation.

This small program changes a neighboring token while holding one token fixed. Its LayerNorm result stays fixed; a batch-normalized feature coordinate changes.

PyTorch implementation

The following code demonstrates a basic PyTorch implementation. It takes an input tensor of shape (batch, seq_len, d_model) and calculates the mean and variance across the final feature dimension. It then applies the learnable scale and shift, returning the normalized tensor.

The illustration above traces the same four numbers through each step: raw input, centered, normalized, and final output. Notice how the bars shrink to a controlled range after normalization.

---

Where you put LayerNorm matters

Now that you know what LayerNorm does, the next question is where to place it inside a transformer block. There are two choices, and the difference is simply whether the normalization happens before or after the residual addition.

The unbounded update problem

• Post-LN normalizes after the residual update has been added. Its residual path therefore crosses a normalization operation in every block.
• Pre-LN normalizes before the sublayer update. Its residual path retains an identity contribution to the backward Jacobian.

Post-LN: the original transformer

In this arrangement (used by the original 2017 Transformer and BERT-style encoders), normalization happens after the residual addition:^[2][3]

$x_{l+1} = \text{LayerNorm}(x_l + \text{Sublayer}(x_l))$

In plain terms: compute the sublayer output (attention or FFN), add it to the original input along the skip connection, then normalize the sum. The problem is that normalization sits directly on the main highway, so gradients through the residual path pass through a LayerNorm Jacobian in each Post-LN block.

Pre-LN: the GPT-2 layout

In this arrangement, used by GPT-2, normalization happens before the sublayer:^[4]

$x_{l+1} = x_l + \text{Sublayer}(\text{LayerNorm}(x_l))$

Normalize the input first, feed it into the sublayer, then add the result to a direct copy of the input. The residual path from x_l to x_{l+1} keeps an identity contribution to the backward pass.

At the full-model level, GPT-2 and the Pre-LN architecture analyzed by Xiong et al. apply a final normalization before prediction.^[4][5] That final operation normalizes the accumulated residual stream before the output projection. Removing it produces a different architecture whose output scale needs evaluation.

The diagram on the left shows Pre-LN: LayerNorm cleans the sublayer input first, and the raw input still jumps straight over to the addition. The diagram on the right shows Post-LN: the sublayer and the skip path meet first, then LayerNorm cleans up the result.

Holding a sublayer update fixed isolates the placement difference. Post-LN centers and rescales the updated stream immediately; Pre-LN lets the updated stream continue along the residual path.

---

A two-layer walkthrough

To see how the placement handles a uniform offset, trace a deliberately simplified stream through two blocks. Assume each block's sublayer produces the same update, 0.5, for every feature. Also set LayerNorm's learned scale to γ = 1 and shift to β = 0. A trained attention or feed-forward sublayer would usually produce nonuniform updates.

Post-LN stack

• Start: x_0 = [1.0, 1.0, 1.0, 1.0]
• After block 1: x_1 = LayerNorm(x_0 + 0.5) = LayerNorm([1.5, 1.5, 1.5, 1.5]) = [0.0, 0.0, 0.0, 0.0] (mean-centered, then scaled)
• After block 2: x_2 = LayerNorm(x_1 + 0.5) = LayerNorm([0.5, 0.5, 0.5, 0.5]) = [0.0, 0.0, 0.0, 0.0]

The constant offset is reset at every layer. This toy is intentionally stark: LayerNorm removes a uniform shift, so the residual stream doesn't preserve that offset.

Pre-LN stack

• Start: x_0 = [1.0, 1.0, 1.0, 1.0]
• After block 1: x_1 = x_0 + Sublayer(LayerNorm(x_0)) = [1.0, 1.0, 1.0, 1.0] + 0.5 = [1.5, 1.5, 1.5, 1.5]
• After block 2: x_2 = x_1 + Sublayer(LayerNorm(x_1)) = [1.5, 1.5, 1.5, 1.5] + 0.5 = [2.0, 2.0, 2.0, 2.0]

The running sum flows through the network because Pre-LN leaves the residual path outside the block normalization. The backward equation below shows why this placement includes an identity gradient route.

---

Why Pre-LN can train more stably

Post-LN gradient flow is like sending a package through a long sorting line where every station re-labels it before passing it on. Each relabeling step slightly changes the signal. After eighty stations, the original priority mark is hard to track. Pre-LN is like keeping a clean bypass label that travels beside the package while each station adds useful updates without disrupting the main path.

Post-LN gradient problem

The key result from Xiong et al. is that Post-LN produces especially large expected gradients near the output layers at initialization under their analysis.^[5] For one block, the backward Jacobian is:

$\frac{\partial x_{l+1}}{\partial x_l} = J_{\text{LN},l} \cdot (I + J_{F,l})$

where J_LN,l is the LayerNorm Jacobian and J_F,l is the sublayer Jacobian. Across many layers, the backward pass multiplies many such terms together. In practice that means:

• Gradient scale becomes uneven across depth. Early layers may receive tiny updates while late layers receive huge ones.
• Top layers can receive very large updates at initialization. A single step with a moderately high learning rate can push weights far from a good basin.
• Learning-rate warmup addresses this risk in the evaluated Post-LN setups. Start the learning rate near zero and ramp it up so early updates are smaller.

Pre-LN gradient advantage

For Pre-LN, the Jacobian is:

$\frac{\partial x_{l+1}}{\partial x_l} = I + J_{F,l} \cdot J_{\text{LN},l}$

That leading identity term I is the key difference. Every block includes a direct residual contribution to the gradient, so the backward signal avoids relying entirely on repeated normalization Jacobians. Xiong et al. report that, in their evaluated tasks:

• Pre-LN trains without learning-rate warmup in experiments where Post-LN needed it for stable optimization.^[5]
• Its initialization-time gradients are better behaved in their analysis.

This scalar toy isn't a training experiment. It only shows how putting a normalization factor on the identity path changes a product across blocks.

---

A reported Pre-LN risk: representation collapse

Pre-LN isn't free of trade-offs. Analyses including ResiDual report that its residual stream can dominate newer sublayer updates as depth increases. Under that behavior, late layers make smaller relative changes to the hidden representation.^[6]

This reported behavior is called representation collapse. Treat it as a design risk to measure, not a guarantee for every Pre-LN model or training run.

Here is the practical summary:

There is no universal winner. Xiong et al. demonstrate the Pre-LN optimization advantage in their evaluated tasks, while DeepNorm later scales a Post-LN-derived design to a 1,000-layer machine-translation experiment using residual scaling and matching initialization.^[5][7]

---

Common mistakes and how to debug them

---

RMSNorm: scale-only normalization

RMSNorm (Root Mean Square Layer Normalization) replaces LayerNorm's centered standard deviation with a root-mean-square scale. Zhang and Sennrich propose it as a cheaper normalization that retains rescaling invariance while removing mean-centering; their evaluated models achieve comparable quality to LayerNorm with run-time reductions that vary by setup.^[8] Qwen2.5 and Gemma 2 reports provide later architecture examples using RMSNorm.^[9][10]

$\text{RMSNorm}(x) = \gamma \odot \frac{x}{\sqrt{\frac{1}{d}\sum_{i=1}^d x_i^2 + \epsilon}}$

Instead of centering (subtracting the mean) and then scaling (dividing by standard deviation), RMSNorm just divides by the root-mean-square. It doesn't subtract the mean, and many implementations also omit a learned bias term.

The original RMSNorm paper reports comparable quality to LayerNorm on its evaluated tasks and observed run-time reductions from 7% to 64% across different models and implementations.^[8] Fused kernels, hardware, and the fraction of time spent in normalization determine the result in another stack, so those measurements are evidence from the paper rather than a promised speedup.

---

Reported normalization layouts

Once you understand Pre-LN and Post-LN, advanced variants become easier to read. These entries summarize specific reports; their experiments differ in model size, objective, and optimization setup.

No row establishes a globally superior recipe. Peri-LN names an input-and-output normalization layout and reports experiments up to 3.2B parameters; Rybakov et al. evaluate attention-specific normalization and capping on an 830M-parameter model driven into instability with high learning rates.^[12][11]

This illustration uses RMS rescaling to expose the mechanism: multiplying a query by twenty changes the raw dot-product logit, but rescaling both query and key by their RMS removes that scale-only change.

Peri-LN and hybrids

Kim et al. use Peri-LN for normalization placed both before and after each sublayer, and report more balanced variance growth and steadier gradients than their compared layouts in experiments up to 3.2B parameters.^[12] Gemma 2's technical report states that it applies RMSNorm to both the input and output of each transformer sublayer; this matches an input-and-output normalization layout, although that report uses no Peri-LN label.^[10]

---

Practice: build a toggle-switch transformer block

A practical way to internalize the difference is to write one block that can switch between Post-LN and Pre-LN. Below is a minimal PyTorch module. Run it, initialize stacks of each type, and compare the gradient norms across blocks at initialization.

The profiles differ across depth and seed. This toy is a diagnostic, not a proof of the Xiong et al. result. The real lesson is how to inspect layerwise gradients instead of treating a diverging model as a mystery.

Mini exercise

1. Change depth from 6 to 12. How do the per-block gradient profiles change?
2. Replace final_norm with nn.Identity() in Pre-LN mode and inspect output scale and loss under several seeds.
3. Try replacing nn.LayerNorm with the RMSNorm class from earlier. Does the gradient norm change meaningfully?

---

Mastery check

Key concepts

• LayerNorm mean, variance, scale, and shift over one token's feature dimension
• BatchNorm versus LayerNorm behavior during variable-length and batch-size-1 decoding
• Pre-LN versus Post-LN block equations
• Gradient stability and warmup pressure in Post-LN
• Reported representation-collapse risk in deep Pre-LN stacks
• RMSNorm as scale-only normalization
• QK-Norm and hybrid norm layouts

Evaluation rubric

• Foundational: Derives the LayerNorm formula and computes it by hand on a small vector
• Intermediate: Identifies Pre-LN or Post-LN placement from a diagram or model.py block
• Advanced: Explains why Post-LN gradients near the output layer can be large at initialization and why warmup helps
• Advanced: Explains why Pre-LN includes an identity gradient path and why later work measures representation-collapse risk at depth
• Advanced: Implements RMSNorm and states exactly what it removes from LayerNorm
• Advanced: Explains why QK-Norm controls attention-logit scale before softmax
• Advanced: Compares plain Pre-LN, Post-LN, DeepNorm, and Peri-LN-style hybrid layouts

Follow-up questions

Common pitfalls

"Pre-LN means no final norm is needed"

Symptom: An implementation of a reference Pre-LN model shows unexpected output-scale behavior. Cause: The final normalization specified by that architecture may have been removed. Fix: Match the reference architecture first; GPT-2 and the Pre-LN setup analyzed by Xiong et al. include a final normalization before prediction.

"RMSNorm is just a faster name for LayerNorm"

Symptom: You describe RMSNorm as if it still subtracts the mean and uses the same statistics. Cause: RMSNorm keeps scale control but removes centering. Fix: State the exact difference: LayerNorm centers and rescales; RMSNorm only rescales by root-mean-square magnitude.

"Warmup matters only because models are deep"

Symptom: A Post-LN run diverges and you only blame model size or optimizer choice. Cause: Post-LN can create especially large top-layer gradients at initialization, even before later training dynamics matter. Fix: Tie warmup reasoning back to normalization placement and layerwise gradient scale, not depth alone.

"QK-Norm replaces residual normalization"

Symptom: Attention logits are controlled, but another instability remains. Cause: QK-Norm targets query-key scale before softmax rather than every residual-path dynamic. Fix: Treat QK-Norm as an attention-specific intervention and separately evaluate the broader normalization layout.