Single-Head Attention Modules

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

Programming

Why One Attention Mechanism Is Not the Whole Story

So far, attention has been built as a single causal attention class that processes an input sequence with one set of weights. The final step in this line of work is to extend that class so it covers several heads at once. This extension is known as 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.

This article 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. Along the way you will trace an input sequence through a two-head module step by step, because the two-head case is the smallest version that shows the full structure.

What a Single-Head Module Does

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.

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

What the Word Head Means

Definition

A head is one instance of the self-attention mechanism, with its own set of weights. Multi-head attention extends causal attention across several heads, each operating independently with its own weights.

In practical terms, building multi-head attention means creating several instances of the self-attention mechanism. Each instance has its own weights. After each instance has done its work, the outputs of the instances are combined into a single result. The instances do not share weights, and that is what makes each head an independent view of the same input.

Key point

The difference between a single causal attention module and one head of a multi-head module is not the mechanism itself. Each head is the same kind of attention computation. What differs is that a multi-head module runs several of these computations, each with its own weights, on the same input.

Tracing a Sequence Through Two Heads

The two-head case shows the full structure. The embedded input tokens enter the module and stay as they are at the start. From there, the same input is handled by two separate heads. Each head applies its own query, key, and value weight matrices, and each head produces its own set of context vectors. These are labeled Z1 and Z2. The two sets are then combined into a single context vector matrix Z.

What happens to the input sequence after it enters the module, and how do the two head outputs get merged back into one result?

1
Step 1: The input token embeddings enter the module unchanged.
2
Step 2: Head 1 processes the sequence with its own query, key, and value weight matrices and produces context vectors Z1.
3
Step 3: Head 2 processes the same sequence with a separate set of query, key, and value weight matrices and produces context vectors Z2.
4
Step 4: The results Z1 and Z2 are combined into a single context vector matrix Z, which is the output of the module.

Which Weights Belong to Which Head

Compared with the single-head module, the two-head module changes the weight matrices in the middle. Where a single-head module had one value weight matrix, a two-head module has two, one per head. The same change applies to the query and key weight matrices. Each head therefore has its own query, key, and value matrices, and nothing in the source says that these are shared between heads.

Which weight matrices belong to each head, and what does each head produce?

What do you think happens?

How many value weight matrices does a two-head module have?

One
Two
Four

Reveal answer

Answer:

Two, one per head.

A single-head module has one value weight matrix. A two-head module has two, one for each head, and the same doubling applies to the query and key weight matrices.

Building Up by Stacking Modules

The expansion from causal attention to multi-head attention is approached in two stages. The first stage builds the idea intuitively by stacking multiple causal attention modules, each acting as a single head, and combining their results. 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, and the multi-head module is 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.

Single-Head and Multi-Head Side by Side

The two designs keep the same overall flow: tokens go in, weights transform them, and context vectors come out. The differences lie in how many weight sets sit in the middle, how many outputs are produced before merging, and what the computation costs.

How many weight sets exist, what each design produces, and how much computation it costs.

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 (Z1, Z2, ...), combined into one matrix Z
Processing	The same weights process every token	Each head processes the input independently with its own weights
Cost	One attention computation	Computationally intensive; each added head adds weights and context vectors
Pattern capacity	One attention mechanism doing all of the work	Different heads can attend to the input in different ways

Single-head and multi-head attention compared on weights, outputs, and cost

Why the Extra Cost Pays Off

Multi-head attention is computationally intensive, and that cost is the price of the approach. The reason it is worth paying is 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 attention mechanism must do all of its work with one set of weights, while several heads let the module hold several separate views of the same input at once.

Worked Example: Counting the Pieces in a Two-Head Module

Worked example

Following one input through a two-head module

A module receives an input sequence of embedded tokens. It is configured with two heads. Describe which weights are used, what each head produces, and what leaves the module.

1
Identify the input
The embedded input tokens enter the module and remain as they are at this stage.
2
Identify the weights
Head 1 has its own query, key, and value weight matrices. Head 2 has a separate set of its own. There are two sets in total, rather than the single set a single-head module would use.
3
Identify the head outputs
Head 1 produces context vectors Z1. Head 2 produces context vectors Z2. Each head works on the same input but with its own weights, so Z1 and Z2 can differ.
4
Combine the outputs
Z1 and Z2 are combined into one context vector matrix Z, which is the single result the module returns.

Result

The module takes one input sequence, runs two independent heads with separate weights, and returns one combined context vector matrix Z.

Mistakes Beginners Make With Heads

Common mistake
Thinking a head is a different kind of attention mechanism
Why it fails
Each head is an instance of the same self-attention mechanism. What makes it a separate head is that it has its own weights.
Do this instead
Treat a head as one more copy of the same computation, with its own weight matrices.
Common mistake
Assuming the heads share one set of weights
What it looks like
Believing that Head 1 and Head 2 both use the same query matrix because the input is the same.
Why it fails
The instances do not share weights. That independence is what makes each head a distinct view of the input.
Do this instead
Give each head its own query, key, and value weight matrices.
Common mistake
Expecting the output to be one result per head
Why it fails
The head outputs are combined into a single result. The module does not return Z1 and Z2 separately as its final output.
Do this instead
Remember that Z1 and Z2 are merged into one context vector matrix Z.
Common mistake
Believing that more heads are always better without accounting for cost

Practice: Single-Head or Multi-Head?

EASY

A colleague says that their module has one set of query, key, and value weights, sends every token through those same weights, and returns one set of context vectors. Is this module single-head or multi-head, and what would change if they added a second head?

Hints

Count the weight sets first.
Recall that each head has its own query, key, and value weight matrices.
Think about what happens to the outputs of the two heads before the module returns.

Your attempt (not saved)

MEDIUM

Explain in your own words why multi-head attention costs more than a single-head module, and why a model might still choose it for handling complex patterns.