One-layer Transformer Provably Learns One-nearest Neighbor In Context

Abstract

Transformers have achieved great success in recent years. Interestingly, transformers have shown particularly strong in-context learning capability -- even without fine-tuning, they are still able to solve unseen tasks well purely based on task-specific prompts. In this paper, we study the capability of one-layer transformers in learning one of the most classical nonparametric estimators, the one-nearest neighbor prediction rule. Under a theoretical framework where the prompt contains a sequence of labeled training data and unlabeled test data, we show that, although the loss function is nonconvex when trained with gradient descent, a single softmax attention layer can successfully learn to behave like a one-nearest neighbor classifier. Our result gives a concrete example of how transformers can be trained to implement nonparametric machine learning algorithms, and sheds light on the role of softmax attention in transformer models.

Stats

• citations0
• S2 citations—
• github stars0
• HF likes0
• heat score0.00
• arxiv keyli2024one

Related papers

• Towards In-context Scene Understanding (2023)0.00
• Test-time Training On Nearest Neighbors For Large Language Models (2023)0.00
• With A Little Help From My Friends: Nearest-neighbor Contrastive Learning Of Visual Representations (2021)18.76
• REGTR: End-to-end Point Cloud Correspondences With Transformers (2022)22.67
• Explaining The Success Of Nearest Neighbor Methods In Prediction (2025)18.63
• A Self-encoder For Learning Nearest Neighbors (2023)0.00
• Why Do Nearest Neighbor Language Models Work? (2023)3.56
• Simpleshot: Revisiting Nearest-neighbor Classification For Few-shot Learning (2019)0.00