Continuum Transformers Perform In-Context Learning by Operator Gradient Descent
Abhiti Mishra, Yash Patel, Ambuj Tewari
TL;DR
The paper provides a rigorous theoretical characterization of in-context learning in continuum transformers for operator learning by showing that ICL arises from gradient descent in an operator RKHS. It introduces a generalized continuum attention mechanism, proves a representer-theorem–based form for operator gradient updates, and establishes BLUP/Bayes-optimality in the infinite-depth limit under Gaussian operator kernels. It also shows that pre-training can yield the gradient-descent parameters as fixed points under symmetry assumptions, and supports these results with empirical experiments across multiple operator RKHSs. The work offers practical guidance for PDE meta-learning by aligning continuum-transformer design with the underlying operator-valued kernels and highlights avenues for depth-independent convergence proofs and broader operator-ICL theory.
Abstract
Transformers robustly exhibit the ability to perform in-context learning, whereby their predictive accuracy on a task can increase not by parameter updates but merely with the placement of training samples in their context windows. Recent works have shown that transformers achieve this by implementing gradient descent in their forward passes. Such results, however, are restricted to standard transformer architectures, which handle finite-dimensional inputs. In the space of PDE surrogate modeling, a generalization of transformers to handle infinite-dimensional function inputs, known as "continuum transformers," has been proposed and similarly observed to exhibit in-context learning. Despite impressive empirical performance, such in-context learning has yet to be theoretically characterized. We herein demonstrate that continuum transformers perform in-context operator learning by performing gradient descent in an operator RKHS. We demonstrate this using novel proof strategies that leverage a generalized representer theorem for Hilbert spaces and gradient flows over the space of functionals of a Hilbert space. We additionally show the operator learned in context is the Bayes Optimal Predictor in the infinite depth limit of the transformer. We then provide empirical validations of this optimality result and demonstrate that the parameters under which such gradient descent is performed are recovered through the continuum transformer training.