# Implicit Optimization Bias of Next-Token Prediction in Linear Models

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

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Builds on 72 of its 103 references that are in the map.
- First posted on arXiv 28 Feb 2024.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Rethinking the Inception Architecture for Computer Vision](https://mlsanity.com/paper/1512.00567/) (2015)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Efficient Estimation of Word Representations in Vector Space](https://mlsanity.com/paper/1301.3781/) (2013)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [Understanding deep learning requires rethinking generalization](https://mlsanity.com/paper/1611.03530/) (2016)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (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)
- [Benign Overfitting in Linear Regression](https://mlsanity.com/paper/1906.11300/) (2019)
- [Surprises in High-Dimensional Ridgeless Least Squares Interpolation](https://mlsanity.com/paper/1903.08560/) (2019)
- [Hungry Hungry Hippos: Towards Language Modeling with State Space Models](https://mlsanity.com/paper/2212.14052/) (2022)
- [In Search of the Real Inductive Bias: On the Role of Implicit Regularization in Deep Learning](https://mlsanity.com/paper/1412.6614/) (2014)
- [Gradient Descent Maximizes the Margin of Homogeneous Neural Networks](https://mlsanity.com/paper/1906.05890/) (2019)
- [An Investigation of Why Overparameterization Exacerbates Spurious Correlations](https://mlsanity.com/paper/2005.04345/) (2020)
- [Characterizing Implicit Bias in Terms of Optimization Geometry](https://mlsanity.com/paper/1802.08246/) (2018)
- [Implicit Bias of Gradient Descent on Linear Convolutional Networks](https://mlsanity.com/paper/1806.00468/) (2018)
- [Human uncertainty makes classification more robust](https://mlsanity.com/paper/1908.07086/) (2019)
- [Directional convergence and alignment in deep learning](https://mlsanity.com/paper/2006.06657/) (2020)
- [A mathematical perspective on Transformers](https://mlsanity.com/paper/2312.10794/) (2023)
- [Towards Understanding Sharpness-Aware Minimization](https://mlsanity.com/paper/2206.06232/) (2022)
- [Inductive Biases and Variable Creation in Self-Attention Mechanisms](https://mlsanity.com/paper/2110.10090/) (2021)
- [Fast and Faster Convergence of SGD for Over-Parameterized Models and an Accelerated Perceptron](https://mlsanity.com/paper/1810.07288/) (2018)
- [Neural collapse with unconstrained features](https://mlsanity.com/paper/2011.11619/) (2020)
- [Label Noise SGD Provably Prefers Flat Global Minimizers](https://mlsanity.com/paper/2106.06530/) (2021)
- [The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains](https://mlsanity.com/paper/2402.11004/) (2024)
- [Convergence of Gradient Descent on Separable Data](https://mlsanity.com/paper/1803.01905/) (2018)

## Built on by (most cited first, 14 of 14)

- [Implicit Bias of Spectral Descent and Muon on Multiclass Separable Data](https://mlsanity.com/paper/2502.04664/) (2025)
- [Implicit Geometry of Next-token Prediction: From Language Sparsity Patterns to Model Representations](https://mlsanity.com/paper/2408.15417/) (2024)
- [Trained Transformer Classifiers Generalize and Exhibit Benign Overfitting In-Context](https://mlsanity.com/paper/2410.01774/) (2024)
- [Unifying Low Dimensional Spectra in Deep Learning](https://mlsanity.com/paper/2404.06106/) (2024)
- [Neural Collapse is Globally Optimal in Deep Regularized ResNets and Transformers](https://mlsanity.com/paper/2505.15239/) (2025)
- [Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis](https://mlsanity.com/paper/2410.09605/) (2024)
- [Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selection](https://mlsanity.com/paper/2410.14581/) (2024)
- [Towards Understanding Steering Strength](https://mlsanity.com/paper/2602.02712/) (2026)
- [Towards Understanding Transformers in Learning Random Walks](https://mlsanity.com/paper/2511.23239/) (2025)
- [Structure Before Collapse: Transient semantic geometry in next-token prediction](https://mlsanity.com/paper/2606.26749/) (2026)
- [Ensemble Debiasing Across Class and Sample Levels for Fairer Prompting Accuracy](https://mlsanity.com/paper/2503.05157/) (2025)
- [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)
- [Self-Improving is Often Sudden: Enlightenment-style Finetuning for Large-Scale Models](https://mlsanity.com/paper/2607.13395/) (2026)

## Abstract (arXiv)

We initiate an investigation into the optimization properties of next-token prediction (NTP), the dominant training paradigm for modern language models. Specifically, we study the structural properties of the solutions selected by gradient-based optimizers among the many possible minimizers of the NTP objective. By framing NTP as cross-entropy minimization across distinct contexts, each tied with a sparse conditional probability distribution across a finite vocabulary of tokens, we introduce "NTP-separability conditions" that enable reaching the data-entropy lower bound. With this setup, and focusing on linear models with fixed context embeddings, we characterize the optimization bias of gradient descent (GD): Within the data subspace defined by the sparsity patterns of distinct contexts, GD selects parameters that equate the logits' differences of in-support tokens to their log-odds. In the orthogonal subspace, the GD parameters diverge in norm and select the direction that maximizes a margin specific to NTP. These findings extend previous research on implicit bias in one-hot classification to the NTP setting, highlighting key differences and prompting further research into the optimization and generalization properties of NTP, irrespective of the specific architecture used to generate the context embeddings.