Many Heads, One-Way Street: Multi-Head Attention and Causal Masking

One attention can only ask one question

A single attention operation learns one notion of *"relevant."* But language needs many at once. To resolve it you care about which noun it refers to; to conjugate a verb you care about the subject's number; to close a quote you care about where the opening quote was. Forcing one attention to juggle all of these blurs them together.

The fix is multi-head attention. Instead of one large attention, the model runs several smaller ones in parallel — each is a head with its own independent query, key, and value matrices. Each head computes its own scores, its own softmax, its own value mix, over the same input. A typical model has anywhere from 8 to 128 heads per layer.

Splitting, attending, recombining

1. Split. The word's vector is divided across the heads — with 8 heads, each head works in a smaller slice of the dimensions. The total work stays about the same as one big head.
2. Attend. Each head independently does the full self-attention routine from the last guide, producing its own context-mixed output.
3. Recombine. The heads' outputs are concatenated and passed through one more learned matrix, which lets the model blend what the heads found into a single updated vector per word.

What the heads end up doing

Nobody tells the heads what to specialize in — they discover roles during training. When researchers inspect a trained model, recurring patterns appear. Some heads are positional, mostly attending to the previous one or two words. Some are syntactic, linking verbs to their subjects or determiners to their nouns. A famous kind is the induction head, which spots a repeated pattern (*"...Dumbledore said... Dumbledore"*) and copies what came after it last time — a key ingredient of in-context learning.

We can see these roles directly by drawing each head's attention weights as a grid, source words on one axis, target words on the other, brightness for weight. This attention pattern visualization is one of the main windows researchers have into what a model is actually doing inside.

The one-way street: causal masking

There's a problem hiding in everything so far. A language model is trained to do next-token prediction — given the words so far, guess the next one. But plain self-attention lets every word see every other word, including the ones after it. If the model can peek at the answer, it learns nothing useful. We have to blindfold it.

Causal masking does exactly that. Before the softmax, every score where a word would attend to a future word is set to negative infinity. After softmax those positions become zero weight — invisible. The result is a strict one-way street: position 5 may attend to positions 1 through 5, never to 6 onward. This is what makes the architecture decoder-only and autoregressive: it generates one token at a time, each conditioned only on its past.

Causal masking adds −∞ to every future position before the softmax, so each token can attend only to itself and the past.

# causal mask applied inside attention, before softmax
scores = Q @ K.T / sqrt(d_k)     # (seq_len, seq_len)

# upper triangle = future positions -> forbid them
scores = scores.masked_fill(future_positions, -infinity)

weights = softmax(scores)        # future positions now get weight 0
out     = weights @ V