Why do you transpose dimensions 1 and 2 when computing batched multi-head attention?

Sreedath Panat explains that batched queries arrive with shape batch, tokens, heads, head dimension, for example 2, 6, 2, 3. Naively transposing the last two axes of the keys and multiplying would give 2 by 2 score matrices per token, not the 6 by 6 token-by-token matrix attention needs. Swapping axes 1 and 2 first gives 2, 2, 6, 3, so queries times transposed keys gives 2, 2, 6, 6: one 6 by 6 score matrix per head per batch item. After the values are applied the same transpose is reversed, and the heads are merged back to 2, 6, 6.

• Scaling by the square root of d, causal masking and softmax leave the tensor shape unchanged.
• Panat accidentally ran the reverse transpose cell twice in Colab, which flipped the shape back, and uses it to show why the step matters.

Source: our own transcription of the episode. The full conversation is on The attention family (2026-10-03). Read the whole episode summary →

▸More from this episode