Causal Self-Attention

Implement single-head scaled dot product attention with a causal mask, the attention used in decoder-only language models such as GPT.

A language model predicts each token from the tokens before it. During training it sees the whole sequence at once, so the attention must be blocked from looking ahead. The causal mask sets the score from position i to every later position j > i to -inf before the softmax, which makes those weights exactly 0:

scores  = Q @ K.T / sqrt(d_k)
scores[i, j] = -inf    for j > i
weights = softmax(scores, axis=-1)
out     = weights @ V

Function signature

def causal_attention(Q, K, V):

Example

Input

Q = K = np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])
V = np.array([[1.0, 0.0], [0.0, 2.0], [3.0, 3.0]])
causal_attention(Q, K, V)

Output

array([[1.        , 0.        ],
       [0.33023845, 1.3395231 ],
       [1.75872461, 2.00697969]])

Reasoning

Position 0 can only see itself, so its output is V[0]. Position 1 sees positions 0 and 1 with scores 0 and 1 / sqrt(2) = 0.707, which give weights 0.330 and 0.670. Position 2 sees all three. The output of position 1 does not depend on V[2] at all, which is the property the mask exists to give.

Constraints and notes

• Q and K have shape (T, d_k) and V has shape (T, d_v).
• Subtract each row's maximum before the exponent.
• Changing a later row of K or V must not change any earlier output row.