@inproceedings{
collins2024incontext,
title={In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness},
author={Liam Collins and Advait U Parulekar and Aryan Mokhtari and sujay sanghavi and Sanjay Shakkottai},
booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
year={2024},
url={https://openreview.net/forum?id=lfxIASyLxB}
}
A striking property of transformers is their ability to perform in-context learning (ICL), a machine learning framework in which the learner is presented with a novel context during inference implicitly through some data, and tasked with making a prediction in that context. As such, that learner must adapt to the context without additional training. We explore the role of *softmax* attention in an ICL setting where each context encodes a regression task. We show that an attention unit learns a window that it uses to implement a nearest-neighbors predictor adapted to the landscape of the pretraining tasks. Specifically, we show that this window widens with decreasing Lipschitzness and increasing label noise in the pretraining tasks. We also show that on low-rank, linear problems, the attention unit learns to project onto the appropriate subspace before inference. Further, we show that this adaptivity relies crucially on the softmax activation and thus cannot be replicated by the linear activation often studied in prior theoretical analyses.
In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness
Liam Collins Advait U Parulekar Aryan Mokhtari sujay sanghavi Sanjay Shakkottai