A Short Course for Modern Large Language ModelsGlossary

Lesson 18 of 23 · Transformer block

Hide the future before probabilities exist.

Your win: build one scaled query-key score row, apply a causal mask, and explain why the mask must precede softmax.

14 minutesNeeds: dot productsOutcome: construct a causal score row
Mission link: RoPE changes \(Q\) and \(K\); causal attention then compares \(Q'\) with \(K'\) without allowing a token to read later positions.

Scores have one row per query

\[ S=\frac{Q'K'^{\mathsf T}}{\sqrt{d_h}}, \qquad S_{ij}=\frac{\mathbf q'_i\cdot\mathbf k'_j}{\sqrt{d_h}}. \]

For one head, \(Q'\in\mathbb R^{T\times d_h}\) and \(K'^{\mathsf T}\in\mathbb R^{d_h\times T}\), so \(S\in\mathbb R^{T\times T}\). Row \(i\) asks: “How strongly does query token \(i\) match every key token \(j\)?”

Numbers substituted: one scaled dot product

Let \(d_h=4\), \(\mathbf q'_3=(1,2,0,-1)\), and \(\mathbf k'_2=(2,0,1,1)\). Then

\[ S_{3,2}=\frac{(1)(2)+(2)(0)+(0)(1)+(-1)(1)}{\sqrt4} =\frac{1}{2}=0.5. \]

The factor \(1/\sqrt{d_h}\) keeps dot-product magnitudes from growing merely because a head has more coordinates.

The causal mask is triangular

\[ M_{ij}=\begin{cases}0,&j\le i,\\-\infty,&j>i,\end{cases} \qquad A=S+M. \]

At query position \(i=3\), raw scores \((0.7,1.1,0.2,2.0)\) become \((0.7,1.1,0.2,-\infty)\). The tempting future score \(2.0\) is removed before normalization.

Explore which keys are visible

\(i=\)3
KeyRaw scoreAfter maskSoftmax weightValue

Weighted output:

Order matters: masking after softmax would assign positive probability to future tokens first. Adding \(-\infty\) before softmax makes \(e^{-\infty}=0\), so forbidden weights are exactly zero in the mathematical model.

Code checkpoint · mask scores before softmax

The upper triangle represents future key positions. Replacing those scores with \(-\infty\) guarantees zero probability after softmax.

Show score scaling and the causal mask
scores = q @ k.transpose(-2, -1) / math.sqrt(head_dim)
# scores: [B, H, query_token, key_token]

tokens = scores.size(-1)
future = torch.ones(tokens, tokens, dtype=torch.bool, device=x.device)
future = future.triu(diagonal=1)
scores = scores.masked_fill(future, float("-inf"))

Trace it: In row 2 of a \(4\times4\) mask, which column receives \(-\infty\)?

Retrieval check

Which keys are visible to query position \(i=2\)?

Practice before moving on

  1. For \(T=4\), write all four rows of the causal mask using \(0\) and \(-\infty\).
  2. Mask the row \((3,-1,2,5)\) for query position \(i=2\).
  3. Compute \(S_{1,2}\) when \(d_h=2\), \(\mathbf q'_1=(2,1)\), and \(\mathbf k'_2=(1,-1)\).
  4. If \(Q'\) has shape \(6\times8\), give the shapes of \(K'^{\mathsf T}\) and \(S\).
  5. Why is the mask added before softmax?
  6. Does causal masking remove a key vector from memory, or only prevent particular query-key connections?
Check solutions
  1. \((0,-\infty,-\infty,-\infty)\), \((0,0,-\infty,-\infty)\), \((0,0,0,-\infty)\), and \((0,0,0,0)\).
  2. \((3,-1,-\infty,-\infty)\).
  3. \(((2)(1)+(1)(-1))/\sqrt2=1/\sqrt2\approx0.707\).
  4. \(K'^{\mathsf T}:8\times6\); \(S:6\times6\).
  5. So forbidden entries exponentiate to zero and receive no probability mass.
  6. Only particular connections are blocked; the key still exists for queries at its own or later positions.

Primary source: Vaswani et al., Sections 3.2.1 and 3.2.3.