Layer Normalization: Pre-Norm, Post-Norm, and RMSNorm

Normalize each token’s feature vector before or after a residual update. Placement changes both activations and gradient paths.

The Norm boxes control the scale of the vectors read by attention and FFN blocks. Their placement also changes the path a gradient takes through the residual connections.

Architectures differ in what kind of normalization they use and where they put it in the block.

Prerequisites

You'll want familiarity with the residual stream picture from the self-attention post and how it accumulates updates from attention and FFN blocks. The gradients post covers why exploding and vanishing gradients matter for deep networks. The FFN post describes the second sub-block in each layer.

Why activation scale matters

As we set up in the self-attention post, each block adds its output back to a shared per-token vector — the residual stream. Both attention and the FFN contribute updates at each layer.

So the residual stream is an accumulator. Each layer adds something. And after 32 or 48 or 96 layers of additive updates, the magnitudes of those vectors can drift a lot. Some dimensions grow large. Others stay small. The variance across the vector keeps changing.

This is a problem for two reasons. First, the layers downstream receive inputs with unpredictable scales. The weights in layer 30 were initialized expecting inputs of a certain magnitude, but by the time activations reach layer 30, they could be 10x or 100x larger than the network expected. Second, during backpropagation, gradients flowing through these layers can explode or vanish depending on the local scales. The gradients post describes how these local derivatives combine across layers.

The fix is simple in concept: normalize the activations at regular intervals so they stay in a predictable range.

Why not Batch Normalization?

Batch Normalization, commonly used in CNNs, computes statistics across batch examples for each feature. This makes an example’s normalized activations depend on the other examples included in those statistics.

Layer Normalization instead computes statistics across the feature dimension ($d_\text{model}$) for each token independently. Every token gets its own mean and variance, computed from its own $d_\text{model}$-dimensional vector. No other token contributes to those statistics, and inference needs no running estimate of population statistics.

The LayerNorm equations

Given a single token's activation vector $x$ of dimension $d$:

Compute the mean across features:

Compute the variance:

Normalize:

The epsilon is a tiny constant (usually $10^{-5}$) to prevent division by zero. After this step, the vector has zero mean and approximately unit variance when its variance is large relative to epsilon.

Re-scale and shift with learned parameters:

Here $\gamma$ and $\beta$ are learned vectors of shape $(d)$, one value per dimension. They set a learned gain and offset after normalization. They do not recover the input-dependent mean and scale that normalization removed; their role is to let the network adjust each feature while controlling the scale of its inputs.

RMSNorm: do we even need the mean?

In 2019, Zhang and Sennrich asked a natural question: does the mean-centering step in LayerNorm actually help? They proposed Root Mean Square Normalization (RMSNorm), which drops the mean subtraction entirely:

There is no mean computation, no mean subtraction, and no $\beta$ parameter. You divide by the root-mean-square of the vector and scale by learned gamma weights.

RMSNorm omits some reductions and arithmetic. Zhang and Sennrich reported runtime improvements with comparable task quality in their tested settings. Actual speed depends on the implementation and hardware.

Ignoring epsilon, RMSNorm retains invariance to positive rescaling while giving up invariance to adding a constant to every feature. If a vector already has near-zero mean, the two normalizers give similar outputs. That algebraic case does not show that all trained activations are mean-zero, or that fixed learned gains can reproduce input-dependent centering.

RMSNorm with pre-norm placement is common in open model families, including Llama. Placement and additional normalization sites still vary by architecture.

Post-Norm vs Pre-Norm

Normalization can go before a sub-block or after its residual addition:

Post-Norm is what the original transformer paper describes. The norm goes after the residual addition:

Pre-Norm puts the norm before the sub-block, and the residual addition happens on the raw (unnormalized) stream:

In Post-Norm, normalization sits on the main residual path. Xiong et al. (2020) analyze large gradients near the output at initialization and explain why a large initial learning rate can destabilize that setup. Warmup reduces those early updates.

In Pre-Norm, the derivative has an identity term: $I+J_f J_{Norm}$. The branch derivative still matters and can amplify or cancel components. Xiong et al. found better-behaved initialization gradients in their Pre-Norm setup.

Pre-Norm leaves the residual stream itself unnormalized between additions, so its norm can grow as updates accumulate. The amount and direction of that change depend on the learned updates.

The toy above holds branch updates fixed to isolate placement. Its Pre-Norm-like path adds them directly; its Post-Norm path rescales after addition. In a trained network the updates depend on the current activations and learned parameters, so growth need not follow the toy trajectory.

The original Attention Is All You Need architecture uses Post-Norm. GPT-2 uses Pre-Norm. These choices produce different gradient paths, which is why a reproduction needs to specify normalization placement as well as the number and width of its layers. We will compare the layouts again in the close reading of AIAYN.

Worked example

Let's trace through both normalizations with a concrete vector: $[12.3, -0.5, 8.1, 0.2]$. This is a deliberately wonky vector with a noticeable mean, so the difference between LN and RMSNorm shows up.

LayerNorm:

Mean: $\mu = (12.3 + (-0.5) + 8.1 + 0.2) / 4 = 5.025$

Variance: $\sigma^2 = ((12.3 - 5.025)^2 + (-0.5 - 5.025)^2 + (8.1 - 5.025)^2 + (0.2 - 5.025)^2) / 4 \approx 29.046875$

Normalize each dimension:

• $(12.3 - 5.025) / \sqrt{29.046875} \approx 1.350$
• $(-0.5 - 5.025) / \sqrt{29.046875} \approx -1.025$
• $(8.1 - 5.025) / \sqrt{29.046875} \approx 0.571$
• $(0.2 - 5.025) / \sqrt{29.046875} \approx -0.895$

