# Mechanics of Next Token Prediction with Self-Attention

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

## Key facts

- Cited by 34 papers in the map (data through 14 Oct 2026).
- Builds on 44 of its 64 references that are in the map.
- First posted on arXiv 12 Mar 2024.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (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)
- [The Implicit Bias of Gradient Descent on Separable Data](https://mlsanity.com/paper/1710.10345/) (2017)
- [Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth](https://mlsanity.com/paper/2103.03404/) (2021)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [Using the Output Embedding to Improve Language Models](https://mlsanity.com/paper/1608.05859/) (2016)
- [Characterizing Implicit Bias in Terms of Optimization Geometry](https://mlsanity.com/paper/1802.08246/) (2018)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [Towards Explaining the Regularization Effect of Initial Large Learning Rate in Training Neural Networks](https://mlsanity.com/paper/1907.04595/) (2019)
- [ProGen2: Exploring the Boundaries of Protein Language Models](https://mlsanity.com/paper/2206.13517/) (2022)
- [Inductive Biases and Variable Creation in Self-Attention Mechanisms](https://mlsanity.com/paper/2110.10090/) (2021)
- [Implicit regularization for deep neural networks driven by an Ornstein-Uhlenbeck like process](https://mlsanity.com/paper/1904.09080/) (2019)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [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)
- [Risk and parameter convergence of logistic regression](https://mlsanity.com/paper/1803.07300/) (2018)
- [O(n) Connections are Expressive Enough: Universal Approximability of Sparse Transformers](https://mlsanity.com/paper/2006.04862/) (2020)
- [Generative Pre-Training for Speech with Autoregressive Predictive Coding](https://mlsanity.com/paper/1910.12607/) (2019)
- [A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexity](https://mlsanity.com/paper/2302.06015/) (2023)
- [Label-Imbalanced and Group-Sensitive Classification under Overparameterization](https://mlsanity.com/paper/2103.01550/) (2021)
- [Transformers as Support Vector Machines](https://mlsanity.com/paper/2308.16898/) (2023)

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

- [TS-Reasoner: Domain-Oriented Time Series Inference Agents for Reasoning and Automated Analysis](https://mlsanity.com/paper/2410.04047/) (2024)
- [Do LLMs dream of elephants (when told not to)? Latent concept association and associative memory in transformers](https://mlsanity.com/paper/2406.18400/) (2024)
- [Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysis](https://mlsanity.com/paper/2410.02167/) (2024)
- [Non-asymptotic Convergence of Training Transformers for Next-token Prediction](https://mlsanity.com/paper/2409.17335/) (2024)
- [Gating is Weighting: Understanding Gated Linear Attention through In-context Learning](https://mlsanity.com/paper/2504.04308/) (2025)
- [Next-token prediction capacity: general upper bounds and a lower bound for transformers](https://mlsanity.com/paper/2405.13718/) (2024)
- [How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Bias](https://mlsanity.com/paper/2505.00926/) (2025)
- [Generalization or Hallucination? Understanding Out-of-Context Reasoning in Transformers](https://mlsanity.com/paper/2506.10887/) (2025)
- [Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selection](https://mlsanity.com/paper/2410.14581/) (2024)
- [Sequential keypoint density estimator: an overlooked baseline of skeleton-based video anomaly detection](https://mlsanity.com/paper/2506.18368/) (2025)
- [On the Power of Convolution Augmented Transformer](https://mlsanity.com/paper/2407.05591/) (2024)
- [On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recovery](https://mlsanity.com/paper/2410.13981/) (2024)
- [Decoupling Positional and Symbolic Attention Behavior in Transformers](https://mlsanity.com/paper/2511.11579/) (2025)
- [On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capability](https://mlsanity.com/paper/2405.16845/) (2024)
- [FloorPlan-DeepSeek (FPDS): A multimodal approach to floorplan generation using vector-based next room prediction](https://mlsanity.com/paper/2506.21562/) (2025)
- [Facts in Stats: Impacts of Pretraining Diversity on Language Model Generalization](https://mlsanity.com/paper/2510.16096/) (2025)
- [Dynamics of Spontaneous Topic Changes in Next Token Prediction with Self-Attention](https://mlsanity.com/paper/2501.06382/) (2025)
- [Towards Understanding Transformers in Learning Random Walks](https://mlsanity.com/paper/2511.23239/) (2025)
- [Unlabeled Data Can Provably Enhance In-Context Learning of Transformers](https://mlsanity.com/paper/2601.10058/) (2026)
- [Breaking the Reversal Curse in Autoregressive Language Models via Identity Bridge](https://mlsanity.com/paper/2602.02470/) (2026)
- [Transformers Provably Learn to Internalize Chain-of-Thought](https://mlsanity.com/paper/2605.28600/) (2026)
- [REAL: A Reasoning-Enhanced Graph Framework for Long-Term Memory Management of LLMs](https://mlsanity.com/paper/2606.10694/) (2026)
- [Recent Advances in Federated Learning Driven Large Language Models: A Survey on Architecture, Performance, and Security](https://mlsanity.com/paper/2406.09831/) (2024)
- [Reorganizing attention-space geometry with expressive attention](https://mlsanity.com/paper/2407.18601/) (2024)
- [On Next-Token Prediction in LLMs: How End Goals Determine the Consistency of Decoding Algorithms](https://mlsanity.com/paper/2505.11183/) (2025)
- [How Transformers Learn In-Context Recall Tasks? Optimality, Training Dynamics and Generalization](https://mlsanity.com/paper/2505.15009/) (2025)
- [Attention with Trained Embeddings Provably Selects Important Tokens](https://mlsanity.com/paper/2505.17282/) (2025)
- [Faster Query-Key Learning Sharpens Attention in Self-Attention Models](https://mlsanity.com/paper/2608.06776/) (2026)
- [Large Investment Model](https://mlsanity.com/paper/2408.10255/) (2024)
- [Fighter: Unveiling the Graph Convolutional Nature of Transformers in Time Series Modeling](https://mlsanity.com/paper/2510.17106/) (2025)

## Abstract (arXiv)

Transformer-based language models are trained on large datasets to predict the next token given an input sequence. Despite this simple training objective, they have led to revolutionary advances in natural language processing. Underlying this success is the self-attention mechanism. In this work, we ask: textitWhat textitdoes textita textitsingle textitself-attention textitlayer textitlearn textitfrom textitnext-token textitprediction? We show that training self-attention with gradient descent learns an automaton which generates the next token in two distinct steps: (1) Hard retrieval: Given input sequence, self-attention precisely selects the textithigh-priority textitinput textittokens associated with the last input token. (2) Soft composition: It then creates a convex combination of the high-priority tokens from which the next token can be sampled. Under suitable conditions, we rigorously characterize these mechanics through a directed graph over tokens extracted from the training data. We prove that gradient descent implicitly discovers the strongly-connected components (SCC) of this graph and self-attention learns to retrieve the tokens that belong to the highest-priority SCC available in the context window. Our theory relies on decomposing the model weights into a directional component and a finite component that correspond to hard retrieval and soft composition steps respectively. This also formalizes a related implicit bias formula conjectured in [Tarzanagh et al. 2023]. We hope that these findings shed light on how self-attention processes sequential data and pave the path toward demystifying more complex architectures.