In-Context Operator Learning on the Space of Probability Measures

TL;DR

We introduce in-context operator learning on probability-measure spaces to approximate OT maps $T_{\rho_0\to\rho_1}$ as a single, task-agnostic operator $\mathscr{T}$ that is invoked by few-shot prompts without gradient updates at test time. The framework uses a transformer-based solution operator that maps empirical prompts $(\hat{\rho}_0,\hat{\rho}_1)$ to predictions of the OT map, with formal analysis in both nonparametric (low-intrinsic-dimension task manifolds) and parametric Gaussian settings. In the nonparametric regime, the authors derive generalization bounds that depend on the intrinsic task dimension, prompt size, and model capacity; in the parametric setting, they construct an explicit transformer that recovers the exact OT map in context and provide finite-sample excess-risk bounds. Numerical experiments on synthetic transports and high-dimensional generative-modeling benchmarks corroborate the theory, showing accurate prompt-driven transport and a scaling behavior with prompt length that aligns with the theory.

Abstract

We introduce \emph{in-context operator learning on probability measure spaces} for optimal transport (OT). The goal is to learn a single solution operator that maps a pair of distributions to the OT map, using only few-shot samples from each distribution as a prompt and \emph{without} gradient updates at inference. We parameterize the solution operator and develop scaling-law theory in two regimes. In the \emph{nonparametric} setting, when tasks concentrate on a low-intrinsic-dimension manifold of source--target pairs, we establish generalization bounds that quantify how in-context accuracy scales with prompt size, intrinsic task dimension, and model capacity. In the \emph{parametric} setting (e.g., Gaussian families), we give an explicit architecture that recovers the exact OT map in context and provide finite-sample excess-risk bounds. Our numerical experiments on synthetic transports and generative-modeling benchmarks validate the framework.

Figures (10)

• Figure 1: Illustration of the ICL model architecture.
• Figure 2: Model prediction with prompts sampled from Gaussian with varying target means and fixed covariance as stated in \ref{['mean_shift']}. Blue and red contours show the predicted and target distribution respectively. The segments connecting landmarks to their mapped predictions are nearly parallel, which is consistent with the optimal transport map $T(x)=x+\mu$.
• Figure 3: Model prediction with prompts sampled from different diagonal covariances and fixed mean as stated in \ref{['cov_shift']}. Blue and red contours show the predicted and target distribution respectively. The black segments are learned trajectories for selected landmarks.
• Figure 4: Visualization results of the Gaussian to MNIST. Different task corresponds to mapping from Standard Gaussian to one digit class. Conditional sampling images from one digit class as prompts, the column-wise images are the output of the learned in-context model for various Gaussian inputs.
• Figure 5: Visualization results of the Gaussian to Fashion-MNIST. Each column corresponds to the inference results of each task. Given the prompts (images) sampled from one fashion class as conditioning, the learned model generates images from the same fashion class with consistency.

Theorems & Definitions (72)

• Definition 1
• Theorem 1: Generalization of Few-Shot In-Context Learning
• Theorem 2: Excess loss estimate
• Remark 1
• Remark 2
• Theorem 3: Transport map generalization error
• Lemma 1
• Lemma 2
• Lemma 3
• proof