# Linear attention is (maybe) all you need (to understand transformer optimization)

arXiv: https://arxiv.org/abs/2310.01082 · HTML: https://mlsanity.com/paper/2310.01082/

## Key facts

- Cited by 98 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 31 references that are in the map.
- First posted on arXiv 2 Oct 2023.

## Builds on (most cited first, 24 of 24)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Sparks of Artificial General Intelligence: Early experiments with GPT-4](https://mlsanity.com/paper/2303.12712/) (2023)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Why gradient clipping accelerates training: A theoretical justification for adaptivity](https://mlsanity.com/paper/1905.11881/) (2019)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [The Marginal Value of Adaptive Gradient Methods in Machine Learning](https://mlsanity.com/paper/1705.08292/) (2017)
- [Understanding the Difficulty of Training Transformers](https://mlsanity.com/paper/2004.08249/) (2020)
- [One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention](https://mlsanity.com/paper/2307.03576/) (2023)
- [Noise Is Not the Main Factor Behind the Gap Between SGD and Adam on Transformers, but Sign Descent Might Be](https://mlsanity.com/paper/2304.13960/) (2023)
- [Unveiling Transformers with LEGO: a synthetic reasoning task](https://mlsanity.com/paper/2206.04301/) (2022)
- [How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding](https://mlsanity.com/paper/2303.04245/) (2023)
- [Robustness to Unbounded Smoothness of Generalized SignSGD](https://mlsanity.com/paper/2208.11195/) (2022)
- [A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexity](https://mlsanity.com/paper/2302.06015/) (2023)
- [Toward Understanding Why Adam Converges Faster Than SGD for Transformers](https://mlsanity.com/paper/2306.00204/) (2023)
- [Learning threshold neurons via the "edge of stability"](https://mlsanity.com/paper/2212.07469/) (2022)
- [On the Role of Attention in Prompt-tuning](https://mlsanity.com/paper/2306.03435/) (2023)
- [The Crucial Role of Normalization in Sharpness-Aware Minimization](https://mlsanity.com/paper/2305.15287/) (2023)
- [How Does Adaptive Optimization Impact Local Neural Network Geometry?](https://mlsanity.com/paper/2211.02254/) (2022)
- [A Mechanism for Sample-Efficient In-Context Learning for Sparse Retrieval Tasks](https://mlsanity.com/paper/2305.17040/) (2023)

## Built on by (most cited first, 30 of 98)

- [Adam-mini: Use Fewer Learning Rates To Gain More](https://mlsanity.com/paper/2406.16793/) (2024)
- [Why Transformers Need Adam: A Hessian Perspective](https://mlsanity.com/paper/2402.16788/) (2024)
- [Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Models](https://mlsanity.com/paper/2402.19449/) (2024)
- [Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?](https://mlsanity.com/paper/2410.08292/) (2024)
- [Transformers Implement Functional Gradient Descent to Learn Non-Linear Functions In Context](https://mlsanity.com/paper/2312.06528/) (2023)
- [Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMs](https://mlsanity.com/paper/2410.13835/) (2024)
- [From Gradient Clipping to Normalization for Heavy Tailed SGD](https://mlsanity.com/paper/2410.13849/) (2024)
- [Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape](https://mlsanity.com/paper/2402.01258/) (2024)
- [Training Dynamics of In-Context Learning in Linear Attention](https://mlsanity.com/paper/2501.16265/) (2025)
- [Why Larger Language Models Do In-context Learning Differently?](https://mlsanity.com/paper/2405.19592/) (2024)
- [SAMformer: Unlocking the Potential of Transformers in Time Series Forecasting with Sharpness-Aware Minimization and Channel-Wise Attention](https://mlsanity.com/paper/2402.10198/) (2024)
- [A Theoretical Understanding of Self-Correction through In-context Alignment](https://mlsanity.com/paper/2405.18634/) (2024)
- [Deconstructing What Makes a Good Optimizer for Language Models](https://mlsanity.com/paper/2407.07972/) (2024)
- [Lions and Muons: Optimization via Stochastic Frank-Wolfe under Heavy-Tailed Noise](https://mlsanity.com/paper/2506.04192/) (2025)
- [Does SGD really happen in tiny subspaces?](https://mlsanity.com/paper/2405.16002/) (2024)
- [Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought](https://mlsanity.com/paper/2502.21212/) (2025)
- [Pretrained transformer efficiently learns low-dimensional target functions in-context](https://mlsanity.com/paper/2411.02544/) (2024)
- [Training Tensor Attention Efficiently: From Cubic to Almost Linear Time](https://mlsanity.com/paper/2405.16411/) (2024)
- [Understanding Adam Optimizer via Online Learning of Updates: Adam is FTRL in Disguise](https://mlsanity.com/paper/2402.01567/) (2024)
- [Adam with model exponential moving average is effective for nonconvex optimization](https://mlsanity.com/paper/2405.18199/) (2024)
- [Setting the Record Straight on Transformer Oversmoothing](https://mlsanity.com/paper/2401.04301/) (2024)
- [Superiority of Multi-Head Attention in In-Context Linear Regression](https://mlsanity.com/paper/2401.17426/) (2024)
- [Can SGD Handle Heavy-Tailed Noise?](https://mlsanity.com/paper/2508.04860/) (2025)
- [Is your batch size the problem? Revisiting the Adam-SGD gap in language modeling](https://mlsanity.com/paper/2506.12543/) (2025)
- [A Survey to Recent Progress Towards Understanding In-Context Learning](https://mlsanity.com/paper/2402.02212/) (2024)
- [Short-length Adversarial Training Helps LLMs Defend Long-length Jailbreak Attacks: Theoretical and Empirical Evidence](https://mlsanity.com/paper/2502.04204/) (2025)
- [Towards Understanding Inductive Bias in Transformers: A View From Infinity](https://mlsanity.com/paper/2402.05173/) (2024)
- [Explicitly Encoding Structural Symmetry is Key to Length Generalization in Arithmetic Tasks](https://mlsanity.com/paper/2406.01895/) (2024)
- [On the Role of Depth and Looping for In-Context Learning with Task Diversity](https://mlsanity.com/paper/2410.21698/) (2024)
- [What Does It Mean to Be a Transformer? Insights from a Theoretical Hessian Analysis](https://mlsanity.com/paper/2410.10986/) (2024)

## Abstract (arXiv)

Transformer training is notoriously difficult, requiring a careful design of optimizers and use of various heuristics. We make progress towards understanding the subtleties of training Transformers by carefully studying a simple yet canonical linearized shallow Transformer model. Specifically, we train linear Transformers to solve regression tasks, inspired by J.~von Oswald et al.~(ICML 2023), and K.~Ahn et al.~(NeurIPS 2023). Most importantly, we observe that our proposed linearized models can reproduce several prominent aspects of Transformer training dynamics. Consequently, the results obtained in this paper suggest that a simple linearized Transformer model could actually be a valuable, realistic abstraction for understanding Transformer optimization.