With $\gamma = [1, 1, 1, 1]$ and $\beta = [0, 0, 0, 0]$ (untrained), the output is $[1.350, -1.025, 0.571, -0.895]$. Zero mean and approximately unit variance (before rounding, and allowing for epsilon).

RMSNorm:

$\text{RMS} = \sqrt{(12.3^2 + 0.5^2 + 8.1^2 + 0.2^2) / 4} = \sqrt{(151.29 + 0.25 + 65.61 + 0.04)/4} = \sqrt{54.2975} \approx 7.369$

Divide each dimension by RMS:

• $12.3 / 7.369 \approx 1.669$
• $-0.5 / 7.369 \approx -0.068$
• $8.1 / 7.369 \approx 1.099$
• $0.2 / 7.369 \approx 0.027$

Output: $[1.669, -0.068, 1.099, 0.027]$. The mean of this output is about 0.68, not zero. RMSNorm controls the scale and leaves the offset alone.

Where norms live in the architecture

In a Pre-Norm transformer, the flow through one block looks like this:

1. Normalize the residual stream.
2. Feed the normalized stream into multi-head attention.
3. Add the attention output back to the (unnormalized) residual stream.
4. Normalize the residual stream again.
5. Feed it into the FFN.
6. Add the FFN output back to the residual stream.

In pseudocode: x = x + Attn(Norm(x)), then x = x + FFN(Norm(x)).

Pre-Norm decoders commonly apply a final normalization before the unembedding projection. It normalizes the accumulated residual stream before the output head; learned gains can then change individual feature scales.

Implementation from scratch

Here are LayerNorm and RMSNorm in PyTorch:

import torch
import torch.nn as nn
 
class LayerNorm(nn.Module):
    def __init__(self, d_model, eps=1e-5):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(d_model))
        self.beta = nn.Parameter(torch.zeros(d_model))
        self.eps = eps
 
    def forward(self, x):
        # x shape: (batch, seq_len, d_model)
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        x_norm = (x - mean) / torch.sqrt(var + self.eps)
        return self.gamma * x_norm + self.beta
 
class RMSNorm(nn.Module):
    def __init__(self, d_model, eps=1e-5):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(d_model))
        self.eps = eps
 
    def forward(self, x):
        # x shape: (batch, seq_len, d_model)
        rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return self.gamma * (x / rms)
 
# Sanity check: same input, both norms.
x = torch.tensor([[12.3, -0.5, 8.1, 0.2]])  # the worked-example vector
print("LayerNorm:", LayerNorm(4)(x))
print("RMSNorm:  ", RMSNorm(4)(x))
 
# Parameter count comparison.
d_model = 4096
ln = LayerNorm(d_model)
rn = RMSNorm(d_model)
print(f"LayerNorm params: {sum(p.numel() for p in ln.parameters())}")  # 8192
print(f"RMSNorm params:   {sum(p.numel() for p in rn.parameters())}")  # 4096

LayerNorm has $2 d_\text{model}$ parameters (gamma and beta). RMSNorm has $d_\text{model}$ (just gamma). For a model with $d_\text{model} = 4096$ and 32 layers, each layer has two norms, so that's 64 norm sites. The parameter savings from RMSNorm vs LayerNorm is $64 \times 4096 = 262{,}144$ parameters, which is tiny next to the billions in the attention and FFN weights.

You can verify these implementations match PyTorch's built-in torch.nn.LayerNorm:

Misconceptions

"LayerNorm normalizes across the batch." In these transformer blocks, LayerNorm uses only the $d_\text{model}$ features at one token position. BatchNorm instead combines statistics across examples.

"Pre-Norm and Post-Norm are basically interchangeable." Moving normalization changes the residual derivative and can require a different learning-rate or initialization recipe. Xiong et al. found Pre-Norm easier to train in their comparison; Post-Norm variants such as DeepNorm use other scaling choices to support deep networks.

"RMSNorm necessarily sacrifices quality for speed." The original paper found comparable quality in its tested architectures and tasks. Its reported 7–64% running-time reductions were results for those model experiments, not a universal speedup for a normalization kernel or a modern LLM.

What's next

How do these blocks fit together, and what changes with separate encoder and decoder stacks? The next post compares those arrangements in Decoder-Only vs. Encoder-Decoder: Architecture Trade-offs.

Additional reading (and watching)

• Ba, J. L., Kiros, J. R., & Hinton, G. E. (2016). Layer Normalization. arXiv:1607.06450. The original LayerNorm paper that introduced normalization across the feature dimension as a sequence-friendly alternative to BatchNorm.
• Zhang, B. & Sennrich, R. (2019). Root Mean Square Layer Normalization. NeurIPS 2019. Introduced RMSNorm and showed empirically that the mean-centering step in LayerNorm contributes little.
• Xiong, R., et al. (2020). On Layer Normalization in the Transformer Architecture. ICML 2020. Theoretical and empirical analysis of why Pre-Norm trains more stably than Post-Norm.
• Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS 2017. The figure that started the Pre-Norm vs Post-Norm confusion.
• Touvron, H., et al. (2023). LLaMA: Open and Efficient Foundation Language Models. arXiv:2302.13971. Adopted RMSNorm with Pre-Norm as the default recipe for open-weight decoder-only models.
• Elhage, N., et al. (2021). A Mathematical Framework for Transformer Circuits. Anthropic interpretability writeup. Introduces the residual stream as a shared read/write communication channel and provides the mathematical framework for understanding how attention heads and MLPs interact through it.
• Henry, A., et al. (2020). Query-Key Normalization for Transformers. Findings of EMNLP 2020. The original QK-Norm proposal that several large-scale models later adopted.