# Self-supervised In-context Operator Learning for Stochastic Mean-Field Control

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 42 of its 58 references that are in the map.
- First posted on arXiv 18 Aug 2026.

## Builds on (most cited first, 30 of 42)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Normalizing Flows for Probabilistic Modeling and Inference](https://mlsanity.com/paper/1912.02762/) (2019)
- [NICE: Non-linear Independent Components Estimation](https://mlsanity.com/paper/1410.8516/) (2014)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [DGM: A deep learning algorithm for solving partial differential equations](https://mlsanity.com/paper/1708.07469/) (2017)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [Neural Spline Flows](https://mlsanity.com/paper/1906.04032/) (2019)
- [The Deep Ritz method: A deep learning-based numerical algorithm for solving variational problems](https://mlsanity.com/paper/1710.00211/) (2017)
- [Solving high-dimensional partial differential equations using deep learning](https://mlsanity.com/paper/1707.02568/) (2017)
- [Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion](https://mlsanity.com/paper/2310.02279/) (2023)
- [Conditional Neural Processes](https://mlsanity.com/paper/1807.01613/) (2018)
- [Flow Matching on General Geometries](https://mlsanity.com/paper/2302.03660/) (2023)
- [Diffusion Schr"odinger Bridge Matching](https://mlsanity.com/paper/2303.16852/) (2023)
- [Flow map matching with stochastic interpolants: A mathematical framework for consistency models](https://mlsanity.com/paper/2406.07507/) (2024)
- [In-Context Operator Learning with Data Prompts for Differential Equation Problems](https://mlsanity.com/paper/2304.07993/) (2023)
- [On the relation between optimal transport and Schr"odinger bridges: A stochastic control viewpoint](https://mlsanity.com/paper/1412.4430/) (2014)
- [A Machine Learning Framework for Solving High-Dimensional Mean Field Game and Mean Field Control Problems](https://mlsanity.com/paper/1912.01825/) (2019)
- [The Schr"odinger Bridge between Gaussian Measures has a Closed Form](https://mlsanity.com/paper/2202.05722/) (2022)
- [Probability flow solution of the Fokker-Planck equation](https://mlsanity.com/paper/2206.04642/) (2022)
- [Alternating the Population and Control Neural Networks to Solve High-Dimensional Stochastic Mean-Field Games](https://mlsanity.com/paper/2002.10113/) (2020)
- [Actor-Critic Method for High Dimensional Static Hamilton--Jacobi--Bellman Partial Differential Equations based on Neural Networks](https://mlsanity.com/paper/2102.11379/) (2021)
- [Recent Developments in Machine Learning Methods for Stochastic Control and Games](https://mlsanity.com/paper/2303.10257/) (2023)

## Abstract (arXiv)

Stochastic mean-field control (MFC) provides a fundamental framework for coordinating large populations of interacting agents under uncertainty, with a wide range of applications. Existing numerical and deep-learning methods solve one MFC problem instance at a time and must be re-optimized whenever the task changes. In this work, we formulate stochastic MFC as an operator-learning problem and develop, to the best of our knowledge, the first mesh-free, self-supervised neural operator for stochastic MFC. The main challenge is that the diffusion term in the controlled Fokker--Planck equation precludes deterministic transport-map representations. We address this challenge by combining the probability-flow ODE with an invertible normalizing-flow-based transformer, which recasts the dynamics as a deterministic continuity equation and enables closed-form score evaluation through the exact inverse and analytical log-determinant of the normalizing flow, with O(d) cost per particle for networks of fixed size. Through transformer-based in-context learning, task prompts, represented by compact distribution parameters or raw particle clouds, condition the transport map, enabling a single pretrained operator to solve unseen tasks in one forward pass. The resulting Normalizing Flow Invertible Solution Transformer (NFIST) is trained end-to-end by minimizing the stochastic control objective directly, requiring no precomputed numerical solutions for training. We further prove the consistency of the proposed operator-learning formulation with task-by-task optimization. Numerical experiments on stochastic optimal control, Schr"odinger bridge, systemic-risk control, and obstacle-avoiding path planning demonstrate effective zero-shot generalization while substantially reducing the computational cost of solving large families of stochastic MFC problems.