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

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

Abstract

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.

1. Introduction

Diffusion and flow-based generative models generate novel samples from a data distribution by learning time-dependent dynamics that transport an initial Gaussian distribution p0 to the target p1. Such models have achieved state-of-the-art performance across a wide range of domains, including images, molecules, and videos.

Inference-time steering methods aim to direct a pretrained generative model towards higher-quality samples or samples that fulfil certain task-specific properties, without requiring further training. Most existing inference-time steering methods aim to improve samples with respect to a pointwise reward function r: Rd → R defined on individual generated samples. Early steering approaches generally incorporated the reward gradient ∇r(Xt) into the inference dynamics, to 'push' the generations towards areas of higher reward. While effective at improving sample quality, these heuristic modifications to the dynamics come with a poor theoretical understanding of the sampled distribution.

This limitation has motivated a large number of recent works aiming to place inference-time steering on a firmer theoretical footing. To do so, the dominant paradigm aims to sample from the KL-regularised objective, arg maxµ EX∼µ[r(X)] − KL(µ∥p1), which corresponds to the exponentially-tilted distribution µ∗1(x1) ∝ p1(x1)er(x1). Since gradient-steering alone does not sample from this target tilted measure, recent works have sought to correct for the discrepancy using Feynman-Kac particle reweighting schemes, enabling correct tracking of the desired distributional path via SMC resampling procedures.

Following these recent advances, steering according to pointwise-defined reward functions is now well-understood and based upon a principled theoretical foundation. However, many important control objectives are inherently distributional rather than pointwise. In many applications, rather than aiming for individual samples to score highly, one may wish to shape the overall generated distribution according to global properties, such as calibration to experimental observations, balancing across modes, or encouraging diversity. To formalise such distributional-steering, we will instead consider optimising according to a measure-defined reward function R: P(Rd) → R, arg maxµ R(µ) − KL(µ∥p1).

A number of existing works have incorporated steering terms that operate at the distributional level. For example, Corso et al. (2024); Kirchhof et al. (2025); Lam et al. (2026) steer towards generating diverse samples, Maddipatla et al. (2025) steer according to experimental observations, Saravanan et al. (2026) steer to increase expected information gain, and Sani et al. (2026) guide with MMD gradients. However, as in the pointwise-reward case, such approaches are largely heuristic in nature and it is not clear precisely what distribution these procedures target. Unlike in the pointwise-reward setting, a theoretically-grounded framework and approach for distributional inference-time steering of diffusion models is currently lacking.

Contributions

In this work, we address this gap by reconciling these two currently distinct literatures through a mean-field framework, allowing us to place distributional steering on a similarly principled footing to the pointwise case. In particular, we make the following contributions:

• We formalise distribution-level steering (2) as aiming to sample from a tilted distribution for a measure-defined reward function R, in which the tilt now depends on the measure itself. This gives a distributional analogue of the reward-tilted targets used in the pointwise case.

• Using this mean-field framework, we derive weighted McKean-Vlasov dynamics whose marginals correctly track the desired tilted path. From this, we develop a novel weighted interacting-particle sampling procedure that enables principled distribution-level steering.

• We show that our framework recovers standard pointwise reward steering as a special case, and provides a principled counterpart to existing gradient-steering methods for distributional control. Empirically, we validate that our procedure correctly targets the tilted distribution in verifiable low-dimensional settings, and examine its behaviour on higher-dimensional protein conformation tasks relative to standard gradient-steering approaches.

2. Background: Inference-Time Steering for Pointwise Rewards

Given a reward r defined on the generated samples x1, inference-time steering methods aim to modify the generation procedure to instead sample from the reward-tilted distribution, which is the maximiser of the KL-regularised objective, µ∗(x) ∝ p1(x)er(x). Inference-time steering differs from the fine-tuning paradigm, as there is no additional training. Here, the dynamics are modified for the currently-chosen reward function r, and different choices of r can be used without having to perform extra training.

