One-Layer Transformer Provably Learns One-Nearest Neighbor In Context
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2024
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866917840070639616 |
|---|---|
| author | Li, Zihao Cao, Yuan Gao, Cheng He, Yihan Liu, Han Klusowski, Jason M. Fan, Jianqing Wang, Mengdi |
| author_facet | Li, Zihao Cao, Yuan Gao, Cheng He, Yihan Liu, Han Klusowski, Jason M. Fan, Jianqing Wang, Mengdi |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2411_10830 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | One-Layer Transformer Provably Learns One-Nearest Neighbor In Context Li, Zihao Cao, Yuan Gao, Cheng He, Yihan Liu, Han Klusowski, Jason M. Fan, Jianqing Wang, Mengdi Machine Learning Artificial Intelligence Optimization and Control 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. |
| title | One-Layer Transformer Provably Learns One-Nearest Neighbor In Context |
| topic | Machine Learning Artificial Intelligence Optimization and Control |
| url | https://arxiv.org/abs/2411.10830 |