Single-Head Attention and Stacked Attention Layers

Multi-head attention is often introduced as a theory of several attention heads working side by side. This article takes a different route. It starts with the code of the MultiHeadAttention class, which is a PyTorch module, and follows the tensors through its forward method, where every computation happens. The goal is to understand not only what the method computes, but how the shape of the data changes at each line.

Read the forward method once, from top to bottom. As you read, notice which lines create new tensors and which lines only rearrange tensors that already exist. Most of the method is the familiar attention computation. What makes it multi-head is the reorganization that happens around that computation, and that reorganization is almost entirely a matter of shapes.

Three things in the method matter most. First, the input is projected by three linear layers into keys, queries, and values. Second, three view calls and three transpose calls reorganize those projections so that each head has its own slice. Third, the final view and the out_proj layer turn the per-head results back into a single output.

Before following the tensors, make a prediction. The input x has shape (b, num_tokens, d_in), and the constructor was given d_out as the projection size. Decide what shape the method returns.

The answer is not the same as the input shape. Keep that difference in mind. The method changes the width from d_in to d_out, and it also changes the number of axes temporarily when it splits the features into heads. Tracing those intermediate shapes is the whole job of this article.

The class does not create a separate weight matrix for each head. It creates one projection per role: W_query, W_key, and W_value. Each of these maps d_in to d_out. Each projection therefore produces one wide matrix that contains the information for all heads side by side.

The head structure is not built into the weights. It is imposed afterward, by reshaping that wide matrix. This is the central design decision: a single large projection per role, split internally, rather than many small independent projections.

The constructor enforces this relationship. It asserts that d_out is divisible by num_heads, and then it sets head_dim to d_out divided by num_heads. Written as a formula, head_dim is $\frac{d_{\text{out}}}{\text{num\_heads}}$. The head count must divide the output width evenly, because an uneven split would leave some head without a well-defined slice.

After the projection, each tensor is still one wide matrix per token. The head structure is created by two operations applied in sequence: a view call that regroups the last dimension into heads, and a transpose call that moves the head dimension ahead of the token dimension. The two operations are easiest to follow as a before-and-after pair.

To make the split concrete, imagine a tiny case with d_out equal to 6 and num_heads equal to 2. Then head_dim is 3. The diagram below labels each of the six columns of the projected vector with the head it belongs to. This example is generated for teaching purposes; the same arithmetic applies to any d_out that is divisible by num_heads.

After the view, the first three columns form head 0 and the last three form head 1. The transpose then reorders the axes so that each head's three-column slice is processed as its own matrix, separate from the other head. Nothing in the weights was split; the split exists only in how the same wide vector is indexed.

The next diagram shows the full path of data through the forward method. Each stage corresponds to one statement in the method, so it is useful for deciding which line to inspect first. Shapes are given for a batch of b sequences, each with num_tokens tokens.

Read the diagram as a sequence of one view, then one transpose, for each of the three projections. The score stage is the only one where the token axis appears twice, because each token's query is compared with every token's key. The stage labeled attention scores is an inference from the matrix product queries @ keys.transpose(2, 3), which multiplies the query matrix for each head by the transposed key matrix for that head.

After the attention weights multiply the values, each head's context has shape (b, num_heads, num_tokens, head_dim). The method first moves the head dimension back with transpose(1, 2), which returns the tensor to (b, num_tokens, num_heads, head_dim). It then calls contiguous() and view(b, num_tokens, self.d_out).

Because d_out equals num_heads times head_dim, the final view merges the last two axes back into one feature axis. The per-head results sit side by side again as a single width of d_out.

The source describes the final linear layer as optional. Here out_proj is a Linear layer that takes the merged output and produces the returned context_vec. Its role is to combine the information from all heads. The shape of the returned tensor is (b, num_tokens, d_out), which matches the prediction you made earlier only after every intermediate reshape has been undone.

The earlier MultiHeadAttentionWrapper built multiple heads by stacking complete single-head attention layers, each with its own weight matrices. The MultiHeadAttention class takes an integrated approach instead: it starts with one multi-head projection per role and splits it internally. The diagram below shows the query side of this difference.

The source states that the same integrated idea applies to keys and values, which it leaves out of the figure to reduce clutter. The result is that the class reaches the same conceptual outcome as the wrapper, with fewer separate weight matrices and a single large projection per role.

When a shape error appears, the fastest route is to print the shape after each step and compare it with the expected shape from the pipeline above. The first step whose printed shape differs from the expected one is the step where the bug begins. Everything after that point inherits the wrong shape, so the error message at the end is usually a symptom rather than the cause.