Elamur Research⌕
☰
home / papers / 2023 / Linear attention is (maybe) all you need (to understand transformer optimization)

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

Paper · 2023 · Kwangjun Ahn, Xiang Cheng, Minhak Song et al.

Summary

Training Transformers is notoriously hard. Plain SGD lags well behind adaptive optimizers such as Adam, and the loss landscape shows unusual properties that classical optimization theory does not account for. Full Transformers are too complex to analyze directly, so the authors look for the simplest model that still behaves the same way during training. They study shallow linear Transformers, which have linear attention and no nonlinear activations or feedforward blocks. These models are trained on random linear regression tasks, following the in-context learning setup of von Oswald et al. (ICML 2023) and Ahn et al. (NeurIPS 2023). The toy model reproduces the main training phenomena previously reported for real Transformers on PTB, WikiText-2 and SQuAD:

  • Adam outperforms SGD.
  • Stochastic gradient noise is heavy-tailed.
  • The robust condition number is larger under SGD than under Adam.
  • SGD and Adam differ in directional smoothness.
  • Gradient norm correlates with smoothness.

Because the model is so simple, the authors can vary the data distribution and depth in controlled ways. This makes linear attention a practical testbed for theory and for designing new optimizers.

Key points

  • Setup. The base regression task (Setting 1) uses N = 20 in-context examples in d = 5 dimensions, with Gaussian covariates and Gaussian w⋆. Settings 2 and 3 vary this configuration.
  • Optimizers. Both momentum SGD (momentum 0.9) and Adam (β1 = β2 = 0.9) use global gradient clipping at 1. Learning rates are tuned separately for each setting, ranging from about 0.005 to 5.
  • Conditioning. The ratio of SGD's robust condition number to Adam's ranges from about 1.6 to 9.6 for the linear models. This mirrors the gap Jiang et al. (2022) reported for full Transformers.
  • Controlled variations. The distinctive training features become stronger when covariates are more heavy-tailed (compared with spherical ones) and when depth increases (tested at L = 2, 4, 6, 8 layers).
  • Nonlinear extension. An appendix applies the method to targets distorted by a random ReLU MLP, with an extra MLP layer added before the linear Transformer blocks. The resulting curves look much like those for the linear case.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Massachusetts Institute of TechnologyCambridge · United States274 works · 82,013 citations🎓Korea Advanced Institute of Science and TechnologyDaejeon · South Korea59 works · 3,320 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksArts and Humanities · 17,959 worksLanguage and Linguistics · 17,789 works

Cite

Cite as Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie and Suvrit Sra. 2023. Linear attention is (maybe) all you need (to understand transformer optimization). arXiv:2310.01082.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2310.01082 · arXiv 2310.01082

@misc{ahn2023linear,
  title = {{Linear attention is (maybe) all you need (to understand transformer optimization)}},
  author = {Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie and Suvrit Sra},
  year = {2023},
  doi = {10.48550/arxiv.2310.01082},
  eprint = {2310.01082},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2310.01082}
}

References

29 references, 11 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Jacob Abernethy, Alekh Agarwal, Teodor V Marinov, and Manfred K Warmuth. A mechanism for sample-efficient in-context learning for sparse retrieval tasks. arXiv preprint arXiv:2305.17040, 2023. 9 — [arXiv:2305.17040](https://arxiv.org/abs/2305.17040) find
  2. Kwangjun Ahn, Sébastien Bubeck, Sinho Chewi, Yin Tat Lee, Felipe Suarez, and Yi Zhang. Learning threshold neurons via the “edge of stability”. NeurIPS 2023 (arXiv:2212.07469), 2023a. 9 — [arXiv:2212.07469](https://arxiv.org/abs/2212.07469) find
  3. Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023 (arXiv:2306.00297), 2023b. 5, 6, 7
  4. Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. International Conference on Learning Representations, 2022. 5, 7
  5. Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 1, context-free grammar. arXiv preprint arXiv:2305.13673, 2023. 9
  6. Sébastien Bubeck. Convex optimization: Algorithms and complexity. Foundations and Trends® in Machine Learning, 8(3-4):231–357, 2015. 3 find
  7. Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv:2303.12712, 2023. 1
  8. Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, and Zhenxun Zhuang. Robustness to unbounded smoothness of generalized signsgd. Advances in Neural Information Processing Systems, 35:9955–9968, 2022. 5 find
  9. Yan Dai, Kwangjun Ahn, and Suvrit Sra. The crucial role of normalization in sharpness-aware minimization. NeurIPS 2023 (arXiv:2305.15287), 2023. 9 — [arXiv:2305.15287](https://arxiv.org/abs/2305.15287) find
  10. J Devlin, MW Chang, K Lee, and K Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding in: Proceedings of the 2019 conference of the north american chapter of the association for computational linguistics, 4171–4186.. acl. ACL. DOI: https://doi. org/10.18653/v1,…
  11. Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583–30598, 2022. 5, 7
  12. Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023. 6 — [arXiv:2310.05249](https://arxiv.org/abs/2310.05249) find
  13. Kaiqi Jiang, Dhruv Malik, and Yuanzhi Li. How does adaptive optimization impact local neural network geometry? arXiv preprint arXiv:2211.02254, 2022. 1, 2, 3, 4, 7, 8 find
  14. Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020. 1
  15. Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be. In International Conference on Learning Representations (ICLR) (arXiv:2304.13960), 2023. 1, 2, 3, 4, 7, 8 — [arXiv:2304.13960](https://arxiv.org/abs/2304.13960) find
  16. Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06015, 2023a. 6 — [arXiv:2302.06015](https://arxiv.org/abs/2302.06015) find
  17. Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. International Conference on Machine Learning (ICML) (arXiv:2303.04245), 2023b. 6, 9 find
  18. Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. In 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, pages 5747–5763. Association for Computational Linguistics (ACL), 2020. 1
  19. Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:2307.03576, 2023. 5
  20. Yurii Nesterov et al. Lectures on convex optimization, volume 137. Springer, 2018. 3 find
  21. Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning. arXiv preprint arXiv:2306.03435, 2023. 6 — [arXiv:2306.03435](https://arxiv.org/abs/2306.03435) find
  22. Yan Pan and Yuanzhi Li. Toward understanding why adam converges faster than sgd for transformers. arXiv preprint arXiv:2306.00204, 2023. 1, 2, 3, 4, 5, 7 find
  23. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 2017. 1, 6
  24. Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151–35174. PMLR, 2023. 5, 6, 7
  25. Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. In Advances in Neural Information Processing Systems, pages 4148–4158, 2017. 2 find
  26. Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. In International Conference on Learning Representations (ICLR), 2020a. 1, 3, 5, 7, 13 find
  27. Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33:15383–15393, 2020b. 1, 2, 3, 4, 7, 8 find
  28. Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927, 2023. 5 — [arXiv:2306.09927](https://arxiv.org/abs/2306.09927) find
  29. Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: a synthetic reasoning task. arXiv preprint arXiv:2206.04301, 2022. 9 find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/ahn-2023-linear-attention-maybe-all-need/