GPT-2 › Architectural Changes from GPT-1
Language AI HandbookAll chapters GPT-2 › Architectural Changes from GPT-1One section · Sections in this chapter3 / 11 ↓
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
GPT-2 maintained the decoder-only transformer architecture from GPT-1 but introduced several modifications that proved important for stable training at larger scales. These changes became standard practice for subsequent language models. Understanding each change requires understanding the problem it was designed to solve: as networks grow deeper, gradients become harder to propagate reliably during backpropagation, and small numerical instabilities compound across layers.
Think of training a deep network as a game of telephone with 48 players (the layers). Each player passes a message to the next, and each also passes feedback in the opposite direction to help earlier players improve. If any player distorts the feedback signal, every player upstream receives corrupted information. GPT-2's architectural changes are engineering solutions to this telephone problem. This keeps useful gradient information survives the journey from the output layer all the way back to the first embedding layer.
Pre-Normalization
The most significant architectural change was moving layer normalization before each sub-block rather than after. GPT-1 and the original transformer used post-normalization, applying LayerNorm after the residual connection. GPT-2 switched to pre-normalization, applying LayerNorm before the attention and feed-forward operations.
To understand why this matters, consider what happens to activations inside a deep residual network. At each layer, the residual path adds information to the main stream. In post-normalization, unnormalized values from the residual path are added together and then normalized. The problem is that the unnormalized values can have high and unpredictable variance, which makes the normalization layer's job harder and makes gradients flowing backward through the normalization potentially unstable.
With pre-normalization, the LayerNorm is applied to the input before it enters the attention or feed-forward block. The residual connection then adds the block's output to the original, un-normalized input. This means the inputs to each block are always well-behaved, and the gradient signal traveling backward through the residual connection is clean because it bypasses the normalization entirely. For 12-layer networks, both approaches often work. For 48-layer networks, pre-normalization is nearly needed for avoiding training divergence.
class GPT2Block(nn.Module):
"""GPT-2 transformer block with pre-normalization."""
def __init__(
self, hidden_size: int = 768, num_heads: int = 12, dropout: float = 0.1
):
super().__init__()
self.hidden_size = hidden_size
self.num_heads = num_heads
# Pre-normalization: LayerNorm BEFORE attention and FFN
self.ln_1 = nn.LayerNorm(hidden_size, eps=1e-5)
self.ln_2 = nn.LayerNorm(hidden_size, eps=1e-5)
# Attention (simplified for clarity)
self.c_attn = nn.Linear(
hidden_size, 3 * hidden_size
) # Q, K, V combined
self.c_proj = nn.Linear(hidden_size, hidden_size)
# Feed-forward
self.c_fc = nn.Linear(hidden_size, 4 * hidden_size)
self.c_proj_ffn = nn.Linear(4 * hidden_size, hidden_size)
self.dropout = nn.Dropout(dropout)
def forward(
self, x: torch.Tensor, mask: torch.Tensor = None
) -> torch.Tensor:
# Pre-norm attention
residual = x
x = self.ln_1(x) # Normalize BEFORE attention
x = self._attention(x, mask)
x = residual + self.dropout(x) # Residual connection
# Pre-norm feed-forward
residual = x
x = self.ln_2(x) # Normalize BEFORE FFN
x = self.c_fc(x)
x = F.gelu(x) # GELU activation
x = self.c_proj_ffn(x)
x = residual + self.dropout(x) # Residual connection
return x
def _attention(
self, x: torch.Tensor, mask: torch.Tensor = None
) -> torch.Tensor:
B, T, C = x.size()
head_dim = C // self.num_heads
# Compute Q, K, V
qkv = self.c_attn(x)
q, k, v = qkv.split(C, dim=-1)
# Reshape for multi-head attention
q = q.view(B, T, self.num_heads, head_dim).transpose(1, 2)
k = k.view(B, T, self.num_heads, head_dim).transpose(1, 2)
v = v.view(B, T, self.num_heads, head_dim).transpose(1, 2)
# Scaled dot-product attention with causal mask
# Scale by 1/sqrt(d_k) to prevent dot products from growing too large
# which would push softmax into regions with tiny gradients
scale = 1.0 / (head_dim**0.5)
attn = torch.matmul(q, k.transpose(-2, -1)) * scale
# Apply causal mask
causal_mask = torch.triu(
torch.ones(T, T, device=x.device), diagonal=1
).bool()
attn = attn.masked_fill(causal_mask, float("-inf"))
attn = F.softmax(attn, dim=-1)
attn = self.dropout(attn)
# Apply attention to values
out = torch.matmul(attn, v)
out = out.transpose(1, 2).contiguous().view(B, T, C)
return self.c_proj(out)The causal mask is important for autoregressive generation. It ensures each position can only attend to previous positions (and itself), preventing information leakage from future tokens. Let's visualize what this mask looks like:
The lower triangular pattern means token can only attend to tokens . This is what makes GPT-2 autoregressive: when predicting the next token, the model cannot "peek" at future tokens. During training, this enables efficient parallel computation since the entire sequence can be processed in one forward pass with masking.
Why does this matter? Pre-normalization provides more stable gradients during training. In post-normalization, the residual branch adds unnormalized activations, which can have high variance. The subsequent LayerNorm must handle this variance, but gradients flowing back through the normalization can become unstable at scale. With pre-normalization, the residual connection adds already-normalized values, keeping the residual stream's magnitude more controlled.
GELU Activation
GPT-2 replaced the ReLU activation function with GELU (Gaussian Error Linear Unit). While ReLU simply zeros out negative values, GELU applies a smooth, non-monotonic transformation that weights inputs by their probability under a Gaussian distribution.
The GELU function multiplies each input by its probability of being greater than other inputs from a Gaussian distribution. This creates a smooth gating mechanism where positive values pass through almost unchanged, while negative values are attenuated based on how negative they are.
The exact formulation is:
where:
- : the input value to the activation function
- : the cumulative distribution function (CDF) of the standard normal distribution, representing the probability that a random variable is less than or equal to
- : the error function, a special function that arises in probability and statistics
The intuition is elegant: we multiply each input by the probability that exceeds values drawn from a Gaussian. Large positive values have , so they pass through unchanged. Large negative values have , so they're zeroed out. Values near zero get partially attenuated.
In practice, computing the error function is expensive. GPT-2 and most implementations use this polynomial approximation:
where is the hyperbolic tangent function and the constants and were chosen to minimize approximation error. This formulation is differentiable everywhere and computationally efficient.
Why does this formula make sense? Notice that when is large and positive, approaches 1, so : the activation behaves like the identity. When is large and negative, approaches 0, so : the activation suppresses the signal. But unlike ReLU, the transition is not abrupt. Near , the Gaussian CDF is smooth and differentiable, meaning gradients can flow through even for slightly negative inputs. This smoothness is the key advantage over ReLU.
GELU's smoothness provides better gradient flow compared to ReLU's hard cutoff. To understand why this matters for training, let's compare the derivatives:
The shaded region shows where GELU outperforms ReLU: for slightly negative inputs, ReLU has zero gradient (the "dying ReLU" problem), while GELU allows some gradient to flow. This small difference compounds across billions of training steps and millions of neurons, potentially making the difference between successful and failed optimization at large scales.
Other Modifications
Several additional changes improved GPT-2's training stability. Each reflects a coherent engineering philosophy: keep activations well-behaved throughout the forward pass and keep gradients well-behaved throughout the backward pass.
Increased vocabulary size: GPT-2 uses a vocabulary of 50,257 tokens (versus GPT-1's roughly 40,000), built using byte-level Byte Pair Encoding (BPE). The byte-level encoding is important: it allows the model to represent any text without unknown tokens, because every possible byte can be encoded. This means GPT-2 can handle code, foreign language strings, special characters, and unusual punctuation without discarding information. The odd number 50,257 arose from starting with 256 base bytes, adding 50,000 BPE merge rules, and adding one special end-of-text token.
Extended context length: The context window doubled from 512 to 1024 tokens, letting the model to condition on longer passages. For language modeling, more context is almost always better: a model that can see the last two pages of a document will make better predictions about the next word than one limited to half a page. The doubling here was bounded by memory constraints on the hardware available in 2019.
Modified initialization: Weights in residual layers were scaled by , where is the number of residual layers. Each transformer block has two residual connections (one for attention and one for feed-forward), so for a 12-layer model, , and the scale factor is approximately . Without this scaling, each residual addition increases the variance of activations by a constant amount. After additions, variance grows proportionally to . Scaling by keeps the variance of the residual stream approximately constant regardless of depth, analogous to how the scaling in attention keeps dot-product magnitudes controlled.
Final layer normalization: An additional LayerNorm is applied after the last transformer block, before the output projection. With pre-normalization, the final transformer block's output is added directly to the residual stream without normalization. The final LayerNorm ensures this accumulated residual stream is properly normalized before being projected to logits.
Let's implement a minimal GPT-2 model to see these components together:
class GPT2(nn.Module):
"""Minimal GPT-2 implementation."""
def __init__(
self,
vocab_size: int = 50257,
hidden_size: int = 768,
num_layers: int = 12,
num_heads: int = 12,
max_position: int = 1024,
dropout: float = 0.1,
):
super().__init__()
self.hidden_size = hidden_size
self.num_layers = num_layers
# Token and position embeddings
self.wte = nn.Embedding(vocab_size, hidden_size)
self.wpe = nn.Embedding(max_position, hidden_size)
self.drop = nn.Dropout(dropout)
# Transformer blocks
self.blocks = nn.ModuleList(
[
GPT2Block(hidden_size, num_heads, dropout)
for _ in range(num_layers)
]
)
# Final layer norm (pre-norm style means we need this at the end)
self.ln_f = nn.LayerNorm(hidden_size, eps=1e-5)
# Output projection (tied with token embeddings)
self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)
self.lm_head.weight = self.wte.weight # Weight tying
# Initialize weights
self.apply(self._init_weights)
# Scale residual projections
for block in self.blocks:
block.c_proj.weight.data *= 1.0 / np.sqrt(2 * num_layers)
block.c_proj_ffn.weight.data *= 1.0 / np.sqrt(2 * num_layers)
def _init_weights(self, module):
if isinstance(module, nn.Linear):
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
if module.bias is not None:
torch.nn.init.zeros_(module.bias)
elif isinstance(module, nn.Embedding):
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
B, T = input_ids.size()
device = input_ids.device
# Get embeddings
position_ids = torch.arange(T, device=device).unsqueeze(0)
token_emb = self.wte(input_ids)
pos_emb = self.wpe(position_ids)
x = self.drop(token_emb + pos_emb)
# Apply transformer blocks
for block in self.blocks:
x = block(x)
# Final layer norm and output projection
x = self.ln_f(x)
logits = self.lm_head(x)
return logitsGPT-2 Small: 124,439,808 parameters
GPT-2 Medium: 354,823,168 parameters
GPT-2 Large: 774,030,080 parameters
GPT-2 XL: 1,557,611,200 parameters
Our implementation produces parameter counts closely matching the official GPT-2 specifications. Small discrepancies arise from minor implementation details like bias terms and layer norm parameters, but the overall scaling matches: roughly 3x increase from Small to Medium, 2.2x from Medium to Large, and 2x from Large to XL. Weight tying between the input embeddings and output projection reduces parameter count while maintaining performance, a technique that became standard in language models. The intuition behind weight tying is elegant: the embedding matrix translates tokens into vectors at the input, and the output projection translates vectors back into token scores. These two operations are inverses of each other, so sharing their weights is parameter-efficient and conceptually coherent.