Transformer Forward Pass: End-to-End Walkthrough
A detailed visual guide to how input tokens become output predictions
Overview
┌────────────────────────────────────────────────────────────┐
│ HIGH-LEVEL FLOW │
│ │
│ "The cat sat" │
│ ↓ │
│ ┌─────────────┐ │
│ │ Tokenizer │ │
│ └─────────────┘ │
│ ↓ │
│ [464, 3797, 3332] (token IDs) │
│ ↓ │
│ ┌─────────────┐ │
│ │ Embeddings │ (token + position) │
│ └─────────────┘ │
│ ↓ │
│ ┌─────────────┐ │
│ │ Layer 1 │ │
│ ├─────────────┤ │
│ │ Layer 2 │ │
│ ├─────────────┤ │
│ │ ... │ (N transformer layers) │
│ ├─────────────┤ │
│ │ Layer N │ │
│ └─────────────┘ │
│ ↓ │
│ ┌─────────────┐ │
│ │ Final Norm │ │
│ └─────────────┘ │
│ ↓ │
│ ┌─────────────┐ │
│ │ Unembedding │ (project to vocabulary) │
│ └─────────────┘ │
│ ↓ │
│ [0.01, 0.02, ..., 0.23, ...] (probabilities) │
│ ↓ │
│ "on" (predicted next token) │
│ │
└────────────────────────────────────────────────────────────┘
Model Configuration (GPT-2 Small)
┌────────────────────────────────────────────────────────────┐
│ MODEL PARAMETERS │
│ │
│ d_model (embedding dimension): 768 │
│ n_layers (transformer blocks): 12 │
│ n_heads (attention heads): 12 │
│ d_k (dimension per head): 64 (768 / 12) │
│ d_ff (MLP hidden dimension): 3072 (4 × 768) │
│ vocab_size: 50257 │
│ max_seq_len: 1024 │
│ │
└────────────────────────────────────────────────────────────┘
Step 0: Tokenization
┌────────────────────────────────────────────────────────────┐
│ TOKENIZATION │
│ │
│ Input text: "The cat sat" │
│ │
│ Tokenizer (BPE): │
│ "The" → 464 │
│ " cat" → 3797 (note: space is part of token) │
│ " sat" → 3332 │
│ │
│ Output: [464, 3797, 3332] │
│ │
│ Shape: [3] (3 tokens) │
│ │
└────────────────────────────────────────────────────────────┘
Step 1: Token + Position Embeddings
┌────────────────────────────────────────────────────────────┐
│ EMBEDDING LOOKUP │
│ │
│ Token IDs: [464, 3797, 3332] │
│ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ Token Embedding Table (wte) │ │
│ │ Shape: [50257, 768] │ │
│ │ │ │
│ │ token_id 0: [0.012, -0.034, ..., 0.023] │ │
│ │ token_id 1: [0.045, 0.067, ..., 0.089] │ │
│ │ ... │ │
│ │ token_id 464: [0.123, -0.045, ..., 0.087] ← │ │
│ │ ... │ │
│ │ token_id 3797: [0.456, 0.178, ..., 0.312] ← │ │
│ │ ... │ │
│ │ token_id 3332: [-0.089, 0.267, ..., 0.145] ← │ │
│ │ ... │ │
│ │ token_id 50256: [0.034, -0.078, ..., 0.056] │ │
│ └────────────────────────────────────────────────────┘ │
│ │
│ Token Embeddings [3, 768]: │
│ ┌────────────────────────────────────────────────────┐ │
│ │ The │ 0.123 -0.045 0.234 ... 0.087 │ │
│ │ cat │ 0.456 0.178 -0.123 ... 0.312 │ │
│ │ sat │ -0.089 0.267 0.189 ... 0.145 │ │
│ └────────────────────────────────────────────────────┘ │
│ │
└────────────────────────────────────────────────────────────┘
┌────────────────────────────────────────────────────────────┐
│ POSITION EMBEDDINGS │
│ │
│ Position IDs: [0, 1, 2] │
│ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ Position Embedding Table (wpe) │ │
│ │ Shape: [1024, 768] │ │
│ │ │ │
│ │ position 0: [0.011, -0.023, ..., 0.012] ← │ │
│ │ position 1: [0.022, 0.015, ..., 0.025] ← │ │
│ │ position 2: [-0.015, 0.032, ..., 0.018] ← │ │
│ │ position 3: [0.028, -0.019, ..., 0.031] │ │
│ │ ... │ │
│ │ position 1023: [0.045, 0.012, ..., 0.067] │ │
│ └────────────────────────────────────────────────────┘ │
│ │
│ Position Embeddings [3, 768]: │
│ ┌────────────────────────────────────────────────────┐ │
│ │ p0 │ 0.011 -0.023 0.034 ... 0.012 │ │
│ │ p1 │ 0.022 0.015 -0.018 ... 0.025 │ │
│ │ p2 │ -0.015 0.032 0.021 ... 0.018 │ │
│ └────────────────────────────────────────────────────┘ │
│ │
└────────────────────────────────────────────────────────────┘
┌────────────────────────────────────────────────────────────┐
│ COMBINE EMBEDDINGS │
│ │
│ X = token_embeddings + position_embeddings │
│ │
│ Initial Residual Stream [3, 768]: │
│ ┌────────────────────────────────────────────────────┐ │
│ │ The │ 0.134 -0.068 0.268 ... 0.099 │ │
│ │ cat │ 0.478 0.193 -0.141 ... 0.337 │ │
│ │ sat │ -0.104 0.299 0.210 ... 0.163 │ │
│ └────────────────────────────────────────────────────┘ │
│ │
│ This X is the RESIDUAL STREAM - flows through all layers │
│ │
└────────────────────────────────────────────────────────────┘
Step 2: Transformer Layer (Detailed)
Each of the 12 layers has identical structure. Here's Layer 1 in full detail:
┌────────────────────────────────────────────────────────────┐
│ TRANSFORMER LAYER 1 │
│ │
│ Input: X [3, 768] (residual stream) │
│ │
│ ══════════════════════════════════════════════════════ │
│ ║ ATTENTION BLOCK ║ │
│ ══════════════════════════════════════════════════════ │
│ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ 1. LAYER NORMALIZATION │ │
│ │ │ │
│ │ For each token, normalize across 768 dims: │ │
│ │ X_norm = (X - mean) / sqrt(var + ε) │ │
│ │ X_norm = γ * X_norm + β │ │
│ │ │ │
│ │ X_norm: [3, 768] │ │
│ └────────────────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ 2. Q, K, V PROJECTIONS │ │
│ │ │ │
│ │ Combined weight: W_qkv [768, 2304] (768 × 3) │ │
│ │ │ │
│ │ QKV = X_norm @ W_qkv [3, 768] @ [768, 2304] │ │
│ │ = [3, 2304] │ │
│ │ │ │
│ │ Split into Q, K, V: │ │
│ │ Q = QKV[:, 0:768] [3, 768] │ │
│ │ K = QKV[:, 768:1536] [3, 768] │ │
│ │ V = QKV[:, 1536:2304] [3, 768] │ │
│ └────────────────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ 3. RESHAPE FOR MULTI-HEAD │ │
│ │ │ │
│ │ Reshape [3, 768] → [3, 12, 64] → [12, 3, 64] │ │
│ │ │ │
│ │ Q: [12, 3, 64] (12 heads, 3 pos, 64 dims) │ │
│ │ K: [12, 3, 64] │ │
│ │ V: [12, 3, 64] │ │
│ │ │ │
│ │ Now each head operates independently │ │
│ └────────────────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ 4. ATTENTION SCORES (per head) │ │
│ │ │ │
│ │ scores = Q @ K.transpose(-2, -1) / sqrt(64) │ │
│ │ = [12, 3, 64] @ [12, 64, 3] / 8 │ │
│ │ = [12, 3, 3] │ │
│ │ │ │
│ │ Example for head 0: │ │
│ │ K_The K_cat K_sat │ │
│ │ ┌─────────────────────┐ │ │
│ │ Q_The │ 0.52 0.73 0.41 │ │ │
│ │ Q_cat │ 0.38 0.89 0.56 │ │ │
│ │ Q_sat │ 0.45 0.67 0.62 │ │ │
│ │ └─────────────────────┘ │ │
│ └────────────────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ 5. CAUSAL MASK │ │
│ │ │ │
│ │ Mask future positions with -infinity: │ │
│ │ │ │
│ │ mask = [[0, -∞, -∞], │ │
│ │ [0, 0, -∞], │ │
│ │ [0, 0, 0]] │ │
│ │ │ │
│ │ scores + mask: │ │
│ │ K_The K_cat K_sat │ │
│ │ ┌─────────────────────┐ │ │
│ │ Q_The │ 0.52 -∞ -∞ │ │ │
│ │ Q_cat │ 0.38 0.89 -∞ │ │ │
│ │ Q_sat │ 0.45 0.67 0.62 │ │ │
│ │ └─────────────────────┘ │ │
│ └────────────────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ 6. SOFTMAX (row-wise) │ │
│ │ │ │
│ │ weights = softmax(masked_scores, dim=-1) │ │
│ │ │ │
│ │ Attention weights [12, 3, 3]: │ │
│ │ K_The K_cat K_sat │ │
│ │ ┌─────────────────────┐ │ │
│ │ Q_The │ 1.00 0.00 0.00 │ (sum=1) │ │
│ │ Q_cat │ 0.38 0.62 0.00 │ (sum=1) │ │
│ │ Q_sat │ 0.28 0.40 0.32 │ (sum=1) │ │
│ │ └─────────────────────┘ │ │
│ │ │ │
│ │ "The" only sees itself (causal) │ │
│ │ "cat" sees The + cat │ │
│ │ "sat" sees all three │ │
│ └────────────────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ 7. WEIGHTED SUM OF VALUES │ │
│ │ │ │
│ │ attn_output = weights @ V │ │
│ │ = [12, 3, 3] @ [12, 3, 64] │ │
│ │ = [12, 3, 64] │ │
│ │ │ │
│ │ Each position now has context-aware repr │ │
│ └────────────────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ 8. CONCATENATE HEADS + OUTPUT PROJECTION │ │
│ │ │ │
│ │ Reshape: [12, 3, 64] → [3, 12, 64] → [3, 768] │ │
│ │ │ │
│ │ Project: attn_out = concat @ W_O │ │
│ │ = [3, 768] @ [768, 768] │ │
│ │ = [3, 768] │ │
│ └────────────────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ 9. RESIDUAL CONNECTION │ │
│ │ │ │
│ │ X = X + attn_out │ │
│ │ = [3, 768] + [3, 768] │ │
│ │ = [3, 768] │ │
│ │ │ │
│ │ Original info preserved + attention info added │ │
│ └────────────────────────────────────────────────────┘ │
│ │
│ ══════════════════════════════════════════════════════ │
│ ║ MLP BLOCK ║ │
│ ══════════════════════════════════════════════════════ │
│ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ 10. LAYER NORMALIZATION │ │
│ │ │ │
│ │ X_norm = LayerNorm(X) [3, 768] │ │
│ └────────────────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ 11. MLP: EXPAND │ │
│ │ │ │
│ │ hidden = X_norm @ W_up │ │
│ │ = [3, 768] @ [768, 3072] │ │
│ │ = [3, 3072] │ │
│ │ │ │
│ │ Expand from 768 to 3072 dimensions (4×) │ │
│ └────────────────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ 12. ACTIVATION (GELU) │ │
│ │ │ │
│ │ hidden = GELU(hidden) [3, 3072] │ │
│ │ │ │
│ │ GELU ≈ x * sigmoid(1.702 * x) │ │
│ │ Smooth non-linearity, allows negative values │ │
│ └────────────────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ 13. MLP: CONTRACT │ │
│ │ │ │
│ │ mlp_out = hidden @ W_down │ │
│ │ = [3, 3072] @ [3072, 768] │ │
│ │ = [3, 768] │ │
│ │ │ │
│ │ Contract back from 3072 to 768 dimensions │ │
│ └────────────────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ 14. RESIDUAL CONNECTION │ │
│ │ │ │
│ │ X = X + mlp_out │ │
│ │ = [3, 768] + [3, 768] │ │
│ │ = [3, 768] │ │
│ └────────────────────────────────────────────────────┘ │
│ │
│ Output: X [3, 768] → goes to Layer 2 │
│ │
└────────────────────────────────────────────────────────────┘
Step 3: All Layers
┌────────────────────────────────────────────────────────────┐
│ LAYER PROGRESSION │
│ │
│ X₀ [3, 768] ← initial (token + position embeddings) │
│ │ │
│ ▼ │
│ ┌─────────┐ │
│ │ Layer 1 │ Attention + MLP │
│ └────┬────┘ │
│ │ │
│ ▼ │
│ X₁ [3, 768] ← + local patterns, bigrams │
│ │ │
│ ▼ │
│ ┌─────────┐ │
│ │ Layer 2 │ │
│ └────┬────┘ │
│ │ │
│ ▼ │
│ X₂ [3, 768] ← + syntax, previous token info │
│ │ │
│ ▼ │
│ ┌─────────┐ │
│ │ Layer 3 │ │
│ └────┬────┘ │
│ │ │
│ ▼ │
│ X₃ [3, 768] │
│ │ │
│ ... │
│ │ │
│ ▼ │
│ ┌─────────┐ │
│ │Layer 12 │ │
│ └────┬────┘ │
│ │ │
│ ▼ │
│ X₁₂ [3, 768] ← final representation │
│ │
│ Shape NEVER changes: always [3, 768] │
│ Only the CONTENT evolves through layers │
│ │
└────────────────────────────────────────────────────────────┘
Step 4: Final LayerNorm + Prediction
┌────────────────────────────────────────────────────────────┐
│ FINAL PROCESSING │
│ │
│ X₁₂ [3, 768] (output from layer 12) │
│ │ │
│ ▼ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ FINAL LAYER NORMALIZATION │ │
│ │ │ │
│ │ X_final = LayerNorm(X₁₂) [3, 768] │ │
│ └────────────────────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ SELECT LAST POSITION │ │
│ │ │ │
│ │ For next-token prediction, only need last token: │ │
│ │ │ │
│ │ X_final [3, 768]: │ │
│ │ ┌───────────────────────────────────────┐ │ │
│ │ │ The │ ... 768 values ... │ pos 0 │ │
│ │ │ cat │ ... 768 values ... │ pos 1 │ │
│ │ │ sat │ ... 768 values ... │ pos 2 │ │
│ │ └───────────────────────────────────────┘ ↑ │ │
│ │ │ │ │
│ │ last_hidden = X_final[2] [768] ←──────────┘ │ │
│ └────────────────────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ UNEMBEDDING (Project to Vocabulary) │ │
│ │ │ │
│ │ GPT-2 uses weight tying: W_unembed = W_embed.T │ │
│ │ │ │
│ │ logits = last_hidden @ wte.weight.T │ │
│ │ = [768] @ [768, 50257] │ │
│ │ = [50257] │ │
│ │ │ │
│ │ Each value = "score" for that vocabulary token │ │
│ └────────────────────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ LOGITS (raw scores) │ │
│ │ │ │
│ │ logits [50257]: │ │
│ │ token 0 ("!"): -2.34 │ │
│ │ token 1 ("\""): -3.12 │ │
│ │ token 2 ("#"): -4.56 │ │
│ │ ... │ │
│ │ token 319 ("on"): 4.21 ← HIGH │ │
│ │ ... │ │
│ │ token 866 ("down"): 3.15 │ │
│ │ token 867 ("there"): 2.89 │ │
│ │ ... │ │
│ │ token 50256 ("<|eos|>"): -5.67 │ │
│ └────────────────────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ SOFTMAX → PROBABILITIES │ │
│ │ │ │
│ │ probs = softmax(logits) [50257] │ │
│ │ │ │
│ │ All values now between 0 and 1, sum to 1: │ │
│ │ token 319 ("on"): 0.23 (23%) │ │
│ │ token 866 ("down"): 0.12 (12%) │ │
│ │ token 867 ("there"): 0.08 (8%) │ │
│ │ token 373 ("still"): 0.05 (5%) │ │
│ │ ... │ │
│ └────────────────────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌────────────────────────────────────────────────────┐ │
│ │ SAMPLING / ARGMAX │ │
│ │ │ │
│ │ Greedy (argmax): Select "on" (highest prob) │ │
│ │ Temperature: Adjust distribution sharpness │ │
│ │ Top-k: Sample from top k tokens │ │
│ │ Top-p (nucleus): Sample from top cumulative p │ │
│ │ │ │
│ │ Selected: "on" │ │
│ └────────────────────────────────────────────────────┘ │
│ │
│ Output: "on" │
│ │
│ Full sequence: "The cat sat" → "The cat sat on" │
│ │
└────────────────────────────────────────────────────────────┘
Complete Shape Flow Summary
┌────────────────────────────────────────────────────────────┐
│ COMPLETE SHAPE FLOW │
│ │
│ "The cat sat" │
│ ↓ │
│ Tokenize [3] │
│ ↓ │
│ Token Embedding Lookup [3, 768] │
│ + │
│ Position Embedding Lookup [3, 768] │
│ = │
│ Initial Residual Stream [3, 768] │
│ ↓ │
│ ╔════════════════ LAYER 1-12 (each) ════════════════╗ │
│ ║ ║ │
│ ║ LayerNorm [3, 768] ║ │
│ ║ ↓ ║ │
│ ║ Q, K, V projection [3, 768] each ║ │
│ ║ ↓ ║ │
│ ║ Reshape for heads [12, 3, 64] ║ │
│ ║ ↓ ║ │
│ ║ Q @ K.T (attention scores) [12, 3, 3] ║ │
│ ║ ↓ ║ │
│ ║ + Causal mask [12, 3, 3] ║ │
│ ║ ↓ ║ │
│ ║ Softmax (weights) [12, 3, 3] ║ │
│ ║ ↓ ║ │
│ ║ Weights @ V [12, 3, 64] ║ │
│ ║ ↓ ║ │
│ ║ Concat heads [3, 768] ║ │
│ ║ ↓ ║ │
│ ║ Output projection (W_O) [3, 768] ║ │
│ ║ ↓ ║ │
│ ║ + Residual [3, 768] ║ │
│ ║ ↓ ║ │
│ ║ LayerNorm [3, 768] ║ │
│ ║ ↓ ║ │
│ ║ MLP up [3, 3072] ║ │
│ ║ ↓ ║ │
│ ║ GELU [3, 3072] ║ │
│ ║ ↓ ║ │
│ ║ MLP down [3, 768] ║ │
│ ║ ↓ ║ │
│ ║ + Residual [3, 768] ║ │
│ ║ ║ │
│ ╚═══════════════════════════════════════════════════╝ │
│ ↓ │
│ Final LayerNorm [3, 768] │
│ ↓ │
│ Select last position [768] │
│ ↓ │
│ @ wte.weight.T (unembedding) [50257] │
│ ↓ │
│ Softmax [50257] │
│ ↓ │
│ Sample/Argmax → token ID scalar │
│ ↓ │
│ Decode → "on" │
│ │
└────────────────────────────────────────────────────────────┘
Parameter Count Breakdown
┌────────────────────────────────────────────────────────────┐
│ GPT-2 SMALL PARAMETERS │
│ │
│ EMBEDDINGS │
│ ────────────────────────────────────────────────────── │
│ Token embeddings (wte): 50257 × 768 = 38,597,376 │
│ Position embeddings (wpe): 1024 × 768 = 786,432 │
│ │
│ PER TRANSFORMER LAYER (×12) │
│ ────────────────────────────────────────────────────── │
│ Attention: │
│ W_qkv (combined Q,K,V): 768 × 2304 = 1,769,472 │
│ W_O (output projection): 768 × 768 = 589,824 │
│ LayerNorm (ln_1): 768 × 2 = 1,536 │
│ │
│ MLP: │
│ W_up: 768 × 3072 = 2,359,296 │
│ W_down: 3072 × 768 = 2,359,296 │
│ LayerNorm (ln_2): 768 × 2 = 1,536 │
│ │
│ Layer subtotal: = 7,080,960 │
│ │
│ ALL 12 LAYERS │
│ ────────────────────────────────────────────────────── │
│ 12 × 7,080,960 = 84,971,520 │
│ │
│ FINAL │
│ ────────────────────────────────────────────────────── │
│ Final LayerNorm (ln_f): 768 × 2 = 1,536 │
│ │
│ TOTAL │
│ ────────────────────────────────────────────────────── │
│ Embeddings: 38,597,376 │
│ + Position: 786,432 │
│ + Layers: 84,971,520 │
│ + Final LN: 1,536 │
│ ───────────────────────────────────────────────────── │
│ TOTAL: 124,356,864 │
│ │
│ (Actual GPT-2 small: ~117M with weight tying) │
│ │
└────────────────────────────────────────────────────────────┘
Simplified PyTorch Code
import torch
import torch.nn as nn
import torch.nn.functional as F
def gpt2_forward(token_ids, model):
"""
Simplified GPT-2 forward pass.
Args:
token_ids: [batch_size, seq_len] tensor of token IDs
model: GPT-2 model with wte, wpe, layers, ln_f
Returns:
logits: [batch_size, seq_len, vocab_size]
"""
batch_size, seq_len = token_ids.shape
# Step 1: Embeddings
positions = torch.arange(seq_len, device=token_ids.device)
x = model.wte(token_ids) + model.wpe(positions)
# x: [batch_size, seq_len, 768]
# Step 2: Transformer layers
for layer in model.h: # h = list of transformer blocks
# === Attention block ===
residual = x
x = layer.ln_1(x) # LayerNorm
# Q, K, V projection (combined)
qkv = layer.attn.c_attn(x) # [batch, seq, 2304]
q, k, v = qkv.split(768, dim=-1)
# Reshape for multi-head
# [batch, seq, 768] -> [batch, heads, seq, 64]
q = q.view(batch_size, seq_len, 12, 64).transpose(1, 2)
k = k.view(batch_size, seq_len, 12, 64).transpose(1, 2)
v = v.view(batch_size, seq_len, 12, 64).transpose(1, 2)
# Attention scores
scores = (q @ k.transpose(-2, -1)) / 8.0 # sqrt(64) = 8
# Causal mask
mask = torch.triu(
torch.ones(seq_len, seq_len), diagonal=1
).bool()
scores = scores.masked_fill(mask, float('-inf'))
# Softmax and weighted sum
weights = F.softmax(scores, dim=-1)
attn_out = weights @ v # [batch, heads, seq, 64]
# Reshape back and project
attn_out = attn_out.transpose(1, 2)
attn_out = attn_out.reshape(batch_size, seq_len, 768)
attn_out = layer.attn.c_proj(attn_out) # W_O
x = residual + attn_out # Residual connection
# === MLP block ===
residual = x
x = layer.ln_2(x) # LayerNorm
hidden = layer.mlp.c_fc(x) # Expand: 768 -> 3072
hidden = F.gelu(hidden) # Activation
mlp_out = layer.mlp.c_proj(hidden) # Contract: 3072 -> 768
x = residual + mlp_out # Residual connection
# Step 3: Final LayerNorm
x = model.ln_f(x)
# Step 4: Project to vocabulary (weight tying)
logits = x @ model.wte.weight.T # [batch, seq, vocab_size]
return logits
# Usage:
# logits = gpt2_forward(token_ids, model)
# probs = F.softmax(logits[:, -1, :], dim=-1) # Last position
# next_token = torch.argmax(probs, dim=-1) # Greedy decoding
Information Flow: What Happens Conceptually
┌────────────────────────────────────────────────────────────┐
│ CONCEPTUAL INFORMATION FLOW │
│ │
│ Input: "The cat sat" │
│ │
│ ══════════════════════════════════════════════════════ │
│ EMBEDDINGS │
│ ══════════════════════════════════════════════════════ │
│ │
│ "The" → [word: article, common, sentence-start] │
│ "cat" → [word: noun, animal, singular] │
│ "sat" → [word: verb, past-tense, intransitive] │
│ │
│ ══════════════════════════════════════════════════════ │
│ EARLY LAYERS (1-4): Local Patterns │
│ ══════════════════════════════════════════════════════ │
│ │
│ "The" ← learns: followed by noun │
│ "cat" ← learns: preceded by article, subject position │
│ "sat" ← learns: preceded by noun, verb position │
│ │
│ Previous-token attention, bigram patterns │
│ │
│ ══════════════════════════════════════════════════════ │
│ MIDDLE LAYERS (5-8): Syntactic Structure │
│ ══════════════════════════════════════════════════════ │
│ │
│ "The" ← part of noun phrase "The cat" │
│ "cat" ← SUBJECT of sentence │
│ "sat" ← VERB, agrees with singular subject │
│ │
│ Subject-verb heads, syntactic role assignment │
│ │
│ ══════════════════════════════════════════════════════ │
│ LATER LAYERS (9-11): Semantic Composition │
│ ══════════════════════════════════════════════════════ │
│ │
│ Full meaning: "A cat performed a sitting action" │
│ Context: simple declarative sentence, past event │
│ Expectations: location? manner? continuation? │
│ │
│ ══════════════════════════════════════════════════════ │
│ FINAL LAYER (12): Prediction Preparation │
│ ══════════════════════════════════════════════════════ │
│ │
│ "sat" position now encodes: │
│ - Likely next: preposition (on, down, in) │
│ - Possible: adverb (quietly, still) │
│ - Less likely: noun, verb, punctuation │
│ │
│ ══════════════════════════════════════════════════════ │
│ OUTPUT │
│ ══════════════════════════════════════════════════════ │
│ │
│ Top predictions: │
│ "on" 23% ← selected │
│ "down" 12% │
│ "there" 8% │
│ "still" 5% │
│ ... │
│ │
│ Output: "on" │
│ │
└────────────────────────────────────────────────────────────┘