Causal (Masked) Self-Attention

Radford, A., Narasimhan, K., Salimans, T., & Sutskever, I. (2018). “Improving Language Understanding by Generative Pre-Training.” OpenAI.

Learning Objectives

After completing this chapter, you will be able to:

1. Explain why standard bidirectional attention is illegal for language generation and how causal masking enforces the autoregressive constraint.
2. Write the mathematical formula for causal attention and describe the role of every symbol, including the mask function, $-\infty$, and the lower-triangular structure.
3. Compute causal attention weights and outputs by hand for our shared example “The cat sat on the mat,” and contrast them with the standard (bidirectional) results from Chapter 1.
4. Implement causal attention from scratch in both NumPy and PyTorch, including the production-ready $\texttt{F.scaled\_dot\_product\_attention}$ API with $\texttt{is\_causal=True}$.
5. Describe how causal masking interacts with Flash Attention, KV-cache optimization, and prefix-LM architectures in modern systems like GPT-4, Claude, and Gemini.

Where this appears: Every autoregressive language model in production today — GPT-4, Claude, Gemini, LLaMA, Mistral, Command R — uses causal masking. It is also the basis for code generation (Codex, Copilot), music generation (MusicLM), time series forecasting (TimesFM), and any task where the model must predict the next element in a sequence.

The Real Problem

The Autoregressive Constraint

In Chapter 1, we computed scaled dot-product attention where every token could attend to every other token — including tokens that appear later in the sequence. This bidirectional view is perfect for tasks like sentence classification or named entity recognition, where the model has access to the complete input before making predictions.

But language generation is fundamentally different. When you type a prompt and GPT generates a response, it produces tokens one at a time, left to right. At the moment it is generating the 50th token, the 51st token does not exist yet. The model cannot look into the future because the future has not been written.

This is the autoregressive constraint: at position $i$, the model may only condition on positions $0, 1, \ldots, i$. The probability of a sequence is factored as:

$P(x_1, x_2, \ldots, x_N) = \prod_{i=1}^{N} P(x_i \mid x_1, x_2, \ldots, x_{i-1})$

Each token's probability depends only on the tokens that came before it. This is the chain rule of probability applied to sequential generation, and it is the mathematical foundation of every autoregressive language model.

The Cheating Problem

During training, the model sees the entire sequence at once (for parallelism). If we used standard bidirectional attention, token $i$ could attend to token $i+1$ — the very token it is supposed to predict. This creates a devastating shortcut: the model learns to copy the answer from the future instead of learning the patterns of language.

Consider our example. When predicting what comes after “The cat,” the model should learn that a verb is likely. But if “cat” can attend to “sat” during training, the model simply learns to read off the answer. At inference time, “sat” does not exist yet, and the model collapses — it was trained to cheat, and now the cheat sheet is gone.

Think of it like a student taking an exam. If they can see the answer key during practice tests, they will ace every practice test. But when the real exam arrives without an answer key, they fail completely because they never learned the material — they only learned to copy. Causal masking is the exam proctor that hides the answer key during training, forcing the model to genuinely learn how language works.

The Core Insight: Causal masking does not add new computation. It removes information by blocking future positions. This single constraint is what makes autoregressive language generation possible. Without it, you cannot train a model to generate text.

The Story Behind Causal Masking

GPT and the Decoder-Only Model

The idea of masking future tokens in self-attention was introduced in the original Transformer paper (Vaswani et al., 2017), where the decoder used “masked self-attention” to preserve the autoregressive property. But it was Alec Radford's GPT (2018) that demonstrated the full power of this approach as a standalone architecture.

Radford's key insight was radical in its simplicity: you do not need an encoder at all. While BERT (Devlin et al., 2018) achieved remarkable results with bidirectional attention and masked language modeling, GPT showed that a decoder-only model — using nothing but causal self-attention — could learn powerful language representations through the simple objective of next-token prediction.

The original GPT used 12 transformer decoder layers, each with causal self-attention. The training objective was straightforward: given a sequence of tokens $x_1, \ldots, x_N$, maximize the log-likelihood:

