Layer Normalization: Pre-LN vs Post-LN

Layer normalization stabilizes training by normalizing activations within each layer to have zero mean and unit variance, then applying learned scale and shift. Where you place it in the block matters more than you'd think.

Two placements

Post-LN (original Transformer, 2017): normalize after the residual connection.

x = LayerNorm(x + Sublayer(x))

Trains poorly without learning-rate warmup. Gradient signals at depth tend to explode or vanish unless carefully tuned. Most papers from 2017-2019 use Post-LN; modern models avoid it.

Pre-LN (modern standard): normalize before the sublayer.

x = x + Sublayer(LayerNorm(x))

The residual stream remains un-normalized; gradients can flow back through the skip connection without ever being pushed through layer norms. Training is much more stable. GPT-2, GPT-3, BERT-base, every modern Llama, Mistral, and Claude variant uses Pre-LN.

RMSNorm

Modern Llama and Mistral use RMSNorm instead of LayerNorm. Same shape, but skips the mean-centering step — it normalizes by the root mean square of the activations rather than (activation − mean) / std. Slightly faster, slightly fewer parameters, empirically just as stable.

Code

Pre-LN block (modern)·python

class PreLNBlock(nn.Module):
    def __init__(self, d_model, attn, ffn, norm_cls=nn.LayerNorm):
        super().__init__()
        self.norm1 = norm_cls(d_model)
        self.attn = attn
        self.norm2 = norm_cls(d_model)
        self.ffn = ffn
    def forward(self, x):
        x = x + self.attn(self.norm1(x))    # norm BEFORE attention
        x = x + self.ffn(self.norm2(x))      # norm BEFORE FFN
        return x

RMSNorm in 6 lines·python

class RMSNorm(nn.Module):
    def __init__(self, d_model, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d_model))
        self.eps = eps
    def forward(self, x):
        # Skip mean centering; just normalize by RMS
        rms = x.pow(2).mean(-1, keepdim=True).rsqrt()
        return self.weight * (x * (rms + self.eps))

Layer Normalization: Pre-LN vs Post-LN

Two placements

RMSNorm

Code

External links

Exercise

Progress

Comments 0