# Approximating How Single Head Attention Learns

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

## Key facts

- Cited by 38 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 27 references that are in the map.
- First posted on arXiv 13 Mar 2021.

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

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [Effective Approaches to Attention-based Neural Machine Translation](https://mlsanity.com/paper/1508.04025/) (2015)
- [Attention is not Explanation](https://mlsanity.com/paper/1902.10186/) (2019)
- [Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned](https://mlsanity.com/paper/1905.09418/) (2019)
- [Are Sixteen Heads Really Better than One?](https://mlsanity.com/paper/1905.10650/) (2019)
- [Designing and Interpreting Probes with Control Tasks](https://mlsanity.com/paper/1909.03368/) (2019)
- [Attention is not not Explanation](https://mlsanity.com/paper/1908.04626/) (2019)
- [Is Attention Interpretable?](https://mlsanity.com/paper/1906.03731/) (2019)
- [Multi30K: Multilingual English-German Image Descriptions](https://mlsanity.com/paper/1605.00459/) (2016)
- [A Tail-Index Analysis of Stochastic Gradient Noise in Deep Neural Networks](https://mlsanity.com/paper/1901.06053/) (2019)
- [Algorithmic Regularization in Learning Deep Homogeneous Models: Layers are Automatically Balanced](https://mlsanity.com/paper/1806.00900/) (2018)
- [On Identifiability in Transformers](https://mlsanity.com/paper/1908.04211/) (2019)
- [Sharp Nearby, Fuzzy Far Away: How Neural Language Models Use Context](https://mlsanity.com/paper/1805.04623/) (2018)
- [Learning to Deceive with Attention-Based Explanations](https://mlsanity.com/paper/1909.07913/) (2019)
- [Attention Interpretability Across NLP Tasks](https://mlsanity.com/paper/1909.11218/) (2019)
- [Seq2Seq-Vis: A Visual Debugging Tool for Sequence-to-Sequence Models](https://mlsanity.com/paper/1804.09299/) (2018)
- [Fine-grained Sentiment Analysis with Faithful Attention](https://mlsanity.com/paper/1908.06870/) (2019)
- [Analogs of Linguistic Structure in Deep Representations](https://mlsanity.com/paper/1707.08139/) (2017)

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

- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Birth of a Transformer: A Memory Viewpoint](https://mlsanity.com/paper/2306.00802/) (2023)
- [Inductive Biases and Variable Creation in Self-Attention Mechanisms](https://mlsanity.com/paper/2110.10090/) (2021)
- [Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer](https://mlsanity.com/paper/2305.16380/) (2023)
- [Vision Transformers provably learn spatial structure](https://mlsanity.com/paper/2210.09221/) (2022)
- [How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding](https://mlsanity.com/paper/2303.04245/) (2023)
- [A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexity](https://mlsanity.com/paper/2302.06015/) (2023)
- [Max-Margin Token Selection in Attention Mechanism](https://mlsanity.com/paper/2306.13596/) (2023)
- [JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attention](https://mlsanity.com/paper/2310.00535/) (2023)
- [Understanding Factual Recall in Transformers via Associative Memories](https://mlsanity.com/paper/2412.06538/) (2024)
- [Attention Scheme Inspired Softmax Regression](https://mlsanity.com/paper/2304.10411/) (2023)
- [An Over-parameterized Exponential Regression](https://mlsanity.com/paper/2303.16504/) (2023)
- [Solving Regularized Exp, Cosh and Sinh Regression Problems](https://mlsanity.com/paper/2303.15725/) (2023)
- [The Closeness of In-Context Learning and Weight Shifting for Softmax Regression](https://mlsanity.com/paper/2304.13276/) (2023)
- [What can a Single Attention Layer Learn? A Study Through the Random Features Lens](https://mlsanity.com/paper/2307.11353/) (2023)
- [Attention with Markov: A Framework for Principled Analysis of Transformers via Markov Chains](https://mlsanity.com/paper/2402.04161/) (2024)
- [An Iterative Algorithm for Rescaled Hyperbolic Functions Regression](https://mlsanity.com/paper/2305.00660/) (2023)
- [In-Context Learning for Attention Scheme: from Single Softmax Regression to Multiple Softmax Regression via a Tensor Trick](https://mlsanity.com/paper/2307.02419/) (2023)
- [A Fast Optimization View: Reformulating Single Layer Attention in LLM Based on Tensor and SVM Trick, and Solving It in Matrix Multiplication Time](https://mlsanity.com/paper/2309.07418/) (2023)
- [A Mathematical Abstraction for Balancing the Trade-off Between Creativity and Reality in Large Language Models](https://mlsanity.com/paper/2306.02295/) (2023)
- [Towards a Theoretical Understanding of the 'Reversal Curse' via Training Dynamics](https://mlsanity.com/paper/2405.04669/) (2024)
- [Fast Contextual Adaptation with Neural Associative Memory for On-Device Personalized Speech Recognition](https://mlsanity.com/paper/2110.02220/) (2021)
- [Multi-Layer Transformers Gradient Can be Approximated in Almost Linear Time](https://mlsanity.com/paper/2408.13233/) (2024)
- [Zero-th Order Algorithm for Softmax Attention Optimization](https://mlsanity.com/paper/2307.08352/) (2023)
- [Why Softmax Attention Outperforms Linear Attention](https://mlsanity.com/paper/2310.11685/) (2023)
- [On the Convergence of Encoder-only Shallow Transformers](https://mlsanity.com/paper/2311.01575/) (2023)
- [Unmasking Transformers: A Theoretical Approach to Data Recovery via Attention Weights](https://mlsanity.com/paper/2310.12462/) (2023)
- [Distributional Associations vs In-Context Reasoning: A Study of Feed-forward and Attention Layers](https://mlsanity.com/paper/2406.03068/) (2024)
- [From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics](https://mlsanity.com/paper/2510.06954/) (2025)

## Abstract (arXiv)

Why do models often attend to salient words, and how does this evolve throughout training? We approximate model training as a two stage process: early on in training when the attention weights are uniform, the model learns to translate individual input word `i` to `o` if they co-occur frequently. Later, the model learns to attend to `i` while the correct output is o because it knows `i` translates to `o`. To formalize, we define a model property, Knowledge to Translate Individual Words (KTIW) (e.g. knowing that `i` translates to `o`), and claim that it drives the learning of the attention. This claim is supported by the fact that before the attention mechanism is learned, KTIW can be learned from word co-occurrence statistics, but not the other way around. Particularly, we can construct a training distribution that makes KTIW hard to learn, the learning of the attention fails, and the model cannot even learn the simple task of copying the input words to the output. Our approximation explains why models sometimes attend to salient words, and inspires a toy example where a multi-head attention model can overcome the above hard training distribution by improving learning dynamics rather than expressiveness. We end by discussing the limitation of our approximation framework and suggest future directions.