Causal Attention: Hiding Future Words · Build an LLM From Scratch

The attention of §3.4 has a fatal flaw for text generation: every token can see the whole sequence — including tokens that come after it. A GPT generates left to right; when it predicts the word after “Your journey”, the words “starts with one step” do not exist yet. Training must reflect that. Causal attention (also called masked attention) is the fix — rung 3 of the chapter’s ladder: it restricts the model to only consider previous and current inputs when processing any given token.

Below is the exact matrix from §3.4 (the SelfAttention_v2 weights, seed 789). Toggle to causal-masked α to see what the mask does: everything above the diagonal — each row’s future — is forced to zero, and the surviving entries renormalize so each row still sums to 1:

softmax(ω / 1.41)

	Your	journey	starts	with	one	step
Your	0.19	0.16	0.17	0.15	0.17	0.15
journey	0.20	0.17	0.17	0.15	0.17	0.15
starts	0.20	0.17	0.17	0.15	0.17	0.15
with	0.19	0.17	0.17	0.16	0.17	0.16
one	0.18	0.17	0.17	0.16	0.17	0.16
step	0.19	0.17	0.17	0.15	0.17	0.15

hover a cell — row = query (who is looking), column = key (who is looked at)

Recreation of the book's p. 36 figure with the real §3.4 values. Row «Your» keeps only itself (weight 1.0); row «step», as the last token, keeps everything. Hover a masked cell to see why it's off-limits.

3.5.1 Applying a causal attention mask

Both implementations build the same causal mask — an upper-triangular pattern marking each row’s future. The naive one applies it after the softmax, in three steps; the efficient one before it, in one.

The naive recipe (book p. 37): softmax first, then mask with 0s, then renormalize each row.

Walk through the three steps on one concrete row — “starts”, which may see only “Your”, “journey”, and itself:

Step 1 — the ordinary softmax weights

Start from the §3.4 attention weights: the «starts» row is [0.2036, 0.1659, 0.1662, 0.1498, 0.1664, 0.1480], summing to 1 — but it still assigns weight to "with", "one" and "step", which lie in the future.

1 / 3

The naive masking recipe applied to the row of the query «starts» (row 3 of the matrix).

</> The naive three-step implementation

# [Step 1] the attention weights from the previous section
print(attn_weights)  # 6×6, each row sums to 1

# [Step 2] zero out everything above the diagonal
context_length = attn_scores.shape[0]
mask_simple = torch.tril(torch.ones(context_length, context_length))  #A
masked_simple = attn_weights * mask_simple  #B

# [Step 3] renormalize each row to sum to 1
row_sums = masked_simple.sum(dim=1, keepdim=True)  #C
masked_simple_norm = masked_simple / row_sums
print(masked_simple_norm)
# tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
#         [0.5517, 0.4483, 0.0000, 0.0000, 0.0000, 0.0000],
#         [0.3800, 0.3097, 0.3103, 0.0000, 0.0000, 0.0000],
#         [0.2758, 0.2460, 0.2462, 0.2319, 0.0000, 0.0000],
#         [0.2175, 0.1983, 0.1984, 0.1888, 0.1971, 0.0000],
#         [0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],
#        grad_fn=<DivBackward0>)

#A torch.tril keeps the lower triangle (1s on and below the diagonal), zeroing the rest.
#B Element-wise multiply: future positions become 0, the past is untouched.
#C keepdim=True keeps the sum as a (6×1) column so the division broadcasts row-wise.

No information leakage — softmax’s elegant cancellation

The masked future tokens were part of the original softmax denominator — doesn’t their influence linger? No. Renormalizing after masking is mathematically the same as computing softmax over the smaller, visible subset from the start: the masked positions’ contribution cancels out exactly. They don’t influence the surviving weights in any meaningful way — as if they had never been included.

The −∞ shortcut: mask before the softmax

The three steps collapse into one by exploiting a property of softmax: since $e^{-\infty} \to 0$ , any score set to $-\infty$ receives zero weight, and the rest of the row normalizes automatically — mask and normalization in a single softmax call:

$\mathrm{softmax}\big([\,0.4656,\; 0.1723,\; -\infty,\; \dots\,]/\sqrt{d_k}\big) = [\,0.5517,\; 0.4483,\; 0,\; \dots\,]$

</> The efficient implementation: fill future scores with −inf

mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)  #A
masked = attn_scores.masked_fill(mask.bool(), -torch.inf)  #B
print(masked)
# tensor([[0.2899,   -inf,   -inf,   -inf,   -inf,   -inf],
#         [0.4656, 0.1723,   -inf,   -inf,   -inf,   -inf],
#         [0.4594, 0.1703, 0.1731,   -inf,   -inf,   -inf],
#         [0.2642, 0.1024, 0.1036, 0.0186,   -inf,   -inf],
#         [0.2183, 0.0874, 0.0882, 0.0177, 0.0786,   -inf],
#         [0.3408, 0.1270, 0.1290, 0.0198, 0.1290, 0.0078]],
#        grad_fn=<MaskedFillBackward0>)

attn_weights = torch.softmax(masked / keys.shape[-1]**0.5, dim=1)  #C
# → exactly the same matrix as the three-step version, in one softmax

#A torch.triu with diagonal=1 marks strictly-above-diagonal positions — the future.
#B masked_fill writes −inf wherever the boolean mask is True, leaving the visible scores untouched.
#C One scaled softmax now does masking AND normalization: exp(−inf) = 0 drops the future, the visible entries sum to 1. Fewer operations, same result.

3.5.2 Masking additional weights with dropout

Dropout ignores randomly selected units during training — “dropping them out” — so the model cannot become over-reliant on any specific set of hidden units. Two things to keep straight:

