Stacking Multiple Attention Layers

Up to this point, attention 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 that it covers several heads at once. This extension is known as multi-head attention.

The idea behind the extension is less exotic than its 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 explains what that extension means, how its pieces fit together, and why the approach is worth its added cost. It starts from the single-head case, then looks at how a multi-head module is organized.

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.

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

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. That is what makes each head an independent view of the same input rather than a copy of one view.

The difference from the single-head module is therefore about ownership of weights. A single causal attention module owns one set of query, key, and value weights. A head, in the multi-head sense, owns its own query, key, and value weight matrices and produces its own context vectors. The attention module is the container, and the heads are the parts that each own a separate set of weights.

The two-head case is the smallest version that shows the full structure. Take a three-token sequence as a running example, with tokens labeled 1, 2, and 3. Each token has already been turned into an embedded vector. In a two-head module those embedded vectors enter both heads unchanged.

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. So head 1 runs its own query, key, and value weights over the sequence and produces context vectors $Z_1$. Head 2 runs a separate set of weights over the same embedded input and produces context vectors $Z_2$. The two heads do not see each other's weights, so each one is an independent attempt at the same input.

Each head produces its own set of context vectors. In the two-head diagram these are labeled $Z_1$ and $Z_2$. The module does not return two separate results. Those two sets are combined into a single context vector matrix, written $Z$. The combining step is what turns several independent head outputs back into one result for the rest of the model.

The same pattern extends beyond two heads. Each additional head adds another set of weight matrices and another set of context vectors to be combined. The number of heads changes how many parallel streams feed the merge, but the overall shape of the process stays the same: enter, split into heads, process per head, merge.

The expansion from causal attention to multi-head attention is approached in two stages. The first stage builds the idea intuitively by stacking several causal attention modules, each acting as a single head, and then 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 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. Understanding the stacked version first makes that later step easier to follow.

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.

The reason is that giving the model several independent heads means different heads can, in principle, attend to the input in different ways. Because each head has its own weights, the heads are not forced to learn the same view. The combined output carries those differences forward into the rest of the model.

The clearest way to see the difference is to compare the two configurations on the three dimensions that matter here: weights, outputs, and cost. The single-head module has one set of weights and one output. The two-head module has two sets of weights, two head-level outputs, and one merged output.

Each head in this module is a causal attention module, so the causal constraint applies inside every head. A causal mask decides which positions a token may draw on. The general rule is that each position is limited to itself and earlier positions. Using the same three-token running example, the table of visibility below shows which positions each token may use inside a head.

Because the mask is applied within each head, both head 1 and head 2 in the two-head module obey the same visibility pattern. The heads differ in their weights, not in which earlier positions a token is allowed to draw on.