Asymptotic theory of in-context learning by linear attention

2024-05-20

arXiv.org

AI Summary

This work establishes the theoretical foundations of in-context learning (ICL) in linear-attention Transformers, addressing how pretraining sample size, task diversity, and context length jointly govern ICL performance. Method: Leveraging high-dimensional asymptotic analysis and rigorous scaling-limit derivations, we develop the first asymptotic theory for ICL, characterizing its double-descent learning curve and a phase transition driven by task diversity—where low diversity induces memorization and high diversity enables generalization. Contribution/Results: We derive tight necessary conditions for successful ICL in linear regression tasks and validate them empirically across linear-attention models and fully nonlinear Transformers. Our framework provides the first analytically tractable and empirically verifiable theory of ICL, unifying scaling laws across key architectural and statistical dimensions. The results rigorously connect pretraining data scale, contextual capacity, and task distribution geometry to emergent in-context inference behavior.

Technology Category

• Machine Learning: Learning Theory
• Natural Language Processing: Sentence-level Semantics, Textual Inference, etc.
• Search and Optimization: Learning to Search

Application Category

• Web Mining and Content Analysis: Large pretrained models with web data
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Practical large-scale studies of user experience

Abstract

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.

Problem

• Analyzes in-context learning in linear attention Transformers
• Determines sample complexity and task diversity for ICL
• Explores phase transition between memorization and generalization in ICL

Innovation

• Linear attention Transformers
• Asymptotic learning curve analysis
• Phase transition in task diversity