Research map / Architectures and efficiency / Attention and transformers
Memorization Capacity of Multi-Head Attention in Transformers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proves that a single multi-head attention layer with H heads and context size n<d can memorize Omega(Hmin(n,d_h)) examples under linear-independence assumptions (Kruskal rank of queries, full-rank contexts), and validates the assumptions on ViT and with synthetic experiments.
Problem
How many examples a multi-head attention layer can memorize as a function of number of heads and context length, given that the commonly used General Position assumption fails for transformer inputs.
Approach
Replace General Position with two relaxed linear-independence assumptions (query Kruskal rank at least n; each context matrix full rank n), and show each head can be made responsible for min(n,d_h)-1 examples by exploiting softmax saturation, so ranks add up linearly across heads.
Key results
- Memorization capacity of one attention layer lower-bounded by H(min(n,d_h)-1)+1 examples under the two assumptions, with ~2Hdd_h+Hdd_v+d*dout parameters.
- Synthetic experiments show memorization increases linearly with H (R^2 >= 0.98 on the linear fit) and monotonically with n, with no gain for d_h > n.
- Assumption tests on ImageNet show General Position fails in all tested models while the proposed assumptions hold for random attention, random ViT and trained ViT (Embedding layer violates only the query Kruskal rank assumption).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| increase the number of attention heads H | single-layer MHA · synthetic classification | memorization accuracy | Helped |
| increase context size n | single-layer MHA · synthetic classification | memorization accuracy | Helped |
| increase head dimension d_h while d_h ≤ n | single-layer MHA · synthetic classification | memorization accuracy | Helped |
| increase head dimension d_h beyond n | single-layer MHA · synthetic classification | memorization fraction | No clear effect |
| apply one randomly initialized Self-Attention layer | ViT embedding-layer representations · ImageNet | query-vector Kruskal-rank assumption validity | Helped |
| use the embedding layer without an Attention layer | ViT embedding-layer representations · ImageNet | query-vector Kruskal-rank assumption validity | Hurt |
| increase the number of attention heads H | single-layer MHA · under the paper's linear-independence assumptions | memorization capacity | Helped |
| increase context size n | single-layer MHA · under the paper's linear-independence assumptions | memorization capacity | Helped |
| increase head dimension d_h beyond n | single-layer MHA · under the paper's linear-independence assumptions | memorization capacity | No clear effect |
Surprising findings
- Memorization gains vanish for d_h > n, so larger key/query dimensions do not increase capacity beyond n, though they speed optimization.
- The embedding layer violates the query Kruskal-rank assumption, which only holds after one attention layer with skip connection.
Limitations
- Results cover only a single attention layer and a single prediction per input, not sequence-to-sequence outputs.
- Upper bounds on MHA memorization are not established; the lower bound is only order-wise.
- Experiments rely on gradient descent, which may not find memorizing solutions that exist.
Experimental setup
Datasets: ImageNet synthetic uniform random data
Benchmarks: synthetic memorization tasks (classification with 100 classes, regression) ImageNet assumption tests on ViT variants
Metrics: average accuracy MSE rank test pass rate R^2 of linear fit
Baselines: two-layer ReLU network (upper bound comparison)
Builds on: Vision Transformer Kruskal rank General Position memorization results for FCNs
Closest papers by meaning
- On the Optimal Memorization Capacity of Transformers (ICLR 2025 · similarity 0.89)
- Understanding Factual Recall in Transformers via Associative Memories (ICLR 2025 · similarity 0.89)
- What Can Transformer Learn with Varying Depth? Case Studies on Sequence Learning Tasks (ICML 2024 · similarity 0.89)
- Are Transformers with One Layer Self-Attention Using Low-Rank Weight Matrices Universal Approximators? (ICLR 2024 · similarity 0.89)
- Theory, Analysis, and Best Practices for Sigmoid Self-Attention (ICLR 2025 · similarity 0.89)
- Transformers Learn Faster with Semantic Focus (NeurIPS 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.