Paper
One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context
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 o · Authors: Skanda Athreya, Yutong Wang · Published 2026-09-01 · arXiv:2609.01311
- Published
- 2026-09-01
- Authors
- 2
- arXiv
- 2609.01311
- Category
- cs.LG
Abstract
arXiv:2609.01311
Skanda Athreya, Yutong Wang
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.