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.

A transformer block repeatedly updates one token's hidden-state vector. Each layer adds attention and MLP outputs to the same representation. After fifty layers, 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.

Assume 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. During backpropagation, 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. 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, start with one concrete token hidden state:

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.

If you feed 16 request-latency feature vectors through LayerNorm, which numbers determine the mean and variance for one vector?

Answer

Only that vector's own feature values determine its mean and variance. The other 15 vectors in the batch don't affect it, which is why LayerNorm behaves the same during real-time batch-size-1 prediction and during larger-batch training.

The general formula

For a vector x ∈ ℝ^d (a single token's hidden state), LayerNorm computes:^{[1]Reference 1Layer Normalization.https://arxiv.org/abs/1607.06450}

$\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
• ε: tiny positive constant to prevent division by zero (10^-5 in this example; defaults vary by implementation). In mixed precision (FP16 / BF16), keep ε large enough that the variance floor doesn't underflow to zero; common production floors sit around 1e-5 to 1e-6 rather than aggressive 1e-12 values that can look fine in FP32 and blow up in half precision.

Subtracting the mean centers the values around zero. Dividing by the standard deviation 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. Batch Normalization (BatchNorm) instead computes each feature coordinate's statistics across examples in a mini-batch (and, for some data layouts, additional positions). This makes LayerNorm independent of batch size, which is useful for variable-length sequences and autoregressive generation.

Why is BatchNorm awkward for autoregressive decoding while LayerNorm is natural?

Answer

BatchNorm depends on batch-level statistics or stored running statistics, so behavior can change between training and inference. Autoregressive decoding often uses batch size 1 and variable sequence lengths. LayerNorm normalizes each token independently across its feature dimension, so population statistics are unnecessary.

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

This PyTorch implementation takes an input tensor of shape (batch, seq_len, d_model), calculates the mean and variance across the final feature dimension, then applies the learnable scale and shift.

The colored paths preserve each feature's identity across both operations. Centering translates the whole vector; scaling changes its spread. Learned γ and β can then reshape that standardized vector for the next layer.

---

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 whether the normalization happens before or after the residual addition.

The placement difference

• 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]Reference 2Attention Is All You Need.https://arxiv.org/abs/1706.03762[3]Reference 3BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.https://arxiv.org/abs/1810.04805}

$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. Normalization now 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]Reference 4Language Models are Unsupervised Multitask Learners.https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf}

$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]Reference 4Language Models are Unsupervised Multitask Learners.https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf[5]Reference 5On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745} 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.

In one equation, what is the architectural difference between Post-LN and Pre-LN?

Answer

Post-LN is $x_{l+1} = LayerNorm(x_l + F(x_l))$. Pre-LN is $x_{l+1} = x_l + F(LayerNorm(x_l))$. The important difference is whether normalization sits on the residual highway after the add, or only on the sublayer input before the add.

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 does the Pre-LN architecture analyzed by Xiong et al. include a final normalization before the output head?

Answer

Inside each Pre-LN block, the residual stream bypasses block-output normalization and accumulates updates. A final norm sets its scale before the vocabulary projection. Removing that norm changes the architecture and requires rechecking output-scale and training behavior.

---

Why Pre-LN can train more stably

Post-LN gradient flow is like passing an error signal through a long chain of adapters where every block rescales it before sending it backward. Each rescaling step slightly changes the signal. After eighty blocks, the original direction is hard to track. Pre-LN keeps a cleaner residual path beside the block updates, so useful gradients can travel backward without being repeatedly disrupted.

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]Reference 5On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745} 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]Reference 5On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745}
• 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]Reference 6ResiDual: Transformer with Dual Residual Connections.https://arxiv.org/abs/2304.14802}

