In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness
This provides theoretical insights into transformer mechanisms for in-context learning, addressing a foundational problem in machine learning, though it is incremental as it builds on prior analyses.
The paper investigates how softmax attention in transformers adapts to in-context learning for regression tasks, showing that the attention window widens with decreasing Lipschitzness and increasing label noise, and that softmax activation is crucial for this adaptivity, unlike linear activation.
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.