In-Context Learning of Linear Systems: Generalization Theory and Applications to Operator Learning
This work addresses the theoretical understanding of in-context learning for linear systems, which is incremental as it builds on existing transformer architectures but provides new theoretical insights.
The paper tackles the problem of providing theoretical guarantees for in-context learning of linear systems using linear transformers, establishing neural scaling laws for in-domain generalization and showing that task diversity is crucial for out-of-domain generalization under distribution shifts.
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.