Mean-Field Path-Integral Diffusion: From Samples to Interacting Agents

Michael Chertkov

math.OCFeb 23, 2026

Mean-Field Path-Integral Diffusion: From Samples to Interacting Agents

Authors: Michael Chertkov

Organizations: Graduate Interdisciplinary Program in Applied Mathematics & Department of Mathematics, University of Arizona, Tucson, AZ, USA

Abstract

Independent sample generation is the prevailing paradigm in modern diffusion-based generative models of AI. We ask a different question: can samples coordinate through shared population statistics to transport probability mass more efficiently? We introduce Mean-Field Path-Integral Diffusion (MF-PID), a framework in which samples are promoted to interacting agents whose drift depends self-consistently on the evolving population density. The coupling converts distribution matching into a McKean--Vlasov extension of the stochastic optimal transport problem, unifying generative modeling and multi-agent control under the same Hamilton--Jacobi--Bellman/Kolmogorov--Fokker--Planck duality. We identify two analytically tractable regimes: a Linear--Quadratic--Gaussian (LQG) benchmark in which the infinite-dimensional mean-field system reduces to a finite set of Riccati and linear ODEs, and a Gaussian-mixture regime governed by a piecewise-constant protocol that preserves closed-form solvability. For a quadratic interaction potential with schedule $β_t$ and zero base drift we prove that the self-consistent MF guidance is the exact linear interpolant between initial and target global means -- a result that holds for arbitrary initial and target densities and any $β_t$ . Applied to demand-response control of energy systems, where agents aggregated into an ensemble are energy consumers (e.g.\ thermal zones within a building), MF-PID achieves 19--24% reductions in cumulative control energy over independent-agent baselines while matching the prescribed terminal distribution exactly, and reveals how coordination redistributes actuation effort across heterogeneous sub-populations.

Figures & tables

Panel of Figure 1 — Overview of the MF-PID framework. (a) Independent Agents/Samples: In the baseline, each sample evolves under a pre-computed drift; trajectories carry no information about one another, yielding no inter-agent coupling. (b) Interacting Agents: MF-PID promotes samples to interacting agents whose optimal drift $u^{(*)}_{t}$ depends self-consistently on the evolving population density $\rho_{t}$ through the effective potential $V^{(\mathrm{eff})}_{t}$ . (c) LQG: Two analytically tractable regimes— Linear–Quadratic–Gaussian (LQG), where the infinite-dimensional MF system reduces to Riccati ODEs with variance/mean decoupling; and Gaussian-Mixture (GM) representation for densities with Piece-Wise-Constant (PWC) in time protocol for $\beta_{t}$ results in an explicit – gradient-of-log-of-Gaussian mixture expression for the score function. (d) DR: Application to demand-response control of a building TCL fleet; MF coordination achieves 11.6 % (Scenario A) and 22.6 % (Scenario B) reductions in cumulative control energy $\mathcal{E}$ relative to the unguided IA baseline.

Panel of Figure 2 — Figure 2: Top: Initial and target distributions. Middle/Bottom: Sample trajectories (50 paths) for Scenarios A (middle) and B (bottom). Blue: occupied; orange: unoccupied. Black: ensemble mean; gray band: $\pm 1\sigma$ ; shaded rectangles: target $\pm\sigma_{k}^{(\mathrm{tar})}$ bands.

	Scenario A			Scenario B
Method	$\mathcal{E}_{\mathrm{occ}}$	$\mathcal{E}_{\mathrm{unocc}}$	Total	$\mathcal{E}_{\mathrm{occ}}$	$\mathcal{E}_{\mathrm{unocc}}$	Total
$\nu=0$	13.89	56.92	31.30	3.38	37.40	17.15
$\nu=\bar{m}$	13.43	53.59	29.68	2.63	34.36	15.47
$\nu=\nu^{\mathrm{MF}}$	14.72	46.73	27.67	3.21	28.07	13.27

Table 1: Per-mode control energy. MF guidance ( 11 ) redistributes effort from the occupied to the unoccupied sub-population, yielding net savings.

