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
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 ut(∗) depends self-consistently on the evolving population density ρt through the effective potential Vt(eff) . (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 β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 E relative to the unguided IA baseline.
Figure 1: MF vs. IA controls in the scalar TCL example. Top left: Trajectory ensembles for MF (blue), IA with mˉ=0 (orange), and IA with mˉ=m(tar) (green); thin curves: sample paths, thick: analytic means, shaded band: target variance. Top right: Instantaneous control power P(t) . Bottom left: Cumulative control energy E(t)=∫0tP(u)du . Bottom right: Mean trajectories mt . All curves are obtained analytically.
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: ±1σ ; shaded rectangles: target ±σk(tar) bands.
Scenario A
Scenario B
Method
Eocc
Eunocc
Total
Eocc
Eunocc
Total
ν=0
13.89
56.92
31.30
3.38
37.40
17.15
ν=mˉ
13.43
53.59
29.68
2.63
34.36
15.47
ν=ν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
E/d (MF)
E/d (IA 0 )
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 E(1)/d is invariant across d=1 – 32 , confirming Theorem ( 11 )’s dimension-independence. Wall-clock time (single CPU, B=4,000 particles, nsteps=2,500 ) scales sub-cubically in this range; the asymptotic O(d3) Cholesky cost dominates beyond d≫32 .
Figure 3: Multi-zone scalability ( K=2 ). (a) Per-zone cumulative energy E(1)/d vs. number of zones d . MF (blue diamonds) is flat at ≈13.5 ; IA( ν=0 ) (orange) at ≈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 O(d3) trend (Cholesky of d×d covariance matrices per time step); for d≤32 overhead dominates and scaling is milder.
Figure 4: Spatial mechanism at d=10 , K=2 . Left: Zone × time heatmap of ensemble mean temperature xˉj(t) for MF (top) and IA( ν=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 ( ≈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 nsteps
2,500
PWC intervals M
8
β -schedule (β0,γ)
(12.0,0.65)
MF tolerance ϵ
2×10−4
Target mixture (π,m(tar),σ(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∈{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 ρ(tar) ; black dashed: initial ρ(in) .
Figure 6: Scenario A : trajectory ensembles (50 sample paths, colored by initial component). Blue: occupied; orange: unoccupied. Black curve: ensemble mean; gray band: ±1σ . Shaded rectangles at t=1 : target ±σk(tar) bands.
Figure 7: Scenario A : per-component analysis. (a) Per-component mean trajectories: the unoccupied mode (orange) must travel ∼4.5 units while the occupied mode (blue) travels ∼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)≈−S(t)x−s(t) . (a) Spring constant S(t) : identical across methods, set by β(t) and ρ(tar) . (b) Shift s(t) : differs between methods; ∣sMF∣ is largest at early times. (c) Affinity R2(t) : drops from ∼1.0 to ∼0.03 , marking the onset of genuinely non-affine Gaussian-mixture score structure. Vertical gray lines: β -interval boundaries.
Figure 12: Cross-scenario energy comparison. (Left) Total control energy 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 ν(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∣ν(MF)−νlin∣=0.078 for A, 0.030 for B) are due to the PWC temporal discretisation ( M=8 intervals).
Figure 14: Convergence of the fixed-point iteration to the linear interpolant. Left: Scenario A (12 iterations to tolerance 2×10−4 ). Right: Scenario B (10 iterations). (a) max∣Δν∣ vs. iteration (dashed: tolerance). (b) Residual νi(MF)−νilin at each PWC midpoint, confirming Theorem C.2 in the PWC limit.
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(H2/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≥103).
Deep generative models often produce structured outputs composed of interacting components. Modelling these outputs with a single model requires learning both the component distributions and their interactions. We pursue a modular alternative: reuse independently trained component generators and learn only how to coordinate them to produce coherent structured outputs. Our framework, Coordinated Multi-Agent Diffusion Steering (CMDS), treats frozen pretrained diffusion models as reusable generative primitives and coordinates their reverse processes through a learned control. We formulate coordination as a stochastic optimal control problem, balancing an assembly-level reward that specifies the desired properties of the combined output against deviations from the pretrained dynamics. The learned control amortises this optimisation, allowing reuse across new task instances. Experiments show that CMDS can recover a known target distribution, satisfy different spatial constraints with the same trained control, and recover individual sources from degraded mixtures. Across multi-agent maze navigation, articulated robot planning, and text-conditioned human motion, CMDS turns frozen models into coordinated multi-agent generators.
Diffusion models are increasingly used as controllable samplers, whose generations can be steered at inference time according to a chosen reward function. While such rewards are typically defined on individual samples, for many applications it is desirable to steer according to distribution-level rewards, for example to calibrate with population-level information or to encourage diversity. In both cases, simply incorporating the reward gradient into the dynamics, while often effective, comes with few theoretical guarantees on the sampled distribution. For pointwise rewards, recent work has therefore sought to develop a principled framework for targeting a prescribed tilted distribution using particle reweighting. However, an analogous theoretically-grounded approach for distributional rewards is currently lacking. In this work, we formulate inference-time distributional control as targeting a tilted measure under a mean-field framework, and derive a weighted interacting particle scheme to target it in a principled manner. Our framework recovers pointwise-reward steering as a special case, while providing a theoretical foundation for existing batch-level steering methods. Empirically, we verify that the procedure correctly targets the prescribed distribution in tractable low-dimensional settings, and investigate its behaviour in higher-dimensional protein conformation tasks.
Samuel Howard, Nikolas Nüsken
Department of Statistics, University of Oxford · Department of Mathematics, King’s College London