Masked attention for autoregression

---

The Core Problem: Train-Test Mismatch

The mismatch:

• Training time: All tokens $(y_1, y_2, \ldots, y_T)$ are available in parallel
• Inference time: We have only $(y_1, \ldots, y_{t-1})$ when predicting $y_t$

Without masking, the attention mechanism would let token $t$ attend to token $t+1$, creating information leakage. The model would learn to copy the answer rather than predict it.

---

Deriving the Mask: From First Principles

Standard Attention (Encoder-Style)

Start with the vanilla scaled dot-product attention:

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

For a sequence of length $T$, the attention matrix $A = \frac{QK^T}{\sqrt{d_k}}$ is $T \times T$. Element $A_{ij}$ represents how much position $i$ attends to position $j$.

Problem: Softmax over row $i$ gives all positions equal consideration—position $i$ can attend to position $i+1, i+2, \ldots, T$.

Introducing the Causal Mask

We define a causal mask $M \in \{0, -\infty\}^{T \times T}$:

0 & \text{if } j \leq i \\ -\infty & \text{if } j > i \end{cases}$$ **Why $-\infty$ instead of $0$?** Because we apply the mask **before** softmax: $$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V$$ When $M_{ij} = -\infty$, the pre-softmax score becomes $-\infty$, and $\text{softmax}(-\infty) = 0$. This completely zeros out the attention weight. **Matrix form of $M$ for $T=4$:** $$M = \begin{bmatrix} 0 & -\infty & -\infty & -\infty \\ 0 & 0 & -\infty & -\infty \\ 0 & 0 & 0 & -\infty \\ 0 & 0 & 0 & 0 \end{bmatrix}$$ This is a ==lower-triangular matrix== (including diagonal). Position $i$ can attend to positions $\{1, 2, \ldots, i\}$ but not $\{i+1, \ldots, T\}$. ![[4.1.11-Masked-attention-for-autoregression.png]] --- ## Worked Examples > [!example] Example 1: Predicting "Paris" in "The capital of France is___" **Sequence:** `["The", "capital", "of", "France", "is", "Paris"]` When predicting token 6 ("Paris"), the model has access to: - ✅ Tokens 1-5: "The capital of France is" - ❌ Token 6: "Paris" (the answer we're trying to predict) **Attention matrix at position 6:** | From\To | The | capital | of | France | is | Paris | |------|---------|----|--------|----|----| | Paris | 0.15| 0.20 |0.10| 0.40 |0.15| 0.0 | The last column is **zeroed out** by the mask. The model must infer "Paris" from context, not copy it. **Why this step?** At inference, we only have `["The", "capital", "of", "France", "is"]` and must predict the next token. The mask ensures training mirrors this constraint. --- > [!example] Example 2: Step-by-Step Computation Input: 3 tokens with $d_k = 2$ $$Q = \begin{bmatrix} 1 & 2 \\ 3 & 1 \\ 2 & 2 \end{bmatrix}, \quad K = \begin{bmatrix} 2 & 1 \\ 1 & 3 \\ 2 & 1 \end{bmatrix}$$ **Step 1:** Compute $QK^T$ $$QK^T = \begin{bmatrix} 1 & 2 \\ 3 & 1 \\ 2 & 2 \end{bmatrix} \begin{bmatrix} 2 & 1 & 2 \\ 1 & 3 & 1 \end{bmatrix} = \begin{bmatrix} 4 & 7 & 4 \\ 7 & 6 & 7 \\ 6 & 8 & 6 \end{bmatrix}$$ **Why this step?** Raw similarity scores between queries and keys. **Step 2:** Scale by $\sqrt{d_k} = \sqrt{2} \approx 1.41$ $$\frac{QK^T}{\sqrt{d_k}} = \begin{bmatrix} 2.83 & 4.95 & 2.83 \\ 4.95 & 4.24 & 4.95 \\ 4.24 & 5.66 & 4.24 \end{bmatrix}$$ **Why this step?** Prevent dot products from growing too large (stabilizes gradients). **Step 3:** Apply causal mask $$M = \begin{bmatrix} 0 & -\infty & -\infty \\ 0 & 0 & -\infty \\ 0 & 0 & 0 \end{bmatrix}$$ $$\text{Masked Scores} = \begin{bmatrix} 2.83 & -\infty & -\infty \\ 4.95 & 4.24 & -\infty \\ 4.24 & 5.66 & 4.24 \end{bmatrix}$$ **Why this step?** Enforces causality—no peeking ahead. **Step 4:** Softmax (row-wise) Row 1: $\text{softmax}([2.83, -\infty, -\infty]) = [1.0, 0, 0]$ Row 2: $\text{softmax}([4.95, 4.24, -\infty]) = [0.67, 0.33, 0]$ Row 3: $\text{softmax}([4.24, 5.66, 4.24]) = [0.16, 0.68, 0.16]$ **Why this step?** Convert scores to probability distribution. Note how $-\infty$ becomes $0$ after softmax. *Check for Row 3:* $e^{4.24} \approx 69.4$, $e^{5.66} \approx 287.1$, $e^{4.24} \approx 69.4$. Sum $\approx 425.9$. Weights: $69.4/425.9 \approx 0.16$, $287.1/425.9 \approx 0.68$, $69.4/425.9 \approx 0.16$. ✅ --- ## Common Mistakes & Steel-Manning > [!mistake] Mistake 1: "Just use a binary mask of 0s and 1s" **Why it feels right:** Binary masks are simpler—just multiply attention scores by 0 or 1 after softmax. **The fix:** This doesn't work correctly. If you apply binary mask *after* softmax, the remaining weights don't sum to 1: $$\text{softmax}([2, 3, 4]) \odot [1, 1, 0] = [0.09, 0.24, 0] \quad \text{(sums to 0.33, not 1!)}$$ You'd need to **renormalize**, which is exactly what applying the mask *before* softmax achieves automatically. The $-\infty$ approach is both correct and efficient. --- > [!mistake] Mistake 2: "Masking is only for training" **Why it feels right:** At inference, we generate one token at a time, so there's nothing to mask—future tokens literally don't exist yet. **The fix:** This is *partially* true but misses the point. During inference in ==teacher forcing== or ==scheduled sampling== setups, or when using ==cached keys/values==, we still need the mask to maintain consistency. More importantly, the mask is essential during training to ensure the model learns the same constraints it will face at inference. Without it, train-test mismatch destroys generation quality. --- > [!mistake] Mistake 3: "Diagonal elements should be masked too" **Why it feels right:** Token $i$ technically hasn't been "generated" yet when we're predicting it, so shouldn't it be masked? **The fix:** No. Position $i$ attending to itself is like "using the question to help answer the question"—it provides positional and contextual self-reference. The model needs this self-loop. What we're preventing is attending to *future* tokens ($j > i$), not the current one. --- ## The Math: Why Lower-Triangular? > [!formula] Causal Attention Formula $$\text{CausalAttention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V$$ where $M$ is lower-triangular with: $$M_{ij} = \begin{cases} 0 & j \leq i \\ -\infty & j > i \end{cases}$$ **Properties:** 1. **Autoregressive factorization:** Ensures $P(y_i \mid y_{<i})$ only depends on $y_1, \ldots, y_{i-1}$ 2. **Efficient parallelization:** All positions computed in parallel during training (unlike RNNs) 3. **Inference caching:** Keys and values from previous steps can be reused (==KV cache==) **Derivation of causality:** For position $i$, the attention output is: $$\text{out}_i = \sum_{j=1}^{T} \alpha_{ij} V_j$$ where $\alpha_{ij}$ is the softmax weight. With masking: $$\alpha_{ij} = \begin{cases} \frac{\exp(s_{ij}/\sqrt{d_k})}{\sum_{k=1}^{i} \exp(s_{ik}/\sqrt{d_k})} & j \leq i \\ 0 & j > i \end{cases}$$ Thus $\text{out}_i$ is a weighted sum of only $V_1, \ldots, V_i$, ensuring no information leakage from the future. --- ## Implementation Details > [!example] Example 3: PyTorch Implementation ```python import torch import torch.nn.functional as F def causal_attention(Q, K, V): """ Q, K, V: (batch, seq_len, d_k) Returns: (batch, seq_len, d_k) """ d_k = Q.size(-1) seq_len = Q.size(1) # Step 1: Compute attention scores scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5) # scores: (batch, seq_len, seq_len) # Step 2: Create causal mask mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1) mask = mask.bool() # mask[i, j] = True if j > i (future positions) # Step 3: Apply mask with -inf scores = scores.masked_fill(mask, float('-inf')) # Step 4: Softmax and attend attn_weights = F.softmax(scores, dim=-1) output = torch.matmul(attn_weights, V) return output, attn_weights ``` **Why `triu` with `diagonal=1`?** Upper-triangular with offset 1 gives us the strictly-upper part (future positions). We fill those with $-\infty$. --- ## Multi-Head Extension In practice, masked attention is applied **per head** in multi-head attention: $$\text{head}_h = \text{CausalAttention}(QW^Q_h, KW^K_h, VW^V_h)$$ The mask $M$ is **shared across all heads**—causality is a structural constraint, not a learned one. --- > [!recall]- Explain to a 12-year-old > Imagine you're playing a game where you have to guess the next word in a story, but you can only read what came before, not what comes after. That's exactly what a language model does! > > But here's the problem: when we *train* the model, we already have the full story written out. If the model could peek at future words, it would just memorize them instead of learning to actually *predict*. That's cheating! > > **Masked attention** is like putting a blindfold over the future words. When the model is looking at word #5, we cover up words #6, #7, #8, etc. The model can only see words #1 through #5. This way, even though we have the full story during training, the model has to learn the same way it will work later—one word at a time, no peeking! > > Think of it like learning to ride a bike with training wheels (training), then riding without them (testing). The mask is like practicing without holding the handlebars—you learn the real skill, not just how to balance when someone's helping you. --- > [!mnemonic] Remember CAUSAL > **C**an't **A**ttend to **U**nknown **S**ubsequent (future) **A**ll **L**ower-triangular > > Visualize a staircase 🪜: You can only step on stairs **below or equal** to your current level, never above. --- ## Connections - [[Self-Attention Mechanism]]: Masked attention is self-attention + causality constraint - [[Decoder Architecture]]: GPT-style decoders use masked self-attention exclusively - [[Encoder-Decoder Models]]: Decoders use masked self-attention; encoders use full attention - [[KV Caching]]: Masking enables efficient incremental decoding - [[Teacher Forcing]]: Training strategy that requires masking to prevent leakage - [[Positional Encoding]]: Works with masking to preserve sequence order information - [[Attention Score Scaling]]: $\sqrt{d_k}$ scaling applied before masking --- ## Why This Matters ==Masked attention== is the fundamental mechanism enabling: 1. **Autoregressive generation**: GPT, LLaMA, Claude—all rely on this 2. **Parallel training**: Unlike RNNs, we train on full sequences at once 3. **Inference efficiency**: KV caching exploits the causal structure 4. **Consistent behavior**: Train and test conditions match Without masking, language models would collapse into sophisticated copy machines. --- #flashcards/ai-ml What is the purpose of masked attention in autoregressive models? ::: To prevent positions from attending to future tokens during training, ensuring the model learns to predict based only on past context (matching inference conditions) What value is used in the causal mask for future positions and why? ::: $-\infty$, because it becomes 0 after softmax, completely eliminating attention to future positions while maintaining normalized probability distributions Write the formula for causal attention ::: $\text{CausalAttention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V$ where $M_{ij} = 0$ if $j \leq i$ else $-\infty$ What is the shape of the causal mask matrix? ::: Lower-triangular (including diagonal): element $M_{ij}$ is 0 if $j \leq i$, meaning position $i$ can attend to positions 1 through $i$ but not beyond Why can't we apply a binary mask after softmax? ::: Because multiplying by 0/1 after softmax breaks normalization—the remaining weights won't sum to 1, requiring manual renormalization (which is what applying $-\infty$ before softmax does automatically) What is the train-test mismatch that masking solves? ::: During training all tokens are available in parallel, but at inference we generate one token at a time. Masking ensures training mirrors the sequential constraint of inference Why do we allow attention to the diagonal (position $i$ attending to itself)? ::: The current position needs self-reference for positional and contextual information. We only prevent attending to *future* positions ($j > i$), not the current one How does masking enable KV caching? ::: Because attention is causal, keys and values from previous time steps never change for new predictions, allowing them to be cached and reused rather than recomputed What is the complexity trade-off of masked attention? ::: Training is $O(T^2)$ per layer (parallelizable), but inference is $O(T)$ per new token (sequential). The mask enables this asymmetry In multi-head attention, is the mask head-specific? ::: No, the causal mask is shared across all heads—causality is a structural constraint, not a learned parameter ## 🖼️ Concept Map ```mermaid flowchart TD A[Autoregressive Generation] -->|requires| B[Predict token i from past only] C[Train-Test Mismatch] -->|training sees all tokens| D[Information Leakage] A -->|inference has only past| C D -->|solved by| E[Masked Attention] F[Scaled Dot-Product Attention] -->|allows future access| D E -->|adds| G[Causal Mask M] G -->|is| H[Lower-Triangular Matrix] G -->|uses -inf before softmax| I[softmax of -inf equals 0] I -->|zeros out| J[Future Attention Weights] F -->|modified as softmax QKt plus M| E E -->|enforces| B ``` ## 🔊 Hinglish (regional understanding) > [!intuition]- Hinglish mein samjho > Dekho, autoregressive models jaise GPT ek baar mein ek token generate karte hain—current word predict karne ke liye sirf pehle ke words dekh sakte hain, future ke nahi. Lekin training time pe humein poori sentence already pata hoti hai! Agar model ko future words dekhne denge, to ![[audio/4.1.11-Masked-attention-for-autoregression.mp3]]