In-Context Operator Learning on the Space of Probability Measures

2026-01-15

AI Summary

This work addresses the problem of learning optimal transport maps between probability distributions during inference without gradient updates and using only a few in-context samples. It introduces the first in-context operator learning framework that directly predicts optimal transport maps over measure spaces, applicable to both nonparametric and parametric distribution settings. By integrating in-context learning, optimal transport theory, and the low-dimensional manifold hypothesis, the method employs an explicit neural architecture and provides theoretical guarantees on generalization error and excess risk. Experiments on synthetic transport tasks and generative modeling benchmarks demonstrate the approach’s effectiveness, achieving high-accuracy map prediction without fine-tuning.

Technology Category

• Machine Learning: Learning with Manifolds
• Search and Optimization: Learning to Search
• Computer Vision: Learning & Optimization for CV

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Incentives in network design for Web infrastructures and ecosystems
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs

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.

Problem

• in-context learning
• operator learning
• optimal transport
• probability measures
• few-shot learning

Innovation

• in-context operator learning
• optimal transport
• probability measure spaces
• few-shot learning
• generalization bounds