The aim of inference-time steering methods is to 'push' the particles towards areas of higher reward during the generation, while performing SMC-style reweighting of the particles to ensure they trace out a prescribed path of marginal distributions from p0 to µ∗. To do so, one chooses a time-dependent reward function rt which satisfies r0 = 0 and r1 = r; we aim to trace out the tilted path of marginals µ∗t(x) ∝ pt(x)ert(x), which will interpolate from p0 to the target tilted distribution µ∗ ∝ p1(x)er(x). A common approach is then to modify the dynamics by incorporating a reward-gradient steering term, dXt = [bt(Xt) + εtst(Xt) + εt∇rt(Xt)]dt + √(2εt)dBt. As the original reward function r is defined on data, a natural choice is rt(xt) = t·r(x̂1(t, xt)), where x̂1(t, xt) is an estimate of the terminal point x1 given the current point (t, xt) in the trajectory.

While the additional steering term pushes the particles towards areas of higher reward, the dynamics do not in general follow the desired tilted marginal path µ∗t. To correct for the discrepancy, it is necessary to adjust the measure using particle weights, which can then be used in an SMC-style resampling scheme to target µ∗ in a theoretically-principled manner. The log-weights Ait evolve according to the dynamics dAt = [bt(Xt) · ∇rt(Xt) + ∂trt(Xt)]dt.

3. A Mean-Field Framework for Inference-Time Distributional Control

3.1. Distributional-Reward Measure Tilting

Consider now that we want to tilt towards a desirable distributional property of the generated measure, quantified by a measure-dependent reward function R: P(Rd) → R defined on the terminal distribution µ1. We can formalise this as aiming to sample from µ∗ = arg maxµ R(µ) − KL(µ∥p1), assuming for the moment that the maximiser is unique. To understand the properties of µ∗, let us introduce the functional derivative Ψ(x, µ):= δR/δµ(x, µ), which quantifies the change of the reward R when the measure µ is perturbed. Using this object, we can characterise µ∗ as an exponential tilt, mirroring the pointwise case.

Proposition 3.1 (First-order optimality condition). Assuming mild regularity properties and concavity of R, the maximiser in (2) is unique and satisfies µ∗(dx) = (1/Z)eΨ(x,µ∗)p1(x)dx, for normalising constant Z = ∫Rd eΨ(x,µ∗)p1(x)dx.

The proof, essentially setting the functional derivative of the objective R(µ) − KL(µ∥p1) to zero, can be found in Appendix D.1. In contrast to the pointwise case, the tilt is implicit: the target measure µ∗ appears on both sides of the equation, and the tilting potential Ψ(x, µ∗) is itself measure-dependent. Proposition 3.1 therefore places the problem in a mean-field regime: the optimality condition is a self-consistency relation, in which the target measure is characterised by a potential that depends on the measure itself.

