Layer Normalization

April 2026 · ML Series · Transformers

Why Normalization Matters

Deep neural networks are stacks of transformations. Each layer takes its input, multiplies by a weight matrix, adds a bias, and applies a non-linearity. The output of one layer becomes the input of the next. When the network has 12, 24, or 96 layers, the composition of all these transformations creates an acute problem: activation drift.

Consider what happens without normalization. If the weights of a layer tend to amplify their input slightly - say, by a factor of 1.05 on average - then after 50 layers the activations have grown by 1.05⁵⁰ ≈ 11.5×. After 100 layers, 1.05¹⁰⁰ ≈ 131×. The values explode. Conversely, if each layer shrinks activations by a factor of 0.95, after 100 layers you have 0.95¹⁰⁰ ≈ 0.006 - the signal effectively vanishes.

This phenomenon is sometimes called internal covariate shift (the term from the original Batch Normalization paper, Ioffe & Szegedy, 2015). The idea is that the distribution of inputs to each layer changes during training as the parameters of all preceding layers are updated. Each layer is trying to learn a mapping, but the input distribution it is learning from is constantly shifting under its feet. The optimizer has to chase a moving target.

What Normalization Does

Normalization layers act as stabilizers. They take the activations at a given point in the network, compute their statistics (mean, variance), and re-center and re-scale them to a standard distribution. This provides three concrete benefits:

• Bounded activations: No matter what the preceding layers do, the normalized activations have approximately zero mean and unit variance. The signal neither explodes nor vanishes.
• Smoother loss landscape: Normalization has been shown empirically (Santurkar et al., 2018) to smooth the loss surface, making it more Lipschitz-continuous. This means the gradients change less abruptly, and gradient-based optimization methods can take larger, more confident steps.
• Faster convergence: Networks with normalization layers consistently train faster. Higher learning rates become feasible without divergence, and the model reaches a given loss level in fewer iterations.

Without normalization, training very deep networks (beyond 10-20 layers) is extremely fragile. You need meticulous weight initialization, tiny learning rates, and gradient clipping. With normalization, you can stack hundreds of layers and train with standard hyperparameters.

Batch Normalization: The Predecessor

Before Layer Normalization, the dominant approach was Batch Normalization (BN), introduced by Ioffe and Szegedy in 2015. BN was a breakthrough that made training deep convolutional networks practical and helped enable architectures like ResNet, Inception-v2, and many others.

How Batch Normalization Works

BN normalizes across the batch dimension for each feature independently. Suppose you have a mini-batch of 32 training samples, each represented as a vector of d features. BN takes each feature j (j = 0, 1, ..., d-1), looks at the 32 values that feature takes across all samples in the batch, and normalizes those 32 values to zero mean and unit variance.

Here γ_j and β_j are learned per-feature scale and shift parameters. The key word is across the batch: the statistics depend on what other samples happen to be in the same mini-batch.

Training vs Inference

During training, BN uses the mini-batch statistics (μ_batch, σ²_batch). During inference, you don’t have a batch - you might be processing a single input. So BN maintains running averages of the mean and variance computed during training, and uses those fixed statistics at inference time.

This works well for CNNs because the batch statistics are stable: each feature in a convolutional layer corresponds to a specific spatial filter response, and the distribution of that response is relatively consistent across different batches of images.

Why BN Works for CNNs

• Large batches: Image training typically uses batch sizes of 32, 64, 128, or even larger. The batch statistics are good estimates of the population statistics.
• Fixed input sizes: Images are typically resized to a fixed resolution (224×224, 299×299). Every sample in the batch has the same spatial structure.
• Feature-wise statistics: Each convolutional filter produces one feature map, and normalizing each feature map across the batch ensures that no single filter dominates or vanishes.

Why BatchNorm Fails for Transformers

When researchers attempted to apply Batch Normalization to sequence models - first RNNs, then Transformers - they found that it performed poorly. There are three fundamental reasons why BN is a bad fit for the Transformer architecture:

Layer Normalization: The Formula

Layer Normalization (Ba et al., 2016) takes a fundamentally different approach from Batch Normalization. Instead of normalizing across the batch dimension for each feature, it normalizes across the feature dimension for each sample independently. No batch statistics, no running averages, no dependence on other samples.

Setup

Consider a single token’s representation at some point in the network: a vector x of dimension d (the model dimension, d_model). In a Transformer with d_model = 768, this is a 768-dimensional vector. Layer Normalization operates on this single vector.

Step 1: Compute the Mean

Sum all d elements of the vector and divide by d. This gives the average activation value across the feature dimension. For a 768-dimensional vector, you are averaging 768 numbers.

Step 2: Compute the Variance

Subtract the mean from each element, square the result, and average. This measures how spread out the activations are. A large variance means the values are widely dispersed; a small variance means they are clustered near the mean.

Step 3: Normalize

Subtract the mean (centering: makes the mean zero) and divide by the standard deviation (scaling: makes the variance one). The ε term (typically 10⁻⁵) prevents division by zero when all elements happen to be identical (variance = 0).

After this step, the normalized vector x̂ has mean ≈ 0 and variance ≈ 1 across its features.

Step 4: Scale and Shift (Affine Transform)

The final output is an element-wise affine transformation using learned parameters γ (scale) and β (shift), both vectors of dimension d. These allow the network to undo the normalization if it is beneficial - we will discuss why this matters in the Learned Parameters section.

The Complete Formula

Key Properties

• Sample-independent: Each token is normalized using only its own feature statistics. No dependence on other tokens in the sequence or other samples in the batch.
• Identical at train and test time: Since there are no batch statistics or running averages, the computation is exactly the same during training and inference.
• Sequence-length agnostic: The normalization operates on individual feature vectors. Whether the sequence has 5 tokens or 5,000 tokens, each token is normalized the same way.
• Differentiable: All operations (mean, variance, division, affine transform) are smooth and differentiable, so gradients flow through without issues.

Numerical Walkthrough

Let’s work through a concrete example to make the formula tangible. We’ll use a small 4-dimensional vector (in practice d_model is 768 or larger, but the math is identical).

Full Computation Table

Verify: the normalized values sum to approximately 0 (mean = 0) and their variance is approximately 1. This is the invariant that Layer Normalization enforces at every layer, for every token, at every training step.

BatchNorm vs LayerNorm: Visual Comparison

The key difference between BatchNorm and LayerNorm is the axis of normalization. This SVG diagram makes it crystal clear. Imagine a matrix where rows are samples (or tokens) and columns are features:

BatchNorm (left): For each feature column, compute the mean and variance over all samples in the batch. The highlighted red column shows feature f1 being normalized across all 4 samples. This couples every sample to every other sample in the batch.

LayerNorm (right): For each sample row, compute the mean and variance over all features. The highlighted blue row shows sample 1 being normalized across all 5 features. Each sample is self-contained - no coupling to other samples.

Interactive Animation: LayerNorm Step by Step

The animation below visualizes Layer Normalization applied to a 6-element vector. Each step of the computation is shown as a transformation of the bar chart: the bars shift, compress, and scale as the normalization proceeds. Use the Step button to advance through each stage and the Reset button to start over.

The Learned Parameters γ and β

After normalizing to zero mean and unit variance, why do we immediately apply a learned scale (γ) and shift (β)? Doesn’t this just undo the normalization?

Why They Are Necessary

Normalization constrains the output to have mean 0 and variance 1. But the optimal representation at a given layer might not have mean 0 and variance 1. Perhaps some features should be larger than others. Perhaps the mean should be shifted away from zero for a particular dimension that represents a bias-like signal.

Without γ and β, the network loses expressiveness. The normalization forces all layers to produce outputs in a very narrow band, which can limit the functions the network can represent. In the worst case, it could prevent the network from learning the identity function (outputting its input unchanged), which is important for residual connections.

