Causal Attention

Multi-head attention extends causal attention across several heads, each operating independently with its own weights.

Programming

One Attention Module, Several Views

So far, attention in this course has been built as a single causal attention class. That class processes an input sequence using one set of weights. The final step in this line of work is to extend the class so it covers several heads at once. This extension is called multi-head attention. The idea is less exotic than the name suggests. Instead of asking one attention mechanism to do all of the work, we run several attention mechanisms side by side and gather what they produce.

What do you think happens?

How many sets of query weight matrices does a two-head module contain?

One, shared by both heads
Two, one per head
Four, one for each of query, key, value, and output

Reveal answer

Answer:

Two, one per head

Each head has its own query, key, and value weight matrices. The two-head module therefore has two query matrices, two key matrices, and two value matrices. The instances do not share weights, which is what makes each head an independent view of the same input.

This article follows the same path the course uses. It starts from the single-head case, then looks at how a multi-head module is organized, and finally explains why the approach is worth its added cost.

Tracing a Single-Head Module

A single causal attention module can be regarded as single-head attention. In this arrangement there is exactly one set of attention weights, and that one set processes the input sequence. Every token passes through the same weights, and the module produces one set of context vectors as its result.

Read the single-head flow from left to right. The input tokens enter the module. A single weight configuration transforms them as a sequence. A single context vector matrix leaves the module. Keep this picture in mind, because the multi-head version keeps the same overall flow. What changes is how many weight sets sit in the middle.

What the Causal Mask Allows

The word causal in causal attention describes which positions a token is allowed to use. Each token attends only to itself and to earlier tokens. Future positions are blocked. The map below shows the pattern for a four-token sequence. Each row is a token, and each row lists the positions it can attend to.

Notice the direction of the blocking. Token 2 can see token 1, but token 1 cannot see token 2. The mask removes the future in every row, so the upper right part of the map is always empty. The multi-head module keeps this rule for every head, because each head is a causal attention instance.

Tracing Input Through Two Heads

In practical terms, a multi-head module creates several instances of the self-attention mechanism. The diagram below shows the two-head case. The embedded input tokens stay exactly as they are. Where a single-head module had one value weight matrix, the two-head module has two, one for each head. The same change applies to the query and key weight matrices. Each head therefore produces its own set of context vectors, labeled Z1 and Z2 in the source. Those two sets are then combined into one context vector matrix Z.

Worked example

Following one input sequence through two heads

A causal attention module with two heads receives a sequence of four tokens. Trace what happens to the sequence from entry to the final output.

1
Step 1: The input enters
The four embedded tokens enter the module unchanged. Both heads receive the same input sequence.
2
Step 2: Head 1 processes the sequence
Head 1 uses its own query, key, and value weight matrices. It applies the causal rule, so each token only draws on itself and earlier tokens. The result is the context vectors Z1.
3
Step 3: Head 2 processes the same sequence
Head 2 has a separate set of query, key, and value weights. It works on the same input and applies the same causal rule, but with its own weights. The result is the context vectors Z2, which can differ from Z1.
4
Step 4: The head outputs are combined
The two sets of context vectors, Z1 and Z2, are combined into one context vector matrix Z. The module returns a single result, just as the single-head module does.

Result

The module receives one input sequence and returns one combined context vector matrix Z. Internally, two independent heads each produced their own view of that sequence.

Adding More Heads and Merging Their Outputs

The same pattern extends to more than two heads. Each additional head adds another set of weight matrices and another set of context vectors to be combined. The two-head case is simply the smallest version that shows the full structure. Whatever the number of heads, the module keeps one input and one combined output.

Because the heads do not share weights, each head is an independent view of the same input. The merged matrix Z carries the results of all of those views forward together.

Two Stages of Building the Module

The expansion from causal attention to multi-head attention is approached in two stages. The first stage builds the idea intuitively by stacking several CausalAttention modules. Each one acts as a single head, and their results are combined. The second stage implements the same multi-head module in a more complicated but more computationally efficient way.

Stacking is the right starting point because it makes the structure visible. Each stacked module is a head you already understand. The multi-head module is simply several of them arranged so that their outputs can be merged. The later efficient version computes the same kind of result with a different organization, so understanding the stacked version first makes that later step easier to follow.

Why Multiple Heads Are Worth Their Cost

Multi-head attention is computationally intensive. It is worth that cost because of the kind of pattern recognition it supports. Models based on transformers, which include many large language models, rely on this ability to handle complex patterns in data.

Giving the model several independent heads means that different heads can, in principle, attend to the input in different ways. The combined output carries those differences forward. A single head produces one view of the sequence. Several heads produce several views, and the merged result keeps all of them.

Aspect	Single-head attention	Multi-head attention
Weight sets	One set of query, key, and value weights	One set of query, key, and value weights per head
Outputs	One set of context vectors	One set of context vectors per head, combined into one matrix Z
Input flow	Input passes through one weight configuration	Input passes through several independent heads that share the same input
Cost	Baseline cost for one module	More weights and more computation, growing with each added head
Pattern capacity	One view of the sequence	Several views, which can attend to the input in different ways

Single-head and multi-head causal attention compared

Mistakes Learners Make with Heads

Common mistake
Thinking that a head is a separate kind of attention
Why it fails
A head is an instance of the same self-attention mechanism. Its meaning comes from having its own weights, not from a different algorithm.
Do this instead
Treat each head as one more copy of the causal attention module, with its own query, key, and value weights.
Common mistake
Assuming the heads share one set of weights
Why it fails
The instances do not share weights. That independence is what makes each head a separate view of the input.
Do this instead
In a two-head module, count two query matrices, two key matrices, and two value matrices.
Common mistake
Expecting the output to grow into several separate results
Why it fails
The head outputs are combined into one context vector matrix Z. The module still returns one result for one input sequence.
Do this instead
Follow Z1 and Z2 from each head into the single merged matrix Z.
Common mistake
Believing more heads are always free
Why it fails
Multi-head attention is computationally intensive, and each added head adds weight matrices and context vectors to compute and combine.
Do this instead
Weigh the pattern recognition gained from extra heads against the added computation and memory.
Common mistake
Dropping the causal rule when adding heads
Why it fails
Each head is a causal attention instance, so the rule that a token attends only to itself and earlier tokens applies to every head.
Do this instead
Check that every head blocks future positions, the same way the single-head module does.

Practice Tracing Two Heads

MEDIUM

A two-head causal attention module receives an input sequence. Describe, in order, every object the sequence passes through before the module returns its final result. Name the weight matrices that each head uses, and state what the two head outputs become.

Hints

Start with the embedded input tokens and note whether they change on entry.
Count the query, key, and value weight matrices separately for each head.
Look for the single matrix that the two sets of context vectors merge into.

Your attempt (not saved)