@inproceedings{
li2023transformers,
title={Transformers as Multi-Task Feature Selectors: Generalization Analysis of In-Context Learning},
author={Hongkang Li and Meng Wang and Songtao Lu and Hui Wan and Xiaodong Cui and Pin-Yu Chen},
booktitle={NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning},
year={2023},
url={https://openreview.net/forum?id=BMQ4i2RVbE}
}
Transformer-based large language models have displayed impressive capabilities in the domain of in-context learning, wherein they use multiple input-output pairs to make predictions on unlabeled test data. To lay the theoretical groundwork for in-context learning, we delve into the optimization and generalization of a single-head, one-layer Transformer in the context of multi-task learning for classification. Our investigation uncovers that lower sample complexity is associated with increased training-relevant features and reduced noise in prompts, resulting in improved learning performance.  The trained model exhibits the mechanism to first attend to demonstrations of training-relevant features and then decode the corresponding label embedding. Furthermore, we delineate the necessary conditions for successful out-of-domain generalization for in-context learning, specifically regarding the relationship between training and testing prompts.
Transformers as Multi-Task Feature Selectors: Generalization Analysis of In-Context Learning
Hongkang Li Meng Wang Songtao Lu Hui Wan Xiaodong Cui Pin-Yu Chen