Multi-head attention, one shape at a time (Part 9)

Part 9 of LLMs, the whole thing. In Part 8, I worked through causal masking and dropout. The next section of Sebastian Raschka’s Build a Large Language Model (From Scratch) adds multiple attention heads.

I understood the idea before I understood the code. Two heads meant two sets of query, key and value parameters. Then I reached view(...).transpose(...) and had to stop. What had been computed, and what was merely being rearranged?

The explanation that clicked was to follow one token through those operations. This time I use the shorter sentence The river bank was muddy. The shifted training inputs are The / river / bank / was; their targets are river / bank / was / muddy. That gives us four input positions, with bank at index 2.

Same input, different parameters

A head computes attention weights and uses them to mix value vectors. Multiple heads can compute different mixtures for the same token because each head has its own query, key and value parameters. We do not assign them jobs such as grammar or meaning.

My tiny example uses four input features, two heads and three features per head. Each head reads all four input features. I pack the two heads’ query matrices into one larger linear layer; keys and values get separate layers of the same size.

mha_W_query = nn.Linear(4, 6, bias=False)
mha_W_key   = nn.Linear(4, 6, bias=False)
mha_W_value = nn.Linear(4, 6, bias=False)

packed_queries = mha_W_query(x_batch)
packed_keys    = mha_W_key(x_batch)
packed_values  = mha_W_value(x_batch)

Each weight matrix has shape (6, 4). Rows 0–2 belong to head 1; rows 3–5 belong to head 2. Calling a layer computes x_batch @ layer.weight.T. The six outputs per token already contain both heads’ results.

view: the queries are already there

packed_queries has shape (1, 4, 6): one sentence, four tokens, six query numbers per token. Now I give the heads an axis:

queries_by_token = packed_queries.view(1, 4, 2, 3)
# General form: (batch_size, num_tokens, NUM_HEADS, HEAD_DIM)

Read that as one sentence, four tokens, two heads per token, three query numbers per head. For bank, one six-number row becomes two groups of three. No query is recalculated. No weight changes.

The total stays at 24 elements: 1 × 4 × 6 = 1 × 4 × 2 × 3. This view shares the original storage; its shape must be compatible with that storage layout. Keys and values receive the same transformation. These are computed tensor values, distinct from the trainable weights stored inside the linear layers.

transpose: put all tokens inside each head

head_queries = queries_by_token.transpose(1, 2)
# (B, T, H, D) -> (B, H, T, D)
# (1, 4, 2, 3) -> (1, 2, 4, 3)

Axis 1 held tokens; axis 2 held heads. Swapping them changes how we index the same query vectors. Before, each token contained both heads. Afterwards, each head contains every token.

queries_by_token[0, 2, 0]  # sentence 0, bank, head 1
head_queries[0, 0, 2]      # exactly the same vector

For these dense tensors, transpose also shares storage. Changing a shape directly to (1, 2, 4, 3) with view would not perform the same axis swap: matching dimensions alone does not preserve which vector belongs to which head.

Each head gets its own attention matrix

We arrange heads first because batched matrix multiplication treats the final two axes as matrices. Each batch/head pair independently compares every query with every key:

head_raw_scores = head_queries @ head_keys.transpose(-2, -1)
# (1, 2, 4, 3) @ (1, 2, 3, 4) -> (1, 2, 4, 4)

scaled = head_raw_scores / HEAD_DIM**0.5
masked = scaled.masked_fill(causal_mask, -torch.inf)
head_attention_weights = torch.softmax(masked, dim=-1)

Scale by the square root of three, the key width of one head. The same (4, 4) causal mask broadcasts across the batch and head axes. Both heads let bank use The, river and itself, while blocking was.

Collect values, concatenate, then mix

After attention dropout, each head collects its own value vectors. With p=0.5, surviving attention weights double; individual rows need not sum to one afterwards.

head_contexts = mha_dropout(head_attention_weights) @ head_values
# (1, 2, 4, 3)

joined_contexts = head_contexts.transpose(1, 2).contiguous().view(1, 4, 6)
out_proj = nn.Linear(6, 6, bias=False)
mha_context_vectors = out_proj(joined_contexts)

The transpose brings each token’s heads together. contiguous() makes the layout suitable for merging the head and feature axes with view. Concatenation preserves both three-number results as six features. The trainable output projection can then mix features from both heads, separately at each token position. It creates no path to a future token.

The final shape is (1, 4, 6). This teaching script deliberately keeps a six-feature output from a four-feature input; a full GPT block needs the attention output width to match the stream it adds it to through the residual connection.

The checks that made the shapes trustworthy

The walkthrough checks packed attention against heads calculated separately from their weight rows, verifies normalization before dropout, and checks that future weights stay zero. In a separate check, changing the final input left all earlier outputs exactly unchanged. A batch of two distinct inputs also passed the head-equivalence checks.

All numerical figures come from the full script with seed 123 on PyTorch 2.13.0 CPU. Nothing is trained. The outputs are context features, not vocabulary probabilities. I can now read Step 15 literally: the projections compute the values; view exposes the heads; transpose organizes them for the next multiplication.

The complete runnable gptversion.py prints every stage. Install torch and tiktoken, then run python gptversion.py. Steps 14–20 cover this post. The chapter reference is Raschka’s official Chapter 3 notebook, section 3.6.