Masked attention for autoregression

4.1.11 · AI-ML › Transformer Architecture

---

Core Problem: Train-Test Mismatch

Mismatch yeh hai:

• Training time: Saare tokens $(y_1,y_2,\dots ,y_T)$ parallel mein available hote hain
• Inference time: $y_t$ predict karte waqt hamare paas sirf $(y_1,\dots ,y_{t-1})$ hote hain

Masking ke bina, attention mechanism token $t$ ko token $t+1$ attend karne deta, jisse information leakage hoti. Model answer copy karna seekh leta, predict karna nahi.

---

Mask Derive Karna: First Principles Se

Standard Attention (Encoder-Style)

Vanilla scaled dot-product attention se shuru karo:

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

Length $T$ ki sequence ke liye, attention matrix $A=\frac{Q{K}^T}{\sqrt{d_k}}$ ka size $T\times T$ hota hai. Element $A_{ij}$ represent karta hai ki position $i$ kitna position $j$ ko attend karta hai.

Problem: Row $i$ par softmax saare positions ko equal consideration deta hai—position $i$, position $i+1,i+2,\dots ,T$ ko attend kar sakta hai.

Causal Mask Introduce Karna

Hum ek causal mask $M\in \{0,-\infty {\}}^{T\times T}$ define karte hain:

