In-Context Operator Learning on the Space of Probability Measures

TL;DR

This paper proposes a novel in-context operator learning framework for optimal transport on probability measures, enabling accurate OT map predictions from few samples without gradient updates, supported by theoretical analysis and experiments.

Contribution

It introduces a new in-context learning approach for OT maps on probability measures, with theoretical generalization bounds and explicit architectures for specific families.

Findings

The framework achieves accurate OT map predictions with few samples.

Theoretical bounds relate accuracy to prompt size and task complexity.

Experimental results validate the approach on synthetic and benchmark datasets.

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.