ML Course: Solution Articles for New Problems

articles/code-gpt.md

@@ -4,7 +4,7 @@ Before attempting this problem, you should be comfortable with:
 
 - **Transformer Blocks** - Multi-headed attention, FFN, layer normalization, and residual connections, because the GPT model stacks multiple transformer blocks
 - **Word and Position Embeddings** - Token embeddings map IDs to vectors, position embeddings encode order, and they are added together before entering the transformer stack
-- **Softmax** - The final operation that converts logits into a probability distribution over the vocabulary
+- **Output Projection ($W_O$)** - Each multi-head attention layer projects the concatenated head outputs through a final linear layer
 
 ---
 
@@ -16,10 +16,9 @@ GPT (Generative Pre-trained Transformer) assembles everything from the course in
 2. **Position embeddings**: Add learned position vectors using a second `nn.Embedding`. Unlike the sinusoidal encoding from earlier, GPT uses learned positions.
 3. **$N$ Transformer blocks**: Each block applies multi-headed self-attention (for inter-token communication) and a feed-forward network (for per-token computation), connected by residual paths and layer normalization.
 4. **Final layer normalization**: Stabilizes the output of the last transformer block.
-5. **Vocabulary projection**: A linear layer that maps from $d_{\text{model}}$ to vocabulary size, producing logits for every possible next token.
-6. **Softmax**: Converts logits to probabilities.
+5. **Vocabulary projection**: A linear layer that maps from $d_{\text{model}}$ to vocabulary size, producing **logits** (raw unnormalized scores) for every possible next token.
 
-At each position $t$, the model outputs a probability distribution over the vocabulary, predicting what token should come at position $t+1$. Causal masking inside the attention layers ensures position $t$ only sees tokens $0$ through $t$, so the model can be used autoregressively: generate one token, append it, and repeat.
+At each position $t$, the model outputs logits over the vocabulary, predicting what token should come at position $t+1$. During training, `cross_entropy` applies softmax internally. During generation, you apply softmax yourself to sample the next token. Causal masking inside the attention layers ensures position $t$ only sees tokens $0$ through $t$, so the model can be used autoregressively: generate one token, append it, and repeat.
 
 This architecture scales remarkably well. GPT-2 Small uses $d=768$, 12 blocks, 12 heads. GPT-3 uses $d=12288$, 96 blocks, 96 heads. The structure is identical; only the numbers change.
 
@@ -29,7 +28,7 @@ This architecture scales remarkably well. GPT-2 Small uses $d=768$, 12 blocks, 1
 
 ### Intuition
 
-Compose all previously built components: embedding layers, a sequence of transformer blocks, final normalization, and a linear projection to vocabulary logits. The forward pass adds token and position embeddings, processes through all blocks, normalizes, projects, and applies softmax.
+Compose all previously built components: embedding layers, a sequence of transformer blocks, final normalization, and a linear projection to vocabulary logits. The forward pass adds token and position embeddings, processes through all blocks, normalizes, and projects to logits. Note: the model returns raw logits, not probabilities — this matches how GPT models work in practice, since `cross_entropy` and generation each handle softmax separately.
 
 ### Implementation
 
@@ -63,8 +62,7 @@ class GPT(nn.Module):
         output = self.final_norm(self.transformer_blocks(embedded))
         logits = self.vocab_projection(output)  # (B, T, vocab_size)
 
-        probabilities = nn.functional.softmax(logits, dim=-1)
-        return torch.round(probabilities, decimals=4)
+        return torch.round(logits, decimals=4)
 
     class TransformerBlock(nn.Module):
 
