In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2024
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866910460827140096 |
|---|---|
| author | Collins, Liam Parulekar, Advait Mokhtari, Aryan Sanghavi, Sujay Shakkottai, Sanjay |
| author_facet | Collins, Liam Parulekar, Advait Mokhtari, Aryan Sanghavi, Sujay Shakkottai, Sanjay |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2402_11639 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness Collins, Liam Parulekar, Advait Mokhtari, Aryan Sanghavi, Sujay Shakkottai, Sanjay Machine Learning Artificial Intelligence Computation and Language 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. |
| title | In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness |
| topic | Machine Learning Artificial Intelligence Computation and Language |
| url | https://arxiv.org/abs/2402.11639 |