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

2026-09-01

AI Summary

研究通过单层Transformer与argmax分类头结合，在多分类场景下实现了一最近邻分类器的行为，解决了之前工作中的方法不标准问题。

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.

Problem

• one-layer transformer
• multiclass
• one-nearest-neighbor
• argmax

Innovation

• one-layer transformers
• multiclass one-nearest neighbor
• simplex encoding
• argmax classification head