The ResiDual authors call this reported behavior representation collapse.^{[6]Reference 6ResiDual: Transformer with Dual Residual Connections.https://arxiv.org/abs/2304.14802} Treat it as a design risk to measure, not a guarantee for every Pre-LN model or training run.

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]Reference 5On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745[7]Reference 7DeepNet: Scaling Transformers to 1,000 Layers.https://arxiv.org/abs/2203.00555}

What is the representation-collapse risk reported for Pre-LN?

Answer

The residual stream can dominate newer sublayer updates as depth grows. Late blocks still run, but their relative changes may become small compared with the accumulated stream. This is an observed and analyzed risk, rather than an inevitable outcome for every Pre-LN run.

Residual magnitude growth across depth

A short numeric toy makes the scale story concrete. Treat each block as adding a fixed update vector (size 1 along one axis). Pre-LN lets the stream accumulate without a reset. Post-LN re-normalizes after each add. DeepNorm-style residual scaling multiplies the incoming stream by a constant α before the add, then normalizes (toy form of the paper's update).

In this toy, Pre-LN stream magnitude grows with depth while Post-LN and DeepNorm-scale both re-normalize after the add so ||x|| stays controlled. DeepNorm's distinguishing mechanism is the depth-dependent residual scale α plus matched init on the update branch, which keeps optimization of a Post-LN-shaped stack stable when depth is extreme. Real models use learned sublayers rather than a fixed update, so treat the table as geometry intuition, then measure ||x_l|| on your own stack.

DeepNorm: scaled residual + matched init

DeepNorm (DeepNet) keeps a Post-LN-shaped block but multiplies the residual branch by a depth-dependent constant α and down-scales residual-branch initialization so expected update magnitudes stay controlled at large depth.^{[7]Reference 7DeepNet: Scaling Transformers to 1,000 Layers.https://arxiv.org/abs/2203.00555} The paper's encoder update has the form:

$x_{l+1} = \mathrm{LN}\big(\alpha\, x_l + G_l(x_l)\big)$

Here $G_l$ is the sublayer stack (attention and/or FFN writes). α grows slowly with depth so the residual highway still carries signal, while initialization of $G_l$ is reduced (the paper pairs α with a gain bound β on residual-branch weights). Architecture-specific constants from the paper include encoder-only $\alpha=(2N)^{1/4}$ for $N$ layers, decoder-only $\alpha=(2M)^{1/4}$ for $M$ layers, and encoder-decoder pairs with different α on each stack. Copy constants only from the table matching your architecture.

DeepNorm is therefore not "plain Post-LN with no residual scaling." The residual scale and matched init are the design, and the 1,000-layer machine-translation experiment is evidence under that recipe, not a free upgrade for an arbitrary Post-LN codebase.

What two mechanisms does DeepNorm add on top of a Post-LN-shaped block?

Answer

A depth-dependent residual scale α in the update $x_{l+1}=\mathrm{LN}(\alpha x_l + G_l(x_l))$, and matched down-scaling of residual-branch initialization so expected update size stays controlled as depth grows.

---

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]Reference 8Root Mean Square Layer Normalization.https://arxiv.org/abs/1910.07467} Qwen2.5 and Gemma 2 reports provide later architecture examples using RMSNorm.^{[9]Reference 9Qwen2.5 Technical Reporthttps://arxiv.org/abs/2412.15115[10]Reference 10Gemma 2: Improving Open Language Models at a Practical Sizehttps://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf}

$\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]Reference 8Root Mean Square Layer Normalization.https://arxiv.org/abs/1910.07467} 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.

What does RMSNorm remove from LayerNorm, and what does it keep?

Answer

RMSNorm removes mean-centering and often removes the learned bias. It keeps per-token rescaling by a feature-wise magnitude estimate and keeps a learned scale vector.

---

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]Reference 12Peri-LN: Revisiting Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2502.02732[11]Reference 11QK-Norm: Improving Transformer Training Stability.https://arxiv.org/abs/2410.16682}

QK-Norm targets attention logits, not every residual-stream dynamic.

