Pre-Norm vs Post-Norm › Output Scale: An Important Trade-off

Michael BrenndoerferUpdated June 10, 20254 min read
Language AI HandbookAll chapters
Part 15 of 67Transformer Blocks
Chapter 106Pre-Norm vs Post-Norm › Output Scale: An Important Trade-off
Reading section 7 of 18 · Output Scale: An Important Trade-off · 4 minSection 7 of 18 · 4 min

Book contents

Section 7 of 18 · 4 min read

Pre-Norm vs Post-Norm › Output Scale: An Important Trade-offOne section · Sections in this chapter7 / 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

Pre-norm's stability advantage comes with a trade-off: the output of each block accumulates without normalization. This can lead to growing activation magnitudes as we stack more layers.

Let's measure this effect:

In[14]:
Code
def measure_activation_growth(n_layers, norm_type="pre"):
    """
    Measure how activation magnitudes evolve through stacked blocks.
    """
    np.random.seed(42)
    d_model = 64

    x = np.random.randn(4, d_model)  # 4 tokens
    x = x / np.linalg.norm(x, axis=-1, keepdims=True)  # Normalize input

    gamma = np.ones(d_model)
    beta = np.zeros(d_model)

    magnitudes = [np.mean(np.linalg.norm(x, axis=-1))]

    for i in range(n_layers):
        np.random.seed(42 + i)
        W = np.random.randn(d_model, d_model) * 0.1

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

        magnitudes.append(np.mean(np.linalg.norm(x, axis=-1)))

    return magnitudes


# Compare activation growth
n_layers = 24
pre_magnitudes = measure_activation_growth(n_layers, "pre")
post_magnitudes = measure_activation_growth(n_layers, "post")
Out[15]:
Visualization
Line plot comparing activation magnitudes for pre-norm (growing) and post-norm (constant) across 24 layers.
Activation magnitude across layers. Pre-norm (blue) shows gradual accumulation as residuals add up without normalization. Post-norm (orange) maintains constant magnitude due to normalization after each block.
Out[16]:
Console
Activation Magnitude Summary
========================================
Pre-Norm:  Start=1.00, End=19.48, Growth=19.48x
Post-Norm: Start=1.00, End=8.00, Growth=8.00x

The growth in pre-norm isn't necessarily problematic, as the final layer normalization (applied before the output projection in most architectures) handles the accumulated magnitude. However, it does mean that numerical precision matters more in very deep pre-norm networks.

To understand where this growth comes from, we can decompose the output into contributions from the original input and each sublayer:

Out[17]:
Visualization
Stacked bar chart showing the decomposition of activation magnitude into input and sublayer contributions across 12 layers.
Residual contribution decomposition in pre-norm. Each bar shows how the total activation magnitude is built up from the original input (blue) and successive sublayer contributions. The input signal persists through all layers while sublayer contributions accumulate.

The decomposition reveals the structure of pre-norm's activation growth. The original input signal (dark blue) persists unchanged through all layers, which is the residual connection at work. Each sublayer adds its own contribution on top, and these contributions accumulate. The total magnitude (red line) grows steadily, but this growth is predictable and well-behaved, making it easy to compensate for with a final normalization.