In-Context Operator Learning on the Space of Probability Measures
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2026
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866914256171040768 |
|---|---|
| author | Cole, Frank Wang, Dixi Chen, Yineng Lu, Yulong Lai, Rongjie |
| author_facet | Cole, Frank Wang, Dixi Chen, Yineng Lu, Yulong Lai, Rongjie |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2601_09979 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | In-Context Operator Learning on the Space of Probability Measures Cole, Frank Wang, Dixi Chen, Yineng Lu, Yulong Lai, Rongjie Machine Learning Numerical Analysis 68T07, 49Q22, 65C20, 91G60 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. |
| title | In-Context Operator Learning on the Space of Probability Measures |
| topic | Machine Learning Numerical Analysis 68T07, 49Q22, 65C20, 91G60 |
| url | https://arxiv.org/abs/2601.09979 |