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.

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\)?

Engineering bridge · fast path and explicit fallback

MiniMind uses PyTorch's scaled dot-product attention when the inputs are a plain training prefix. It keeps an explicit score path for cached decoding and padded or otherwise masked batches, where the key length can exceed the current query length.

Show the runtime decision
if self.flash and seq_len > 1 and past_key_value is None \
        and (attention_mask is None or torch.all(attention_mask == 1)):
    output = F.scaled_dot_product_attention(
        q, k, v, dropout_p=self.dropout if self.training else 0.0,
        is_causal=True,
    )
else:
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
    # Only the newest query rows need a new causal triangle.
    scores[:, :, :, -seq_len:] += torch.full(
        (seq_len, seq_len), float("-inf"), device=scores.device
    ).triu(1)
    if attention_mask is not None:
        scores += (1.0 - attention_mask[:, None, None, :]) * -1e9
    output = F.softmax(scores.float(), dim=-1).to(q.dtype) @ v

Shape check: with a cached prefix, \(T_q=\) the new token count but \(T_k=T_{\mathrm{past}}+T_q\). The output still returns only \(T_q\) rows.

Retrieval check

Answer once before reopening the code; use the feedback to correct your explanation.

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

STAGE 4 · PROVE IT

Practice studio · recall → build → verify

  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.