Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Pre-Norm vs Post-Norm › Training Stability: The Gradient Flow Perspective

Michael BrenndoerferUpdated June 10, 202511 min read

Language AI HandbookAll chapters

Part 15 of 67Transformer Blocks

Chapter 106Pre-Norm vs Post-Norm › Training Stability: The Gradient Flow Perspective

Reading section 5 of 18 · Training Stability: The Gradient Flow Perspective · 11 minSection 5 of 18 · 11 min

Section 5 of 18 · 11 min read

Pre-Norm vs Post-Norm › Training Stability: The Gradient Flow PerspectiveOne section · Sections in this chapter5 / 18 ↓

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

The practical motivation for pre-norm comes from training stability. While the forward pass differences we've seen are instructive, the real story unfolds during training, when gradients must flow backward through potentially hundreds of layers. Understanding this gradient flow explains why pre-norm enables stable training of very deep networks while post-norm struggles.

The Vanishing Gradient Problem in Deep NetworksLink Copied

Training neural networks requires computing how changes in each parameter affect the final loss. We do this through backpropagation: starting from the loss, we work backward through the network, computing gradients layer by layer using the chain rule.

Here's the challenge. The chain rule multiplies gradients together. If you have 100 layers and each layer slightly shrinks the gradient (say, by a factor of 0.99), the gradient reaching the first layer is $0.99^{100} \approx 0.37$ of what it started as. If each layer shrinks it by 0.95, you get $0.95^{100} \approx 0.006$ , and the earliest layers barely learn anything. This is the vanishing gradient problem.

The converse, exploding gradients, happens when each layer slightly amplifies the gradient. Even small amplification factors compound exponentially, leading to numerical overflow and unstable training.

Out[10]:

Visualization

Line plot showing exponential decay of gradient magnitude for different per-layer scaling factors (0.99, 0.98, 0.95) across 100 layers. — Exponential gradient decay across network depth. Even small per-layer decay factors compound dramatically: a 5% reduction per layer leaves only 0.6% of the original gradient after 100 layers. The dashed lines mark typical transformer depths (12 and 24 layers).

