Abstract:Foundation models for natural language processing, powered by the transformer architecture, exhibit remarkable in-context learning (ICL) capabilities, allowing pre-trained models to adapt to downstream tasks using few-shot prompts without updating their weights. Recently, transformer-based foundation models have also emerged as versatile tools for solving scientific problems, particularly in the realm of partial differential equations (PDEs). However, the theoretical foundations of the ICL capabilities in these scientific models remain largely unexplored. This work develops a rigorous error analysis for transformer-based ICL applied to solution operators associated with a family of linear elliptic PDEs. We first demonstrate that a linear transformer, defined by a linear self-attention layer, can provably learn in-context to invert linear systems arising from the spatial discretization of PDEs. This is achieved by deriving theoretical scaling laws for the prediction risk of the proposed linear transformers in terms of spatial discretization size, the number of training tasks, and the lengths of prompts used during training and inference. These scaling laws also enable us to establish quantitative error bounds for learning PDE solutions. Furthermore, we quantify the adaptability of the pre-trained transformer on downstream PDE tasks that experience distribution shifts in both tasks (represented by PDE coefficients) and input covariates (represented by the source term). To analyze task distribution shifts, we introduce a novel concept of task diversity and characterize the transformer's prediction error in terms of the magnitude of task shift, assuming sufficient diversity in the pre-training tasks. We also establish sufficient conditions to ensure task diversity. Finally, we validate the ICL-capabilities of transformers through extensive numerical experiments.

Learning advanced mathematical computations from examples

Transformers discover an elementary calculation system exploiting local attention and grid-like problem representation

Attending to Mathematical Language with Transformers

Transformers Can Do Arithmetic with the Right Embeddings

Understanding In-Context Learning in Transformers and LLMs by Learning to Learn Discrete Functions

Transformers as Neural Operators for Solutions of Differential Equations with Finite Regularity

What is my math transformer doing? -- Three results on interpretability and generalization

How Well Can Transformers Emulate In-context Newton's Method?

Fourier Circuits in Neural Networks and Transformers: A Case Study of Modular Arithmetic with Multiple Inputs

Learning stochastic dynamics and predicting emergent behavior using transformers

Pretrained Language Models are Symbolic Mathematics Solvers too!

Multi-scale Time-stepping of Partial Differential Equations with Transformers

Can Transformers Learn to Solve Problems Recursively?

Learning Elementary Cellular Automata with Transformers

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

Volume-Preserving Transformers for Learning Time Series Data with Structure

Learning integral operators via neural integral equations

A mathematical perspective on Transformers

Global Convergence in Training Large-Scale Transformers

Beyond Transformers for Function Learning