mlsanityData to 7 Oct 2026

On the Benefits of Rank in Attention Layers

arXiv 2407.16153 ↗First posted 23 Jul 2024cs.LG stat.MLcited by 4

Not on the map. Why?

Abstract (arXiv)

Attention-based mechanisms are widely used in machine learning, most prominently in transformers. However, hyperparameters such as the rank of the attention matrices and the number of heads are scaled nearly the same way in all realizations of this architecture, without theoretical justification. In this work we show that there are dramatic trade-offs between the rank and number of heads of the attention mechanism. Specifically, we present a simple and natural target function that can be represented using a single full-rank attention head for any context length, but that cannot be approximated by low-rank attention unless the number of heads is exponential in the embedding dimension, even for short context lengths. Moreover, we prove that, for short context lengths, adding depth allows the target to be approximated by low-rank attention. For long contexts, we conjecture that full-rank attention is necessary. Finally, we present experiments with off-the-shelf transformers that validate our theoretical findings.

Key facts. Cited by 4 papers in the map (data through 7 Oct 2026). Builds on 42 of its 69 references that are in the map. First posted on arXiv 23 Jul 2024.

Builds on

42 of 69 references.

References · 42

  1. 2017Attention Is All You Needcited by 59,932
  2. 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,204
  3. 2021Learning Transferable Visual Models From Natural Language Supervisioncited by 33,456
  4. 2020Language Models are Few-Shot Learnerscited by 32,199
  5. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
  6. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,160
  7. 2021LoRA: Low-Rank Adaptation of Large Language Modelscited by 14,084
  8. 2023LLaMA: Open and Efficient Foundation Language Modelscited by 13,958
  9. 2023Llama 2: Open Foundation and Fine-Tuned Chat Modelscited by 12,865
  10. 2022PaLM: Scaling Language Modeling with Pathwayscited by 5,304
  11. 2014Neural Machine Translation by Jointly Learning to Align and Translatecited by 4,530
  12. 2019Root Mean Square Layer Normalizationcited by 1,425
30 more
  1. 2021Scaling Language Models: Methods, Analysis & Insights from Training Gophercited by 1,199
  2. 2021Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolationcited by 1,085
  3. 2024OLMo: Accelerating the Science of Language Modelscited by 599
  4. 2019Are Transformers universal approximators of sequence-to-sequence functions?cited by 404
  5. 2023Trained Transformers Learn Linear Models In-Contextcited by 331
  6. 2014Breaking the Curse of Dimensionality with Convex Neural Networkscited by 329
  7. 2023Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selectioncited by 291
  8. 2022Transformers Learn Shortcuts to Automatacited by 260
  9. 2016Benefits of depth in neural networkscited by 248
  10. 2024Repeat After Me: Transformers are Better than State Space Models at Copyingcited by 208
  11. 2023Birth of a Transformer: A Memory Viewpointcited by 154
  12. 2021Inductive Biases and Variable Creation in Self-Attention Mechanismscited by 144
  13. 2015On the Equivalence between Kernel Quadrature Rules and Random Feature Expansionscited by 136
  14. 2023Representational Strengths and Limitations of Transformerscited by 132
  15. 2023The Truth is in There: Improving Reasoning in Language Models with Layer-Selective Rank Reductioncited by 120
  16. 2021Statistically Meaningful Approximation: a Case Study on Approximating Turing Machines with Transformerscited by 101
  17. 2023Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformercited by 101
  18. 2020Low-Rank Bottleneck in Multi-head Attention Modelscited by 88
  19. 2024Transformers, parallel computation, and logarithmic depthcited by 75
  20. 2024Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimalitycited by 71
  21. 2016Depth-Width Tradeoffs in Approximating Natural Functions with Neural Networkscited by 58
  22. 2023The Expressive Power of Transformers with Chain of Thoughtcited by 47
  23. 2023Are Transformers with One Layer Self-Attention Using Low-Rank Weight Matrices Universal Approximators?cited by 46
  24. 2023Memorization Capacity of Multi-Head Attention in Transformerscited by 39
  25. 2023What can a Single Attention Layer Learn? A Study Through the Random Features Lenscited by 39
  26. 2023On the Optimization and Generalization of Multi-head Attentioncited by 37
  27. 2024Separations in the Representational Capabilities of Transformers and Recurrent Architecturescited by 29
  28. 2020Approximate is Good Enough: Probabilistic Variants of Dimensional and Margin Complexitycited by 25
  29. 2023Six Lectures on Linearized Neural Networkscited by 19
  30. 2019Depth-Width Trade-offs for ReLU Networks via Sharkovsky's Theoremcited by 15

Built on by

Cited by year

Citing papers per year: 2025: 4.

4 papers.

  1. 2025Compositional Reasoning with Transformers, RNNs, and Chain of Thoughtcited by 6
  2. 2025Achilles' Heel of Mamba: Essential difficulties of the Mamba architecture demonstrated by synthetic datacited by 3
  3. 2025When Do Transformers Outperform Feedforward and Recurrent Networks? A Statistical Perspectivecited by 2
  4. 2025Customizing the Inductive Biases of Softmax Attention using Structured Matricescited by 0