# One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 11 of its 15 references that are in the map.
- First posted on arXiv 1 Sep 2026.

## Builds on (most cited first, 11 of 11)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Generalization through Memorization: Nearest Neighbor Language Models](https://mlsanity.com/paper/1911.00172/) (2019)
- [Prevalence of Neural Collapse during the terminal phase of deep learning training](https://mlsanity.com/paper/2008.08186/) (2020)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel](https://mlsanity.com/paper/1908.11775/) (2019)
- [Mimetic Initialization of Self-Attention Layers](https://mlsanity.com/paper/2305.09828/) (2023)
- [Rates of Convergence for Nearest Neighbor Classification](https://mlsanity.com/paper/1407.0067/) (2014)
- [Why do Nearest Neighbor Language Models Work?](https://mlsanity.com/paper/2301.02828/) (2023)
- [One-Layer Transformer Provably Learns One-Nearest Neighbor In Context](https://mlsanity.com/paper/2411.10830/) (2024)
- [Cut your Losses with Squentropy](https://mlsanity.com/paper/2302.03952/) (2023)
- [Fast and Bayes-consistent nearest neighbors](https://mlsanity.com/paper/1910.05270/) (2019)

## Abstract (arXiv)

We extend recent work establishing an equivalence between one-layer transformers and nearest-neighbor classifiers in the binary setting to the multiclass case. By leveraging the simplex encoding, we show that one-layer transformers with an argmax classification head behave identically to a one-nearest-neighbor classifier in the multiclass setting. This closes a gap left by prior work, whose multiclass result relied on a non-standard rounding-based approach rather than the typical argmax head used in practice.