Masked Attention — Module 10.2

How This Lesson Fits the Module & Volume

MHSA lets every position see every other. For autoregressive decoding that would leak the future. Masked attention (causal masking) zeros out illegal links so position i may attend only to positions ≤ i.

This is essential in the decoder block, in GPT-style LMs, and connects Module 10.1 attention scores to Vol. 11 next-token prediction.

Learning Objectives

By the end of this lesson, students should be able to:

Define causal (look-ahead) masking for autoregressive Transformers.
Distinguish causal masks from padding masks.
Build an upper-triangular −∞ mask and apply it before softmax.
Explain why teacher forcing still needs causal masks.
Implement masked MHSA in PyTorch.
Relate masking to valid training of language models in Vol. 11.

Definition

Masked (causal) attention modifies attention scores so disallowed positions receive −∞ before softmax (hence weight 0). In the standard causal case, query position i cannot attend to key positions j > i, enforcing the autoregressive factorization p(x_t | x_{<t}).

Mask Types You Will Meet

Scores

QK^T / √d_k.

Add mask

0 allowed; −∞ blocked.

Softmax

Blocked positions → 0.

Weighted V

Only legal context mixes.

Mask	Purpose	Typical shape
Causal / look-ahead	Block future tokens	(T, T) upper triangle
Padding	Block PAD keys	(B, L) or (B, 1, 1, L)
Combined	Both constraints	Broadcast-sum of masks
Custom / sparse	Local windows, etc.	Task-specific pattern

Causal vs Padding

Causal mask

About time / order.
Same pattern every batch (for fixed T).
Required for autoregressive decoders.

Padding mask

About invalid tokens.
Varies per example length.
Used in encoders and decoders.

Together

Add both to scores.
Softmax ignores all blocked cells.
Common bug: forgetting one of them.

Code: Causal Mask + Attention

import math
import torch
import torch.nn.functional as F

def causal_attn(q, k, v):
    # q,k,v: (B, h, T, d_k)
    T = q.size(-2)
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(q.size(-1))
    causal = torch.triu(torch.ones(T, T, device=q.device), diagonal=1).bool()
    scores = scores.masked_fill(causal, float("-inf"))
    weights = scores.softmax(dim=-1)
    return weights @ v, weights

B, h, T, d_k = 1, 1, 4, 8
q = k = v = torch.randn(B, h, T, d_k)
out, w = causal_attn(q, k, v)
print(out.shape)           # torch.Size([1, 1, 4, 8])
print(w[0, 0].round(decimals=2))
# Row i should be ~0 for columns j > i
# Example:
# tensor([[1.00, 0.00, 0.00, 0.00],
#         [0.48, 0.52, 0.00, 0.00],
#         [...],
#         [...]])

Strengths and Tradeoffs

Strengths

Enables parallel teacher-forced training without future leak.
Matches left-to-right generation at inference.
Simple additive implementation on scores.

Tradeoffs

Blocks useful bidirectional context (by design).
Easy to mis-wire float mask vs bool mask APIs.
Inference still sequential unless optimized (KV cache).

Common Misconception

“Masking deletes tokens from the batch.” Masking does not remove tokens from the tensor; it sets their attention logits to −∞ so softmax weight is zero. The sequence length stays T; illegal positions simply contribute nothing to the weighted sum of values.

Knowledge Check

Short Answer: What does a causal mask prevent? Answer: Attending to future positions (j > i).
True/False: Padding masks and causal masks solve the same problem. Answer: False—pads vs future tokens.
Multiple Choice: Illegal scores are usually set to: (a) +1, (b) −∞, (c) the mean. Answer: (b).
Short Answer: Why is causal masking needed with teacher forcing? Answer: So the model cannot read future gold tokens while predicting the current one.
True/False: After softmax, masked positions should have weight ~0. Answer: True.
Multiple Choice: Causal masking is standard in: (a) BERT encoders, (b) GPT-style decoders, (c) bag-of-words. Answer: (b).
Short Answer: Which Module 10.1 topic produces the scores you mask? Answer: Attention scores / scaled dot-product attention.
Short Answer: How do you combine pad + causal masks? Answer: Add both (broadcast) to the score matrix before softmax.
Multiple Choice: Masking changes sequence length: (a) always, (b) never by itself, (c) only on CPU. Answer: (b).
True/False: Autoregressive LMs in Vol. 11 rely on causal masked attention. Answer: True.

Key Takeaways

Causal masks enforce left-to-right information flow in decoders.
Implement by adding −∞ to forbidden score entries before softmax.
Distinct from padding masks; often used together.
Required for honest autoregressive training and generation.
Next: Positional Embedding.

Trainer’s Guide

Hands-on idea: Train a tiny character LM for 50 steps with and without the causal mask; show the unmasked model’s suspiciously low loss from cheating.

Discussion prompt: When would you intentionally use a non-causal bidirectional mask instead (e.g., BERT MLM)?

Recap: Masked attention blocks illegal links—especially the future—so decoders train and generate fairly. Continue with Positional Embedding.