The Identity Initialization Trick

γ is initialized to all 1s and β is initialized to all 0s. This means that at initialization:

The affine transform is the identity - it does nothing. The output is just the normalized vector. During training, the optimizer adjusts γ and β to find the optimal scale and shift for each feature dimension. Some dimensions might end up with γ > 1 (amplified), others with γ < 1 (suppressed). Some might have positive β, others negative.

Parameter Count

Each LayerNorm has 2 × d_model learnable parameters (γ and β, each of size d_model). For a Transformer with d_model = 768, that’s 1,536 parameters per LayerNorm. A standard GPT-2 Small has about 25 LayerNorm layers (2 per Transformer block, plus 1 final), so the total LayerNorm parameters are roughly 25 × 1,536 = 38,400 - a tiny fraction of the model’s 117M total parameters.

What Does the Network Learn?

In practice, after training, the γ values tend to cluster around 1 (many features remain at roughly unit scale) but with meaningful deviations. Some features consistently get amplified (γ ≈ 1.5–2.0), indicating they carry important signals. Others get suppressed (γ ≈ 0.3–0.5). The β values remain near 0 on average but develop per-feature biases that reflect the learned data distribution.

Researchers have even used the learned γ values as a form of feature importance analysis: features with large γ are features the model considers important.

RMSNorm: A Simpler Alternative

In 2019, Zhang and Sennrich proposed Root Mean Square Layer Normalization (RMSNorm), a simplified version of LayerNorm that has since been adopted by many high-profile models including LLaMA, LLaMA 2, T5, PaLM, and Gemma.

The Key Insight

Standard LayerNorm performs two operations: re-centering (subtracting the mean) and re-scaling (dividing by the standard deviation). Zhang and Sennrich hypothesized that the re-centering step is less important than the re-scaling step. Their experiments confirmed this: removing mean subtraction had minimal impact on model quality, but removing the variance normalization was catastrophic.

The RMSNorm Formula

That’s it. No mean subtraction, no β parameter. Just divide each element by the root-mean-square of the vector, then scale by γ.

Why It’s Faster

RMSNorm skips two operations that LayerNorm must perform:

1. Mean computation: LayerNorm computes μ = (1/d)∑x_i, which requires a full reduction over the feature dimension.
2. Mean subtraction: LayerNorm subtracts μ from every element, which requires a full pass over the vector.

RMSNorm replaces both with a single computation: ∑x_i², which is just the sum of squares - one reduction, no subtraction pass. In practice, this yields a 10–15% speedup in the normalization operation, which matters at scale because LayerNorm is called many times per forward pass.

RMSNorm Numerical Example

Using the same vector x = [2.0, 4.0, −1.0, 3.0] with γ = [1, 1, 1, 1]:

LayerNorm vs RMSNorm Comparison

Notice that the RMSNorm outputs are not zero-centered - the mean of the RMSNorm output is 0.7303 rather than 0. This is exactly what we expect, since RMSNorm does not subtract the mean. In practice, the learned γ parameters and subsequent layers compensate for this shift, and the model performs equally well.

Pre-Norm vs Post-Norm

Where you place the LayerNorm in relation to the sublayer (attention or FFN) and the residual connection has a profound impact on training stability and final model quality. There are two main configurations:

Post-Norm (Original Transformer)

The original “Attention Is All You Need” paper (Vaswani et al., 2017) used Post-Norm. The residual and the sublayer output are summed first, and then the sum is normalized. The data flow is:

1. Input x passes through the sublayer (e.g., multi-head attention): produces Sublayer(x).
2. Add the residual: x + Sublayer(x).
3. Normalize the sum: LayerNorm(x + Sublayer(x)).

The problem: every path from input to output passes through LayerNorm. During backpropagation, the gradient must flow through the LayerNorm Jacobian at every layer. The Jacobian of LayerNorm is a dense matrix that modifies gradient direction and magnitude. Across many layers, these Jacobians compound, and the gradient signal can degrade.

