Faster Query-Key Learning Sharpens Attention in Self-Attention Models
Not on the map. Why?
Abstract (arXiv)
A standard self-attention layer consists of two interacting circuits: the query-key circuit that governs attention allocation, and the output-value circuit that maps attended representations to predictions. Collapsed and factorized parameterizations of the query-key and output-value circuits lead to qualitatively different attention patterns. In particular, some parameterizations give sharper attention to task-relevant tokens, at a similar training loss. We analyze how the parameterizations of these circuits shape the parameter trajectories in single-layer self-attention models trained for next-token prediction. Through gradient-flow analysis, we show that factorization induces implicit rescaling of the two circuits' learning rates. We derive closed-form dynamics showing that output-value and query-key parameters move along a line, with relative speeds determined by their learning rates. Faster query-key learning relative to output-value learning thus produces sharper attention, as the model compensates for slower output-value learning by increasing attention mass on relevant tokens. Experiments show that differences in the relative learning rates of the two circuits govern attention concentration. This improves attention interpretability proxies while maintaining comparable predictive performance.
Key facts. Cited by 1 paper in the map (data through 14 Oct 2026). Builds on 27 of its 31 references that are in the map. First posted on arXiv 7 Aug 2026.
Builds on
27 of 31 references.
References · 27
- 2017Attention Is All You Needcited by 60,361
- 2020Language Models are Few-Shot Learnerscited by 32,419
- 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 30,186
- 2016SQuAD: 100,000+ Questions for Machine Comprehension of Textcited by 4,148
- 2022In-context Learning and Induction Headscited by 964
- 2019Attention is not Explanationcited by 929
- 2020Quantifying Attention Flow in Transformerscited by 775
- 2018Universal Transformerscited by 647
- 2019Attention is not not Explanationcited by 636
- 2019ERASER: A Benchmark to Evaluate Rationalized NLP Modelscited by 530
- 2020HateXplain: A Benchmark Dataset for Explainable Hate Speech Detectioncited by 407
- 2019Is Attention Interpretable?cited by 389
15 more
- 2016Assessing the Ability of LSTMs to Learn Syntax-Sensitive Dependenciescited by 354
- 2018On the Optimization of Deep Networks: Implicit Acceleration by Overparameterizationcited by 284
- 2018Algorithmic Regularization in Learning Deep Homogeneous Models: Layers are Automatically Balancedcited by 177
- 2023Birth of a Transformer: A Memory Viewpointcited by 161
- 2023The emergence of clusters in self-attention dynamicscited by 140
- 2023Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformercited by 102
- 2023How Do Transformers Learn Topic Structure: Towards a Mechanistic Understandingcited by 86
- 2020On the Computational Power of Transformers and its Implications in Sequence Modelingcited by 85
- 2023On the Optimization and Generalization of Multi-head Attentioncited by 37
- 2024Mechanics of Next Token Prediction with Self-Attentioncited by 34
- 2024Clustering in Causal Attention Maskingcited by 32
- 2024Self-attention Networks Localize When QK-eigenspectrum Concentratescited by 19
- 2020On the Dynamics of Training Attention Modelscited by 8
- 2022On the Interpretability of Attention Networkscited by 5
- 2023On the Learning Dynamics of Attention Networkscited by 2
Built on by
Citing papers per year: 2026: 1.
1 paper.