Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Pre-Norm vs Post-Norm › The Final Layer Norm

Michael BrenndoerferUpdated June 10, 20252 min read

Language AI HandbookAll chapters

Part 15 of 67Transformer Blocks

Chapter 106Pre-Norm vs Post-Norm › The Final Layer Norm

Reading section 9 of 18 · The Final Layer Norm · 2 minSection 9 of 18 · 2 min

Section 9 of 18 · 2 min read

Pre-Norm vs Post-Norm › The Final Layer NormOne section · Sections in this chapter9 / 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

Pre-norm architectures typically add a final layer normalization after all transformer blocks but before the output projection. This serves two purposes: it normalizes the accumulated activations, and it provides a consistent interface for the output layer.

In[18]:

Code

def pre_norm_transformer(x, n_layers, d_model, d_ff):
    """
    Complete pre-norm transformer stack with final normalization.

    Args:
        x: Input embeddings of shape (seq_len, d_model)
        n_layers: Number of transformer blocks
        d_model: Model dimension
        d_ff: Feed-forward hidden dimension

    Returns:
        Output representations of shape (seq_len, d_model)
    """
    np.random.seed(42)

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

    for i in range(n_layers):
        # Pre-norm attention block
        x_norm = layer_norm(x, gamma, beta)
        # Simplified attention (just a projection for demonstration)
        W_attn = np.random.randn(d_model, d_model) * 0.02
        attn_out = x_norm @ W_attn
        x = x + attn_out

        # Pre-norm feed-forward block
        x_norm = layer_norm(x, gamma, beta)
        W1 = np.random.randn(d_model, d_ff) * 0.02
        W2 = np.random.randn(d_ff, d_model) * 0.02
        ff_out = np.maximum(0, x_norm @ W1) @ W2  # ReLU activation
        x = x + ff_out

    # Final layer normalization
    x = layer_norm(x, gamma, beta)

    return x

In[19]:

Code

# Test the complete stack
np.random.seed(42)
x_input = np.random.randn(4, 64)
x_output = pre_norm_transformer(x_input, n_layers=12, d_model=64, d_ff=256)

Out[20]:

Console

Pre-Norm Transformer Stack
========================================
Input shape:  (4, 64)
Output shape: (4, 64)
Input magnitude:  7.78
Output magnitude: 8.00

Despite passing through 12 transformer blocks (each adding residual contributions), the final output magnitude remains controlled. The final layer normalization ensures that regardless of how activations accumulated through the network, the output has well-controlled statistics suitable for the output projection. This design pattern is why pre-norm architectures include a final LayerNorm before the language modeling head.