Pre-Norm (Modern Standard)

Pre-Norm applies the normalization before the sublayer. The residual connection bypasses both the normalization and the sublayer. The data flow is:

1. Normalize the input: LayerNorm(x).
2. Pass the normalized input through the sublayer: Sublayer(LayerNorm(x)).
3. Add the residual (the original, unnormalized input): x + Sublayer(LayerNorm(x)).

The crucial difference: the residual path is a direct, unmodified connection from input to output. During backpropagation, gradients can flow straight through the residual path without passing through any LayerNorm or sublayer Jacobian. This is similar to how skip connections in ResNets provide a gradient highway.

Why Pre-Norm Is More Stable

For shallow models (6–12 layers), both Post-Norm and Pre-Norm work fine. For deep models (32+ layers), Pre-Norm is significantly more stable and often the only configuration that converges without extensive warmup and careful learning rate tuning. GPT-2, GPT-3, LLaMA, and most modern LLMs use Pre-Norm.

The Performance Trade-off

There is a subtle trade-off. Some research (Xiong et al., 2020) has found that Post-Norm can achieve slightly better final performance when training is stable. The hypothesis is that passing gradients through LayerNorm at every layer provides a form of implicit regularization. Pre-Norm’s “clean” gradient highway makes optimization easier but may allow the model to overfit slightly more.

In practice, the stability advantage of Pre-Norm far outweighs the marginal quality difference, especially at scale. When training a 175B parameter model, you cannot afford to have training diverge after days of computation.

Gradient Flow Analysis

Understanding how LayerNorm affects gradient flow is important for understanding why it stabilizes training. The gradient dynamics are more nuanced than simply “keeping activations bounded.”

The LayerNorm Jacobian

For a vector x of dimension d, the Jacobian of the normalization step (ignoring γ and β) is:

This Jacobian has three components:

• 1/σ · I (scaling): Scales the gradient by 1/σ. If activations have large variance, gradients are compressed. If activations have small variance, gradients are amplified. This is the core stabilization mechanism.
• −(1/d) · 11ᵀ (centering): Subtracts the mean of the gradient vector. This ensures the gradients are zero-centered, preventing any consistent directional bias from accumulating across layers.
• −x̂ x̂ᵀ / d (decorrelation): Removes the component of the gradient that is aligned with the normalized activations. This prevents the gradient from simply reinforcing the current activation pattern.

Why This Prevents Explosion and Vanishing

The 1/σ scaling factor is the key. Consider two scenarios:

• Activations are exploding (σ is large): The Jacobian scales gradients down by 1/σ, compressing them. Large forward-pass activations produce small backward-pass gradients. This self-regulates the gradient magnitude.
• Activations are vanishing (σ is small): The Jacobian scales gradients up by 1/σ, amplifying them. Small forward-pass activations produce large backward-pass gradients. Again, self-regulating.

Combined with residual connections, this creates a powerful stabilization loop. The residual connection provides a gradient highway (gradients can bypass sublayers entirely). LayerNorm ensures that gradients passing through sublayers are neither too large nor too small. Together, gradient magnitude stays well-bounded across dozens or hundreds of layers.

Empirical Evidence

Measurements on trained Transformers confirm this analysis. The gradient norm per layer remains remarkably flat across all layers in a Pre-Norm Transformer, varying by less than a factor of 2 from the first layer to the last. Without LayerNorm, the gradient norm can vary by orders of magnitude, with later layers receiving vanishingly small gradients.

Implementation

Clean pseudo-code for both LayerNorm and RMSNorm, following the conventions used in production frameworks.

Layer Normalization

