English

In-Context Operator Learning on the Space of Probability Measures

Machine Learning 2026-01-16 v1 Numerical Analysis Numerical Analysis

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.

Keywords

Cite

@article{arxiv.2601.09979,
  title  = {In-Context Operator Learning on the Space of Probability Measures},
  author = {Frank Cole and Dixi Wang and Yineng Chen and Yulong Lu and Rongjie Lai},
  journal= {arXiv preprint arXiv:2601.09979},
  year   = {2026}
}
R2 v1 2026-07-01T09:05:08.585Z