dropout is used only during training and disabled afterward;
to keep the expected magnitude unchanged, the surviving values are scaled up by $1/(1-p)$ — with rate $p = 0.5$ , survivors are doubled.

In transformer blocks, dropout on the attention mechanism is typically applied at one of two points: after computing the attention scores (before the softmax) or — as GPT implementations and the book’s code do — after the softmax, on the attention weights:

</> Where dropout goes (pseudocode) — the book uses the after-softmax variant

function scaledDotProductAttention(Q, K, V, dropoutRate):
  scores = matmul(Q, transpose(K)) / sqrt(d_k)

  attentionWeights = softmax(scores)

  # applied AFTER softmax, on the attention weights  <-- this book / GPT
  attentionWeights = applyDropout(attentionWeights, dropoutRate)

  output = matmul(attentionWeights, V)
  return output

</> A 50% dropout layer on a matrix of ones

torch.manual_seed(123)
dropout = torch.nn.Dropout(0.5) #A
example = torch.ones(6, 6) #B
print(dropout(example))
# tensor([[2., 2., 0., 2., 2., 0.],
#         [0., 0., 0., 2., 0., 2.],
#         [2., 2., 2., 2., 0., 2.],
#         [0., 2., 2., 0., 0., 2.],
#         [0., 2., 0., 2., 0., 2.],
#         [0., 2., 2., 2., 2., 0.]])

#A Rate 0.5 for the demonstration; when training the GPT later the book uses 10–20%.
#B A 6×6 matrix of 1s: after dropout, ≈half are 0 and the SURVIVORS print as 2 — the 1/(1−0.5) = 2× rescaling in action.

Try it on the real causal attention weights. Drag the rate and re-roll the random mask — watch survivors grow by $1/(1-p)$ as their neighbors vanish:

dropout rate p = 0.50survivors × 1/(1−p) = ×2.00 · kept 20/36

	Your	journey	starts	with	one	step
Your	2.00	0	0	0	0	0.00
journey	1.10	0.90	0.00	0	0.00	0.00
starts	0	0	0.62	0	0.00	0.00
with	0	0.49	0.49	0	0	0
one	0.43	0.40	0	0.38	0.39	0.00
step	0.39	0	0.33	0	0.33	0

Dropout applied to the causal attention-weight matrix (training only). Note rows no longer sum to 1 after dropout — that's expected; the 1/(1−p) rescaling keeps the EXPECTED sum right, and at inference dropout is off entirely.

3.5.3 A compact causal attention class

Listing 3.3 packages causal masking + dropout into a module that also handles batched inputs of shape (b, num_tokens, d_in) — like the 2-sentence batch chapter 2’s data loader produced:

</> Listing 3.3 — A compact causal attention class

import torch.nn as nn
class CausalAttention(nn.Module):
  def __init__(self, d_in, d_out, context_length,
               dropout, qkv_bias=False):
      super().__init__()
      self.d_out = d_out
      self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
      self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)
      self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
      self.dropout = nn.Dropout(dropout) #A
      self.register_buffer(
          'mask',
          torch.triu(
              torch.ones(context_length, context_length),
              diagonal=1
          )
      ) #B

  def forward(self, x):
      b, num_tokens, d_in = x.shape #C
      keys    = self.W_key(x)
      queries = self.W_query(x)
      values  = self.W_value(x)

      attn_scores = queries @ keys.transpose(1, 2) #C
      attn_scores.masked_fill_( #D
          self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)
      attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
      attn_weights = self.dropout(attn_weights)

      context_vec = attn_weights @ values
      return context_vec

#A Dropout on the attention weights, exactly as in the pseudocode above — active in training, disabled in eval mode.
#B register_buffer stores the triu mask on the module: it travels with the model to CPU/GPU like a parameter, but the optimizer never trains it.
#C Inputs are now batched: (b, num_tokens, d_in). transpose(1, 2) swaps the token and feature dims so the matmul runs per batch element — keys.T would wrongly flip the batch dim too.
#D The trailing underscore = in-place. Slicing [:num_tokens, :num_tokens] lets one pre-built context_length-sized mask serve shorter sequences.

</> Using CausalAttention on a batch of two sequences

batch = torch.stack((inputs, inputs), dim=0)  #A

torch.manual_seed(123)
context_length = batch.shape[1]
d_in, d_out = 3, 2
ca = CausalAttention(d_in, d_out, context_length, 0.0)
context_vecs = ca(batch)
print("context_vecs.shape:", context_vecs.shape)
# context_vecs.shape: torch.Size([2, 6, 2])

#A Stacking the same 6-token example twice simulates a batch of 2 texts: shape (2, 6, 3). The output keeps the batch: (2, 6, 2) — for each of 2 sequences, 6 context vectors of dimension 2.

The mask travels with the module because of register_buffer : when you later call model.to("cuda"), buffers move along with parameters — but unlike an nn.Parameter, the optimizer never updates the mask.

Key idea — causal attention in one sentence

Fill the strictly-upper triangle of the score matrix with

-\infty

before the scaled softmax, then (during training) apply dropout to the resulting weights: each token now builds its context vector only from itself and the past, one token at a time — exactly how a GPT must read while generating. One rung remains: running several of these heads in parallel — multi-head attention.

📝 Check yourself: causal attention & dropout

0 / 5

1.What restriction does causal (masked) attention add to standard self-attention?
2.In the naive mask-then-renormalize recipe, masked future tokens were part of the ORIGINAL softmax denominator. Why is there still no information leakage?
3.Why does replacing future positions' SCORES with −∞ before the softmax implement masking in one step?
4.Predict: a dropout layer with rate p = 0.5 is applied to a matrix of 1.0s during training. What do the SURVIVING entries print as?
5.Why does CausalAttention store its mask with register_buffer instead of a plain attribute or an nn.Parameter?