Multi-head attention is a sum of low-rank maps

ProofAdvanced⏱ 30 minDeep Learningmulti-headrankloracircuits

Theory: The Mathematics of Transformers →

Problem

Show that multi-head attention equals $\sum_m A^mXW_V^mW_O^m$ with each $W_V^mW_O^m$ of rank at most $d/h$ , and that head $m$ 's logits depend on $W_Q^m(W_K^m)^\top$ only. Discuss the consequences for interpretability and for LoRA.

Write a complete argument. Every inequality you use should be one you could state precisely.

Solution

Opening this before you have a real attempt on paper is the fastest way to feel like you understand something you cannot reproduce.

Related problems

All problems

Problem	Field	Type	Difficulty	Time
Adapter parameter counts	Fine-Tuning	Numerical	Foundation	15 min
LoRA parameter accounting	LLMs	Numerical	Beginner	15 min
When does attention dominate the MLP?	Deep Learning	Numerical	Beginner	30 min
Will this fine-tune fit?	Fine-Tuning	Numerical	Beginner	30 min
Attention is permutation equivariant	Deep Learning	Proof	Intermediate	30 min