# One-Layer Transformer Provably Learns One-Nearest Neighbor In Context

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

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 43 references that are in the map.
- First posted on arXiv 16 Nov 2024.

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

- [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)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?](https://mlsanity.com/paper/2202.12837/) (2022)
- [MaskGIT: Masked Generative Image Transformer](https://mlsanity.com/paper/2202.04200/) (2022)
- [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)
- [Spectrally-normalized margin bounds for neural networks](https://mlsanity.com/paper/1706.08498/) (2017)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [Linear Convergence of Gradient and Proximal-Gradient Methods Under the Polyak-Lojasiewicz Condition](https://mlsanity.com/paper/1608.04636/) (2016)
- [A Mean Field View of the Landscape of Two-Layers Neural Networks](https://mlsanity.com/paper/1804.06561/) (2018)
- [Larger language models do in-context learning differently](https://mlsanity.com/paper/2303.03846/) (2023)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Stabilizing Transformers for Reinforcement Learning](https://mlsanity.com/paper/1910.06764/) (2019)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [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)
- [MAE-AST: Masked Autoencoding Audio Spectrogram Transformer](https://mlsanity.com/paper/2203.16691/) (2022)
- [Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality](https://mlsanity.com/paper/2402.19442/) (2024)
- [A Mean-field Analysis of Deep ResNet and Beyond: Towards Provable Optimization Via Overparameterization From Depth](https://mlsanity.com/paper/2003.05508/) (2020)
- [Rates of Convergence for Nearest Neighbor Classification](https://mlsanity.com/paper/1407.0067/) (2014)

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

- [Lost in Cultural Translation: Do LLMs Struggle with Math Across Cultural Contexts?](https://mlsanity.com/paper/2503.18018/) (2025)
- [Decoupling Positional and Symbolic Attention Behavior in Transformers](https://mlsanity.com/paper/2511.11579/) (2025)
- [Transformers versus the EM Algorithm in Multi-class Clustering](https://mlsanity.com/paper/2502.06007/) (2025)
- [In-Context Learning of Linear Dynamical Systems with Transformers: Approximation Bounds and Depth-Separation](https://mlsanity.com/paper/2502.08136/) (2025)
- [Softmax ≥ Linear: Transformers may learn to classify in-context by kernel gradient descent](https://mlsanity.com/paper/2510.10425/) (2025)
- [Attention-based clustering](https://mlsanity.com/paper/2505.13112/) (2025)
- [Bilevel Data Curation for LLM Fine-tuning: Offline Selection and Online Self-Refining Generation](https://mlsanity.com/paper/2511.21056/) (2025)
- [Towards Understanding Transformers in Learning Random Walks](https://mlsanity.com/paper/2511.23239/) (2025)
- [Attention-Based Sampler for Diffusion Language Models](https://mlsanity.com/paper/2604.08564/) (2026)
- [Transformers Can Implement Preconditioned Richardson Iteration for In-Context Gaussian Kernel Regression](https://mlsanity.com/paper/2605.08475/) (2026)
- [A Theoretical Analysis of Why Masked Diffusion Models Mitigate the Reversal Curse](https://mlsanity.com/paper/2602.02133/) (2026)
- [Attention-based PCA](https://mlsanity.com/paper/2605.18315/) (2026)
- [One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context](https://mlsanity.com/paper/2609.01311/) (2026)
- [Theory on Attention Dynamics for Out-of-Distribution In-Context Learning](https://mlsanity.com/paper/2609.36448/) (2026)
- [Retrieval-Based In-Context Learning: A Domain Adaptation Framework](https://mlsanity.com/paper/2610.05717/) (2026)
- [Adaptive Mean Estimation by In-Context Learning: A Gradient-Flow Analysis](https://mlsanity.com/paper/2610.07804/) (2026)

## Abstract (arXiv)

Transformers have achieved great success in recent years. Interestingly, transformers have shown particularly strong in-context learning capability -- even without fine-tuning, they are still able to solve unseen tasks well purely based on task-specific prompts. In this paper, we study the capability of one-layer transformers in learning one of the most classical nonparametric estimators, the one-nearest neighbor prediction rule. Under a theoretical framework where the prompt contains a sequence of labeled training data and unlabeled test data, we show that, although the loss function is nonconvex when trained with gradient descent, a single softmax attention layer can successfully learn to behave like a one-nearest neighbor classifier. Our result gives a concrete example of how transformers can be trained to implement nonparametric machine learning algorithms, and sheds light on the role of softmax attention in transformer models.