GPT-2 › Worked Example: Tracing a Forward Pass

To solidify your understanding of how GPT-2 processes text, let's trace a minimal forward pass step by step with concrete dimensions. We'll use GPT-2 Small as our example, with hidden size $d_{model} = 768$, 12 layers, and 12 attention heads.

Suppose the input is the sequence "The cat sat", which tokenizes to 3 tokens with indices, say, [464, 3797, 3332].

Step 1: Token Embedding Lookup. We look up each token index in the token embedding table, which has shape $(50257, 768)$. For 3 tokens, we retrieve a matrix of shape $(3, 768)$. Each row is the 768-dimensional representation of that token.

Step 2: Position Embedding Lookup. We look up positions 0, 1, 2 in the position embedding table, also of shape $(1024, 768)$. We retrieve another $(3, 768)$ matrix. Adding token embeddings and position embeddings gives the initial hidden state $x_0$ of shape $(3, 768)$.

Step 3: Pre-Norm Attention in Layer 1. We apply LayerNorm to $x_0$, creating $\hat{x}_0$ of shape $(3, 768)$ with mean 0 and variance 1 per position. We then project $\hat{x}_0$ through the combined $QKV$ projection matrix of shape $(768, 2304)$ to get a $(3, 2304)$ matrix, which we split into three $(3, 768)$ matrices for $Q$, $K$, and $V$.

Step 4: Multi-Head Attention Split. We reshape each of $Q$, $K$, and $V$ from $(3, 768)$ into $(12, 3, 64)$ by splitting the 768 dimensions across 12 heads of 64 dimensions each. For head $h$, the query matrix is $Q_h$ of shape $(3, 64)$, and similarly for $K_h$ and $V_h$.

Step 5: Scaled Dot-Product Attention. For each head $h$, we compute attention scores and weighted values:

$$A_h = \text{softmax}\left(\frac{Q_h K_h^T}{\sqrt{64}}\right) V_h$$

where:

• $Q_h K_h^T$ has shape $(3, 3)$: a score matrix where entry $(i,j)$ measures how much position $i$ should attend to position $j$
• $\sqrt{64} = 8$: the scaling factor preventing large dot products from saturating the softmax
• The causal mask zeroes out the upper triangle: position 0 only attends to itself, position 1 attends to positions 0 and 1, position 2 attends to all three

The output $A_h$ has shape $(3, 64)$.

Step 6: Concatenate and Project. We concatenate the 12 head outputs to get a $(3, 768)$ matrix, then project through the output projection of shape $(768, 768)$, giving a $(3, 768)$ attention output.

Step 7: First Residual Addition. We add the attention output to the original $x_0$: $x_1 = x_0 + \text{attn\_output}$. Both have shape $(3, 768)$. The residual connection preserves the original information while the attention output adds what the attention mechanism learned.

Step 8: Pre-Norm Feed-Forward. We apply LayerNorm to $x_1$, then pass through the feed-forward network. The first linear layer expands from 768 to 3072 dimensions, GELU activation is applied, and the second linear layer compresses back to 768. The result has shape $(3, 768)$.

Step 9: Second Residual Addition. We add the feed-forward output to $x_1$: $x_2 = x_1 + \text{ffn\_output}$. This completes one transformer block.

Steps 10 through 21: Repeat for layers 2 through 12. The $(3, 768)$ hidden state passes through 11 more identical blocks. Each adds its learned transformations via residual connections, progressively enriching each token's representation with broader context.

Step 22: Final LayerNorm. After all 12 layers, we apply the final layer normalization to the $(3, 768)$ output.

Step 23: Logit Projection. We multiply by the transposed token embedding matrix of shape $(768, 50257)$, creating logits of shape $(3, 50257)$. For next-token prediction, we take the last position's logits (shape $(50257,)$) and sample from this distribution to generate the 4th token.

This entire forward pass computes all three positions simultaneously during training, with the causal mask so each position only uses information from prior positions. The efficiency comes from this parallel computation: rather than running the model three times sequentially to train on positions 1, 2, and 3, we train on all positions in a single forward pass.