📌 Multi-Head Attention Shape Transformations (Cheat Sheet)

[talk2jaydip]

Step	Before Shape	Operation	After Shape	Major Change
1️⃣ Input X	(batch_size, num_tokens, d_in)	-	(batch_size, num_tokens, d_in)	No change
2️⃣ Q, K, V projections	(batch_size, num_tokens, d_in)	Linear (W_q, W_k, W_v)	(batch_size, num_tokens, d_out)	d_in → d_out
3️⃣ Split into heads	(batch_size, num_tokens, d_out)	Reshape into heads	(batch_size, num_tokens, num_heads, head_dim)	d_out → num_heads × head_dim
4️⃣ Transpose for attention	(batch_size, num_tokens, num_heads, head_dim)	Swap axes (1,2)	(batch_size, num_heads, num_tokens, head_dim)	num_tokens ↔ num_heads
5️⃣ Compute Attention Scores	(batch_size, num_heads, num_tokens, head_dim) and (batch_size, num_heads, head_dim, num_tokens)	Q @ Kᵀ	(batch_size, num_heads, num_tokens, num_tokens)	head_dim → num_tokens
6️⃣ Apply Mask & Softmax	(batch_size, num_heads, num_tokens, num_tokens)	Mask + Softmax	(batch_size, num_heads, num_tokens, num_tokens)	No shape change
7️⃣ Multiply with Values (V)	(batch_size, num_heads, num_tokens, num_tokens) and (batch_size, num_heads, num_tokens, head_dim)	Attention @ V	(batch_size, num_heads, num_tokens, head_dim)	num_tokens ↔ head_dim
8️⃣ Merge Heads	(batch_size, num_heads, num_tokens, head_dim)	Reshape back	(batch_size, num_tokens, d_out)	num_heads × head_dim → d_out
9️⃣ Final Linear Layer	(batch_size, num_tokens, d_out)	Linear layer	(batch_size, num_tokens, d_out)	No change

[rasbt]

This is a really nice and organized figure/cheatsheet! Thanks a lot for sharing!

[talk2jaydip]

Thanks @rasbt

[yashgadbail]

Helpful!