0 & \text{if } j \leq i \\ -\infty & \text{if } j > i \end{cases}$$ **$0$ ki jagah $-\infty$ kyun?** Kyunki hum mask **softmax se pehle** apply karte hain: $$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V$$ Jab $M_{ij} = -\infty$ hota hai, pre-softmax score $-\infty$ ban jaata hai, aur $\text{softmax}(-\infty) = 0$. Isse attention weight completely zero ho jaata hai. **$T=4$ ke liye $M$ ki matrix form:** $$M = \begin{bmatrix} 0 & -\infty & -\infty & -\infty \\ 0 & 0 & -\infty & -\infty \\ 0 & 0 & 0 & -\infty \\ 0 & 0 & 0 & 0 \end{bmatrix}$$ Yeh ek ==lower-triangular matrix== hai (diagonal including). Position $i$ sirf positions $\{1, 2, \ldots, i\}$ ko attend kar sakta hai, $\{i+1, \ldots, T\}$ ko nahi. ![[4.1.11-Masked-attention-for-autoregression.png]] --- ## Worked Examples > [!example] Example 1: "The capital of France is___" mein "Paris" Predict Karna **Sequence:** `["The", "capital", "of", "France", "is", "Paris"]` Token 6 ("Paris") predict karte waqt, model ko access hai: - ✅ Tokens 1-5: "The capital of France is" - ❌ Token 6: "Paris" (wo answer jo hum predict karne ki koshish kar rahe hain) **Position 6 par Attention matrix:** | From\To | The | capital | of | France | is | Paris | |------|---------|----|--------|----|----| | Paris | 0.15| 0.20 |0.10| 0.40 |0.15| 0.0 | Last column mask se **zero out** ho jaata hai. Model ko "Paris" context se infer karna padta hai, copy nahi kar sakta. **Yeh step kyun?** Inference par, hamare paas sirf `["The", "capital", "of", "France", "is"]` hota hai aur humein next token predict karna hota hai. Mask ensure karta hai ki training mein yahi constraint mirror ho. --- > [!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:** $QK^T$ compute karo $$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}$$ **Yeh step kyun?** Queries aur keys ke beech raw similarity scores. **Step 2:** $\sqrt{d_k} = \sqrt{2} \approx 1.41$ se scale karo $$\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}$$ **Yeh step kyun?** Dot products ko bahut bada hone se rokna (gradients stabilize hote hain). **Step 3:** Causal mask apply karo $$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}$$ **Yeh step kyun?** Causality enforce karta hai—aage dekhna allowed nahi. **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]$ **Yeh step kyun?** Scores ko probability distribution mein convert karo. Dekho ki $-\infty$ softmax ke baad $0$ ban jaata hai. *Row 3 ka check:* $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: "Sirf 0s aur 1s ka binary mask use karo" **Kyun sahi lagta hai:** Binary masks simple hote hain—softmax ke baad attention scores ko 0 ya 1 se multiply karo. **Fix:** Yeh sahi kaam nahi karta. Agar binary mask softmax ke *baad* apply karo, to baaki bache weights 1 sum nahi karte: $$\text{softmax}([2, 3, 4]) \odot [1, 1, 0] = [0.09, 0.24, 0] \quad \text{(sums to 0.33, not 1!)}$$ Tumhe **renormalize** karna padega, jo ki exactly wahi hai jo softmax se *pehle* mask lagane par automatically ho jaata hai. $-\infty$ approach dono sahi hai aur efficient bhi. --- > [!mistake] Mistake 2: "Masking sirf training ke liye hai" **Kyun sahi lagta hai:** Inference mein, hum ek waqt mein ek token generate karte hain, to mask karne ko kuch hai nahi—future tokens literally abhi exist nahi karte. **Fix:** Yeh *partially* sahi hai lekin poori baat miss kar raha hai. ==Teacher forcing== ya ==scheduled sampling== setups mein inference ke dauran, ya jab ==cached keys/values== use ho rahe hoon, tab bhi consistency maintain karne ke liye mask ki zaroorat hoti hai. Aur bhi important yeh hai ki, mask training ke dauran essential hai taaki model wohi constraints seekhe jo use inference mein face karni hongi. Iske bina, train-test mismatch generation quality barbad kar deta hai. --- > [!mistake] Mistake 3: "Diagonal elements ko bhi mask karna chahiye" **Kyun sahi lagta hai:** Token $i$ technically abhi "generate" nahi hua jab hum use predict kar rahe hain, to kya use mask nahi karna chahiye? **Fix:** Nahi. Position $i$ ka apne aap ko attend karna aise hai jaise "sawaal ko answer mein help karne dena"—yeh positional aur contextual self-reference provide karta hai. Model ko yeh self-loop chahiye. Hum jo rok rahe hain wo hai *future* tokens attend karna ($j > i$), current token nahi. --- ## Math: Lower-Triangular Kyun? > [!formula] Causal Attention Formula $$\text{CausalAttention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V$$ jahan $M$ lower-triangular hai with: $$M_{ij} = \begin{cases} 0 & j \leq i \\ -\infty & j > i \end{cases}$$ **Properties:** 1. **Autoregressive factorization:** Ensure karta hai ki $P(y_i \mid y_{<i})$ sirf $y_1, \ldots, y_{i-1}$ par depend kare 2. **Efficient parallelization:** Training ke dauran saare positions parallel compute hote hain (RNNs ke unlike) 3. **Inference caching:** Pichle steps ki keys aur values reuse ho sakti hain (==KV cache==) **Causality ki derivation:** Position $i$ ke liye, attention output hai: $$\text{out}_i = \sum_{j=1}^{T} \alpha_{ij} V_j$$ jahan $\alpha_{ij}$ softmax weight hai. Masking ke saath: $$\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}$$ Is tarah $\text{out}_i$ sirf $V_1, \ldots, V_i$ ka weighted sum hai, jo ensure karta hai ki future se koi information leakage na ho. --- ## 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: Attention scores compute karo scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5) # scores: (batch, seq_len, seq_len) # Step 2: Causal mask banao 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: -inf ke saath mask apply karo scores = scores.masked_fill(mask, float('-inf')) # Step 4: Softmax aur attend karo attn_weights = F.softmax(scores, dim=-1) output = torch.matmul(attn_weights, V) return output, attn_weights ``` **`triu` with `diagonal=1` kyun?** Offset 1 ke saath upper-triangular humein strictly-upper part (future positions) deta hai. Unhe hum $-\infty$ se fill karte hain. --- ## Multi-Head Extension Practice mein, masked attention multi-head attention mein **per head** apply hota hai: $$\text{head}_h = \text{CausalAttention}(QW^Q_h, KW^K_h, VW^V_h)$$ Mask $M$ **saare heads mein shared** hota hai—causality ek structural constraint hai, koi learned parameter nahi. --- > [!recall]- Ek 12-saal ke bachche ko explain karo > Socho tum ek aisa game khel rahe ho jahan tumhe story ka agla word guess karna hai, lekin tum sirf wahi padh sakte ho jo pehle aaya, aage kya hai nahi. Exactly yahi ek language model karta hai! > > Lekin yahan ek problem hai: jab hum model ko *train* karte hain, poori story pehle se likhi hoti hai. Agar model future words peek kar sake, to wo unhe yaad kar lega instead of actually *predict* karna seekhne ke. Yeh cheating hai! > > **Masked attention** aise hai jaise future words par blindfold laga dena. Jab model word #5 dekh raha hota hai, hum words #6, #7, #8, etc. chhupa dete hain. Model sirf words #1 se #5 tak dekh sakta hai. Is tarah, training mein poori story hone ke bawajood, model waise hi seekhna padta hai jaise wo baad mein kaam karega—ek word at a time, koi peeking nahi! > > Aise socho jaise training wheels ke saath cycle chalana seekhna (training), phir unke bina chalana (testing). Mask aise hai jaise handlebars pakde bina practice karna—tum asli skill seekhte ho, sirf yeh nahi ki koi help karte waqt balance kaise karte hain. --- > [!mnemonic] CAUSAL yaad rakho > **C**an't **A**ttend to **U**nknown **S**ubsequent (future) **A**ll **L**ower-triangular > > Ek staircase 🪜 visualize karo: tum sirf woh seedhiyaan chadh sakte ho jo tumhare **current level se neeche ya barabar** hain, kabhi upar nahi. --- ## Connections - [[Self-Attention Mechanism|Self-Attention Mechanism]]: Masked attention = self-attention + causality constraint - [[Decoder Architecture|Decoder Architecture]]: GPT-style decoders exclusively masked self-attention use karte hain - [[Encoder-Decoder Models|Encoder-Decoder Models]]: Decoders masked self-attention use karte hain; encoders full attention use karte hain - [[KV Caching|KV Caching]]: Masking efficient incremental decoding enable karta hai - [[Teacher Forcing|Teacher Forcing]]: Training strategy jo leakage rokne ke liye masking require karta hai - [[Positional Encoding|Positional Encoding]]: Masking ke saath kaam karta hai sequence order information preserve karne ke liye - [[Attention Score Scaling|Attention Score Scaling]]: $\sqrt{d_k}$ scaling masking se pehle apply hoti hai --- ## Yeh Kyun Zaroori Hai ==Masked attention== yeh fundamental mechanism hai jo enable karta hai: 1. **Autoregressive generation**: GPT, LLaMA, Claude—sab isi par rely karte hain 2. **Parallel training**: RNNs ke unlike, hum ek saath full sequences par train karte hain 3. **Inference efficiency**: KV caching causal structure ko exploit karta hai 4. **Consistent behavior**: Train aur test conditions match karte hain Masking ke bina, language models sophisticated copy machines mein collapse ho jaate. --- #flashcards/ai-ml Autoregressive models mein masked attention ka kya purpose hai? ::: Future tokens ko attend karne se positions ko rokna training ke dauran, taaki model sirf past context ke basis par predict karna seekhe (inference conditions se match kare) Causal mask mein future positions ke liye kaunsi value use hoti hai aur kyun? ::: $-\infty$, kyunki softmax ke baad yeh 0 ban jaata hai, future positions par attention completely eliminate ho jaata hai jabki normalized probability distributions maintain rehti hain Causal attention ka formula likhो ::: $\text{CausalAttention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V$ jahan $M_{ij} = 0$ if $j \leq i$ else $-\infty$ Causal mask matrix ki shape kya hoti hai? ::: Lower-triangular (diagonal including): element $M_{ij}$ 0 hota hai agar $j \leq i$, matlab position $i$ sirf positions 1 se $i$ tak attend kar sakta hai, aage nahi Softmax ke baad binary mask apply karna kyun kaam nahi karta? ::: Kyunki softmax ke baad 0/1 se multiply karna normalization tod deta hai—baaki bache weights 1 sum nahi karenge, manual renormalization ki zaroorat padti hai (jo ki softmax se pehle $-\infty$ apply karna automatically karta hai) Masking kis train-test mismatch ko solve karta hai? ::: Training ke dauran saare tokens parallel mein available hote hain, lekin inference mein hum ek waqt mein ek token generate karte hain. Masking ensure karta hai ki training inference ke sequential constraint ko mirror kare Diagonal par attention (position $i$ apne aap ko attend karta hai) allow karna kyun zaroori hai? ::: Current position ko positional aur contextual information ke liye self-reference chahiye. Hum sirf *future* positions attend karna rokते hain ($j > i$), current one nahi Masking KV caching kaise enable karta hai? ::: Kyunki attention causal hai, pichle time steps ki keys aur values naye predictions ke liye kabhi nahi badalte, isliye unhe cache karke reuse kiya ja sakta hai, baar baar compute karne ki zaroorat nahi Masked attention ki complexity trade-off kya hai? ::: Training har layer ke liye $O(T^2)$ hai (parallelizable), lekin inference har naye token ke liye $O(T)$ hai (sequential). Mask yeh asymmetry enable karta hai Multi-head attention mein, kya mask head-specific hota hai? ::: Nahi, causal mask saare heads mein shared hota hai—causality ek structural constraint hai, koi learned parameter nahi ## 🖼️ 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 ```