$\mathcal{L} = \sum_{i=1}^{N} \log P(x_i \mid x_1, \ldots, x_{i-1}; \theta)$

This was the birth of the paradigm that would scale to GPT-2 (1.5B parameters), GPT-3 (175B), GPT-4, and eventually Claude, Gemini, and LLaMA. The architecture is identical in all of them: stacked layers of causal self-attention followed by feedforward networks. The only things that changed were scale, data, and training recipes.

The remarkable consequence is that the single most important mechanism in modern AI — the one powering every chatbot, code assistant, and reasoning engine — is not a complex invention. It is standard attention with a single modification: set the upper triangle to $-\infty$ before softmax.

The Mathematical Definition

Causal attention is identical to the scaled dot-product attention from Chapter 1, with one additional step: a mask is applied to the score matrix before softmax.

$\text{CausalAttn}(Q, K, V) = \text{softmax}\!\left(\text{mask}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)\right) V$

where the mask function is defined as:

$\text{mask}(S)_{ij} = \begin{cases} S_{ij} & \text{if } j \leq i \quad \text{(past or present)} \\ -\infty & \text{if } j > i \quad \text{(future --- blocked)} \end{cases}$

Symbol-by-Symbol Breakdown

The Mask Matrix

The causal mask is a binary lower-triangular matrix. Entry $(i, j)$ is 1 if token $i$ is allowed to attend to token $j$, and 0 otherwise. The rule is simple: a token can see itself and everything before it.

This is a lower-triangular matrix: everything on or below the diagonal is 1 (visible), and everything above the diagonal is 0 (blocked). In implementation, we represent the blocked positions with a boolean mask and set those positions to $-\infty$ in the score matrix. The softmax then naturally assigns zero weight to those positions.

Notice that each successive row gains one more visible token. “The” can only see itself. “cat” can see “The” and itself. “mat” (the last token) can see the entire sequence — for the last token, causal attention is identical to standard attention.

Interactive: Causal Mask Explorer

The interactive visualization below lets you explore the causal mask. Click on any token to see what it can attend to and the corresponding attention weights. Use the step slider to watch the mask build token by token. Toggle “Weights” to see the actual attention weight values instead of the binary mask.

Step-by-Step Calculation

We use the same Q, K, V matrices from Chapter 1. The first two steps (dot products and scaling) are identical. Causal masking changes only Step 3 — everything else follows mechanically.

Step 1: Raw Dot Products ($QK^\top$)

Compute the $5 \times 5$ matrix of all pairwise dot products between queries and keys. This is identical to Chapter 1:

Step 2: Scaling ($\div \sqrt{d_k} = \div 2.0$)

Divide every entry by $\sqrt{4} = 2.0$:

Step 3: Apply Causal Mask

This is the new step. Set every entry where $j > i$ (future positions) to $-\infty$:

The lower triangle and diagonal retain their original scaled values. The upper triangle is now $-\infty$. When we apply softmax next, $\exp(-\infty) = 0$, so these positions will receive exactly zero attention weight.

Step 4: Softmax Over Visible Tokens Only

Softmax is applied row-by-row. Because the $-\infty$ entries contribute zero to both numerator and denominator, the softmax effectively operates over only the visible tokens in each row:

Key observations:

• “The” gets weight 1.0 on itself because it is the only visible token. Its representation is entirely self-referential.
• “cat” allocates 81.8% of its attention to “The” and only 18.2% to itself. This is because $Q_{\text{cat}} \cdot K_{\text{The}}$ has a high score (1.500 after scaling).
• “mat” (the last token) sees the entire sequence, so its causal weights are identical to its standard attention weights from Chapter 1.

Step 5: Weighted Sum of Values

Each output row is a weighted sum of the visible value vectors:

$O_{\text{cat}} = 0.8176 \cdot V_{\text{The}} + 0.1824 \cdot V_{\text{cat}}$

$= 0.8176 \cdot [1, 0, 0, 0] + 0.1824 \cdot [0, 1, 0, 0] = [0.8176, 0.1824, 0.0, 0.0]$

