Efficient Implementation of Multi-Head Attention

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 focuses on what the extension means and how its pieces fit together. 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.

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 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 only how many weight sets sit in the middle.

In practical terms, a head is one independent attention mechanism inside a multi-head module. It is a full self-attention instance with its own weights. A head does not borrow weights from another head. Each head therefore reads the same input sequence but does so through its own query, key, and value weight matrices, and it produces its own context vectors.

The contrast with a single causal attention module is the key idea. A single module has one value weight matrix, one query weight matrix, and one key weight matrix. A head is one of several such complete sets. Calling each set a head emphasizes that it is one view of the input, and that several views can be held at once.

The extension is approached in two stages. The first stage builds the idea intuitively by stacking multiple CausalAttention 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 comes first because it makes the structure visible: each stacked module is a head you already understand. The later efficient version computes the same kind of result with a different organization.

Consider the two-head case, the smallest version that shows the full structure. The embedded input tokens stay exactly as they are at the start. 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 then runs its own attention computation on the input and produces its own set of context vectors, labeled $Z_1$ and $Z_2$. Those two sets are combined into a single context vector matrix $Z$.

Read the diagram as a sequence of stages. First, the input enters the module once, and it reaches both heads. Second, Head 1 and Head 2 work on that input independently, each through its own weights. Third, each head emits its own context vectors, $Z_1$ from the first head and $Z_2$ from the second. Finally, the two outputs are combined into one matrix $Z$. The merge is the only point where the two paths meet, so the two heads are fully separate until that step.

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 flow stays the same: one input, several independent paths, one merge.

The clearest way to see the difference is to compare the two arrangements on three dimensions: weights, outputs, and cost. The diagram places the single-head and two-head cases in parallel columns.

On weights, the single-head module has one set. The two-head module has two sets of query, key, and value weight matrices, and the instances do not share them. On outputs, the single-head module produces one set of context vectors. The two-head module produces two sets, $Z_1$ and $Z_2$, which are then combined into one matrix $Z$. On cost, the source describes multi-head attention as computationally intensive. Each added head brings another set of weight matrices and another set of context vectors, so the work grows with the number of heads.

Multi-head attention is worth its 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, and the combined output carries those differences forward.