The visualization makes the exponential decay visceral. Even at the modest depth of 12 layers (BERT's configuration), a 2% per-layer reduction leaves only about 78% of the gradient. At 100 layers, a 5% reduction per layer is catastrophic, leaving less than 1% of the original gradient signal.

Residual connections were designed to address this problem by providing a direct path for gradients. But as we'll see, where you place the layer normalization determines whether this path remains truly direct.

Deriving the Gradient FlowLink Copied

To understand the stability difference, we need to trace how gradients flow through each block type. Let's denote:

$L$ : the loss function we're minimizing
$\mathbf{x}$ : the input to a transformer block
$\mathbf{x}'$ : the output of the block
$\frac{\partial L}{\partial \mathbf{x}'}$ : the gradient arriving from subsequent layers (we receive this during backpropagation)
$\frac{\partial L}{\partial \mathbf{x}}$ : the gradient we need to compute and pass to earlier layers

The chain rule tells us that $\frac{\partial L}{\partial \mathbf{x}} = \frac{\partial L}{\partial \mathbf{x}'} \cdot \frac{\partial \mathbf{x}'}{\partial \mathbf{x}}$ . The key question is: what does $\frac{\partial \mathbf{x}'}{\partial \mathbf{x}}$ look like for each block type?

Gradient Flow in Post-NormLink Copied

Recall the post-norm formulation: $\mathbf{x}' = \text{LayerNorm}(\mathbf{x} + \text{Sublayer}(\mathbf{x}))$ . To compute the gradient, we apply the chain rule from outside to inside:

First, the gradient passes through LayerNorm
Then, it splits into two paths: through the residual ( $+\mathbf{x}$ ) and through the sublayer

Mathematically, this gives:

\frac{\partial L}{\partial \mathbf{x}} = \frac{\partial L}{\partial \mathbf{x}'} \cdot \frac{\partial \text{LayerNorm}}{\partial (\mathbf{x} + \text{Sublayer}(\mathbf{x}))} \cdot \left(1 + \frac{\partial \text{Sublayer}}{\partial \mathbf{x}}\right)

where:

$\frac{\partial L}{\partial \mathbf{x}}$ : the gradient of the loss with respect to the block input, which we want to compute
$\frac{\partial L}{\partial \mathbf{x}'}$ : the gradient of the loss with respect to the block output (provided by the next layer during backpropagation)
$\frac{\partial \text{LayerNorm}}{\partial (\cdot)}$ : the Jacobian of the layer normalization operation, a matrix that describes how small changes in the input to LayerNorm affect its output
$1 + \frac{\partial \text{Sublayer}}{\partial \mathbf{x}}$ : the combined contribution from the residual path (the $1$ ) and the sublayer path

Notice the structure: the LayerNorm Jacobian appears as a multiplicative factor in the main gradient path. This is the crux of the instability problem. Even though LayerNorm has reasonably well-behaved gradients for a single layer (typically with eigenvalues close to 1), stacking many post-norm blocks means multiplying many of these Jacobians together. Small deviations from 1 compound across layers, leading to vanishing or exploding gradients in deep networks.

Gradient Flow in Pre-NormLink Copied

Now consider pre-norm: $\mathbf{x}' = \mathbf{x} + \text{Sublayer}(\text{LayerNorm}(\mathbf{x}))$ . The structural difference determines how the gradient splits. Applying the chain rule:

First, the gradient splits at the addition: one path goes through the residual ( $+\mathbf{x}$ ), another through the sublayer
Only the sublayer path passes through LayerNorm

This yields:

\frac{\partial L}{\partial \mathbf{x}} = \frac{\partial L}{\partial \mathbf{x}'} \cdot \left(1 + \frac{\partial \text{Sublayer}}{\partial \text{LayerNorm}(\mathbf{x})} \cdot \frac{\partial \text{LayerNorm}}{\partial \mathbf{x}}\right)

where:

$\frac{\partial L}{\partial \mathbf{x}}$ : the gradient of the loss with respect to the block input
$\frac{\partial L}{\partial \mathbf{x}'}$ : the gradient of the loss with respect to the block output
$1$ : the gradient contribution from the direct residual path (since $\frac{\partial \mathbf{x}}{\partial \mathbf{x}} = 1$ )
$\frac{\partial \text{Sublayer}}{\partial \text{LayerNorm}(\mathbf{x})}$ : the Jacobian of the sublayer with respect to its (normalized) input
$\frac{\partial \text{LayerNorm}}{\partial \mathbf{x}}$ : the Jacobian of layer normalization

The Gradient HighwayLink Copied

The key insight lies in that $+1$ term inside the parentheses. In pre-norm, the gradient from the output $\frac{\partial L}{\partial \mathbf{x}'}$ has a direct, unmodified path back to the input. This path bypasses both the sublayer and the layer normalization entirely.

Think of it as a gradient highway: no matter what happens in the sublayer (vanishing activations, saturated nonlinearities, poorly conditioned weight matrices), the gradient can always flow back through the residual connection. The $+1$ term ensures that at minimum, the full gradient $\frac{\partial L}{\partial \mathbf{x}'}$ reaches the input.

In post-norm, this highway doesn't exist in the same form. The LayerNorm Jacobian gates all gradient flow, including the residual path. If these Jacobians systematically shrink or grow gradients even slightly, the effects compound across many layers.

This mathematical property explains the empirical observations:

Pre-norm networks train stably even at 100+ layers
Post-norm networks require careful warmup schedules and smaller learning rates
Pre-norm tolerates larger learning rates without diverging

Out[11]:

Visualization

Heatmap showing gradient flow paths in post-norm, with LayerNorm Jacobian gating all paths. — Post-norm gradient paths: All gradients must pass through the LayerNorm Jacobian (shown in pink). There is no direct path from output to input, so every gradient signal is modulated.

Heatmap showing gradient flow paths in pre-norm, with a clear diagonal identity path. — Pre-norm gradient paths: The identity path (dark blue diagonal) provides a direct gradient highway. Gradients can flow unimpeded from any layer''s output to any earlier layer''s input.

The heatmaps reveal the structural difference in gradient flow. In post-norm, gradient strength decays as it travels backward through layers, with no escape from the cascading LayerNorm Jacobians. In pre-norm, the bright diagonal represents the gradient highway: full-strength gradients flowing directly from any layer's output to its input, bypassing all intermediate transformations.

Let's simulate this difference with a numerical experiment:

In[12]:

Code

def simulate_gradient_flow(n_layers, norm_type="pre"):
    """
    Simulate gradient magnitude through stacked transformer blocks.

    Args:
        n_layers: Number of transformer blocks to stack
        norm_type: 'pre' or 'post'

    Returns:
        List of gradient magnitudes at each layer
    """
    np.random.seed(42)
    d_model = 64

    # Initialize parameters for each layer
    sublayer_weights = [
        np.random.randn(d_model, d_model) * 0.02 for _ in range(n_layers)
    ]
    gamma = np.ones(d_model)
    beta = np.zeros(d_model)

    # Forward pass: track activations
    x = np.random.randn(1, d_model)
    activations = [x.copy()]

    for i in range(n_layers):
        W = sublayer_weights[i]

        if norm_type == "pre":
            x_norm = layer_norm(x, gamma, beta)
            sublayer_out = np.tanh(x_norm @ W)
            x = x + sublayer_out
        else:  # post
            sublayer_out = np.tanh(x @ W)
            x = layer_norm(x + sublayer_out, gamma, beta)

        activations.append(x.copy())

    # Backward pass: track gradient magnitudes
    # Start with unit gradient at output
    grad = np.ones_like(x)
    # Express magnitudes relative to the unit-gradient baseline. A raw vector
    # norm starts at sqrt(d_model), which would hide the curves on this chart.
    gradient_baseline = np.sqrt(d_model)
    gradient_magnitudes = [np.linalg.norm(grad) / gradient_baseline]

    for i in range(n_layers - 1, -1, -1):
        W = sublayer_weights[i]
        x_prev = activations[i]

        if norm_type == "pre":
            # Simplified gradient (actual would involve LN Jacobian)
            # The key point is the +1 from residual
            x_norm = layer_norm(x_prev, gamma, beta)
            tanh_deriv = 1 - np.tanh(x_norm @ W) ** 2
            sublayer_grad = (grad * tanh_deriv) @ W.T

            # Gradient through residual (direct path)
            residual_grad = grad.copy()

            # Combine (simplified: ignoring LN Jacobian details)
            grad = residual_grad + sublayer_grad * 0.5  # LN scaling effect
        else:
            # Post-norm: gradient must go through LN first
            # Simplified: LN tends to have gradient magnitude ~1
            tanh_deriv = 1 - np.tanh(x_prev @ W) ** 2
            sublayer_grad = (grad * tanh_deriv) @ W.T
            grad = grad + sublayer_grad
            grad = (
                grad
                / (np.linalg.norm(grad) + 1e-6)
                * np.linalg.norm(grad)
                * 0.95
            )

        gradient_magnitudes.append(np.linalg.norm(grad) / gradient_baseline)

    return gradient_magnitudes[::-1]


# Compare gradient flow for different depths
depths = [6, 12, 24, 48]
pre_norm_grads = {d: simulate_gradient_flow(d, "pre") for d in depths}
post_norm_grads = {d: simulate_gradient_flow(d, "post") for d in depths}

Out[13]:

Visualization

Line plot showing stable gradient magnitudes across layers for pre-norm at various depths. — Pre-norm gradient magnitudes remain stable across layers and depths. The direct residual path ensures gradients don't vanish even in very deep networks.

Line plot showing declining gradient magnitudes in deeper post-norm networks. — Post-norm gradient magnitudes can decay over many layers. Deeper networks show more pronounced gradient reduction, potentially causing training difficulties.

The visualization illustrates the stability difference. Pre-norm maintains consistent gradient magnitudes across all depths because the residual connection provides a direct gradient path. Post-norm gradients tend to decay in deeper networks, making optimization more challenging.