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

arXiv:2609.01311v1 Announce Type: new 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

• Multi-layer Cross-Attention is Provably Optimal for Multi-modal In-context Learning
• Sparse Attention as Compact Kernel Regression
• Limit Theorems for Stochastic Gradient Descent in High-Dimensional Single-Layer Networks
• RetrievalFormer: A Dual-Encoder Transformer for Efficient Approximate Nearest Neighbor Retrieval and Cold-Item Recommendation

Source: arXiv cs.LG | 2026-09-02