Notice that dimensions 2 and 3 of “cat”'s output are exactly zero. In standard attention (Chapter 1), they were 0.3018 and 0.2058 respectively. Those values came from “sat” and “on” through the value vectors — but causal masking blocked “cat” from seeing those tokens entirely.

Worked Example: What “sat” Sees

Let us trace through the full calculation for “sat” (position 2), which can see “The” (position 0), “cat” (position 1), and itself (position 2).

$Q_{\text{sat}} = [1.0, 1.0, 1.0, 0.0]$

Softmax over the three visible entries $[0.500, 1.000, 1.000]$:

$\exp(0.5) = 1.6487, \quad \exp(1.0) = 2.7183, \quad \exp(1.0) = 2.7183$

$\text{Sum} = 1.6487 + 2.7183 + 2.7183 = 7.0853$

$A_{\text{sat}} = \left[\frac{1.6487}{7.0853}, \frac{2.7183}{7.0853}, \frac{2.7183}{7.0853}, 0, 0\right] = [0.2327, 0.3837, 0.3837, 0, 0]$

Output:

$O_{\text{sat}} = 0.2327 \cdot [1,0,0,0] + 0.3837 \cdot [0,1,0,0] + 0.3837 \cdot [0,0,1,0] = [0.2327, 0.3837, 0.3837, 0.0]$

Compare to standard attention (Chapter 1): $O_{\text{sat}}^{\text{std}} = [0.2495, 0.3481, 0.3481, 0.2495]$. Dimension 3 is 0.2495 in standard attention (contributed by “on” and “mat” through their value vectors) but exactly 0.0 in causal attention because those future tokens are blocked.

Full Attention Weights and Output

Causal Attention Weight Matrix ($5 \times 5$)

Causal Output Matrix ($5 \times 4$)

Interactive: Standard vs Causal Attention

The visualization below shows the standard (bidirectional) and causal attention weights side by side. Hover over any cell to see exact values and the difference between the two modes. Switch to the “Output Vectors” view to see how the masked information loss propagates to the output.

Interactive: Autoregressive Generation

Watch a language model generate our sentence token by token. At each step, the model can only see the tokens it has already generated. Toggle “Cheat Mode” to see what would happen if the model could peek at future tokens — and why that breaks everything.

Applications Across Domains

Language Generation

Every modern large language model — GPT-4, Claude, Gemini, LLaMA, Mistral, Qwen — uses causal self-attention as its core mechanism. The entire forward pass of these models is a stack of causal attention layers followed by feedforward networks. When Claude generates a response to your prompt, each new token is produced by attending causally to all previous tokens.

Code Generation

Code is a sequence, and code generation is autoregressive. Tools like GitHub Copilot, Cursor, and Claude's code generation use causal attention to predict the next token of code given the preceding context. The causal structure naturally captures the sequential dependencies in code: a variable must be declared before it is used, a function must be defined before it is called, and an import must appear before the module is referenced.

Time Series Forecasting

In time series modeling, the future is literally unknown at prediction time. Models like TimesFM (Google, 2024) and Lag-Llama apply causal attention to temporal sequences, ensuring that predictions at time $t$ depend only on observations at times $1, 2, \ldots, t-1$. The causal mask prevents the temporal leakage that would make forecasting trivially (and falsely) accurate during training.

Music and Audio Generation

MusicLM (Google, 2023), Jukebox (OpenAI, 2020), and AudioLM use causal attention to generate audio tokens sequentially. Music has a natural temporal structure: a note at beat 16 should be influenced by the melody at beats 1–15, but not by what comes at beat 17. Causal masking enforces this temporal causality, producing musically coherent output.

Connection to Modern Systems

Flash Attention with Causal Masks

Flash Attention (Dao et al., 2022) computes attention using IO-aware tiling to minimize GPU memory transfers. It provides a dedicated causal mode ($\texttt{is\_causal=True}$) that is even faster than the general masked version because it can skip entire tiles in the upper triangle that would be masked to zero. This makes causal attention cheaper than standard attention in practice — roughly 50% fewer operations since the upper triangle is never computed at all.

