Attention: Which Position Is Comparing With Which?

Fri, 28 Aug 2026 18: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 assertion anyone would think to write passes.

Now the same split written the other way:

Masks on Programmer.ie

Attention: Which Position Is Comparing With Which?