As in the standard case, let us fix a family of measure-defined reward functions Rt: P(Rd) → R interpolating from R0 = 0 to R1 = R. Extending the above, we can define the curve of distributions µ∗t(dx) = (1/Zt)eΨt(x,µ∗t)pt(x)dx, t ∈ [0, 1], tracing from p0 to µ∗1. As in the pointwise case, the measures µ∗t are defined implicitly and self-consistently through the corresponding measure-dependent potentials Ψt(x, µt), similarly defined as functional derivatives of the time-dependent rewards Rt. We note that these intermediate measures µ∗t satisfy analogous variational principles, namely µ∗t = arg maxµ Rt(µ) − KL(µ∥pt). Following the pointwise case, in this work we will take Rt = t·R(x̂1(·, t)#µ), using the denoiser estimate x̂1.

3.2. Idealised Dynamics

We will use this mean-field framework to develop a probabilistically-principled inference-time steering procedure that traces out this distributional path µ∗t during the generation process. As in the standard setting, it is again intuitive to include a gradient-steering term to 'push' the generated distribution in the direction of higher reward during the inference procedure. To do so, one can consider the modified dynamics, dXt = [bt(Xt) + εtst(Xt) + εt∇xΨt(Xt; µt)]dt + √(2εt)dBt. In fact, such gradient-steering resembles many existing works for batch-level reward steering that incorporate gradient terms into the inference dynamics, for example for repulsive potentials, ensemble likelihoods given experimental data, expected information gain, and maximum mean discrepancy to a reference dataset.

Just as in the standard case, merely incorporating a gradient steering term into the dynamics does not in general follow the target marginals µ∗t, nor come with any strong theoretical understanding of the generated distribution. It is therefore again necessary to incorporate a Feynman-Kac style reweighting scheme to correct for this discrepancy. This setting differs from the standard pointwise case, as the particles now interact during the inference dynamics due to the dependence of the steering term on the evolving distribution. Fortunately, by appealing to our mean-field framework we can derive the corrective potentials under this interacting particle behaviour. This yields our main result, in which we construct mean-field type, or McKean-Vlasov, dynamics whose coefficients are chosen so that the corresponding law follows the desired path µ∗t.

Theorem 3.2 (Weighted McKean-Vlasov Dynamics). Consider the dynamics dXt = [bt(Xt) + εtst(Xt) + εt∇xΨt(Xt; µt)]dt + √(2εt)dBt, X0 ∼ p0, and dAt = [bt(Xt) · ∇xΨt(Xt; µt) + Ψ̇t(Xt)]dt, A0 = 0, with the weighted measure µt ∝ Law(Xt)eAt, and Ψ̇t(x):= d/dt Ψt(x; µt). Then, under mild conditions, µt = µ∗t, for all t ∈ [0, 1].

Note that these dynamics strongly resemble the pointwise setting. However, a key difference is that now the time-derivative term ∂tr(·) has been replaced with the term Ψ̇t(·) = d/dt Ψt(·, µt); here, the position x is kept fixed but the time-dependence now arises through both Ψt and the measure evolution µt. Again in contrast to the standard case, this induces implicit behaviour in the log-weight updates—their dynamics involve the change in the measure µt, but the measure evolution is itself dependent on the change in the weights. We remark that the term Ψ̇t can be equivalently characterised as solving the following implicit expression, which we will later use in our inference procedure.

Proposition 3.3. Define the second variation of the reward functional, Φt(x, y, µ) = δ2Rt/δµ2(x, y, µ), and a mean-centred version Φ̃t(x, y, µ) = Φt(x, y, µ) − ∫Φt(x, z, µ)µ(dz). Then, under the target flow (10), the time derivatives Ψ̇t satisfy the following implicit equality: Ψ̇t(x) = ∂tΨt(x, µ∗t) + ∫[Φ̃t(x, y, µ∗t)Ψ̇t(y) + bt(y) · ∇yΦ̃t(x, y, µ∗t) + Φ̃t(x, y, µ∗t)bt(y) · ∇yΨt(y, µ∗t)]dµ∗t(y).

3.3. Method

To turn the idealised dynamics of Theorem 3.2 into a finite-particle algorithm, we approximate the McKean-Vlasov system with an ensemble (Xti, wti)Ni=1, where the weighted empirical measure µ̂t = ΣNj=1 wtjδXjt approximates µt. To do so, we need the following terms.

Gradient-Steering Term ∇xΨt(x; µt): The gradient-steering term ∇xΨt(x; µt) in the positional updates can be easily computed with autodifferentiation, evaluated on the current empirical measure µ̂t with weights wtj ∝ eAjt. When one can analytically compute the first variation (as is often the case; see Appendix A.2), then one can use a finite-sample approximation Ψt(Xti, µ̂t) and then differentiate with respect to the first argument. Alternatively, one could approximate the steering term using (wti)−1∇XtiRt(µ̂t), by directly evaluating the reward on the empirical measure and differentiating with respect to the positions Xti.

Time-Derivative Term Ψ̇t(x): The more challenging term in the dynamics is the time-derivative Ψ̇t(x) = d/dt Ψt(x, µt) due to its implicit behaviour. We provide two different approaches for computing this term. The first involves solving a linear system based on the implicit expression in (13), which involves only tractable terms, and the second solves for Ψ̇t(x) through a lightweight Picard iteration scheme; we describe these in Algorithms 3 and 4 respectively. We remark that the cost of such procedures is low, and the computational cost is often dominated by the network evaluations rather than these lightweight solvers. We found that the choice of solver for Ψ̇t(x) had little effect on the empirical performance of the method and we therefore primarily used the fixed-point iterative solver for simplicity, but we remark that the first approach may be useful in certain settings.

Generalising the Pointwise Setting: We remark that the pointwise-reward setting is recovered from our framework by choosing R(µ) = ∫r(x)dµ(x). In this case, the first variation is Ψ(x, µ) = r(x), so our mean-field steering procedure recovers the standard procedure from Section 2.

SMC Resampling: As in the pointwise-reward Feynman-Kac steering approaches, we use an SMC-style resampling scheme to avoid weight degeneracy. Any such scheme could be used; in our experiments we use residual resampling at fixed intervals.

Algorithm: We outline the inference procedure for a single batch of N particles in Algorithm 1, which will approximate the target tilted measure µ∗ as the batch size N goes to infinity. In practice, the batch size is limited by computational constraints; we therefore also provide a multi-batch procedure in Algorithm 2 in Appendix B.

Convergence in N: Under mild assumptions on the dynamics and the Ψ̇t solver, we can show convergence of the finite-particle system to the idealised dynamics (12). In particular, for δN quantifying the approximation error when solving for…

Based on the paper, here are the specific improvements you can make to AI systems and what the improved system can do:

1. Add a distributional-reward steering module to diffusion-based generative models

• Implement the weighted McKean-Vlasov dynamics (Theorem 3.2) as a plug-in inference-time controller.

• Use the interacting-particle scheme (Algorithm 1) with SMC resampling to correct for distributional drift.

• Include a fixed-point or Picard solver for the implicit time-derivative term Ψ̇t (Algorithms 3 and 4).

2. Replace heuristic batch-level steering with a principled mean-field correction

• For existing methods that use gradient steering (e.g., MMD, repulsive potentials, ensemble likelihoods), augment them with the Feynman-Kac reweighting term derived in the paper.

• This ensures the generated distribution converges to the prescribed tilted measure µ∗, rather than an ad-hoc distribution.

3. Enable calibration to population-level or experimental observations

• Use distributional rewards (e.g., MMD to a target ensemble, log-likelihood of crystallographic densities, or DEER-derived state proportions) to steer generations.

• The system will explicitly target the distribution that balances the reward and the KL divergence to the base model, not just push individual samples.

4. Support diversity and mode-balancing objectives

• Define rewards that penalize mode collapse or encourage coverage (e.g., negative MMD to a uniform reference, or entropy-based rewards).

• The mean-field framework ensures the final ensemble matches the desired diversity profile, not just the mean or mode.

5. Provide theoretical guarantees for inference-time control

• Use the convergence result (Theorem 3.4) to certify that the finite-particle approximation approaches the target distribution as batch size grows.

• This makes the system suitable for safety-critical or scientific applications where distributional accuracy is essential.

• Generate protein conformational ensembles that match experimentally observed distributions (e.g., HIV-1 protease flap states, adenylate kinase bond angles, crystallographic electron densities) with a quantifiable accuracy guarantee.

• Steer diffusion models toward diverse, balanced outputs without retraining, by specifying a distributional reward (e.g., match this reference distribution or maximize entropy).

• Calibrate generative models to population-level statistics (e.g., from DEER, X-ray, or MD simulations) in a principled way, ensuring the output distribution is the exact tilted target, not a heuristic approximation.

• Provide a drop-in replacement for existing gradient-steering methods (e.g., MMD guidance, ensemble likelihood steering) that corrects the distributional bias, making those methods theoretically sound.

• Operate in high-dimensional settings (e.g., proteins, molecules) where pointwise rewards are insufficient, by handling measure-dependent objectives with interacting particles and resampling.

• Offer a tunable trade-off between matching the reward and staying close to the base distribution, via the λ parameter, with explicit knowledge of the resulting distribution.

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.