KV-Cache and Causal Structure

The causal mask creates a critical property: the attention computation for position $i$ depends only on keys and values at positions $0, \ldots, i$. This means that once we have computed K and V for a token, they never change — future tokens do not affect past computations.

This enables the KV-cache optimization: during autoregressive generation, we store the K and V tensors for all previously generated tokens. When generating token $i+1$, we only need to compute $Q_{i+1}$, $K_{i+1}$, and $V_{i+1}$ for the new token and append them to the cache. The attention for the new token is computed against all cached keys and values. Without causal masking, the KV-cache would be invalid because earlier tokens' representations could change when new tokens are added.

Prefix-LM: A Hybrid Approach

Some models (e.g., T5, UL2, PaLM 2 in certain modes) use a prefix-LM strategy: the prompt (prefix) uses bidirectional attention (no mask), while the generated continuation uses causal attention. This gives the model full context understanding of the input while maintaining the autoregressive constraint for generation. Mathematically, the mask becomes:

$M_{ij} = \begin{cases} 1 & \text{if } i < P \text{ or } j \leq i \\ 0 & \text{otherwise} \end{cases}$

where $P$ is the length of the prefix. This is a strictly more general mask than pure causal — it reduces to standard attention when $P = N$ and to pure causal when $P = 0$.

Complexity Analysis

The theoretical complexity is unchanged because the mask is $O(N^2)$ to construct. But in practice, Flash Attention's causal mode achieves close to a 2x speedup by never computing the masked tiles. The mask itself requires zero extra parameters — it is purely a function of the sequence positions.

Python Implementation

The complete NumPy implementation below includes the $\texttt{CausalAttention}$ class with all five steps separated into individual methods, plus a detailed $\texttt{explain()}$ method that traces the computation for any query token. Click on any line of code to see the exact values flowing through memory at that point.

       The     cat     sat      on     mat
The  0.000    -inf    -inf    -inf    -inf
cat  1.500   0.000    -inf    -inf    -inf
sat  0.500   1.000   1.000    -inf    -inf
on   0.500   0.500   0.000   1.000    -inf
mat  0.500   0.500   0.500   0.500   0.750

       The     cat     sat      on     mat
The  1.0000  0.0000  0.0000  0.0000  0.0000
cat  0.8176  0.1824  0.0000  0.0000  0.0000
sat  0.2327  0.3837  0.3837  0.0000  0.0000
on   0.2350  0.2350  0.1425  0.3875  0.0000
mat  0.1892  0.1892  0.1892  0.1892  0.2430

       The    cat    sat     on    mat
The  False   True   True   True   True
cat  False  False   True   True   True
sat  False  False  False   True   True
on   False  False  False  False   True
mat  False  False  False  False  False

       The     cat     sat      on     mat
The  0.000   1.000   0.500   0.500   0.750
cat  1.500   0.000   1.000   0.500   0.250
sat  0.500   1.000   1.000   0.500   0.750
on   0.500   0.500   0.000   1.000   0.500
mat  0.500   0.500   0.500   0.500   0.750

       The     cat     sat      on     mat
The  0.000    -inf    -inf    -inf    -inf
cat  1.500   0.000    -inf    -inf    -inf
sat  0.500   1.000   1.000    -inf    -inf
on   0.500   0.500   0.000   1.000    -inf
mat  0.500   0.500   0.500   0.500   0.750

       d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

       d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

       d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

      The   cat   sat    on   mat
The  0.00  2.00  1.00  1.00  1.50
cat  3.00  0.00  2.00  1.00  0.50
sat  1.00  2.00  2.00  1.00  1.50
on   1.00  1.00  0.00  2.00  1.00
mat  1.00  1.00  1.00  1.00  1.50

       The     cat     sat      on     mat