@@ -100,13 +98,14 @@ class GPT(nn.Module):
                 self.att_heads = nn.ModuleList()
                 for i in range(num_heads):
                     self.att_heads.append(self.SingleHeadAttention(model_dim, model_dim // num_heads))
+                self.output_proj = nn.Linear(model_dim, model_dim, bias=False)
 
             def forward(self, embedded: TensorType[float]) -> TensorType[float]:
                 head_outputs = []
                 for head in self.att_heads:
                     head_outputs.append(head(embedded))
                 concatenated = torch.cat(head_outputs, dim = 2)
-                return concatenated
+                return self.output_proj(concatenated)
 
         class VanillaNeuralNetwork(nn.Module):
 
@@ -152,9 +151,8 @@ For `vocab_size = 100`, `context_length = 8`, `model_dim = 16`, `num_blocks = 2`
 | Block 2 | Same architecture, further refining | $(1, 5, 16)$ |
 | Final LN | LayerNorm across dim 16 | $(1, 5, 16)$ |
 | Vocab proj | Linear $16 \to 100$ | $(1, 5, 100)$ |
-| Softmax | Probabilities over vocabulary | $(1, 5, 100)$ |
 
-Each of the 5 positions outputs a distribution over 100 tokens, predicting the next token.
+Each of the 5 positions outputs logits over 100 tokens, predicting the next token.
 
 ### Time & Space Complexity
 
@@ -211,6 +209,6 @@ This becomes `model/gpt.py`. This is the culmination of the entire course: every
 
 ## Key Takeaways
 
-- GPT composes token embeddings, position embeddings, a stack of transformer blocks, final normalization, and a vocabulary projection into a complete autoregressive language model.
+- GPT composes token embeddings, position embeddings, a stack of transformer blocks (each with $W^O$ output projection in multi-head attention), final normalization, and a vocabulary projection into raw logits.
 - Learned position embeddings (rather than sinusoidal) let the model discover its own positional representation during training.
 - The same architecture scales from tiny models (this problem) to GPT-3 (175 billion parameters) by increasing the model dimension, number of blocks, and number of heads.

articles/dead-relu-detector.md

@@ -0,0 +1,158 @@
+## Prerequisites
+
+Before attempting this problem, you should be comfortable with:
+
+- **ReLU Activation** - A neuron is "dead" because ReLU outputs exactly 0 for all negative inputs, and its gradient is also 0, so the weights can never update
+- **Training Diagnostics** - This problem builds on the diagnostic mindset: inspecting internal network state to find problems before they waste training time
+- **PyTorch Module Iteration** - You need to iterate through `model.children()` and check `isinstance(module, nn.ReLU)` to find the right layers
+
+---
+
+## Concept
+
+Dead ReLU neurons output zero for every sample in the batch. Because ReLU's gradient is zero for negative inputs, these neurons receive no gradient updates and are permanently stuck. The `detect_dead_neurons` method measures this per ReLU layer, and `suggest_fix` maps the severity pattern to the most appropriate intervention.
+
+The fix priority reflects real debugging experience: severe death (> 50%) means the activation function itself is the problem (switch to LeakyReLU). Early-layer death (> 30% in layer 1) means initialization is bad (re-init). Depth-increasing death means the learning rate is too aggressive (reduce it).
+
+---
+
+## Solution
+
+### Intuition
+
+Forward the input through the model layer by layer. After each ReLU, check which neurons output zero for all samples in the batch -- those are dead. For `suggest_fix`, apply the priority-ordered rules based on the pattern of dead fractions.
+
+### Implementation
+
+::tabs-start
+```python
+import torch
+import torch.nn as nn
+from typing import List
+
+
+class Solution:
+
+    def detect_dead_neurons(self, model: nn.Module, x: torch.Tensor) -> List[float]:
+        dead_fractions = []
+        with torch.no_grad():
+            for module in model.children():
+                x = module(x)
+                if isinstance(module, nn.ReLU):
+                    # A neuron is dead if it outputs 0 for ALL samples in the batch
+                    dead = (x == 0).all(dim=0).float().mean().item()
+                    dead_fractions.append(round(dead, 4))
+        return dead_fractions
+
+    def suggest_fix(self, dead_fractions: List[float]) -> str:
+        if len(dead_fractions) == 0:
+            return 'healthy'
+
+        max_frac = max(dead_fractions)
+
+        # Any layer > 0.5 dead -> use LeakyReLU
+        if max_frac > 0.5:
+            return 'use_leaky_relu'
+
+        # First layer > 0.3 dead -> reinitialize weights
+        if dead_fractions[0] > 0.3:
+            return 'reinitialize'
+
+        # Dead fraction increases with depth -> reduce learning rate
+        if len(dead_fractions) >= 2:
+            increasing = all(
+                dead_fractions[i] < dead_fractions[i + 1]
+                for i in range(len(dead_fractions) - 1)
+            )
+            if increasing and dead_fractions[-1] > 0.1:
+                return 'reduce_learning_rate'
+
+        # All layers < 0.1 dead -> healthy
+        if max_frac < 0.1:
+            return 'healthy'
+
+        return 'healthy'
+```
+::tabs-end
+
+
+### Walkthrough
+
+**Healthy network** (Kaiming init): dead fractions = `[0.0, 0.0312, 0.0156]`. Max is 0.0312 < 0.1, so `suggest_fix` returns `'healthy'`.
+
+**Broken network** (biases pushed very negative): dead fractions = `[0.5312, 0.8828, 0.9688]`. The first layer already has 0.5312 > 0.5, so the check `max_frac > 0.5` triggers and `suggest_fix` returns `'use_leaky_relu'`.
+
+**Decision logic for `suggest_fix`:**
+
+| Dead Fractions | Rule Triggered | Fix |
+|---|---|---|
+| `[0.02, 0.03, 0.01]` | All < 0.1 | `'healthy'` |
+| `[0.1, 0.4, 0.65]` | Layer 3 has 0.65 > 0.5 | `'use_leaky_relu'` |
+| `[0.35, 0.1, 0.05]` | Layer 1 has 0.35 > 0.3 | `'reinitialize'` |
+| `[0.05, 0.08, 0.15]` | Strictly increasing, last > 0.1 | `'reduce_learning_rate'` |
+
+### Time & Space Complexity
+
+- Time: $O(N \cdot d \cdot L)$ where $N$ is batch size, $d$ is layer width, $L$ is number of ReLU layers
+- Space: $O(d)$ per layer for the boolean dead-neuron mask
+
+---
+
+## Common Pitfalls
+
+### Checking After Linear Instead of After ReLU
+
+Dead neurons are a ReLU-specific concept. The Linear layer's output can be negative (that's fine -- it hasn't been activated yet). You need to check after the ReLU to see which neurons are stuck at exactly zero.
+
+::tabs-start
+```python
+# Wrong: checking after Linear
+if isinstance(module, nn.Linear):
+    dead = (x == 0).all(dim=0).float().mean().item()
+
+# Correct: checking after ReLU
+if isinstance(module, nn.ReLU):
+    dead = (x == 0).all(dim=0).float().mean().item()
+```
+::tabs-end
+
+
+### Using `<= 0` Instead of `== 0` for ReLU Output
+
+After ReLU, the output is either 0 (dead) or positive. Using `<= 0` is equivalent to `== 0` in this context, but conceptually we are checking for exactly zero output, which is what ReLU produces for negative inputs.
+
+### Missing the "Strictly Increasing" Check
+
+The `reduce_learning_rate` fix requires dead fractions to be **strictly** increasing across all layers. If any layer has the same or lower dead fraction as the previous one, the pattern doesn't match.
+
+::tabs-start
+```python
+# Wrong: non-strict comparison
+increasing = all(
+    dead_fractions[i] <= dead_fractions[i + 1]  # allows equal!
+    for i in range(len(dead_fractions) - 1)
+)
+
+# Correct: strictly increasing
+increasing = all(
+    dead_fractions[i] < dead_fractions[i + 1]
+    for i in range(len(dead_fractions) - 1)
+)
+```
+::tabs-end
+
+
+---
+
+## In the GPT Project
+
+GPT uses GELU activation, not ReLU, so the dead neuron problem doesn't apply directly. But the diagnostic technique is universal: inspecting per-layer activation patterns to find silent failures. The same approach works for detecting saturated sigmoids, collapsed layer norms, or any other "the model trains but doesn't learn" scenario.
+
+---
+
+## Key Takeaways
+
+- A dead ReLU neuron outputs zero for every sample in the batch. It receives zero gradient and can never recover. This is a permanent failure mode.
+- The severity pattern determines the fix: widespread death needs a new activation function, early-layer death needs re-initialization, and depth-correlated death needs a lower learning rate.
+- Detection requires checking after the ReLU layer, not after the Linear layer. The Linear output being negative is expected; it's the ReLU output being zero for all samples that indicates death.
+- LeakyReLU, PReLU, ELU, and GELU all avoid this problem by having non-zero gradients for negative inputs.

articles/grouped-query-attention.md

@@ -0,0 +1,156 @@
+## Prerequisites
+
+Before attempting this problem, you should be comfortable with:
+
+- **Multi-Head Self-Attention** - GQA modifies how heads share K and V, so understanding standard MHA with independent KV heads per query head is essential
+- **KV-Cache** - GQA's primary motivation is reducing the KV-Cache memory footprint during inference, so understanding the cache size problem comes first
+- **Tensor Reshaping** - The implementation requires `view`, `transpose`, and `repeat_interleave` to reshape between `(B, T, D)` and `(B, heads, T, head_dim)` formats
+
+---
+
+## Concept
+
+Standard Multi-Head Attention gives every query head its own K and V projections. During inference with KV-Cache, all those K and V tensors must be stored, and the memory cost scales linearly with the number of heads. Grouped Query Attention reduces this by sharing K, V across groups of query heads.
+
+With $h$ query heads and $g$ KV heads, each KV head serves $h/g$ query heads. The key operation is `repeat_interleave`: it expands the $g$ KV heads to $h$ by repeating each one $h/g$ times, making the shapes match for standard attention math. When $g = h$, GQA is identical to MHA. When $g = 1$, it becomes Multi-Query Attention.
+
+---
+
+## Solution
+
+### Intuition
+
+Project x into Q with `num_heads` heads and K, V with `num_kv_heads` heads. Reshape into the multi-head format. Expand K and V by repeating each KV head to match Q's head count. From here, it's standard scaled dot-product attention with a causal mask. Concatenate the heads and apply the output projection.
+
+### Implementation
+
+::tabs-start
+```python
+import torch
+import torch.nn as nn
+from torchtyping import TensorType
+
+class GroupedQueryAttention(nn.Module):
+    def __init__(self, model_dim: int, num_heads: int, num_kv_heads: int):
+        super().__init__()
+        torch.manual_seed(0)
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = model_dim // num_heads
+
+        self.q_proj = nn.Linear(model_dim, num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(model_dim, num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(model_dim, num_kv_heads * self.head_dim, bias=False)
+        self.output_proj = nn.Linear(num_heads * self.head_dim, model_dim, bias=False)
+
+    def forward(self, x: TensorType[float]) -> TensorType[float]:
+        B, T, D = x.shape
+
+        # Project to Q, K, V and reshape into heads
+        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        # Expand K, V to match Q's num_heads by repeating each KV head
+        repeats = self.num_heads // self.num_kv_heads
+        k = k.repeat_interleave(repeats, dim=1)
+        v = v.repeat_interleave(repeats, dim=1)
+
+        # Scaled dot-product attention with causal mask
+        scores = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5)
+        mask = torch.tril(torch.ones(T, T, device=x.device))
+        scores = scores.masked_fill(mask == 0, float('-inf'))
+        weights = torch.softmax(scores, dim=-1)
+
+        out = (weights @ v).transpose(1, 2).contiguous().view(B, T, -1)
+        return torch.round(self.output_proj(out), decimals=4)
+```
+::tabs-end
+
+
+### Walkthrough
+
+For `model_dim=8`, `num_heads=4`, `num_kv_heads=2`:
+
+| Step | Operation | Shape |
+|---|---|---|
+| Q projection | `q_proj(x)` then reshape | `(B, 4, T, 2)` -- 4 query heads, head_dim=2 |
+| K projection | `k_proj(x)` then reshape | `(B, 2, T, 2)` -- only 2 KV heads |
+| V projection | `v_proj(x)` then reshape | `(B, 2, T, 2)` -- only 2 KV heads |
+| Expand K | `repeat_interleave(2, dim=1)` | `(B, 4, T, 2)` -- each KV head repeated 2x |
+| Expand V | `repeat_interleave(2, dim=1)` | `(B, 4, T, 2)` -- now matches Q |
+| Attention | Standard scaled dot-product | `(B, 4, T, 2)` per head |
+| Concat + project | Merge heads, output projection | `(B, T, 8)` |
+
+**Memory savings:** K and V projections produce `num_kv_heads * head_dim = 2 * 2 = 4` values per token instead of `num_heads * head_dim = 4 * 2 = 8`. In the KV-Cache, this halves the memory per layer. For Llama 2 70B (64 query heads, 8 KV heads), the savings are 8x.
+
+### Time & Space Complexity
+
+- Time: $O(T^2 \cdot d)$ for attention (same as MHA, since K, V are expanded before the attention computation)
+- Space: $O(g \cdot T \cdot d_h)$ for KV-Cache storage, where $g$ is the number of KV heads and $d_h$ is the head dimension. This is $g/h$ of standard MHA.
+
+---
+
+## Common Pitfalls
+
+### Using `repeat` Instead of `repeat_interleave`
+
+`repeat` tiles the entire tensor, while `repeat_interleave` repeats each element individually. With 2 KV heads and 4 query heads, you want `[KV0, KV0, KV1, KV1]`, not `[KV0, KV1, KV0, KV1]`.
+
+::tabs-start
+```python
+# Wrong: repeat tiles the whole tensor
+k = k.repeat(1, repeats, 1, 1)  # [KV0, KV1, KV0, KV1] -- wrong grouping!
+
+# Correct: repeat_interleave repeats each element
+k = k.repeat_interleave(repeats, dim=1)  # [KV0, KV0, KV1, KV1] -- correct groups
+```
+::tabs-end
+
+
+### Forgetting the Causal Mask
+
+GQA is used in decoder-only models (GPT, Llama) where tokens must not attend to future positions. Without the lower-triangular mask, the model breaks autoregressive generation.
+
+::tabs-start
+```python
+# Wrong: no causal mask
+weights = torch.softmax(scores, dim=-1)
+
+# Correct: apply causal mask before softmax
+mask = torch.tril(torch.ones(T, T, device=x.device))
+scores = scores.masked_fill(mask == 0, float('-inf'))
+weights = torch.softmax(scores, dim=-1)
+```
+::tabs-end
+
+
+### Wrong Reshape Order
+
+The view/transpose sequence matters. Q has `num_heads` heads, but K and V have `num_kv_heads` heads. Using `num_heads` for the K reshape would produce incorrect head dimensions.
+
+::tabs-start
+```python
+# Wrong: reshaping K with num_heads
+k = self.k_proj(x).view(B, T, self.num_heads, self.head_dim)  # Shape mismatch!
+
+# Correct: reshaping K with num_kv_heads
+k = self.k_proj(x).view(B, T, self.num_kv_heads, self.head_dim)
+```
+::tabs-end
+
+
+---
+
+## In the GPT Project
+
+GQA is the attention variant used by Llama 2/3, Mistral, and Gemma. In a full GPT implementation, you would combine GQA with KV-Cache from the previous problem: the cache stores only `num_kv_heads` worth of K, V per layer (not `num_heads`), giving a direct memory reduction proportional to the head ratio.
+
+---
+
+## Key Takeaways
+
+- GQA shares K and V across groups of query heads, reducing KV-Cache memory by a factor of `num_heads / num_kv_heads` with minimal quality loss.
+- `repeat_interleave` is the key operation: it expands each KV head to serve its assigned group of query heads, making the shapes compatible for standard attention.
+- GQA generalizes both MHA ($g = h$) and MQA ($g = 1$). Most production models use an intermediate value (e.g., Llama 2 70B uses 64 query heads with 8 KV heads).
+- The attention computation after expansion is identical to standard MHA. The savings come entirely from the smaller projection layers and smaller cache.