Causal multi-head attention in NumPy

This is the single most commonly reported from-scratch task in ML engineer and research loops, and it is usually asked without a framework: a blank editor, NumPy, forty minutes. The mathematics is four lines. What actually goes wrong is shapes — a transpose on the wrong axis, a mask broadcast into the wrong dimension, a softmax over the wrong one — so this section states the shape of every intermediate and ends with two assertions that catch most of it.

One head, one sequence#

Start with a single sequence of T positions, each a vector of d_model features. Three learned matrices turn that one input into a query, a key and a value for every position. They are the same input three times, projected differently; there is no separate source for the keys in self-attention.

import numpy as np

rng = np.random.default_rng(0)

T, d_model, d_k = 5, 8, 8
x = rng.normal(size=(T, d_model))                 # (T, d_model)

W_q = rng.normal(size=(d_model, d_k)) * 0.1       # (d_model, d_k)
W_k = rng.normal(size=(d_model, d_k)) * 0.1
W_v = rng.normal(size=(d_model, d_k)) * 0.1

q = x @ W_q                                       # (T, d_k)
k = x @ W_k                                       # (T, d_k)
v = x @ W_v                                       # (T, d_k)

Each row of q asks a question; each row of k advertises what that position offers. Comparing every query with every key is one matrix multiplication, and it produces a square matrix of T by T: how much position i should attend to position j.

scores = q @ k.T / np.sqrt(d_k)                   # (T, d_k) @ (d_k, T) -> (T, T)

The division is not decoration. A dot product of two d_k-dimensional vectors with unit-variance entries has variance d_k, so without scaling the scores grow with the feature dimension, the softmax saturates into a near one-hot row, and the gradient through it goes to nothing. Dividing by the square root of d_k puts the variance back to roughly one.

The causal mask#

A decoder may not look ahead. Position i may attend to positions 0 through i and no further, which is exactly the lower triangle of the score matrix including the diagonal. The mask is applied to the scores, before the softmax, by setting the forbidden entries to negative infinity — so that after exponentiating they are exactly zero and the remaining weights still sum to one.

# True strictly above the diagonal: the future, which must not be attended to
mask = np.triu(np.ones((T, T), dtype=bool), k=1)
scores = np.where(mask, -np.inf, scores)          # (T, T), still

print(np.round(scores, 2))
# [[-0.03  -inf  -inf  -inf  -inf]
#  [ 0.11 -0.05  -inf  -inf  -inf]
#  [ 0.03  0.    0.    -inf  -inf]
#  [ 0.02  0.    0.02 -0.02  -inf]
#  [ 0.05 -0.1   0.02 -0.05  0.03]]

Softmax, stably#

Write it yourself rather than reaching for a library, because the stable form is part of what is being asked. Subtracting the row maximum before exponentiating changes nothing mathematically — the constant cancels between numerator and denominator — and it means the largest exponent is always exp(0), so nothing overflows.

def softmax(z, axis=-1):
    # subtracting the max is a no-op mathematically and stops exp overflowing
    z = z - z.max(axis=axis, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=axis, keepdims=True)

w = softmax(scores, axis=-1)                      # (T, T), each row sums to 1
out = w @ v                                       # (T, T) @ (T, d_k) -> (T, d_k)

The masked entries survive this correctly: a row's maximum is always a real number, because the diagonal is never masked, so a masked entry becomes negative infinity minus a finite number, and its exponential is zero. No row is entirely masked, so there is no division by zero to worry about here.

The two assertions to run#

Nearly every bug in an implementation of this shows up in one of two checks. Run both before claiming it works, and say out loud that you are running them: an interviewer watching you write your own test is watching the thing they are actually grading.

# every row is a probability distribution over the positions it can see
assert np.allclose(w.sum(axis=-1), 1.0)

# nothing attends to the future: the strict upper triangle is exactly zero
assert np.all(np.triu(w, k=1) == 0.0)

Several heads at once#

Multiple heads are not multiple loops. The projections stay one matrix each of shape (d_model, d_model); the head split is a reshape of the feature axis into n_heads chunks of d_head, followed by a transpose that moves heads in front of positions so the score matrix is computed per head by ordinary batched matrix multiplication.

def causal_mha(x, W_q, W_k, W_v, W_o, n_heads):
    B, T, d_model = x.shape                       # (B, T, d_model)
    d_head = d_model // n_heads

    q = x @ W_q                                   # (B, T, d_model)
    k = x @ W_k
    v = x @ W_v

    def split(t):
        # (B, T, d_model) -> (B, T, n_heads, d_head) -> (B, n_heads, T, d_head)
        return t.reshape(B, T, n_heads, d_head).transpose(0, 2, 1, 3)

    q, k, v = split(q), split(k), split(v)

    # transpose the last two axes only: (B, n_heads, d_head, T)
    scores = q @ k.transpose(0, 1, 3, 2) / np.sqrt(d_head)   # (B, n_heads, T, T)

    mask = np.triu(np.ones((T, T), dtype=bool), k=1)         # broadcasts over B and heads
    scores = np.where(mask, -np.inf, scores)
    w = softmax(scores, axis=-1)                             # (B, n_heads, T, T)

    ctx = w @ v                                              # (B, n_heads, T, d_head)
    # heads back behind positions, then flatten them into the feature axis
    ctx = ctx.transpose(0, 2, 1, 3).reshape(B, T, d_model)   # (B, T, d_model)
    return ctx @ W_o, w                                      # (B, T, d_model)

Two details in that function are where people lose marks. The transpose before the score matmul swaps the last two axes and leaves the batch and head axes alone; writing k.T would reverse all four. And the final reshape must follow the transpose back to (B, T, n_heads, d_head), because reshaping straight from (B, n_heads, T, d_head) interleaves positions and heads into nonsense that still has the right shape and will not raise an error.

B, T, d_model, n_heads = 2, 6, 16, 4
x = rng.normal(size=(B, T, d_model))
W_q, W_k, W_v, W_o = (rng.normal(size=(d_model, d_model)) * 0.1 for _ in range(4))

y, w = causal_mha(x, W_q, W_k, W_v, W_o, n_heads)
print(y.shape, w.shape)                           # (2, 6, 16) (2, 4, 6, 6)

assert np.allclose(w.sum(axis=-1), 1.0)
assert np.all(np.triu(w, k=1) == 0.0)

# a third check worth knowing: causality means a prefix is unaffected by what
# follows it, so truncating the input must not change the earlier outputs
y_short, _ = causal_mha(x[:, :4], W_q, W_k, W_v, W_o, n_heads)
assert np.allclose(y[:, :4], y_short)

What is deliberately not here#

This is attention, not a transformer block. A block wraps it in a residual connection and a normalisation, then a feed-forward layer with its own residual, and a real model adds positional information — rotary embeddings, applied to q and k before the scores, in every current model. Modern attention also shares keys and values across heads rather than giving each its own. Those are covered elsewhere in this topic; if an interviewer asks for attention, give them this and then say what would surround it.