$d$	$\mathcal{E}/d$ (MF)	$\mathcal{E}/d$ (IA 0 )	$\mathcal{E}/d$ (IA m )	Saving	Time (s)
1	12.26	16.17	14.34	24.2%	88
2	13.07	16.96	15.13	23.0%	93
4	13.37	17.22	15.41	22.4%	105
8	13.57	17.40	15.59	22.0%	106
16	13.50	17.32	15.51	22.1%	152
32	13.57	17.19	15.42	21.1%	353

Table 2: Dimension sweep (

K=2

, Scenario B parameters). Per-zone energy

\mathcal{E}(1)/d

is invariant across

d=1

–

32

, confirming Theorem ( 11 )’s dimension-independence. Wall-clock time (single CPU,

B=4{,}000

particles,

n_{\mathrm{steps}}=2{,}500

) scales sub-cubically in this range; the asymptotic

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

Cholesky cost dominates beyond

d\gg 32

Figure 3: Multi-zone scalability ( $K=2$ ). (a) Per-zone cumulative energy $\mathcal{E}(1)/d$ vs. number of zones $d$ . MF (blue diamonds) is flat at $\approx 13.5$ ; IA( $\nu{=}0$ ) (orange) at $\approx 17.2$ . (b) MF energy saving (%) vs. $d$ : stable at 21–24% across the full range. (c) Wall-clock time on a single CPU core. The dashed line shows the asymptotic $\mathcal{O}(d^{3})$ trend (Cholesky of $d\times d$ covariance matrices per time step); for $d\leq 32$ overhead dominates and scaling is milder.

Panel of Figure 4 — Figure 4: Spatial mechanism at $d=10$ , $K=2$ . Left: Zone $\times$ time heatmap of ensemble mean temperature $\bar{x}_{j}(t)$ for MF (top) and IA( $\nu{=}0$ ) (bottom). Under MF each zone mean follows a near-linear arc; under IA high-displacement zones (top/bottom rows) show stronger curvature. Right: Per-zone cumulative energy. MF redistributes effort: high-displacement zones (e.g. zone 5) see the largest reduction ( $\approx 14\%$ ) while low-displacement zones change negligibly. Total MF saving: 22.3%.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Parameter	Value
Number of particles $B$	8,000
Euler–Maruyama steps $n_{\mathrm{steps}}$	2,500
PWC intervals $M$	8
$\beta$ -schedule $(\beta_{0},\gamma)$	$(12.0,\;0.65)$
MF tolerance $\epsilon$	$2\times 10^{-4}$
Target mixture $(\pi,m^{(\mathrm{tar})},\sigma^{(\mathrm{tar})})$	$(0.60,0.40)$ ; $(0.0,1.5)$ ; $(0.20,0.30)$

Appendix

Table 3: Numerical parameters for the demand-response simulations.

Figure 5: Scenario A : density snapshots at $t\in\{0.10,0.30,0.50,0.70,1.0\}$ for the three methods. The initially broad, overlapping distribution contracts and splits into the two target modes. Black solid: target $\rho^{(\mathrm{tar})}$ ; black dashed: initial $\rho^{(\mathrm{in})}$ .

Figure 6: Scenario A : trajectory ensembles (50 sample paths, colored by initial component). Blue: occupied; orange: unoccupied. Black curve: ensemble mean; gray band: $\pm 1\sigma$ . Shaded rectangles at $t{=}1$ : target $\pm\sigma_{k}^{(\mathrm{tar})}$ bands.

Figure 7: Scenario A : per-component analysis. (a) Per-component mean trajectories: the unoccupied mode (orange) must travel $\sim 4.5$ units while the occupied mode (blue) travels $\sim 1.0$ . (b) Cumulative per-component energy: the unoccupied mode absorbs most of the control cost; MF coordination reduces this asymmetry.

Figure 8: Scenario B (narrow initial). The MF energy saving increases to 22.6%.

Figure 9: Scenario B : density snapshots. Unlike Scenario A, the initial law is already bimodal; the two peaks translate and sharpen simultaneously.

Figure 10: Scenario B : trajectory ensembles (50 sample paths, colored by component). The two clusters remain well separated throughout the bridge.

Figure 11: Scenario A : score-field affine decomposition $u^{*}(t,x)\approx-S(t)\,x-s(t)$ . (a) Spring constant $S(t)$ : identical across methods, set by $\beta(t)$ and $\rho^{(\mathrm{tar})}$ . (b) Shift $s(t)$ : differs between methods; $|s^{\mathrm{MF}}|$ is largest at early times. (c) Affinity $R^{2}(t)$ : drops from $\sim 1.0$ to $\sim 0.03$ , marking the onset of genuinely non-affine Gaussian-mixture score structure. Vertical gray lines: $\beta$ -interval boundaries.

