# 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 ```python 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" │ │ │ └────────────────────────────────────────────────────────────┘ ```