One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context
Not on the map. Why?
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.
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
11 of 15 references.
References · 11
- 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,305
- 2019Generalization through Memorization: Nearest Neighbor Language Modelscited by 800
- 2020Prevalence of Neural Collapse during the terminal phase of deep learning trainingcited by 719
- 2023Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selectioncited by 298
- 2019Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernelcited by 213
- 2023Mimetic Initialization of Self-Attention Layerscited by 54
- 2014Rates of Convergence for Nearest Neighbor Classificationcited by 46
- 2023Why do Nearest Neighbor Language Models Work?cited by 29
- 2024One-Layer Transformer Provably Learns One-Nearest Neighbor In Contextcited by 16
- 2023Cut your Losses with Squentropycited by 6
- 2019Fast and Bayes-consistent nearest neighborscited by 1
Built on by
Not cited yet.