Panel of Figure 12 — Figure 12: Cross-scenario energy comparison. (Left) Total control energy $\mathcal{E}(1)$ for all three methods in Scenarios A and B. (Right) Cumulative energy curves. The MF saving increases from 11.6% (A) to 22.6% (B).

Figure 13: Self-consistent MF guidance $\nu^{(\mathrm{MF})}(t)$ for Scenarios A (blue) and B (orange), with their respective linear interpolants (dotted). By Theorem C.2 , the two curves coincide exactly in the continuous-time limit; the small residuals ( $\max|\nu^{(\mathrm{MF})}-\nu_{\mathrm{lin}}|=0.078$ for A, $0.030$ for B) are due to the PWC temporal discretisation ( $M=8$ intervals).

Panel of Figure 14 — Figure 14: Convergence of the fixed-point iteration to the linear interpolant. Left: Scenario A (12 iterations to tolerance $2\times 10^{-4}$ ). Right: Scenario B (10 iterations). (a) $\max|\Delta\nu|$ vs. iteration (dashed: tolerance). (b) Residual $\nu_{i}^{(\mathrm{MF})}-\nu_{i}^{\mathrm{lin}}$ at each PWC midpoint, confirming Theorem C.2 in the PWC limit.

Diffusion Models

Diffusion models learn to reverse gradual data corruption, enabling flexible generation, editing, forecasting, and inference.

See all 1,163 papers

Mean-Field Diffuser: Scaling Offline MARL to Thousands of Agents
May 28, 2026Wenhao Li, Xiangfeng Wang, Bo JinMulti-Agent System Optimization Mean Field Games
Tongji University · East China Normal University
One for All, All for One: Coordinated Multi-Agent Diffusion Steering via Stochastic Optimal Control
Oct 6, 2026Riccardo Barbano, Vincent Pauline, Runchang Li +6Diffusion Model Guidance Multi-Agent Coordination
UCL · TUM & MCML · CUHK +5
A Mean-Field Framework for Inference-Time Distributional Control of Diffusion Models
Aug 9, 2026Samuel Howard, Nikolas NüskenInference-Time Guidance Diffusion Models
Department of Statistics, University of Oxford · Department of Mathematics, King’s College London

May 28, 2026J/K move · Enter open · S save

Mean-Field Diffuser: Scaling Offline MARL to Thousands of Agents

Wenhao Li, Xiangfeng Wang, Bo Jin

Tongji University · East China Normal University

Multi-Agent System Optimization Mean Field Games Diffusion-Based Planning Mean-Field Theory Mean-Field MARL Multi-Agent Reinforcement Learning+1

Open paper AskPDF arXiv

Diffusion-based planning has achieved strong results in single-agent offline reinforcement learning, yet scaling to many-agent systems remains intractable due to the curse of dimensionality in the joint trajectory space. We introduce MF-Diffuser, a framework that lifts trajectory planning to the Wasserstein space of trajectory distributions, where the propagation of chaos ensures a small representative subset of agents captures the full population dynamics. Our approach features a value-weighted chaotic entropy objective that reconciles generative fidelity with return maximization, and a hierarchical coarse-to-fine strategy that progressively grows the agent population during denoising. We establish end-to-end suboptimality bounds with four interpretable terms, revealing that mean-field approximation error scales as

O(H^2/\sqrt{N})

while offline distribution shift provably does not grow with population size

N

, and prove the generated policy is an approximate mean-field Nash equilibrium with explicit convergence guarantees. Experiments on three mean-field RL benchmarks -- spanning stage games, sequential dynamics, and adversarial team competition -- show MF-Diffuser achieves the best return in the majority of settings, with the largest gains on suboptimal offline data and at extreme scales (

N \geq 10^3

Abstract

Diffusion Models

Explore similar work

Mean-Field Diffuser: Scaling Offline MARL to Thousands of Agents

One for All, All for One: Coordinated Multi-Agent Diffusion Steering via Stochastic Optimal Control

A Mean-Field Framework for Inference-Time Distributional Control of Diffusion Models