The  0.000   1.000   0.500   0.500   0.750
cat  1.500   0.000   1.000   0.500   0.250
sat  0.500   1.000   1.000   0.500   0.750
on   0.500   0.500   0.000   1.000   0.500
mat  0.500   0.500   0.500   0.500   0.750

       The    cat    sat     on    mat
The  False   True   True   True   True
cat  False  False   True   True   True
sat  False  False  False   True   True
on   False  False  False  False   True
mat  False  False  False  False  False

       The     cat     sat      on     mat
The  0.000    -inf    -inf    -inf    -inf
cat  1.500   0.000    -inf    -inf    -inf
sat  0.500   1.000   1.000    -inf    -inf
on   0.500   0.500   0.000   1.000    -inf
mat  0.500   0.500   0.500   0.500   0.750

       The     cat     sat      on     mat
The  1.0000  0.0000  0.0000  0.0000  0.0000
cat  0.8176  0.1824  0.0000  0.0000  0.0000
sat  0.2327  0.3837  0.3837  0.0000  0.0000
on   0.2350  0.2350  0.1425  0.3875  0.0000
mat  0.1892  0.1892  0.1892  0.1892  0.2430

       d0      d1      d2      d3
The  1.0000  0.0000  0.0000  0.0000
cat  0.8176  0.1824  0.0000  0.0000
sat  0.2327  0.3837  0.3837  0.0000
on   0.2350  0.2350  0.1425  0.3875
mat  0.3108  0.3108  0.3108  0.3108

       d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

       d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

       d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

       The     cat     sat      on     mat
The  1.0000  0.0000  0.0000  0.0000  0.0000
cat  0.8176  0.1824  0.0000  0.0000  0.0000
sat  0.2327  0.3837  0.3837  0.0000  0.0000
on   0.2350  0.2350  0.1425  0.3875  0.0000
mat  0.1892  0.1892  0.1892  0.1892  0.2430

       d0      d1      d2      d3
The  1.0000  0.0000  0.0000  0.0000
cat  0.8176  0.1824  0.0000  0.0000
sat  0.2327  0.3837  0.3837  0.0000
on   0.2350  0.2350  0.1425  0.3875
mat  0.3108  0.3108  0.3108  0.3108

