In-Context Learning of Linear Systems: Generalization Theory and Applications to Operator Learning

2024/09/18 by Cole, Frank, Lu, Yulong, Xu, Wuzhe +1
#FOS: Computer and information sciences #FOS: Mathematics #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Numerical Analysis (math.NA)

paper · doi:10.48550/arxiv.2409.12293

Abstract

We study theoretical guarantees for solving linear systems in-context using a linear transformer architecture. For in-domain generalization, we provide neural scaling laws that bound the generalization error in terms of the number of tasks and sizes of samples used in training and inference. For out-of-domain generalization, we find that the behavior of trained transformers under task distribution shifts depends crucially on the distribution of the tasks seen during training. We introduce a novel notion of task diversity and show that it defines a necessary and sufficient condition for pre-trained transformers generalize under task distribution shifts. We also explore applications of learning linear systems in-context, such as to in-context operator learning for PDEs. Finally, we provide some numerical experiments to validate the established theory.

Related

• In-Context Learning of Linear Dynamical Systems with Transformers: Approximation Bounds and Depth-Separation
• When can in-context learning generalize out of task distribution?
• Re-examining learning linear functions in context
• Transformers as Algorithms: Generalization and Stability in In-context Learning
• Continuum Transformers Perform In-Context Learning by Operator Gradient Descent