L4.6 · LESSON 6 OF 6

Masking and multiple heads

Calculate query-key scores, scaled softmax weights and the weighted value output for one attention head.

THE IDEA

Hiding the future, and splitting into heads

When a model predicts the next token, a position must not see the tokens after it. A causal mask sets the scores of future positions to −∞ before the softmax, so they get weight exactly 0 and the rest still add up to 1. In the worked example, if only the first key is allowed the output becomes [2, 0].

One head looks at one kind of relationship. Multi-head attention splits the model width into h slices, runs attention separately in each (each with its own projected subspace), joins the results and multiplies by an output matrix.

causal mask (a token cannot look at later tokens)

query ↓ / key →	token 1	token 2	token 3	token 4
token 1	1	masked	masked	masked
token 2	0.27	0.73	masked	masked
token 3	0.67	0.09	0.24	masked
token 4	0.17	0.17	0.17	0.48

The mask is applied to the scores as −∞ before the softmax, so each row still sums to 1 over the positions that remain.

number of heads2

model width 8 split into 2 heads of width 4; each head attends on its own slice, then the head outputs are joined and multiplied by an output matrix Wo

Toggle the mask on a 4-token example, then choose a number of heads for a model of width 8.

Predict first. Should forbidden positions keep a small positive softmax weight?

Chapter context

Attention computes a data-dependent mixture of value vectors. A query scores candidate keys, a normalization turns scores into weights, and those weights combine values. Keys decide how to look; values supply what is gathered.

TRY IT YOURSELF

Part 1. Do it hands-on

Work out masking and head widths by hand, then write masked attention and a full multi-head layer and prove that the future cannot leak.

causal mask (a token cannot look at later tokens)

query ↓ / key →	token 1	token 2	token 3	token 4
token 1	1	masked	masked	masked
token 2	0.27	0.73	masked	masked
token 3	0.67	0.09	0.24	masked
token 4	0.17	0.17	0.17	0.48

The mask is applied to the scores as −∞ before the softmax, so each row still sums to 1 over the positions that remain.

number of heads2

model width 8 split into 2 heads of width 4; each head attends on its own slice, then the head outputs are joined and multiplied by an output matrix Wo

positions the 3rd token may see weight of a masked position head width (model 8, 2 heads)

chapter example with the second key masked: output =

To do: Causal mask in a 4-token sequence: how many positions (including itself) may the THIRD token look at?
To do: A forbidden position gets a score of −∞ before the softmax. What weight does it end up with?
To do: Chapter example again, but only the first key is allowed (the second is masked). What is the output vector?
To do: A model of width 8 uses 2 heads. How wide is each head?

Part 2. Do it in code

causal_mask(n): an (n, n) boolean array where entry [i, j] is True when position i may look at position j (j <= i).
masked_attention(Q, K, V, mask=None): scaled dot-product attention. Where mask is False put −inf in the scores BEFORE the softmax. Return (weights, output).
split_heads(X, h): reshape (n, d) into (h, n, d // h). d must be divisible by h (ValueError).
merge_heads(Xh): the inverse, (h, n, d_head) back to (n, h * d_head).
multi_head_attention(X, Wq, Wk, Wv, Wo, h, causal=True): project X, split into h heads, run masked_attention per head (causal mask if causal), merge the heads and multiply by Wo. Return the (n, d_model) output.

import numpy as np

def causal_mask(n):
    pass

def masked_attention(Q, K, V, mask=None):
    pass

def split_heads(X, h):
    pass

def merge_heads(Xh):
    pass

def multi_head_attention(X, Wq, Wk, Wv, Wo, h, causal=True):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Should forbidden positions retain a positive softmax weight?

AYesBNo

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

Q\in\mathbb R^{n_q\times d_k},\ K\in\mathbb R^{n_k\times d_k},\ V\in\mathbb R^{n_k\times d_v}.

S=QK^T/\sqrt{d_k}+M,\quad A=\operatorname{softmax}_{rows}(S),\quad O=AV.

M is zero at allowed positions and conceptually negative infinity at excluded positions. Every query needs a defined policy if all its keys are masked.

Worked example

Use q=[1,0], K=[[1,0],[0,1]], V=[[2,0],[0,2]], and d_k=2. Dot scores [1,0] become [1/sqrt(2),0]. Exponentiate and normalize to approximately [0.669762,0.330238]. The output is [2×0.669762,2×0.330238]≈[1.339523,0.660477]. If only the first key is causally allowed, the output becomes [2,0].

Reference code

The following example uses the Python standard library and does not call a remote service.

from math import sqrt,exp,isclose
q=[1,0]; keys=[[1,0],[0,1]]; values=[[2,0],[0,2]]
scores=[sum(a*b for a,b in zip(q,k))/sqrt(2) for k in keys]
weights=[exp(s-max(scores)) for s in scores]
weights=[w/sum(weights) for w in weights]
out=[sum(weights[i]*values[i][j] for i in range(2)) for j in range(2)]
assert isclose(sum(weights),1)
assert isclose(weights[0],.6697615493266569)
assert isclose(out[0],1.3395230986533138)
assert isclose(out[1],.6604769013466862)

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Attention scores are not weights before normalization. Causal masking happens before softmax. Attention weights alone are not a full causal explanation of a model decision.