Learning In-context n-grams with Transformers: Sub-n-grams Are Near-Stationary Points
Aditya Varre, Gizem Yüce, Nicolas Flammarion
Abstract
Motivated by empirical observations of prolonged plateaus and stage-wise progression during training, we investigate the loss landscape of transformer models trained on in-context next-token prediction tasks. In particular, we focus on learning in-context n-gram language models under cross-entropy loss, and establish a sufficient condition for parameter configurations to be stationary points. We then construct a set of parameter configurations for a simplified transformer model that represent k-gram estimators (for k ⩽ n), and show that the gradient of the population loss at these solutions vanishes in the limit of infinite sequence length and parameter norm. This reveals a key property of the loss landscape: sub-ngrams are near-stationary points of the population cross-entropy loss, offering theoretical insight into widely observed phenomena such as stagewise learning dynamics and emergent phase transitions. These insights are further supported by numerical experiments that illustrate the learning dynamics of n-grams, characterized by discrete transitions between near-stationary solutions.
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.
Cited by top-tier papers6
- Saddle-to-Saddle Dynamics Explains A Simplicity Bias Across Neural Network ArchitecturesYedi Zhang, Andrew M. Saxe, Peter E. LathamICLR 2026 · 15 citations
- What Happens During the Loss Plateau? Understanding Abrupt Learning in TransformersPulkit Gopalani, Wei HuNeurIPS 2025 · 6 citations
- Transformers Learn Latent Mixture Models In-Context via Mirror DescentFrancesco D'Angelo, Nicolas FlammarionICLR 2026 · 2 citations
- Length Generalization Bounds for TransformersAndy Yang, Pascal Bergsträßer, Georg Zetzsche, David Chiang et al.ICML 2026
- Induction Heads Interpolate N-GramsFrancesco D'Angelo, Oğuz Yüksel, Swathi Narashiman, Nicolas FlammarionICML 2026
Builds on21
- Language Models are Few-Shot LearnersTom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah et al.NeurIPS 2020 · 64,255 citations
- What Can Transformers Learn In-Context? A Case Study of Simple Function ClassesShivam Garg, Dimitris Tsipras, Percy Liang, Gregory ValiantNeurIPS 2022 · 883 citations
- Transformers Learn In-Context by Gradient DescentJohannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento et al.ICML 2023 · 729 citations
- Transformers learn to implement preconditioned gradient descent for in-context learningKwangjun Ahn, Xiang Cheng, Hadi Daneshmand, Suvrit SraNeurIPS 2023 · 324 citations
- Birth of a Transformer: A Memory ViewpointAlberto Bietti, Vivien Cabannes, Diane Bouchacourt, Hervé Jégou et al.NeurIPS 2023 · 182 citations
Related papers
- The Evolution of Statistical Induction Heads: In-Context Learning Markov ChainsEzra Edelman, Nikolaos Tsilivis, Benjamin L. Edelman, Eran Malach et al.NeurIPS 2024 · 140 citations
- Optimality and NP-Hardness of Transformers in Learning Markovian Dynamical FunctionsYanna Ding, Songtao Lu, Yingdong Lu, Tomasz Nowicki et al.NeurIPS 2025 · 1 citation
- Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention LandscapeJuno Kim, Taiji SuzukiICML 2024 · 42 citations
- Attention with Markov: A Curious Case of Single-layer TransformersAshok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle et al.ICLR 2025
- In-context Convergence of TransformersYu Huang, Yuan Cheng, Yingbin LiangICML 2024 · 114 citations
