Year: 2024
Type: conference-paper
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.
DOI: 10.52202/079017-2612
Authors: Zihao Li, Yuan Cao, Cheng Gao, Yihan He, Han Liu, Jason Klusowski, Jianqing Fan, Mengdi Wang
Language: English
Subfield: Artificial Intelligence
Field: Computer Science
Domain: Physical Sciences
Open Access status: green
Awards: G12034353, CPS: Medium: Collaborative Research: Provably Safe and Robust Multi-Agent Reinforcement Learning with Applications in Urban Air Mobility, CAREER: Statistical Learning with Recursive Partitioning: Algorithms, Accuracy, and Applications, Deep Learning and Random Forests for High-Dimensional Regression, G4720003262, G5429545050, G6671297155, DMS/NIGMS 2: Collaborative Research: Developing Statistical Learning Methods for Revealing the Molecular Signatures of Microvascular Changes in Neural Injury, Interface of Statistical Learning and Optimal Decisions, HDR TRIPODS: Data Science Principles of the Human-Machine Convergence, G7219397789, G7712150002, RI: Medium: Provable Reinforcement Learning with Function Approximation and Neural Networks, G8876996369
OpenAlex ID: W4415796594