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

Suyi Gao; Mo Zhou; Rongjie Lai

math.OCAug 18, 2026

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

Authors: Suyi Gao, Mo Zhou, Rongjie Lai

Organizations: Department of Mathematics, Purdue University, West Lafayette, USA · Department of Mathematics, University of California, Los Angeles, USA

Abstract

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 $\mathcal{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ödinger 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.

Figures & tables

Panel of Figure 1 — Figure 1: Architecture of NFIST. Top: The solution operator is realized as a stack of affine coupling blocks. Middle: Each coupling block uses a transformer to produce the affine parameters $(\alpha,\beta)$ that define the coupling transformation. Bottom: The transformer conditions the flow on the task prompt. Query particles $(x,t)$ are processed by self-attention and cross-attention with a context $c$ constructed from either non-parametric particle-cloud prompts or parametric prompt tokens.

Figure 2: Single-well one-dimensional SOC problem, one solution operator netowrk (

\gamma=0.1

; random initial mean/variance and well location), on three tasks (strips). In both rows the solid curve is the operator output and the dashed curve the exact reference, drawn on top, so exact agreement appears as a dashed curve riding on a solid one.

Figure 3: One SOC solution operator evaluated zero-shot over the plane of well centers (

d=2

\gamma=0.1

). The SDE trajectories are simulated according to section 2.4 .

		Inference Time (s)		Relative error
dim	$\bm{\gamma}$	transport	$+$ score	trajectory	score
$1$	$0.01$	0.013	0.752	3.60e-4 $\pm$ 2.78e-4	1.11e-3 $\pm$ 7.99e-4
$1$	$0.1$	0.013	1.066	1.11e-3 $\pm$ 6.39e-4	3.29e-3 $\pm$ 1.74e-3
$1$	$1.0$	0.013	0.650	2.57e-3 $\pm$ 1.51e-3	6.84e-3 $\pm$ 3.40e-3
$2$	$0.01$	0.012	0.486	2.45e-3 $\pm$ 1.11e-3	9.67e-3 $\pm$ 4.13e-3
$2$	$0.1$	0.010	0.539	1.07e-3 $\pm$ 4.19e-4	4.83e-3 $\pm$ 1.50e-3

Table 1: Single-well SOC vs. the closed-form Gaussian reference: per-task inference time and the relative errors ( 33 ) and ( 34 ). The

d=1

family additionally randomizes the initial mean and variance.

Figure 4: Single-well SOC with the non-parametric particle-cloud prompt (

d=1

\gamma=0.1

): the operator is conditioned on raw

100

-point sample clouds

X_{0}

rather than on

(m_{0},\sigma_{0})

. Same three tasks and panel layout as the parametric Figure 2 .

Figure 5: One non-parametric SOC solution operator evaluated zero-shot on sampled tasks (

d=2

\gamma=0.1

). The SDE trajectories are simulated according to section 2.4 .

		Inference Time (s)		Relative error
dim	$\bm{\gamma}$	transport	$+$ score	trajectory	score
$1$	$0.01$	0.014	1.370	1.73e-2 $\pm$ 1.66e-2	4.87e-2 $\pm$ 4.59e-2
$1$	$0.1$	0.014	0.651	1.88e-2 $\pm$ 1.66e-2	4.93e-2 $\pm$ 4.00e-2
$1$	$1.0$	0.014	1.373	3.17e-2 $\pm$ 2.26e-2	8.25e-2 $\pm$ 5.72e-2
$2$	$0.01$	0.014	1.346	1.36e-2 $\pm$ 8.21e-3	4.72e-2 $\pm$ 2.42e-2
$2$	$0.1$	0.014	0.664	1.24e-2 $\pm$ 7.47e-3	4.00e-2 $\pm$ 2.02e-2

Table 2: Non-parametric (particle-cloud) single-well SOC vs. the exact Gaussian reference.

Figure 6: Multi-well SOC in the sharply multimodal

\gamma=0.5

regime of three tasks. Each strip samples the score along three high-density slices of the terminal law (as in the 2-D plot on the left).

		Inference Time (s)		Relative error
prompt type	dim	transport	$+$ score	trajectory	score
Parametric	$2$	0.012	1.331	2.16e-2 $\pm$ 4.32e-3	7.25e-2 $\pm$ 1.23e-2
Parametric	$5$	0.012	1.330	4.29e-2 $\pm$ 3.00e-3	1.10e-1 $\pm$ 7.47e-3
Parametric	$8$	0.012	1.364	4.44e-2 $\pm$ 2.51e-3	1.07e-1 $\pm$ 4.58e-3
Parametric	$11$	0.012	1.380	4.56e-2 $\pm$ 1.61e-3	1.02e-1 $\pm$ 2.78e-3
Parametric	$14$	0.012	1.311	4.48e-2 $\pm$ 1.97e-3	9.57e-2 $\pm$ 3.05e-3

Table 3: Multi-well SOC with wells spanning

\mathbb{R}^{d}

vs. the Monte-Carlo Hopf–Cole reference of B.2 : per-task inference time and the relative errors ( 33 ) and ( 34 ), the trajectory scored against the reference probability flow ( 39 ).

Figure 7: Schrödinger bridge, one-dimensional parametric solution operator (

\gamma=1.0

; random mean and variance at both endpoints).

Figure 8: One bridge solution operator (

d=2

\gamma=5.0

), solved zero-shot; each panel framed to its own task. The prescribed terminal law

P_{1}

(green shading, dashed

3\sigma

) is the bridge’s exact terminal marginal; trajectories (blue) carry each initial population (purple) onto it under fixed weights. The SDE trajectories are simulated according to section 2.4 .

		Inference Time (s)		Relative error
dim	$\bm{\gamma}$	transport	$+$ score	trajectory	score
$1$	$1.0$	0.014	0.714	3.27e-2 $\pm$ 2.47e-2	6.57e-2 $\pm$ 4.67e-2
$1$	$5.0$	0.013	1.362	3.13e-2 $\pm$ 1.83e-2	6.90e-2 $\pm$ 2.93e-2
$2$	$1.0$	0.012	1.379	2.12e-2 $\pm$ 1.16e-2	6.83e-2 $\pm$ 3.42e-2
$2$	$5.0$	0.012	1.316	2.16e-2 $\pm$ 1.04e-2	6.14e-2 $\pm$ 2.45e-2
$5$	$1.0$	0.012	1.322	2.28e-2 $\pm$ 7.57e-3	6.15e-2 $\pm$ 1.92e-2

Table 4: Schrödinger bridge vs. the closed-form Gaussian reference: per-task inference time and the relative errors ( 33 ) and ( 34 ), The

d=1

family randomizes mean and variance at both endpoints.

Figure 9: Schrödinger bridge with the non-parametric particle-cloud prompt (

d=1

\gamma=1.0

; endpoints prompted by

100

-point sample clouds).

Figure 10: One non-parametric Schrödinger bridge solution operator (

d=2

\gamma=5.0

), solved zero-shot from particle-cloud prompts; panels as in Figure 8 : prescribed terminal law

P_{1}

(green shading, dashed

3\sigma

), trajectories (blue), initial populations (purple). The SDE trajectories are simulated according to section 2.4 .

		Inference Time (s)		Relative error
dim	$\bm{\gamma}$	transport	$+$ score	trajectory	score
$1$	$1.0$	0.015	1.320	6.70e-2 $\pm$ 4.35e-2	1.53e-1 $\pm$ 7.97e-2
$1$	$5.0$	0.014	0.656	7.52e-2 $\pm$ 4.84e-2	1.70e-1 $\pm$ 1.01e-1
$2$	$1.0$	0.014	0.865	4.15e-2 $\pm$ 2.34e-2	1.31e-1 $\pm$ 5.48e-2
$2$	$5.0$	0.014	1.308	3.93e-2 $\pm$ 2.13e-2	9.66e-2 $\pm$ 3.90e-2
$5$	$1.0$	0.014	0.992	4.60e-2 $\pm$ 1.77e-2	1.24e-1 $\pm$ 3.76e-2

Table 5: Non-parametric (particle-cloud) Schrödinger bridge vs. the exact Gaussian reference,

N_{\rm pt}=100

prompt samples.

	Inference Time (s)		Relative error
Parameters	transport	$+$ score	trajectory	score
$a\in[0.033,\,0.300]$	0.012	0.901	6.36e-3 $\pm$ 2.15e-3	5.49e-2 $\pm$ 1.76e-2
$q\in[0.167,\,0.950]$	0.012	0.901	6.44e-3 $\pm$ 2.33e-3	5.56e-2 $\pm$ 1.76e-2
$\varepsilon\in[0.333,\,3.000]$	0.012	0.901	6.60e-3 $\pm$ 2.57e-3	5.56e-2 $\pm$ 1.85e-2
$c\in[0.333,\,3.000]$	0.012	0.901	1.29e-2 $\pm$ 8.19e-3	1.17e-1 $\pm$ 5.18e-2
$\sigma\in[0.167,\,1.500]$	0.012	0.901	6.76e-3 $\pm$ 3.23e-3	5.23e-2 $\pm$ 1.88e-2

Table 6: Systemic risk vs. the closed-form Riccati/Gaussian reference: per-task inference time and the relative errors ( 33 ) and ( 34 ). Each row sweeps one coordinate of the cost vector

\theta=(a,q,\varepsilon,c,\sigma)

log-uniformly over its training range, with the other four held at

\theta_{\rm base}=\{a=0.1,\,q=0.5,\,\varepsilon=1,\,c=1,\,\sigma=0.5\}

; All results are from the same trained solution operator network, so inference time is constant.

Figure 11: Systemic risk at the base cost

\theta_{\rm base}=\{a=0.1,\,q=0.5,\,\varepsilon=1,\,c=1,\,\sigma=0.5\}

, three tasks with initial mean

m_{0}\in\{-3.8,\,0.1,\,2.5\}

(top/middle/bottom strips).

Figure 12: The same operator prompted zero-shot with the low-volatility cost

\sigma=0.25

, on the tasks and panels of Figure 11 : in every task the density concentrates strongly and the score steepens. Contrast the high-volatility prompt of Figure 13 .

Figure 13: The same operator prompted with the high-volatility cost

\sigma=1.2

, on the tasks and panels of Figure 11 : with no retraining, the density stays markedly wider and the score gentler, exactly as the variance ODE prescribes.

Panel of Figure 14 — Figure 14: One path-planning solution operator, zero-shot on configurations: initial distribution $P_{0}$ (purple), target distribution $P_{T}$ (green, dashed $3\sigma$ ), obstacles (orange rings). Blue curves are Euler–Maruyama sample paths of the controlled diffusion ( 1 ) at the trained $\gamma=0.1$ ; the quadratic programming projection enforces the $3\sigma$ keep-out.

	Inference Time (s)		Peak memory	Parameters
dim	transport	score	(GB)	( $10^{6}$ )
$2$	0.0095	0.6302	0.06	4.55
$4$	0.0095	0.6177	0.06	4.56
$8$	0.0095	0.6192	0.06	4.56
$16$	0.0095	0.6184	0.06	4.58
$32$	0.0095	0.6197	0.06	4.61

Table 7: Per-task inference cost of the multi-well operator versus state dimension, on untrained (randomly initialized) models. “transport” is one batched rollout over all

n_{\rm time}

nodes and “score” is one analytic-score evaluation at each node. Everything but

d

is fixed:

B=1

N_{\rm pt}=100

n_{\rm time}=32

Figure 15: Per-task inference cost versus dimension for the multi-well operator, at

B=1

N_{\rm pt}=100

and

n_{\rm time}=32

. The grey guides are

\mathcal{O}(d)

\mathcal{O}(d^{2})

and

\mathcal{O}(d^{3})

, anchored at the transport measurement at

d=2048

(circled), where growth begins, so that each guide is compared with the data over the range in which an asymptotic rate is claimed.

Explore similar work

Oct 8, 2026

Randomized Transport Maps for Model-Free Policy-Gradient Mean-Field Control

We develop a model-free policy gradient method for discrete-time mean-field control (MFC). In MFC, the policy affects the objective both through the controlled dynamics and through the population distribution. Standard REINFORCE estimators capture the first effect but not the second. We introduce Transport REINFORCE, a transport map-based approach that perturbs a suitable transformation of the population distribution to estimate this missing mean-field contribution. The method applies to both finite and continuous state spaces. In finite state spaces, we perturb the population distribution directly on the probability simplex through a convex combination of the current population weights and random weights. In continuous state spaces, we project the population distribution onto the manifold of Gaussian mixtures, and then randomize it via a transport map that ensures the perturbed law remains within this manifold. We prove consistency of the perturbed objective and gradient as the perturbation vanishes, and derive bias and mean-square error bounds for the resulting sample-based gradient estimator. Numerical experiments on several MFC benchmarks show that Transport REINFORCE improves over standard REINFORCE.

Adonis Jamal, Samy Mekkaoui, Yadh Hafsi +1

École Polytechnique

Policy Gradient Methods Model-Free RL

Ask PDF

Oct 11, 2026

Mesh-Free Convolution: Learned Spectral Attenuation and Transport

Learning solution operators for partial differential equations (PDEs) requires models that can represent spatial structure across grids and scales. This becomes especially important in time-dependent problems, where the learned operator is often used autoregressively: each predicted state is fed back as input to predict future states. In this setting, small spatial errors can be repeatedly propagated through the learned dynamics, leading to instability and drift, vitiating long-term solution fidelity. We propose Mesh-Free Convolutions (MFCs), filters defined in continuous space and discretized only at evaluation time. This allows the same learned filters to be applied across resolutions. Rather than imposing a grid-dependent stencil or a fixed spectral cutoff, MFCs learn smooth scale-dependent filtering and include directional terms for transport-like behavior. Across fluid-dynamics and weather-forecasting benchmarks, MFC-based models produce stable long-horizon rollouts, transfer reliably across resolutions, and better preserve long-time statistics. These results suggest that continuous, mesh-free filtering is a useful inductive bias for PDE operator learning, especially in autoregressive prediction.

Niloufar Zakariaei, Manav Manav, Eldad Haber

University of British Columbia · Computational Geoscience Inc.

PDE Operator Learning Spectral Methods Weather Forecasting

Ask PDF

Jul 13, 2026

Actor-Critic Learning for Extended Mean Field Control with Deterministic Policies

This paper develops a model-free reinforcement learning framework for continuous--time extended mean field control problems, where both the dynamics and reward may depend on the joint distribution of states and controls. We adopt deterministic feedback policies, under which the state--action distribution is induced directly as a push--forward of the state law. This avoids optimization over stochastic kernels and bypasses key limitations of existing approaches in extended mean field settings. We first establish a model--free sensitivity formula for parameterized McKean--Vlasov dynamics and use it to derive a deterministic policy gradient formula expressed through an advantage--rate function on the Wasserstein space. We then refine this formula by introducing local value and advantage--rate representations that depend on the state, action, and joint state--action distribution, yielding a policy gradient that includes both action derivatives and measure--derivative terms with respect to the control distribution. These characterizations lead to a martingale--based learning principle and motivate a continuous--time deep deterministic policy gradient algorithm combining particle approximations, measure--dependent neural networks, temporal--difference learning, and exploration in either action or parameter space. Numerical experiments on stochastic Cucker--Smale consensus control and optimal liquidation with trade crowding demonstrate the efficiency, stability, and robustness of the proposed method, including problems with explicit dependence on the control distribution.

Ziheng Cheng, Xin Guo, Huyên Pham +1

University of California, Berkeley · Ecole Polytechnique, CMAP · Imperial College London

Policy Gradient Continuous-Time RL Model-Free RL Deep Deterministic Policy Gradient Actor-Critic Methods Policy Learning

Ask PDF