CSE DSI Machine Learning Seminar with Yue Lu (SEAS, Harvard)

Asymptotic theory of in-context learning by linear attention

Transformers have a remarkable ability to learn and execute tasks based on examples provided within the input itself, without explicit prior training. It has been argued that this capability, known as in-context learning (ICL), is a cornerstone of Transformers' success, yet questions about the necessary sample complexity, pretraining task diversity, and context length for successful ICL remain unresolved. Here, we provide a precise answer to these questions in an exactly solvable model of ICL of a linear regression task by linear attention. We derive sharp asymptotics for the learning curve in a phenomenologically-rich scaling regime where the token dimension is taken to infinity; the context length and pretraining task diversity scale proportionally with the token dimension; and the number of pretraining examples scales quadratically. We demonstrate a double-descent learning curve with increasing pretraining examples, and uncover a phase transition in the model's behavior between low and high task diversity regimes: In the low diversity regime, the model tends toward memorization of training tasks, whereas in the high diversity regime, it achieves genuine in-context learning and generalization beyond the scope of pretrained tasks. These theoretical insights are empirically validated through experiments with both linear attention and full nonlinear Transformer architectures.

Joint work with Mary Letey, Jacob Zavatone-Veth, Anindita Maiti, and Cengiz Pehlevan.

Yue M. Lu attended the University of Illinois at Urbana-Champaign, where he received both the M.Sc. degree in Mathematics and the Ph.D. degree in Electrical Engineering in 2007.  He is currently a Harvard College Professor and Gordon McKay Professor of Electrical Engineering and of Applied Mathematics at Harvard University. His research interests lie in the mathematical foundations of statistical signal processing and machine learning in high dimensions. His research contributions have been recognized with several best paper awards (from the IEEE ICIP, ICASSP, GlobalSIP), the ECE Illinois Young Alumni Achievement Award (2015), and the IEEE Signal Processing Society Distinguished Lectureship (2022). He is a Fellow of the IEEE (class of 2024).