PreprintCC BY-NC-SA 4.0arXiv
One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context
Abstract
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.
Related
- One-Layer Transformers Provably Learn In-Context K-Nearest Neighbor Prediction with Chain-of-Thought
- Stochastic dominance of first return times for nearest-neighbor random walks on ℤd
- The nearest neighbor interchange proximity measure for the tree-basedness of rooted binary phylogenetic networks
- Vectorized Generalized Nearest Neighbor Decoding for In-block Memory Channel
- Orthogonal Inference for Conditional Z-Estimation without Cross-Fitting: A Distributional Nearest-Neighbor Approach
BibTeX
@misc{athreya2026onelayer,
title = {{One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context}},
author = {Skanda Athreya and Yutong Wang},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.01311},
eprint = {2609.01311},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.01311},
note = {CC BY-NC-SA 4.0},
}
RIS
TY - UNPB
TI - One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context
AU - Skanda Athreya
AU - Yutong Wang
PY - 2026
DA - 2026/09/01
PB - arXiv
DO - 10.48550/arxiv.2609.01311
UR - https://arxiv.org/abs/2609.01311
AB - 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.
ER -