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

• URL: http://arxiv.org/abs/2608.18282v1
• Date: Tue, 18 Aug 2026 20:03:48 GMT
• Title: Self-supervised In-context Operator Learning for Stochastic Mean-Field Control
• Abstract summary: Mean-compute-field control (MFC) provides a framework for coordinating large populations of interacting agents under uncertainty.<n>Existing numerical and deep-learning methods solve one MFC problem instance at a time.<n>We develop the first mesh-free, self-learning neural operator for MFC.
• Score: 17.382073199976062
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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 \emph{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.

Related papers

• Frozen CLIP Priors for Robust Self-Supervised Poisson Inverse Problems [19.52206071875139]
  This paper proposes an ADMM-inspired unrolled plug-and-play solver for Poisson inverse problems that decouples a closed-form data-consistency update from a parameter-efficient prior.<n> Experiments on Poisson CFA demosaicing and deblurring show competitive quality, improved robustness under shifts, and self-supervised performance approaching supervised training.
  arXiv  Detail & Related papers  (2026-08-20T19:37:52Z)
• Relaxation-Informed Training of Neural Network Surrogate Models [3.9722979176564763]
  ReLU neural networks trained surrogate models can be embedded exactly in mixed-integer linear programs (MILPs)<n>The tractability of the resulting MILP depends on structural properties of the network.<n>We show that the proposed regularizers can reduce MILP times by up to four orders of magnitude to an unregularized baseline.
  arXiv  Detail & Related papers  (2026-04-24T17:46:55Z)
• Diffusion Controller: Framework, Algorithms and Parameterization [54.82539154511621]
  We present a unified control-theoretic view that casts reverse diffusion sampling as state-only control within (generalized) linearly-solvable Markov Decision Processes.<n>Under this framework, control acts by reweighting the pretrained reverse-time transition kernels, balancing terminal objectives against an $f$-divergence cost.<n>Experiments on Stable Diffusion v1.4 show consistent gains in preference-alignment win rates and improved quality-efficiency trade-offs.
  arXiv  Detail & Related papers  (2026-03-07T01:49:59Z)
• Flow matching Operators for Residual-Augmented Probabilistic Learning of Partial Differential Equations [0.5729426778193397]
  We formulate flow matching in an infinite-dimensional function space to learn a probabilistic transport.<n>We develop a conditional neural operator architecture based on feature-wise linear modulation for flow matching vector fields.<n>We show that the proposed method can accurately learn solution operators across different resolutions and fidelities.
  arXiv  Detail & Related papers  (2025-12-14T16:06:10Z)
• Optimal Control Meets Flow Matching: A Principled Route to Multi-Subject Fidelity [35.95129874095729]
  Text-to-image (T2I) models excel on single-entity prompts but struggle with multi-subject descriptions.<n>We introduce the first theoretical framework with principled optimizable objective for steering sampling dynamics toward multi-subject fidelity.
  arXiv  Detail & Related papers  (2025-10-02T17:59:58Z)
• Stochastic Optimal Control Matching [53.156277491861985]
  Our work introduces Optimal Control Matching (SOCM), a novel Iterative Diffusion Optimization (IDO) technique for optimal control. The control is learned via a least squares problem by trying to fit a matching vector field. Experimentally, our algorithm achieves lower error than all the existing IDO techniques for optimal control.
  arXiv  Detail & Related papers  (2023-12-04T16:49:43Z)
• A Multi-Head Ensemble Multi-Task Learning Approach for Dynamical Computation Offloading [62.34538208323411]
  We propose a multi-head ensemble multi-task learning (MEMTL) approach with a shared backbone and multiple prediction heads (PHs) MEMTL outperforms benchmark methods in both the inference accuracy and mean square error without requiring additional training data.
  arXiv  Detail & Related papers  (2023-09-02T11:01:16Z)
• Efficient Model-Based Multi-Agent Mean-Field Reinforcement Learning [89.31889875864599]
  We propose an efficient model-based reinforcement learning algorithm for learning in multi-agent systems. Our main theoretical contributions are the first general regret bounds for model-based reinforcement learning for MFC. We provide a practical parametrization of the core optimization problem.
  arXiv  Detail & Related papers  (2021-07-08T18:01:02Z)
• Offline Model-Based Optimization via Normalized Maximum Likelihood Estimation [101.22379613810881]
  We consider data-driven optimization problems where one must maximize a function given only queries at a fixed set of points. This problem setting emerges in many domains where function evaluation is a complex and expensive process. We propose a tractable approximation that allows us to scale our method to high-capacity neural network models.
  arXiv  Detail & Related papers  (2021-02-16T06:04:27Z)
• Control of Stochastic Quantum Dynamics with Differentiable Programming [0.0]
  We propose a framework for the automated design of control schemes based on differentiable programming. We apply this approach to state preparation and stabilization of a qubit subjected to homodyne detection. Despite the resulting poor signal-to-noise ratio, we can train our controller to prepare and stabilize the qubit to a target state with a mean fidelity around 85%.
  arXiv  Detail & Related papers  (2021-01-04T19:00:03Z)
• Gaussian Process-based Min-norm Stabilizing Controller for Control-Affine Systems with Uncertain Input Effects and Dynamics [90.81186513537777]
  We propose a novel compound kernel that captures the control-affine nature of the problem. We show that this resulting optimization problem is convex, and we call it Gaussian Process-based Control Lyapunov Function Second-Order Cone Program (GP-CLF-SOCP)
  arXiv  Detail & Related papers  (2020-11-14T01:27:32Z)
• Learning Likelihoods with Conditional Normalizing Flows [54.60456010771409]
  Conditional normalizing flows (CNFs) are efficient in sampling and inference. We present a study of CNFs where the base density to output space mapping is conditioned on an input x, to model conditional densities p(y|x)
  arXiv  Detail & Related papers  (2019-11-29T19:17:58Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.