How does QK-Norm reduce attention instability?

Answer

QK-Norm normalizes query and key vectors before their dot product, limiting the contribution from their vector magnitude to attention logits. Rybakov et al. study a LayerNorm-based form and report that combining QK-Norm with softmax capping increases the stable learning-rate range in their evaluated setup. The normalizer and its exact placement relative to positional operations remain architecture choices.

This small calculation 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.

Why do hybrid layouts such as Peri-LN or Pre + Post norm exist if Pre-LN is already stable?

Answer

Pre-LN includes a direct identity gradient route, but residual-stream scale can still become a concern. Hybrid layouts add output-side normalization while preserving an input-side norm, then test whether that changes activation and optimization behavior.

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]Reference 12Peri-LN: Revisiting Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2502.02732} 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]Reference 10Gemma 2: Improving Open Language Models at a Practical Sizehttps://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf}

---

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. The minimal PyTorch module below lets you 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?

---

What to check before moving on

• Compute LayerNorm by hand, including mean, population variance, epsilon, learned scale, and learned shift.
• Explain why LayerNorm uses one token's feature values while BatchNorm uses statistics across examples or positions.
• Identify Pre-LN and Post-LN from code, then derive $I + J_FJ_{LN}$ versus $J_{LN}(I + J_F)$.
• Connect Post-LN's initialization-time gradient profile to learning-rate warmup without claiming the result applies to every training setup.
• Explain the reported deep Pre-LN representation-collapse risk and name measurements that would reveal it.
• Implement RMSNorm and state exactly which LayerNorm operations it removes.
• Separate residual-stream normalization, QK-Norm, DeepNorm, and input-output hybrid layouts by the mechanism each changes.

Practice checkpoints

Your 48-layer decoder diverges during first few hundred steps, but the same optimizer works on a 12-layer version. What should you inspect first?

Answer

Start with norm placement and residual-path statistics. If the stack is Post-LN, inspect initialization-time layerwise gradients and check whether warmup or initial learning rate differs from a stable setup. If the model follows a Pre-LN architecture with a final norm, confirm that norm remains present; then inspect attention-logit statistics and candidate controls such as QK-Norm or softmax capping.

Why should you avoid treating Pre-LN as a universal replacement for Post-LN?

Answer

Published comparisons are tied to architecture, depth, initialization, schedule, and task. Xiong et al. establish a Pre-LN initialization and warmup advantage in evaluated setups, while DeepNorm demonstrates a stable Post-LN-derived alternative under a different design. Measure the candidate layout under matched conditions.

When does RMSNorm change the model's behavior less than replacing Pre-LN with Post-LN?

Answer

RMSNorm changes the normalization formula inside a broadly similar residual layout. Switching Pre-LN to Post-LN changes where normalization sits relative to the skip path, which directly changes gradient flow and optimization stability. So RMSNorm is a smaller architectural change than moving the norm across the residual addition.

Common failures and fixes

• Symptom: A reference Pre-LN implementation has unexpected output scale. Cause: Its final normalization was removed because each block already has an input-side norm. Fix: Match the architecture first; GPT-2 and the Pre-LN setup analyzed by Xiong et al. normalize the accumulated stream before prediction.
• Symptom: RMSNorm is described as a faster spelling of LayerNorm. Cause: Mean-centering and RMS rescaling were collapsed into one operation. Fix: State the formulas: LayerNorm centers and rescales; RMSNorm rescales without subtracting the mean.
• Symptom: A Post-LN run diverges and depth alone gets blamed. Cause: Initialization-time gradient scale and early learning rate weren't inspected. Fix: Log layerwise gradient norms, then compare warmup, initial rate, and norm placement under matched conditions.
• Symptom: Attention logits are controlled but another instability remains. Cause: QK-Norm was treated as a replacement for residual normalization. Fix: Treat query-key scale and residual-stream dynamics as separate measurements and interventions.