dorsal/arxiv
View SchemaIn-Context Operator Learning on the Space of Probability Measures
| Authors | Frank Cole, Dixi Wang, Yineng Chen, Yulong Lu, Rongjie Lai |
|---|---|
| Categories | |
| ArXiv ID | 2601.09979vv1 |
| URL | https://arxiv.org/abs/2601.09979 |
| License | http://arxiv.org/licenses/nonexclusive-distrib/1.0/ |
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.
{
"annotation_id": "5620c54f-5620-4c14-a770-9c9a4daa89a0",
"date_created": "2026-02-17T05:53:24.225000Z",
"date_modified": "2026-02-17T05:53:24.225000Z",
"file_hash": "4b5af759397a27aabd6e2b12810a23071de12674fab4bb4fa3eb0129e5ee9127",
"private": false,
"record": {
"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.",
"arxiv_id": "2601.09979",
"authors": [
"Frank Cole",
"Dixi Wang",
"Yineng Chen",
"Yulong Lu",
"Rongjie Lai"
],
"categories": [
"cs.LG",
"cs.NA",
"math.NA"
],
"license": "http://arxiv.org/licenses/nonexclusive-distrib/1.0/",
"title": "In-Context Operator Learning on the Space of Probability Measures",
"url": "https://arxiv.org/abs/2601.09979",
"version": "v1"
},
"schema_id": "dorsal/arxiv",
"source": {
"execution_id": "7cbc8dd1-7d13-47db-9335-2f92975ccf71",
"id": "arXiv Dataset",
"type": "Model",
"variant": "snapshot-2026-01-17",
"version": "0.1.0"
},
"user_id": 1000002
}