Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Pre-Norm vs Post-Norm › Implementation Comparison

Michael BrenndoerferUpdated June 10, 20252 min read

Language AI HandbookAll chapters

Part 15 of 67Transformer Blocks

Chapter 106Pre-Norm vs Post-Norm › Implementation Comparison

Reading section 14 of 18 · Implementation Comparison · 2 minSection 14 of 18 · 2 min

Section 14 of 18 · 2 min read

Pre-Norm vs Post-Norm › Implementation ComparisonOne section · Sections in this chapter14 / 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

Let's implement both variants in a clean, comparable format that highlights their differences:

In[28]:

Code

class TransformerBlock:
    """
    Unified transformer block supporting both pre-norm and post-norm.
    """

    def __init__(self, d_model, d_ff, norm_type="pre"):
        """
        Args:
            d_model: Model dimension
            d_ff: Feed-forward hidden dimension
            norm_type: 'pre' or 'post'
        """
        self.d_model = d_model
        self.d_ff = d_ff
        self.norm_type = norm_type

        np.random.seed(42)

        # Attention weights (simplified)
        self.W_attn = np.random.randn(d_model, d_model) * 0.02

        # Feed-forward weights
        self.W1 = np.random.randn(d_model, d_ff) * 0.02
        self.W2 = np.random.randn(d_ff, d_model) * 0.02

        # Layer norm parameters (2 norms per block)
        self.gamma1 = np.ones(d_model)
        self.beta1 = np.zeros(d_model)
        self.gamma2 = np.ones(d_model)
        self.beta2 = np.zeros(d_model)

    def attention(self, x):
        """Simplified attention (just projection for demo)."""
        return x @ self.W_attn

    def feed_forward(self, x):
        """Two-layer feed-forward with ReLU."""
        hidden = np.maximum(0, x @ self.W1)
        return hidden @ self.W2

    def forward(self, x):
        """Forward pass with specified normalization strategy."""
        if self.norm_type == "pre":
            # Pre-norm attention
            x_norm = layer_norm(x, self.gamma1, self.beta1)
            x = x + self.attention(x_norm)

            # Pre-norm feed-forward
            x_norm = layer_norm(x, self.gamma2, self.beta2)
            x = x + self.feed_forward(x_norm)
        else:
            # Post-norm attention
            x = layer_norm(x + self.attention(x), self.gamma1, self.beta1)

            # Post-norm feed-forward
            x = layer_norm(x + self.feed_forward(x), self.gamma2, self.beta2)

        return x

In[29]:

Code

# Compare forward pass behavior
np.random.seed(42)
x = np.random.randn(4, 64)

pre_block = TransformerBlock(64, 256, norm_type="pre")
post_block = TransformerBlock(64, 256, norm_type="post")

pre_output = pre_block.forward(x)
post_output = post_block.forward(x)

Out[30]:

Console

Single Block Comparison
========================================
Input magnitude:       7.7756
Pre-norm output mag:   7.8477
Post-norm output mag:  8.0000

Even after just one block, the difference is visible. The post-norm output magnitude is close to $\sqrt{d}$ (the expected magnitude for a normalized $d$ -dimensional vector), while pre-norm's output reflects the accumulation of the input and sublayer contributions. This difference compounds across many layers, which is why the architectural choice matters significantly for deep networks.