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

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.

• The paper establishes provable error bounds for transformer-based in-context learning applied to linear systems and linear elliptic PDEs.
• It reduces infinite-dimensional PDEs to finite linear systems and shows that longer prompts and more training tasks significantly lower prediction error.
• It validates theoretical predictions with extensive numerical experiments, demonstrating robustness even under significant task and covariate shifts.

Provable In-Context Learning of Linear Systems and Linear Elliptic PDEs with Transformers

The research presented in this paper develops a rigorous error analysis of in-context learning (ICL) using transformers, specifically targeting linear systems and linear elliptic partial differential equations (PDEs). This study addresses the gap in understanding the theoretical foundations of ICL capabilities in scientific modeling, leveraging transformers for PDE solutions.

Problem Context and Objectives

In recent advancements, foundation models (FMs) based on transformer architectures have demonstrated impressive adaptability in NLP through few-shot prompts without updating model weights. This adaptability is attributed to their in-context learning abilities. In the scientific domain, transformers have shown potential in solving complex PDEs, but theoretical underpinnings of such capabilities remain underexplored. This paper focuses on establishing the theoretical error bounds for transformers applied to a class of linear elliptic PDEs through ICL frameworks.

Core Contributions

1. In-context Learning of Linear Systems:
   • The authors reduce the infinite-dimensional PDE problem to a finite-dimensional linear system problem. They establish that transformers with linear self-attention layers can invert these linear systems in context.
   • They derive error bounds for the prediction risk of generalization from ICL, considering discretization size, number of training tasks, and prompt lengths.
2. Task Distribution Shifts:
   • The study introduces the concept of task diversity to analyze shifts in task distribution. The authors demonstrate conditions ensuring sufficient task diversity, which allows for robust generalization even under significant distribution shifts.
   • They establish bounds quantifying the prediction error under distribution shifts, showcasing the robustness of transformers owing to their adaptability to task diversity.
3. Numerical Validation:
   • Extensive numerical experiments validate the theoretical results. These experiments highlight that the prediction risk decays predictably with increased sample sizes and prompt lengths during both training and inference phases.

Theoretical Results and Implications

In-Domain Generalization

The study presents rigorous bounds for in-domain generalization. For transformers trained to learn linear systems from spatial discretization of PDEs, the generalization error scales as $O(1/m + 1/n^2 + 1/\sqrt{N})$, where $m$ is the prompt length during testing, $n$ the prompt length during training, and $N$ the number of pre-training tasks. This scaling indicates that longer prompts and more training tasks significantly reduce generalization error.

PDE Solution Error Estimation

Theoretical results extend to learning elliptic PDE solutions. The error bound on PDE solutions is expressed in terms of the discretization error and statistical error from learning linear systems. This yields an explicit error structure encompassing the trade-off between increasing discretization (basis functions) and the amount of training data, crucial for practical applications. An illustrative example using FEM discretization in one dimension exemplifies these bounds.

Out-of-Domain Generalization

Out-of-domain generalization is analyzed under two scenarios:

1. Task Distribution Shifts:
   • The paper shows that task diversity ensures OOD generalization by quantifying the prediction error again as $O(1/m)$.
   • Task distributions with centralizer conditions or simultaneously diagonalizable matrices are explored to ensure task diversity.
2. Covariate Distribution Shifts:
   • The error bounds under covariate shifts reveal that transformers are less robust when the data distribution shifts beyond the pre-trained distribution scope. The mismatch error scales with the operator norm difference of covariance matrices.

Future Directions

This paper sets a foundation for theoretically grounded ICL in scientific problems using transformers. Future research calls for exploring more complex settings, including nonlinear and time-dependent PDEs, where depth and nonlinearity play critical roles in approximating solutions. Moreover, a deeper understanding of task diversity in broader contexts beyond linear systems is essential to harness the full potential of transformers in diverse scientific applications.

This research paves the way for leveraging sophisticated in-context learning models in scientific computing, offering pathways to efficient problem-solving without extensive retraining. The blend of theoretical rigor and practical validation in this work underscores the transformative potential of transformer-based methods in scientific modeling and computation.