Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Pre-Norm vs Post-Norm › Initialization Considerations

Michael BrenndoerferUpdated June 10, 20252 min read

Language AI HandbookAll chapters

Part 15 of 67Transformer Blocks

Chapter 106Pre-Norm vs Post-Norm › Initialization Considerations

Reading section 10 of 18 · Initialization Considerations · 2 minSection 10 of 18 · 2 min

Section 10 of 18 · 2 min read

Pre-Norm vs Post-Norm › Initialization ConsiderationsOne section · Sections in this chapter10 / 18 ↓

Sections in this chapter18

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 choice between pre-norm and post-norm affects how we should initialize the network. Post-norm requires more careful initialization because gradients must flow through the normalization layers.

The intuition behind initialization scaling is that in a pre-norm transformer, each block adds its sublayer output to the residual stream. If there are $n$ layers and each sublayer contributes a vector with variance $\sigma^2$ , the total residual stream after all layers has variance roughly $n \cdot \sigma^2$ (since the contributions add). To keep the final residual stream variance similar to the input variance, we want $n \cdot \sigma^2 \approx 1$ , which suggests initializing with $\sigma \approx 1/\sqrt{n}$ . GPT-2 implemented a version of this by scaling the output projections of attention and feed-forward sublayers by $1/\sqrt{2n}$ , where the factor of 2 accounts for having two sublayers (attention and feed-forward) per block.

For pre-norm, a common approach is to scale the output projection of each sublayer by a factor related to the network depth:

In[21]:

Code

def initialize_pre_norm_weights(d_model, n_layers):
    """
    Initialize weights for a pre-norm transformer with proper scaling.

    Following the GPT-2 approach: scale residual connections by 1/sqrt(2*n_layers)
    """
    # Standard initialization
    std = 0.02

    # Residual scaling factor
    residual_scale = 1 / np.sqrt(2 * n_layers)

    weights = {
        "attention_proj": np.random.randn(d_model, d_model) * std,
        "ff_proj": np.random.randn(d_model, d_model) * std * residual_scale,
    }

    return weights, residual_scale


weights, scale = initialize_pre_norm_weights(512, 24)

Out[22]:

Console

Initialization for 24-layer Pre-Norm Transformer
=============================================
Standard weight std:    0.02
Residual scale factor:  0.1443
Scaled projection std:  0.002887

The residual scale factor of approximately 0.14 significantly reduces the contribution of each sublayer's output. With 24 layers (and 2 sublayers per block, hence 48 residual additions), this scaling prevents the accumulated sum from exploding. The GPT-2 paper introduced this technique specifically to enable training of deeper pre-norm networks.