class LayerNorm:
    """Layer Normalization (Ba et al., 2016)"""

    def __init__(self, d_model, eps=1e-5):
        self.gamma = ones(d_model)      # learnable scale, shape (d,)
        self.beta  = zeros(d_model)     # learnable shift, shape (d,)
        self.eps   = eps

    def forward(self, x):
        # x shape: (..., d_model)
        # Compute statistics over the LAST dimension

        mu    = mean(x, dim=-1, keepdim=True)    # mean across features
        var   = var(x, dim=-1, keepdim=True)     # variance across features

        # Normalize: zero mean, unit variance
        x_hat = (x - mu) / sqrt(var + self.eps)

        # Affine transform: learned scale and shift
        return self.gamma * x_hat + self.beta

The key detail: the mean and variance are computed over dim=-1, which is the feature (d_model) dimension. The keepdim=True ensures broadcasting works correctly when subtracting the mean and dividing by the standard deviation.

RMS Normalization

class RMSNorm:
    """Root Mean Square Layer Normalization (Zhang & Sennrich, 2019)
    Used in LLaMA, T5, PaLM, Gemma, Mistral"""

    def __init__(self, d_model, eps=1e-5):
        self.gamma = ones(d_model)      # learnable scale, shape (d,)
        self.eps   = eps               # no beta parameter!

    def forward(self, x):
        # x shape: (..., d_model)

        # Root mean square: sqrt(mean(x^2))
        rms = sqrt(mean(x ** 2, dim=-1, keepdim=True) + self.eps)

        # Normalize and scale (no shift)
        return self.gamma * (x / rms)

Notice how much simpler RMSNorm is. One reduction (mean of squares), one sqrt, one division, one element-wise multiply. No mean subtraction, no separate variance computation, no β addition. This simplicity is why it’s faster and why it’s becoming the default.

Where LayerNorm Appears in a Transformer Block

class TransformerBlock:
    """Pre-Norm Transformer block (GPT-2 style)"""

    def __init__(self, d_model, n_heads, d_ff):
        self.ln1  = LayerNorm(d_model)
        self.attn = MultiHeadAttention(d_model, n_heads)
        self.ln2  = LayerNorm(d_model)
        self.ffn  = FeedForward(d_model, d_ff)

    def forward(self, x):
        # Pre-Norm: normalize BEFORE sublayer
        # Residual: add ORIGINAL input (not normalized)

        x = x + self.attn(self.ln1(x))   # attention sub-block
        x = x + self.ffn(self.ln2(x))    # FFN sub-block
        return x

Each Transformer block has two LayerNorm instances: one before attention, one before the feed-forward network. In a 12-layer Transformer, that’s 24 LayerNorm operations plus typically one final LayerNorm after the last block, for a total of 25.

Summary

• Normalization prevents activation drift in deep networks by keeping activations in a well-behaved range, enabling faster convergence and stable training of very deep architectures.
• Batch Normalization normalizes across the batch dimension for each feature. It works well for CNNs with large, fixed-size batches but fails for Transformers due to variable sequence lengths, small batches, and autoregressive decoding.
• Layer Normalization normalizes across the feature dimension for each sample independently: compute mean and variance of the d_model-dimensional feature vector, normalize, then apply learned γ (scale) and β (shift).
• LayerNorm is sample-independent, identical at train and test time, and agnostic to sequence length - exactly the properties Transformers need.
• The learned parameters γ and β (2 × d_model total) let the model recover any representation, acting as a “pressure release valve” on the normalization constraint.
• RMSNorm simplifies LayerNorm by removing mean subtraction and the β parameter. It divides by the root-mean-square instead of the standard deviation. It is 10–15% faster and is used in LLaMA, Mistral, T5, PaLM, and Gemma.
• Pre-Norm (x + Sublayer(LN(x))) places normalization before the sublayer, providing a clean gradient highway through the residual. It is more stable than Post-Norm for deep networks and is the modern default.
• Gradient flow: LayerNorm’s Jacobian acts as an adaptive gradient normalizer - compressing large gradients, amplifying small ones, and centering them - which keeps gradient magnitudes well-bounded across many layers.