# In-Context Operator Learning on the Space of Probability Measures

arXiv: https://arxiv.org/abs/2601.09979 · HTML: https://mlsanity.com/paper/2601.09979/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 15 Jan 2026.

## Built on by (most cited first, 11 of 11)

- [Learning Probabilistic Filters with Strictly Proper Scoring Rules](https://mlsanity.com/paper/2606.26497/) (2026)
- [One Operator for Many Densities: Amortized Approximation of Conditioning by Neural Operators](https://mlsanity.com/paper/2605.06873/) (2026)
- [Graph In-Context Operator Networks for Generalizable Spatiotemporal Prediction](https://mlsanity.com/paper/2603.12725/) (2026)
- [Chain of Operators: An Inference-Time Harness for In-Context Operator Learning](https://mlsanity.com/paper/2606.12318/) (2026)
- [Transformers as Cross-Task Learners: Shared Structure Drives Sample Efficiency in In-Context Learning](https://mlsanity.com/paper/2609.29060/) (2026)
- [Universal Approximation of Measure-to-Measure Operators by Pushforwards](https://mlsanity.com/paper/2609.35483/) (2026)
- [Attention Kernels for Learning Maps Between Heavy-Tailed Measures](https://mlsanity.com/paper/2610.00564/) (2026)
- [Neural Operator Learning for Collision-Aware Trajectory Planning of Spacecraft Swarms](https://mlsanity.com/paper/2608.00320/) (2026)
- [Self-supervised In-context Operator Learning for Stochastic Mean-Field Control](https://mlsanity.com/paper/2608.18282/) (2026)
- [Stability of Measure-to-Measure Transformers on Sub-Gaussian Data](https://mlsanity.com/paper/2610.07717/) (2026)
- [Transformer-based In-Context Score Operator Learning for Diffusion Models](https://mlsanity.com/paper/2610.13690/) (2026)

## Abstract (arXiv)

We introduce 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 without gradient updates at inference. We parameterize the solution operator and develop scaling-law theory in two regimes. In the 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 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.