1import numpy as np
2import math
3
4class CausalAttention:
5    """
6    Causal (Masked) Self-Attention (Radford et al., 2018)
7
8    Attention(Q, K, V) = softmax(mask(Q @ K^T / sqrt(d_k))) @ V
9
10    The upper triangle of the score matrix is set to -inf
11    before softmax, preventing tokens from attending to
12    future positions in the sequence.
13    """
14
15    def __init__(self, d_k: int):
16        """
17        Args:
18            d_k: Dimension of query/key vectors (used for scaling)
19        """
20        self.d_k = d_k
21        self.scale = math.sqrt(d_k)
22
23    def _softmax(self, x: np.ndarray) -> np.ndarray:
24        """Numerically stable softmax that handles -inf values."""
25        x_safe = np.where(np.isfinite(x), x, -1e9)
26        x_shifted = x_safe - np.max(x_safe, axis=-1, keepdims=True)
27        exp_x = np.exp(x_shifted)
28        return exp_x / np.sum(exp_x, axis=-1, keepdims=True)
29
30    def build_causal_mask(self, N: int) -> np.ndarray:
31        """Build upper-triangular boolean mask (True = blocked)."""
32        return np.triu(np.ones((N, N), dtype=bool), k=1)
33
34    def apply_mask(self, scores: np.ndarray, mask: np.ndarray) -> np.ndarray:
35        """Set masked (future) positions to -inf."""
36        masked_scores = scores.copy()
37        masked_scores[mask] = -np.inf
38        return masked_scores
39
40    def forward(self, Q: np.ndarray, K: np.ndarray, V: np.ndarray):
41        """
42        Full forward pass of causal attention.
43
44        Args:
45            Q: Query matrix  (N, d_k)
46            K: Key matrix    (N, d_k)
47            V: Value matrix  (N, d_v)
48
49        Returns:
50            weights: Causal attention weight matrix (N, N)
51            output:  Context-enriched output         (N, d_v)
52        """
53        N = Q.shape[0]
54        raw_scores = Q @ K.T
55        scaled_scores = raw_scores / self.scale
56        mask = self.build_causal_mask(N)
57        masked_scores = self.apply_mask(scaled_scores, mask)
58        weights = self._softmax(masked_scores)
59        output = weights @ V
60        return weights, output
61
62    def explain(self, Q, K, V, tokens, query_idx=0):
63        """Print a detailed trace for a specific query token."""
64        N = Q.shape[0]
65        raw = Q @ K.T
66        scaled = raw / self.scale
67        mask = self.build_causal_mask(N)
68        masked = self.apply_mask(scaled, mask)
69        weights = self._softmax(masked)
70        output = weights @ V
71
72        t = tokens[query_idx]
73        print(f"\n=== Causal attention trace for '{t}' (row {query_idx}) ===")
74        print(f"Q[{query_idx}] = {Q[query_idx]}")
75        print(f"\nVisible tokens: {tokens[:query_idx+1]}")
76        print(f"Blocked tokens: {tokens[query_idx+1:]}")
77
78        print(f"\nStep 1-2: Raw scores -> scaled (/ {self.scale:.1f}):")
79        for j, tk in enumerate(tokens):
80            s = "BLOCKED" if j > query_idx else f"{scaled[query_idx, j]:.4f}"
81            print(f"  S[{t},{tk}] = {s}")
82
83        print(f"\nStep 3-4: Softmax weights:")
84        for j, tk in enumerate(tokens):
85            w = weights[query_idx, j]
86            bar = '#' * int(w * 40)
87            print(f"  A[{t},{tk}] = {w:.4f} |{bar}|")
88
89        print(f"\nStep 5: Output = weighted sum of V:")
90        print(f"  O[{t}] = {output[query_idx]}")
91
92
93# ── Shared Example (same Q, K, V as every chapter) ──
94tokens = ["The", "cat", "sat", "on", "mat"]
95
96Q = np.array([
97    [1.0, 0.0, 1.0, 0.0],   # The
98    [0.0, 2.0, 0.0, 1.0],   # cat
99    [1.0, 1.0, 1.0, 0.0],   # sat
100    [0.0, 0.0, 1.0, 1.0],   # on
101    [1.0, 0.0, 0.0, 1.0],   # mat
102])
103
104K = np.array([
105    [0.0, 1.0, 0.0, 1.0],   # The
106    [1.0, 0.0, 1.0, 0.0],   # cat
107    [1.0, 1.0, 0.0, 0.0],   # sat
108    [0.0, 0.0, 1.0, 1.0],   # on
109    [1.0, 0.0, 0.5, 0.5],   # mat
110])
111
112V = np.array([
113    [1.0, 0.0, 0.0, 0.0],   # The
114    [0.0, 1.0, 0.0, 0.0],   # cat
115    [0.0, 0.0, 1.0, 0.0],   # sat
116    [0.0, 0.0, 0.0, 1.0],   # on
117    [0.5, 0.5, 0.5, 0.5],   # mat
118])
119
120# ── Run ──
121attn = CausalAttention(d_k=4)
122weights, output = attn.forward(Q, K, V)
123
124print("Causal Attention Weights (5x5):")
125print(np.round(weights, 4))
126
127print("\nCausal Output (5x4):")
128print(np.round(output, 4))
129
130# Detailed trace for "sat"
131attn.explain(Q, K, V, tokens, query_idx=2)

PyTorch Implementation

The PyTorch implementation below includes both a trainable module (with learned $W_Q, W_K, W_V$ projections) and a static method for our manual simulation. The final section demonstrates the production-ready$\texttt{F.scaled\_dot\_product\_attention}$ API with$\texttt{is\_causal=True}$, which automatically uses Flash Attention when available.

       The    cat    sat     on    mat
The  False   True   True   True   True
cat  False  False   True   True   True
sat  False  False  False   True   True
on   False  False  False  False   True
mat  False  False  False  False  False

       The     cat     sat      on     mat
