cs.LG · 2609.01311 Copy arXiv ID · Sep 1, 2026 Save One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context Authors: Skanda Athreya , Yutong Wang
Organizations: James B. Conant High School · Hoffman Estates, IL · Illinois Institute of Technology · Chicago, IL
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.
Explore similar work May 18, 2026 · Penghao Yu, Haotian Jiang, Zeyu Bao +1 Transformer Layers Softmax Attention
May 27, 2026 · Yixiao Huang, Hanlin Zhu, Zixuan Wang +4 Transformer Layers Transformer Architectures
May 7, 2026 · cs.LG J/K move · Enter open · S save
Chenyang Zhang, Yuan Cao
School of Computing & Data Science, The University of Hong Kong
Transformers have demonstrated remarkable in-context learning (ICL) capabilities. The strong ICL performance of transformers is commonly believed to arise from their ability to implicitly execute certain algorithms on the context, thereby enhancing prediction and generation. In this work, we investigate how transformers with softmax attention perform in-context learning on linear classification data. We first construct a class of multi-layer transformers that can perform in-context logistic regression, with each layer exactly performing one step of normalized gradient descent on an in-context loss. Then, we show that our constructed transformer can be obtained through (i) training a single self-attention layer supervised by one-step gradient descent, and (ii) recurrently applying the trained layer to obtain a looped model. Training convergence guarantees of the self-attention layer and out-of-distribution generalization guarantees of the looped model are provided. Our results advance the theoretical understanding of ICL mechanism by showcasing how softmax transformers can effectively act as in-context learners.