Implementing Multi-Head Attention

A useful way to understand a module is to configure it to match a real model. The smallest GPT-2 model fixes three attention settings: 12 heads, 768-dimensional input and output embeddings, and a context length of 1,024 tokens. The task is to supply those settings to MultiHeadAttention and then follow the tensors through its forward pass.

The module arguments connect the model settings to the dimensions used inside the attention computation. Set d_in to 768 so the module accepts GPT-2’s input embedding width. Set d_out to 768 so its projected representations and returned output use that width. Set num_heads to 12. Set context_length to 1,024 to match the model’s supported token context.

The class requires d_out to divide evenly by num_heads. With these settings, 768 divided by 12 gives a head_dim of 64. This is the width assigned to each head. Check this divisibility requirement rather than assuming that every possible output width can be split evenly.

The forward pass uses one projection for each role: query, key, and value. Each projection maps from d_in to d_out, so each projected tensor has shape (b, num_tokens, d_out). The method then reshapes and transposes those tensors to make the head dimension explicit. The same shape changes apply to queries, keys, and values.

After the query, key, and value projections, each tensor is reshaped with view so its d_out axis becomes num_heads and head_dim. The next transpose(1, 2) moves the head axis ahead of the token axis. Queries and keys can then be used to form scores per head: queries are multiplied by the transposed keys. With query and key tensors shaped (b, num_heads, num_tokens, head_dim), the score tensor has shape (b, num_heads, num_tokens, num_tokens).

After attention weights are applied to the values, each head’s context has shape (b, num_heads, num_tokens, head_dim). The method transposes the head and token axes back, calls contiguous(), and uses view to merge num_heads and head_dim into d_out. The output projection then produces the returned context tensor with shape (b, num_tokens, d_out).

The projected tensor initially keeps all head features together in one d_out axis. The view operation reorganizes that axis into two axes: num_heads and head_dim. For the smallest GPT-2 configuration, the 768-wide projection is represented as 12 heads with 64 features per head. The following transpose(1, 2) changes the axis order from (b, num_tokens, num_heads, head_dim) to (b, num_heads, num_tokens, head_dim). It rearranges the axes; it does not change the head count or per-head width.

MultiHeadAttentionWrapper builds multiple heads by stacking complete single-head attention layers. Each of those layers has its own weight matrices. The integrated MultiHeadAttention design instead uses one projection per role—query, key, and value—and splits each projected width into heads internally. In other words, the head structure is created by reshaping a shared wide projection, rather than by giving every head a separate single-head layer.

Both designs organize computation into multiple heads, but their parameters live in different arrangements. In the wrapper, each complete single-head layer owns its matrices. In the integrated design, each role has one wide projection that is divided into heads afterward. The same integrated approach applies to keys and values as well as queries.

When the final output has an unexpected shape, inspect the forward pass in order rather than guessing at the last line. Check the input, then the projected tensors, the result of view, the result of transpose, the attention scores, the head contexts, the merged tensor, and finally out_proj. The first point where the observed shape differs from the expected shape is where to focus your investigation.