The  0.000   1.000   0.500   0.500   0.750
cat  1.500   0.000   1.000   0.500   0.250
sat  0.500   1.000   1.000   0.500   0.750
on   0.500   0.500   0.000   1.000   0.500
mat  0.500   0.500   0.500   0.500   0.750

       The     cat     sat      on     mat
The  1.0000  0.0000  0.0000  0.0000  0.0000
cat  0.8176  0.1824  0.0000  0.0000  0.0000
sat  0.2327  0.3837  0.3837  0.0000  0.0000
on   0.2350  0.2350  0.1425  0.3875  0.0000
mat  0.1892  0.1892  0.1892  0.1892  0.2430

       d0      d1      d2      d3
The  1.0000  0.0000  0.0000  0.0000
cat  0.8176  0.1824  0.0000  0.0000
sat  0.2327  0.3837  0.3837  0.0000
on   0.2350  0.2350  0.1425  0.3875
mat  0.3108  0.3108  0.3108  0.3108

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4import math
5
6class CausalAttention(nn.Module):
7    """
8    Causal (Masked) Self-Attention in PyTorch.
9
10    Supports both:
11    1. Learnable projections (W_Q, W_K, W_V) for training
12    2. Pre-computed Q, K, V for manual simulation
13    """
14
15    def __init__(self, d_model: int, d_k: int):
16        super().__init__()
17        self.d_k = d_k
18        self.W_Q = nn.Linear(d_model, d_k, bias=False)
19        self.W_K = nn.Linear(d_model, d_k, bias=False)
20        self.W_V = nn.Linear(d_model, d_k, bias=False)
21
22    def forward(self, x: torch.Tensor) -> tuple:
23        """Forward pass with learned projections."""
24        N = x.size(0)
25        Q = self.W_Q(x)
26        K = self.W_K(x)
27        V = self.W_V(x)
28        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_k)
29        mask = torch.triu(
30            torch.ones(N, N, dtype=torch.bool, device=x.device),
31            diagonal=1
32        )
33        scores = scores.masked_fill(mask, float('-inf'))
34        weights = F.softmax(scores, dim=-1)
35        output = weights @ V
36        return output, weights
37
38    @staticmethod
39    def manual(Q, K, V):
40        """Run with pre-computed Q, K, V (no learned weights)."""
41        d_k = K.size(-1)
42        N = Q.size(0)
43        scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k)
44        mask = torch.triu(torch.ones(N, N, dtype=torch.bool), diagonal=1)
45        scores = scores.masked_fill(mask, float('-inf'))
46        weights = F.softmax(scores, dim=-1)
47        output = weights @ V
48        return output, weights
49
50
51# ── Manual simulation with shared example ──
52tokens = ["The", "cat", "sat", "on", "mat"]
53
54Q = torch.tensor([
55    [1.0, 0.0, 1.0, 0.0],
56    [0.0, 2.0, 0.0, 1.0],
57    [1.0, 1.0, 1.0, 0.0],
58    [0.0, 0.0, 1.0, 1.0],
59    [1.0, 0.0, 0.0, 1.0],
60])
61
62K = torch.tensor([
63    [0.0, 1.0, 0.0, 1.0],
64    [1.0, 0.0, 1.0, 0.0],
65    [1.0, 1.0, 0.0, 0.0],
66    [0.0, 0.0, 1.0, 1.0],
67    [1.0, 0.0, 0.5, 0.5],
68])
69
70V = torch.tensor([
71    [1.0, 0.0, 0.0, 0.0],
72    [0.0, 1.0, 0.0, 0.0],
73    [0.0, 0.0, 1.0, 0.0],
74    [0.0, 0.0, 0.0, 1.0],
75    [0.5, 0.5, 0.5, 0.5],
76])
77
78output, weights = CausalAttention.manual(Q, K, V)
79print("Causal Weights:\n", weights.round(decimals=4))
80print("Causal Output:\n", output.round(decimals=4))
81
82# Using PyTorch built-in (recommended for production)
83with torch.no_grad():
84    builtin_output = F.scaled_dot_product_attention(
85        Q.unsqueeze(0), K.unsqueeze(0), V.unsqueeze(0),
86        is_causal=True
87    ).squeeze(0)
88print("\nBuilt-in Output:\n", builtin_output.round(decimals=4))

