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

Normalize activations per token, then rescale with learned gain so the residual stream stays stable across layers.

Why does every transformer diagram have those little "Norm" boxes in it? I glossed over them the first time I read the Attention Is All You Need paper. The attention mechanism was the star, the FFN was the supporting actor, and the norms seemed like plumbing.

If you remove them, the model doesn't train. It diverges within the first few hundred steps, and no amount of learning rate tuning brings it back. So they're holding something together.

The other thing worth understanding is the choices modern architectures make about what kind of normalization to use and where to 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. And if you've read the FFN post, you already know the two sub-blocks that normalization sits between.

Why you can't just stack layers

Think about what happens inside a transformer. As we set up in the self-attention post, each block adds its output back to a shared per-token vector — the residual stream. Attention computes an update, adds it. FFN computes an update, adds it. Layer after 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. If you've read the gradients post, you know this is the fundamental challenge of training deep networks.

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

Why not Batch Normalization?

If you've worked with CNNs, your first instinct might be Batch Normalization. It was the breakthrough normalization technique for image models. BatchNorm computes the mean and variance across the batch dimension for each feature, then normalizes.

In language modeling, the batch dimension is a mess. Each sequence in the batch might be a completely different length, from a completely different domain, with completely different statistical properties. One sequence is a code snippet, the next is poetry, the next is a legal contract. Computing statistics across these and using them to normalize each other doesn't make much sense. Worse, at inference time you often have a batch of one, and BatchNorm has nothing to normalize against.

Layer Normalization sidesteps this entirely. Instead of normalizing across the batch, it normalizes 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 cross-token, no cross-sequence contamination. Each token is normalized in isolation.

That last sentence is easy to miss when you first encounter LayerNorm: the "layer" in LayerNorm is the layer of activations within a single token, not a batch axis or a sequence axis. The batch and sequence dimensions are along for the ride.

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 unit variance.

Re-scale and shift with learned parameters:

Here $\gamma$ and $\beta$ are learned vectors of shape $(d)$, one value per dimension. They let the model recover any scale and shift it actually needs. You might wonder: if we just normalized to zero mean and unit variance, why let the model undo that with learned parameters? Because the normalization stabilizes training dynamics, and the learned parameters let the model express whatever representation it wants within that stable regime. The normalization constrains the optimization landscape, not the representational capacity.

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.

This is faster because it requires fewer operations, and empirically the quality is comparable to full LayerNorm. But "comparable" feels suspicious. We added the mean-centering step for a reason, didn't we?

The intuition is that on real activations, the mean is usually small, so subtracting it doesn't move the vector very much. The dominant action of LayerNorm is the re-scaling. The mean step is mostly cleaning up a tiny shift that, if it mattered, the learned $\gamma$ and $\beta$ could absorb anyway. Once the network has been trained for a while, activations stabilize into roughly mean-zero distributions on their own.

That's the intuition behind why RMSNorm works as well as it does. Most of the value of LayerNorm comes from controlling the magnitude. The centering is a finishing touch that the optimizer can do without.

RMSNorm has been adopted by essentially every major architecture since 2023: LLaMA, Mistral, Gemma, Qwen, DeepSeek, and most others. If you read about a model from 2024 or later, it almost certainly uses RMSNorm with Pre-Norm placement.

Post-Norm vs Pre-Norm

We know what normalization does. But where does it go? This is one of the most consequential architectural decisions in modern transformers, and it has a messy history.

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:

The difference looks minor, but it changes how gradients flow through the entire network.

In Post-Norm, the normalization sits on the main residual path. Every gradient flowing backward has to pass through the normalization operation. Xiong et al. (2020) showed that at initialization, the gradient norms near the output layer are much larger than near the input layer. This creates an imbalanced optimization landscape. The later layers update aggressively while the earlier layers barely move. The standard fix is learning rate warmup: start with a very small learning rate and gradually increase it, giving the network time to find a more balanced gradient regime.

In Pre-Norm, the residual path is completely clean. Gradients flow through the addition operations unmodified. The normalization only affects the branch that goes into each sub-block, which means gradient magnitudes are much more uniform across layers at initialization. No warmup required, or at least much less of it. Training is more stable from the start.

There's a tradeoff, though. The clean residual path that makes Pre-Norm so easy to train also lets the residual stream grow as you stack layers. Each block adds a small update to a spine that nothing ever rescales. After thirty or sixty layers, that spine has a noticeably larger norm than it started with.

This "norm growth" effect is real and shows up in trained models. Interpretability work has documented residual stream norms drifting up by several-fold from the first layer to the last in models like GPT-2 and LLaMA. The model copes by learning $\gamma$ values that compensate, and the final norm before the unembedding cleans up whatever's left. The takeaway is that Pre-Norm trades stable training for a residual stream whose statistics drift across depth, and the field has decided that trade is worth it.

The historical confusion. There's a genuinely funny piece of transformer lore here. The original Attention Is All You Need paper shows Post-Norm in Figure 1. But the released tensor2tensor code was eventually updated to use Pre-Norm. People who implemented their own transformers by looking at the figure ended up with Post-Norm. Others read the code and got Pre-Norm. This caused years of inconsistent reproductions and confusion in the community. GPT-2 used Pre-Norm, and once that worked well at scale, the debate was effectively settled. We'll come back to this 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.047$

Normalize each dimension:

• $(12.3 - 5.025) / \sqrt{29.047} \approx 1.350$
• $(-0.5 - 5.025) / \sqrt{29.047} \approx -1.025$
• $(8.1 - 5.025) / \sqrt{29.047} \approx 0.571$
• $(0.2 - 5.025) / \sqrt{29.047} \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, unit variance.

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 (which is what you'll encounter in practice), 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)).

There's also a final normalization after the last layer, right before the unembedding projection (the linear layer that maps from $d_\text{model}$ to vocabulary size). This final norm catches whatever drift the residual stream has accumulated and presents the output projection with a clean, unit-scale input regardless of how many layers the model has.

Implementation from scratch

Both LayerNorm and RMSNorm are simple to implement. Here they are in PyTorch, so you can see there's no magic:

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. The real win is computational: fewer ops per forward pass, which adds up at every layer of every token of every sequence.

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

Misconceptions

"LayerNorm normalizes across the batch." It doesn't. That's BatchNorm. LayerNorm normalizes within a single token's feature vector, using statistics computed only from the $d_\text{model}$ values for that one token. Each token's normalization is completely independent of every other token in the batch. This is exactly what makes it a good fit for sequences: the next-token-prediction setup means tokens in the same batch can be wildly different, and you don't want them contaminating each other's statistics.

"Pre-Norm and Post-Norm are basically interchangeable." They are not interchangeable at scale. Post-Norm typically requires careful learning rate warmup, and even with warmup it can be less stable as depth grows. The gradient flow properties are fundamentally different. You can make Post-Norm work — DeepNorm and a few variants do — but it takes more effort and more hyperparameter tuning. Pre-Norm is more forgiving, which is why it became the default.

"RMSNorm sacrifices quality for speed." The original RMSNorm paper and the LLaMA technical report both showed that on language modeling tasks, the quality difference between LN and RMSNorm is in the noise. The mean-centering step does very little work on its own once the model is trained, so dropping it costs almost nothing. The speedup is roughly 7-64% on the norm operation itself (depending on the model), which is small per call but adds up across hundreds of norm sites per forward pass.

What's next

We now have all the pieces of a single transformer block: attention, the FFN, residual connections, and normalization. The next post steps back and asks a bigger architectural question. Why did the field converge on decoder-only models, and what do you give up compared to encoder-decoder designs? That's the subject of 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.