Layer Normalization › Layer Normalization in Transformers
Language AI HandbookAll chapters Layer Normalization › Layer Normalization in TransformersOne section · Sections in this chapter7 / 16 ↓
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
In transformer architectures, layer normalization appears in two key locations: after the attention mechanism and after the feed-forward network. The original transformer used "post-norm" placement, where normalization comes after the residual connection:
class PostNormTransformerBlock(nn.Module):
"""Transformer block with post-normalization (original architecture)."""
def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
super().__init__()
self.attention = nn.MultiheadAttention(
d_model, n_heads, batch_first=True
)
self.ff = nn.Sequential(
nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model)
)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
# Attention with residual, then normalize
attn_out, _ = self.attention(x, x, x)
x = self.norm1(x + self.dropout(attn_out))
# FFN with residual, then normalize
ff_out = self.ff(x)
x = self.norm2(x + self.dropout(ff_out))
return xModern architectures like GPT-2, GPT-3, and LLaMA use "pre-norm" placement, where normalization comes before the sublayer:
class PreNormTransformerBlock(nn.Module):
"""Transformer block with pre-normalization (modern architecture)."""
def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
super().__init__()
self.attention = nn.MultiheadAttention(
d_model, n_heads, batch_first=True
)
self.ff = nn.Sequential(
nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model)
)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
# Normalize, then attention with residual
normed = self.norm1(x)
attn_out, _ = self.attention(normed, normed, normed)
x = x + self.dropout(attn_out)
# Normalize, then FFN with residual
ff_out = self.ff(self.norm2(x))
x = x + self.dropout(ff_out)
return xPre-Norm vs. Post-Norm: Why the Placement Matters
The difference in placement has clear implications for gradient flow and training dynamics. Understanding why requires tracing the residual stream through a transformer.
In the post-norm design, the residual connection adds the sublayer output to its input, and the sum is then normalized. The residual stream at any depth carries unnormalized activations. When the gradients flow backward, they pass through the normalization operation and then split between the residual path and the sublayer path. The residual path provides a direct gradient highway from the loss to the input embeddings. However, the normalization after each residual addition can interfere with this highway: the gradient must pass through the normalization Jacobian, which can compress or distort the signal.
In the pre-norm design, normalization happens before the sublayer, and the residual connection bypasses the normalization entirely. The residual stream is the raw, unnormalized activations accumulating across layers. The gradient flowing through the residual path does not pass through any normalization operations, giving it a clean highway all the way back to the first layer. This is particularly beneficial in very deep networks, where the clean gradient signal allows earlier layers to receive meaningful updates even during the early stages of training.
Think of the residual stream as a highway and the sublayers as on-ramps. In post-norm, there is a tollbooth (normalization) on the main highway after each on-ramp merge. The tollbooth ensures the merged traffic stays within bounds, but it also introduces bottlenecks and can disrupt the flow of fast-moving traffic (gradients). In pre-norm, the tollbooth sits on the on-ramp before the merge, so the main highway remains unimpeded. The traffic entering from the on-ramp gets regulated, but the existing highway flow is never disrupted.
This explains a well-documented empirical observation: post-norm transformers typically require careful learning rate warmup schedules, while pre-norm transformers can be trained with constant or cosine-decayed learning rates from the start. The warmup in post-norm is essentially giving the network time to settle into a state where the normalized residual stream is stable enough for larger gradient updates. Pre-norm avoids this need by maintaining a stable gradient highway from the beginning of training.
The tradeoff is representation quality at the final layer. In post-norm, the output of the last layer is normalized, putting all outputs in a consistent scale before they reach the final projection layer. In pre-norm, the output is the raw residual accumulation, which may have a non-standard scale. Practitioners have found various ways to address this, such as adding a final layer normalization before the output projection, a practice adopted by GPT-2 and later models.
