Skip to main content
Intermediate9 min read12 of 59

Layer Normalization: Keeping Transformer Activations Stable

How LayerNorm normalizes each token's features, why it stabilizes deep Transformers, and how Pre-LN and RMSNorm change modern LLM training.

Layer Normalization: Keeping Transformer Activations Stable

[Definition] Layer Normalization (LayerNorm) normalizes the features of one token representation, then learns a scale and offset. In Transformers it keeps activations and gradients in a trainable range as residual blocks stack deeply.

Why deep residual networks drift

A Transformer block repeatedly updates a residual stream:

text
xₗ₊₁ = xₗ + attention(xₗ)
xₗ₊₂ = xₗ₊₁ + MLP(xₗ₊₁)

Every addition can change the magnitude and distribution of values flowing to the next block. Across dozens or hundreds of layers, activations can become too large, too small, or differently scaled from one batch to another. That makes optimization brittle:

  • large activations can make gradients explode;
  • small activations can make learning signals vanish;
  • changing scales force later layers to keep adapting;
  • mixed-precision arithmetic becomes less forgiving.

Normalization gives each sublayer a predictable input distribution while preserving useful learned differences between features.

The LayerNorm formula

For one token vector x with d features, LayerNorm computes statistics across the feature dimension:

text
μ = (1 / d) Σ xᵢ
σ² = (1 / d) Σ (xᵢ - μ)²

LayerNorm(x) = γ ⊙ (x - μ) / √(σ² + ε) + β

Where:

  • μ is that token vector's feature mean;
  • σ² is its feature variance;
  • ε prevents division by zero;
  • γ and β are learned vectors, one value per feature;
  • means element-wise multiplication.

The normalized middle term has approximately zero mean and unit variance. The model then learns whether each feature should be amplified, reduced, or shifted back through γ and β.

A tiny example

text
x = [2, 4, 6, 8]
μ = 5
σ² = 5
σ ≈ 2.236

normalized x ≈ [-1.34, -0.45, 0.45, 1.34]

The vector changes scale, but its relative feature pattern remains. With initial γ = 1 and β = 0, this normalized vector is the output. Training learns a feature-specific scale and offset only where they help.

LayerNorm is not BatchNorm

PropertyBatch NormalizationLayer Normalization
Statistics computed acrossexamples in a batchfeatures in one token
Depends on batch sizeyesno
Natural forCNN-style image batchessequences and autoregressive decoding
Train/inference behavioruses running batch statisticssame calculation at both stages

Language-model batch sizes, sequence lengths, and one-token decoding make batch-dependent normalization inconvenient. LayerNorm works for a single sequence, a single token, or a large batch without changing its rule.

Pre-LN versus Post-LN

The original Transformer placed normalization after each residual addition:

text
Post-LN:  xₗ₊₁ = LayerNorm(xₗ + Sublayer(xₗ))

Many modern LLMs normalize before attention or the MLP:

text
Pre-LN:   xₗ₊₁ = xₗ + Sublayer(LayerNorm(xₗ))

Pre-LN leaves a direct residual path from xₗ to xₗ₊₁. That path makes gradients easier to propagate through deep stacks, which is why Pre-LN is common in large decoder-only models.

DesignStrengthTradeoff
Post-LNstrongly normalizes every block outputdeep training can be less stable
Pre-LNreliable gradient path and easier optimizationfinal output often needs an extra normalization

[Key Insight] LayerNorm does not fix a bad architecture. It makes optimization less fragile so attention, MLPs, residual connections, data, and learning-rate schedules can work together at scale.

RMSNorm: a simpler modern alternative

Many current LLM families use RMSNorm rather than full LayerNorm. RMSNorm omits mean subtraction and often omits the learned bias:

text
RMS(x) = √((1 / d) Σ xᵢ² + ε)
RMSNorm(x) = γ ⊙ x / RMS(x)

It normalizes root-mean-square magnitude rather than variance around the mean. This is cheaper and has worked well in practice for decoder-only models such as LLaMA-family architectures.

  • LayerNorm controls mean and variance.
  • RMSNorm controls vector magnitude.
  • Both provide a stable scale for the next sublayer.

Where normalization sits in a modern block

text
residual stream x
       │
       ├── RMSNorm / LayerNorm ──> attention ──┐
       │                                        │
       └────────────────────────────────────────+──> x'
                                                │
       ├── RMSNorm / LayerNorm ──> MLP ─────────┤
       │                                        │
       └────────────────────────────────────────+──> next residual stream

The normalization is local to a token's feature vector. It does not normalize attention weights, token positions, or the whole batch.

Production and debugging signals

LayerNorm problems are usually training or conversion problems, not user-facing runtime errors. Check these when a Transformer train or fine-tune is unstable:

  1. NaN/Inf activations — inspect activation and gradient norms per layer.
  2. Wrong epsilon or dtype — very small ε values can underflow in low precision.
  3. Pre-LN/Post-LN mismatch — loading weights into a differently ordered architecture silently changes behaviour.
  4. Fused kernel differences — verify numerical tolerance when swapping implementations.
  5. Residual scaling changes — normalization interacts with initialization and residual scaling choices.

For inference, the normalization parameters are fixed learned weights. They remain part of every forward pass, so quantization and kernel fusion must preserve their numerical behaviour.

Do not confuse these topics

  • L1/L2 norms measure the size or distance of vectors.
  • LayerNorm rescales activation features inside a neural network.
  • Softmax converts scores to a probability distribution.

They share familiar language but solve different problems.

Continue the path

This is the public conceptual anchor. The Pro Transformer Architecture sequence applies the formula inside a full residual block and shows what breaks when the stabilizing step is removed. Continue with Transformer Architecture, then use the learning roadmap for the structured path.