On the Dynamics of Training Attention Models
Haoye Lu, Yongyi Mao, Amiya Nayak
Abstract
The attention mechanism has been widely used in deep neural networks as a model component. By now, it has become a critical building block in many state-of-the-art natural language models. Despite its great success established empirically, the working mechanism of attention has not been investigated at a sufficient theoretical depth to date. In this paper, we set up a simple text classification task and study the dynamics of training a simple attention-based classification model using gradient descent. In this setting, we show that, for the discriminative words that the model should attend to, a persisting identity exists relating its embedding and the inner product of its key and the query. This allows us to prove that training must converge to attending to the discriminative words when the attention output is classified by a linear classifier. Experiments are performed, which validates our theoretical analysis and provides further insights.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 07d47550-bac5-44f4-ae62-192d8c756959Cited by top-tier papers8
- How Transformers Learn Causal Structure with Gradient DescentEshaan Nichani, Alex Damian, Jason D. LeeICML 2024 · 117 citations
- Transformers are uninterpretable with myopic methods: a case study with bounded Dyck grammarsKaiyue Wen, Yuchen Li, Bingbin Liu, Andrej RisteskiNeurIPS 2023 · 32 citations
- From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training DynamicsZheng-An Chen, Tao LuoNeurIPS 2025 · 13 citations
- Promises and Pitfalls of Generative Masked Language Modeling: Theoretical Framework and Practical GuidelinesYuchen Li, Alexandre Kirchmeyer, Aashay Mehta, Yilong Qin et al.ICML 2024 · 5 citations
- Strong Correlations Induce Cause Only Predictions in Transformer TrainingHaihan Zhang, Yimu Zhang, Cong FangICLR 2026
Builds on2
Related papers
- Attention with Trained Embeddings Provably Selects Important TokensDiyuan Wu, Aleksandr Shevchenko, Samet Oymak, Marco MondelliNeurIPS 2025
- Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow AnalysisHongru Yang, Bhavya Kailkhura, Zhangyang Wang, Yingbin LiangNeurIPS 2024 · 14 citations
- Focus and Dilution: The Multi-stage Learning Process of AttentionZheng-An Chen, Pengxiao Lin, Zhi-Qin John Xu, Tao LuoICML 2026
- Not All Attention Is Needed: Gated Attention Network for Sequence DataLanqing Xue, Xiaopeng Li, Nevin L. ZhangAAAI 2020 · 47 citations
- Max-Margin Token Selection in Attention MechanismDavoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, Samet OymakNeurIPS 2023 · 67 citations
