# Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer

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

## Key facts

- Cited by 102 papers in the map (data through 14 Oct 2026).
- Builds on 72 of its 88 references that are in the map.
- First posted on arXiv 25 May 2023.

## Builds on (most cited first, 30 of 72)

- [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)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [wav2vec 2.0: A Framework for Self-Supervised Learning of Speech Representations](https://mlsanity.com/paper/2006.11477/) (2020)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Scaling Instruction-Finetuned Language Models](https://mlsanity.com/paper/2210.11416/) (2022)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Learning to summarize from human feedback](https://mlsanity.com/paper/2009.01325/) (2020)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [data2vec: A General Framework for Self-supervised Learning in Speech, Vision and Language](https://mlsanity.com/paper/2202.03555/) (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)
- [On Lazy Training in Differentiable Programming](https://mlsanity.com/paper/1812.07956/) (2018)
- [Gradient Descent Provably Optimizes Over-parameterized Neural Networks](https://mlsanity.com/paper/1810.02054/) (2018)

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

- [KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache](https://mlsanity.com/paper/2402.02750/) (2024)
- [Birth of a Transformer: A Memory Viewpoint](https://mlsanity.com/paper/2306.00802/) (2023)
- [How Transformers Learn Causal Structure with Gradient Descent](https://mlsanity.com/paper/2402.14735/) (2024)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [A Primer on the Inner Workings of Transformer-based Language Models](https://mlsanity.com/paper/2405.00208/) (2024)
- [Transformers as Support Vector Machines](https://mlsanity.com/paper/2308.16898/) (2023)
- [Attention Satisfies: A Constraint-Satisfaction Lens on Factual Errors of Language Models](https://mlsanity.com/paper/2309.15098/) (2023)
- [How Large Language Models Encode Context Knowledge? A Layer-Wise Probing Study](https://mlsanity.com/paper/2402.16061/) (2024)
- [RNNs are not Transformers (Yet): The Key Bottleneck on In-context Retrieval](https://mlsanity.com/paper/2402.18510/) (2024)
- [Max-Margin Token Selection in Attention Mechanism](https://mlsanity.com/paper/2306.13596/) (2023)
- [From Words to Numbers: Your Large Language Model Is Secretly A Capable Regressor When Given In-Context Examples](https://mlsanity.com/paper/2404.07544/) (2024)
- [Demystifying Verbatim Memorization in Large Language Models](https://mlsanity.com/paper/2407.17817/) (2024)
- [Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMs](https://mlsanity.com/paper/2410.13835/) (2024)
- [JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attention](https://mlsanity.com/paper/2310.00535/) (2023)
- [On the Role of Attention Masks and LayerNorm in Transformers](https://mlsanity.com/paper/2405.18781/) (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)
- [An Information-Theoretic Analysis of In-Context Learning](https://mlsanity.com/paper/2401.15530/) (2024)
- [Attention with Markov: A Framework for Principled Analysis of Transformers via Markov Chains](https://mlsanity.com/paper/2402.04161/) (2024)
- [On the Optimization and Generalization of Multi-head Attention](https://mlsanity.com/paper/2310.12680/) (2023)
- [How Do Nonlinear Transformers Learn and Generalize in In-Context Learning?](https://mlsanity.com/paper/2402.15607/) (2024)
- [From Self-Attention to Markov Models: Unveiling the Dynamics of Generative Transformers](https://mlsanity.com/paper/2402.13512/) (2024)
- [Mechanics of Next Token Prediction with Self-Attention](https://mlsanity.com/paper/2403.08081/) (2024)
- [Backward Lens: Projecting Language Model Gradients into the Vocabulary Space](https://mlsanity.com/paper/2402.12865/) (2024)
- [Understanding Finetuning for Factual Knowledge Extraction](https://mlsanity.com/paper/2406.14785/) (2024)
- [Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought](https://mlsanity.com/paper/2502.21212/) (2025)
- [Implicit Bias and Fast Convergence Rates for Self-attention](https://mlsanity.com/paper/2402.05738/) (2024)
- [A phase transition between positional and semantic learning in a solvable model of dot-product attention](https://mlsanity.com/paper/2402.03902/) (2024)
- [Linear Transformers are Versatile In-Context Learners](https://mlsanity.com/paper/2402.14180/) (2024)
- [Towards a Theoretical Understanding of the 'Reversal Curse' via Training Dynamics](https://mlsanity.com/paper/2405.04669/) (2024)

## Abstract (arXiv)

Transformer architecture has shown impressive performance in multiple research domains and has become the backbone of many neural network models. However, there is limited understanding on how it works. In particular, with a simple predictive loss, how the representation emerges from the gradient training dynamics remains a mystery. In this paper, for 1-layer transformer with one self-attention layer plus one decoder layer, we analyze its SGD training dynamics for the task of next token prediction in a mathematically rigorous manner. We open the black box of the dynamic process of how the self-attention layer combines input tokens, and reveal the nature of underlying inductive bias. More specifically, with the assumption (a) no positional encoding, (b) long input sequence, and (c) the decoder layer learns faster than the self-attention layer, we prove that self-attention acts as a discriminative scanning algorithm: starting from uniform attention, it gradually attends more to distinct key tokens for a specific next token to be predicted, and pays less attention to common key tokens that occur across different next tokens. Among distinct tokens, it progressively drops attention weights, following the order of low to high co-occurrence between the key and the query token in the training set. Interestingly, this procedure does not lead to winner-takes-all, but decelerates due to a phase transition that is controllable by the learning rates of the two layers, leaving (almost) fixed token combination. We verify this textbfscan and snap dynamics on synthetic and real-world data (WikiText).