Key Takeaways

1. Causal masking is the only difference from standard attention. It sets the upper triangle of the score matrix to $-\infty$ before softmax. No new parameters, no new computation — just information removal.
2. It enforces the autoregressive constraint. Token $i$ can only attend to tokens $0, 1, \ldots, i$. This prevents the model from “cheating” by looking at future tokens during training.
3. Training and inference are consistent. Because the mask blocks future tokens during training (where the full sequence is available), the model learns to operate the same way it must at inference time (where future tokens do not exist).
4. The last token sees everything. For the last position in the sequence, causal attention is identical to standard attention. Earlier tokens are progressively more constrained.
5. It enables the KV-cache. Because past computations never change, keys and values can be cached and reused during generation, turning each step from $O(N^2)$ to $O(N)$.
6. Flash Attention makes it faster, not slower. Despite the mask, Flash Attention's causal mode skips the upper triangle entirely, achieving ~50% fewer FLOPs than standard attention.

Exercises

Exercise 1: Compute “on”'s Output by Hand

Token “on” (position 3) can see “The”, “cat”, “sat”, and itself. Using the causal attention weights from this chapter, compute $O_{\text{on}}$ by taking the weighted sum of the four visible value vectors. Verify that your result matches $[0.2350, 0.2350, 0.1425, 0.3875]$.

Exercise 2: What If “The” Could See “cat”?

In causal attention, “The” gets output $[1, 0, 0, 0]$ (it only sees itself). In standard attention, it gets $[0.2254, 0.4135, 0.2964, 0.2964]$. Explain conceptually why this difference matters for a language model during training.

Exercise 3: Bidirectional Mask for Position 2

Suppose you wanted token “sat” to use standard (bidirectional) attention while all other tokens remain causal. Modify the mask matrix for this scenario and recompute the weights for row 2. How does it differ from pure causal?

Exercise 4: Implement Prefix-LM Masking

Write a function $\texttt{prefix\_lm\_mask(N, P)}$ that returns a mask where positions $0, \ldots, P-1$ (the prefix) use bidirectional attention and positions $P, \ldots, N-1$ use causal attention. Apply it to our example with $P = 2$ (so “The” and “cat” are bidirectional, and “sat”, “on”, “mat” are causal).

Exercise 5: KV-Cache Simulation

Simulate autoregressive generation with KV-caching. Starting with only “The,” compute the output for each position as you add tokens one at a time. At each step, append the new K and V vectors to the cache and compute attention for only the new query against all cached keys. Verify that your results match the full causal attention output.

References

1. Vaswani, A. et al. (2017). “Attention Is All You Need.” Advances in NeurIPS, 30. arXiv:1706.03762.
2. Radford, A., Narasimhan, K., Salimans, T., & Sutskever, I. (2018). “Improving Language Understanding by Generative Pre-Training.” OpenAI Technical Report.
3. Radford, A. et al. (2019). “Language Models are Unsupervised Multitask Learners.” OpenAI Technical Report (GPT-2).
4. Brown, T. B. et al. (2020). “Language Models are Few-Shot Learners.” Advances in NeurIPS, 33 (GPT-3). arXiv:2005.14165.
5. Devlin, J., Chang, M., Lee, K., & Toutanova, K. (2018). “BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.” arXiv:1810.04805.
6. Dao, T., Fu, D., Ermon, S., Rudra, A., & Re, C. (2022). “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” Advances in NeurIPS, 35. arXiv:2205.14135.
7. Tay, Y. et al. (2023). “UL2: Unifying Language Learning Paradigms.” ICLR 2023. arXiv:2205.05131.
8. Das, A. et al. (2024). “A Decoder-only Foundation Model for Time-Series Forecasting.”ICML 2024 (TimesFM). arXiv:2310.10688.