Attention: Which Position Is Comparing With Which?

Sat, 29 Aug 2026 10:00:00 +0100

Here is a tensor of sequence representations and the line that splits it into heads.

B, T, E, Nh = 1, 4, 8, 2
Dh = E // Nh # 4

x = torch.arange(B * T * E).reshape(B, T, E)
heads = x.reshape(B, Nh, T, Dh)

x: (1, 4, 8)
heads: (1, 2, 4, 4)

That is exactly the shape multi-head attention wants: batch, heads, positions, head dimension. Nothing raised. Every shape assertion passes.

Now the same split written the other way:

PyTorch Attention Shapes: Q, K, V, Multi-Head Attention Masks and Transformer Dimension Errors

Sat, 08 Aug 2026 13:30:00 +0100

PyTorch: Zero to Hero — Step 07

Attention code is where tensor-shape mistakes stop being annoying and start becoming architectural.

A CNN usually makes its dimensional assumptions fairly obvious. Attention does not.

A tensor that starts as:

(batch, sequence, embedding)

is projected into Q, K and V, split into heads, transposed, multiplied, masked, normalized, multiplied again, transposed again, concatenated and projected back to the embedding dimension.

A single bad view, transpose, mask shape or head calculation can produce anything from an immediate runtime error to a model that trains while attending to the wrong tokens.

Multihead Attention on AIBussin — AI applications, systems and books

Attention: Which Position Is Comparing With Which?

PyTorch Attention Shapes: Q, K, V, Multi-Head Attention Masks and Transformer Dimension Errors

PyTorch: Zero to Hero — Step 07