Multi-Head Attention

When you write a self-attention module from scratch, you often end up with more than one version of it. An early version might store its trainable projection weights in one layout, while a later, cleaner version organizes the same mathematical operation differently. This article follows one practical question through the whole topic: if you have a trained or carefully set-up instance of one version, how do you move its weights into an instance of the other version so that both produce the same outputs?

The goal is not to make the two classes look alike. The goal is to make them behave alike. Once the weights are assigned correctly, feeding the same input into each instance should give the same output, and that shared output is the check that tells you the transfer worked. The rest of this article uses that check to build up to multi-head attention, the causal masking that a GPT-style model depends on, and the transformer blocks assembled into the full architecture.

The transfer runs from an instance of the second version, called SelfAttention_v2, into an instance of the first version, called SelfAttention_v1. Each class holds trainable weight matrices for its query, key, and value projections. Copying weights is therefore a matter of identifying which matrix in one class corresponds to which matrix in the other, and then assigning values pairwise: the query matrix to the query matrix, the key matrix to the key matrix, and the value matrix to the value matrix.

Pairing the matrices is the easy part. The difficulty is that a matrix you can see in one class may not be stored in the same orientation in the other, and that is where the transfer usually fails. The next section explains the storage layout that causes this.

The step that most often causes trouble is the storage layout. nn.Linear stores its weight matrix in a transposed form. A matrix that you would write with rows for input features and columns for output features is held the other way around inside nn.Linear. If you assign a weight matrix directly without accounting for this, the shapes may still fit by accident in some cases. The values, however, end up arranged in the wrong orientation, and the two instances will disagree even though the code ran without error.

Successful transfer is confirmed by running the same input through both instances and checking that they give the same outputs. You do not need to inspect the classes to decide whether the copy worked. If the outputs match for a shared input, the weights are behaving the same way, and that is the only standard the transfer needs to meet.

The next enhancement adds a causal aspect to self-attention. Causal attention modifies the mechanism so that the model cannot access future information in the sequence. This matters for language modeling, where each prediction should depend only on the words that came before it. If a model could peek at upcoming words while it predicts the current one, it would learn a shortcut that does not exist at generation time.

In GPT-style models, each token is processed in turn. For the current token, the tokens that come after it in the input are masked out. Concretely, the attention weights above the diagonal of the attention matrix are masked out, and the remaining weights in each row are normalized so that they sum to 1. The row for a given token therefore only distributes its attention over itself and the tokens before it.

One way to obtain the masked attention weight matrix is to work in three stages. First, apply the softmax function to the attention scores. Second, zero out the elements above the diagonal so that future positions contribute nothing. Third, normalize the resulting matrix row by row so that each row again sums to 1. The result is a matrix in which every row describes a valid distribution over the tokens that are allowed to be seen.

The diagram above shows the row for the first word. Its only visible position is itself, so after normalization it receives the full weight of 1.0. Rows for later words see more positions. The second-row word, for example, keeps the weight for the word before it and for itself, and nothing beyond that.

The ordering in the three-stage method matters. Zeroing the upper triangle after the softmax leaves rows that no longer sum to 1, which is why the normalization stage follows the masking stage. The masked matrix is unnormalized right after the zeroing and becomes normalized only at the end.

The third enhancement is the multi-head component. Here the attention mechanism is split into several heads. Each head learns different aspects of the data, and the heads together let the model attend to information from different representation subspaces at different positions at the same time. The source presents this as a way to improve performance on complex tasks.

Causal masking and multi-head splitting are presented together because a practical GPT-like model uses both. The mask controls which positions a token may see, while the heads let it look at those positions in several different ways. Each head in the diagram has its own weights and its own attention computation, and the heads' results are brought back together into one output.

The attention modules in transformer-based LLMs involve multiple instances of causal attention, and this arrangement is called multi-head attention. A single causal attention module processes its input once. Multi-head attention runs several such instances side by side, so the model's attention is organized as a group of parallel causal attention units rather than one unit.

The implementation idea is equally direct. A multi-head attention module is created by stacking multiple instances of causal attention modules. You do not need a new kind of attention to get the multi-head version. You reuse the causal attention module you already built and stack several copies of it into one larger module, with each copy acting as one head.

Multi-head attention is one of the core components of an LLM, but it is not the whole model. The next part of the work is to code the other building blocks of an LLM and assemble them into a GPT-like model. Those assembled units are transformer blocks. The source names transformer blocks as the building blocks used in LLM architecture, and implementing them is necessary for creating GPT models of various sizes.

The phrase of various sizes is worth taking seriously. A transformer block is a reusable unit, so the same kind of block can be repeated or configured to produce models of different scales. The learning objective here is to understand how to implement transformer blocks, and the source treats that implementation as a required step rather than an optional refinement.

The LLM architecture is introduced top-down as several building blocks before the individual components are examined. The roadmap that places this chapter inside a larger process has three stages. Read it as a sequence rather than a checklist of unrelated tasks.

The chapter ends at the architecture. The source states that the next chapter adds the training loop and pretrains the LLM, so the model assembled here is not yet trained. Its purpose at this point is to exist as a complete, correctly structured GPT-like model that can later be trained to generate human-like text.