Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Residual Connections › Pre-Norm vs Post-Norm Residuals

Michael BrenndoerferUpdated June 8, 20258 min read

Language AI HandbookAll chapters

Part 15 of 67Transformer Blocks

Chapter 103Residual Connections › Pre-Norm vs Post-Norm Residuals

Reading section 7 of 12 · Pre-Norm vs Post-Norm Residuals · 8 minSection 7 of 12 · 8 min

Section 7 of 12 · 8 min read

Residual Connections › Pre-Norm vs Post-Norm ResidualsOne section · Sections in this chapter7 / 12 ↓

Sections in this chapter12

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

We've established that residual connections enable deep network training through gradient highways. But there's a subtle design choice that significantly affects training dynamics: where do you place layer normalization relative to the residual addition?

Layer normalization is essential for stable training. It prevents activations from exploding or vanishing by normalizing to zero mean and unit variance. But its placement interacts with the skip connection in ways that matter for very deep networks. The two configurations are post-norm (the original transformer design) and pre-norm (now the dominant choice for large models).

Post-Norm: The Original ConfigurationLink Copied

The original "Attention Is All You Need" paper uses post-norm, where normalization comes after the residual addition:

y = \text{LayerNorm}(x + F(x))

where:

$x$ : the input to the block
$F(x)$ : the sub-layer output (e.g., attention or FFN)
$x + F(x)$ : the residual sum
$\text{LayerNorm}(\cdot)$ : layer normalization applied to the combined result
$y$ : the final block output, normalized

In[22]:

Code

class PostNormBlock(nn.Module):
    """Post-norm configuration: normalize after residual."""

    def __init__(self, dim):
        super().__init__()
        self.linear = nn.Linear(dim, dim)
        self.relu = nn.ReLU()
        self.norm = nn.LayerNorm(dim)

    def forward(self, x):
        residual = self.linear(x)
        residual = self.relu(residual)
        return self.norm(x + residual)  # Norm AFTER addition


class PreNormBlock(nn.Module):
    """Pre-norm configuration: normalize before sub-layer."""

    def __init__(self, dim):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.linear = nn.Linear(dim, dim)
        self.relu = nn.ReLU()

    def forward(self, x):
        normalized = self.norm(x)  # Norm BEFORE sub-layer
        residual = self.linear(normalized)
        residual = self.relu(residual)
        return x + residual  # Clean residual path

Pre-Norm: The Modern DefaultLink Copied

Pre-norm rearranges the components, applying normalization before the sub-layer rather than after the residual addition:

y = x + F(\text{LayerNorm}(x))

where:

$\text{LayerNorm}(x)$ : normalization applied to the input before the sub-layer
$F(\text{LayerNorm}(x))$ : the sub-layer operates on normalized input
$x$ : the original input, added directly (not normalized in the skip path)
$y$ : the output, combining raw input with transformed normalized input

This rearrangement might seem cosmetic, but it has a direct consequence for gradient flow. In post-norm, gradients flowing through the skip connection must pass through the LayerNorm operation, where they're multiplied by the normalization's derivative. In pre-norm, the skip path is completely clean: the $+x$ term passes gradients backward without any modification.

The pure gradient highway we derived earlier, where the direct path contributes exactly 1 to the gradient, only holds perfectly for pre-norm. Post-norm introduces the normalization's Jacobian into the skip path, which can amplify or attenuate gradients depending on the input statistics.

In[23]:

Code

def analyze_gradient_path(block_type, depth=20):
    """Analyze gradient flow through stacked blocks."""
    dim = 64

    if block_type == "post_norm":
        blocks = nn.ModuleList([PostNormBlock(dim) for _ in range(depth)])
    else:
        blocks = nn.ModuleList([PreNormBlock(dim) for _ in range(depth)])

    # Forward pass
    x = torch.randn(4, dim, requires_grad=True)
    h = x
    for block in blocks:
        h = block(h)

    # Backward pass
    loss = h.sum()
    loss.backward()

    return x.grad.norm().item()


# Compare gradient norms
torch.manual_seed(42)
post_norm_grad = analyze_gradient_path("post_norm", depth=20)
torch.manual_seed(42)
pre_norm_grad = analyze_gradient_path("pre_norm", depth=20)

Out[24]:

Console

Gradient Flow Comparison (20-layer networks)
--------------------------------------------------
Post-norm input gradient norm: 0.0000
Pre-norm input gradient norm:  35.7888
Ratio (pre/post): 113825718.14x

The gradient norm comparison quantifies what the theory predicts. Pre-norm maintains substantially stronger gradients at the input layer after passing through 20 blocks. The ratio between them indicates that post-norm's normalization operations in the skip path attenuate the gradient signal.

For a 20-layer network, this difference is noticeable but perhaps manageable. For a 96-layer model like GPT-3, the gap compounds dramatically. Pre-norm's clean gradient path becomes essential for stable training at that scale.

Out[25]:

Visualization

Line plot comparing gradient norms for pre-norm and post-norm configurations from 5 to 50 layers. — Gradient magnitude at the input layer for pre-norm and post-norm configurations across increasing network depths. Pre-norm maintains stable gradients regardless of depth, while post-norm shows progressive attenuation.

Training Stability in PracticeLink Copied

The gradient analysis translates directly to practical training differences. Pre-norm transformers are more stable and forgiving during training, especially in challenging regimes:

Very deep networks (50+ layers): The cumulative effect of normalization Jacobians in post-norm's skip path compounds with depth. Pre-norm avoids this entirely.
Large learning rates: Post-norm requires more careful learning rate tuning because gradient magnitudes are less predictable. Pre-norm tolerates larger learning rates.
Reduced warmup schedules: Post-norm typically needs extended warmup periods to avoid early training instabilities. Pre-norm can often skip or shorten warmup.

The mathematical reason is clear: with pre-norm, the skip connection contributes exactly $\frac{\partial \mathcal{L}}{\partial x_L}$ to the gradient at any earlier layer, with no modification and no layer-dependent scaling. With post-norm, this contribution is modulated by the chain of normalization derivatives, introducing variance that can destabilize training.

Out[26]:

Visualization

Diagram showing post-norm block with LayerNorm after residual addition. — Schematic of post-norm configuration. Normalization sits in the residual path, affecting gradient flow through skip connections.

Diagram showing pre-norm block with LayerNorm before the sub-layer. — Schematic of pre-norm configuration. The residual path is clean, allowing unimpeded gradient flow to early layers.

When to Use Each ConfigurationLink Copied

The choice between pre-norm and post-norm depends on your specific situation:

Use Pre-Norm when:

Training very deep transformers (24+ layers)
You want faster training convergence
Stability is more important than squeezing out final performance
Using larger learning rates or shorter warmup schedules

Use Post-Norm when:

Fine-tuning pre-trained models that used post-norm
Willing to use careful learning rate scheduling and warmup
Optimizing for final quality over training convenience
Working with shallower networks where stability is less critical

Modern large language models predominantly use pre-norm due to its training stability at scale. GPT-2, GPT-3, LLaMA, and many other recent models adopt pre-norm configurations.