Composing and Steering Diffusion Models with Generalized Feynman-Kac Correctors

Wael Mattar; Daniel Freedman

cs.LGOct 29, 2026

Composing and Steering Diffusion Models with Generalized Feynman-Kac Correctors

Authors: Wael Mattar, Daniel Freedman

Organizations: School of Mathematical Sciences, Tel Aviv University, Israel.

Abstract

Diffusion models have become the leading paradigm in generative modeling, yet many applications require sampling from distributions that are not those of any single pretrained model, but are instead defined through combinations of several models, possibly tilted toward high values of a downstream reward. Naively combining the scores of the models yields incorrect samples, while existing corrections are restricted to specific families of targets. In this work, we introduce the generalized Feynman-Kac correctors (GFKC), a principled framework for sampling from an arbitrary target function of several pretrained diffusion models. We show that the geometric interpolation between an arbitrary initial function and the target function of the models satisfies a Feynman--Kac partial differential equation, and derive its potential in closed form. To evaluate this potential during generation, we further generalize the log-probability ODE to track the density of each model along trajectories driven by a different velocity field. The resulting algorithm operates entirely in the log domain, incorporates reward-based steering at no additional conceptual cost, and corrects the generation on the fly, without retraining or fine-tuning the models. We validate the method on analytic Gaussian models, where it closely matches the ground-truth targets, and on diffusion models trained on MNIST, where it samples from non-trivial combinations of the models and enables effective steering.

Figures & tables

Panel of Figure 1 — Figure 1: Sampling with the GFKC algorithm in the bivariate Gaussian setting. Each row corresponds to one choice of the terminal function $h_{1}$ : the asymmetric ratio of ( 18 ) on top and its symmetrized counterpart ( 19 ) on bottom. Left column: the ground-truth target densities, obtained by applying $h_{1}$ directly to the analytic densities ( 17 ). Middle column: the empirical densities generated by the GFKC algorithm, using $100{,}000$ particles per agent and simulating up to time $t=1$ . Right column: the relative pointwise error $|h_{1}-\hat{h}_{1}|/\max_{x}h_{1}$ between the ground-truth and the generated densities. Note the close agreement between the left and middle columns.

Figure 2: Convergence of the GFKC algorithm as a function of the number of particles per agent, measured by

\mathrm{KL}(h_{1}\,\|\,\hat{h}_{1})

and the

\ell_{2}

distance, for both choices of

h_{1}

Figure 3: Effect of steering on the symmetric target

h_{1}

of ( 19 ), with reward

r(x)=-\|x\|

and

\kappa=1

. Left: the analytic unsteered target density

h_{1}

. Middle: the empirical density generated by the GFKC algorithm with steering, using

100{,}000

particles per agent and simulating up to time

t=1

. Right: the pointwise difference between the steered empirical density and the unsteered target. Note the shift of mass toward the mode closer to the origin.

Figure 4: Effect of the steering parameter

\kappa

on the symmetric target

h_{1}

of ( 19 ), with reward

r(x)=-\|x\|

. Each panel shows the density generated by the GFKC algorithm for a different value of

\kappa

. As

\kappa

decreases, the steering effect becomes increasingly pronounced, and mass is progressively transferred toward the mode closer to the origin. Conversely, for very large

\kappa

, here

\kappa=10^{6}

, the steering term is negligible and the generated density coincides with the unsteered target

h_{1}

, as expected.

Panel of Figure 5 — Figure 5: Samples generated independently by each agent, without the GFKC algorithm. Left: $100$ images generated by the first agent, trained on the digits $\mathcal{S}_{1}$ . Right: $100$ images generated by the second agent, trained on the digits $\mathcal{S}_{2}$ .

Panel of Figure 6 — Figure 6: Results of the GFKC algorithm with the ratio target ( 18 ). Left: $h_{1}(p_{1},p_{2})$ . Right: $h_{1}(p_{2},p_{1})$ . Each panel displays the $100$ generated particles, $50$ per agent, sorted by their final log-weights in decreasing order, from the top-left to the bottom-right corner. Particles from the first agent (digits $1$ and $6$ ) are shown in cyan, and particles from the second agent (digits $1$ and $7$ ) in magenta. In the left panel, the highest weights are assigned to the digit $6$ , exclusive to the first agent, followed by the shared digit $1$ , while the digit $7$ receives the lowest weights. The right panel exhibits the symmetric behavior, with the digit $7$ ranked highest and the digit $6$ lowest.

Panel of Figure 7 — Figure 7: Histograms of the final log-weights of the particles generated by the GFKC algorithm, with the ratio target $h_{1}(p_{1},p_{2})$ (left) and the swapped target $h_{1}(p_{2},p_{1})$ (right). Colors correspond to the agents as in Figure 6 : cyan for the first agent (digits $1$ and $6$ ) and magenta for the second agent (digits $1$ and $7$ ). Solid curves show kernel density estimates of each histogram. In both cases, the log-weights of the two agents are clearly separated, with the agent holding the favored exclusive digit receiving the higher weights. Since the histograms are over log-weights, differences of several thousands translate into ratios of order $e^{1000}$ or more between the weights themselves, so that, after normalization, essentially all the mass is carried by the particles of the favored agent.

Figure 8: Results of the GFKC algorithm with the annealed sum target $h_{1}=(p_{1}+p_{2})^{2}$ (i.e., $\gamma=2$ ). The $100$ generated particles, $50$ per agent, are sorted by their final log-weights in decreasing order, from the top-left to the bottom-right corner, and colored by agent as in Figure 6 : cyan for the first agent (digits $1$ and $6$ ) and magenta for the second agent (digits $1$ and $7$ ). The highest weights are assigned to the shared digit $1$ , generated by both agents, while the exclusive digits $6$ and $7$ are interleaved among the lower weights.

Figure 9: Histograms of the final log-weights for the sharpened sum target

h_{1}=(p_{1}+p_{2})^{\gamma}

, with

\gamma\in\{1,2,3,10\}

, colored by agent as in Figure 6 . As

\gamma

increases, the range of the log-weights widens and the peak associated with the shared digit

1

shifts toward larger values, so that the separation between the highest-weighted particles and the rest becomes increasingly pronounced. In contrast to the ratio target, the log-weights of the two agents overlap throughout, and the highest-weighted particles originate sometimes from the first agent and sometimes from the second, reflecting that both agents contribute to the shared digit.

Figure 10: Results of the steered GFKC algorithm with the sum target $h_{1}$ of ( 20 ) with $\gamma=1$ , the reward $r$ of ( 21 ) and a small steering parameter $\kappa=10^{-6}$ . The $100$ generated particles are sorted by their final log-weights in decreasing order, from the top-left to the bottom-right corner, and colored by agent as in Figure 6 . Although $h_{1}$ favors the shared digit $1$ , the steering term dominates, and the highest weights are assigned to images of the digit $6$ , all generated by the first agent. As the weights decrease, the generated $6$ ’s become increasingly deformed, and their $\pi$ radians rotations bear progressively less resemblance to the mean image $\boxed{9}$ . They are followed by the digit $1$ , while the particles of the second agent, which cannot generate the digit $6$ , receive the lowest weights.

Figure 11: Left: the highest-weighted particle of the steered GFKC algorithm, a digit $6$ generated by the first agent, rotated by $\pi$ radians. Right: the mean image $\boxed{9}$ of the digit $9$ over the MNIST training set, which defines the reward $r$ of ( 21 ). The rotated sample closely resembles the mean image, illustrating how the steering guides the generation toward high values of the reward.

Figure 12: The reward values against the final log-weights for the

100

particles of Figure 10 , colored by agent. The two quantities are strongly positively correlated, reflecting the dominance of the steering term for the small

\kappa=10^{-6}

. The cluster of roughly

30

points in the top-right corner, combining the highest rewards with the highest log-weights, corresponds to the images of the digit

6

, all generated by the first agent. The particles of the second agent, which cannot generate the digit

6

, are concentrated in the lower-left region.

Figure 13: Highest-weighted image of the steered GFKC algorithm for decreasing values of $\kappa$ . For large $\kappa$ , the top image depicts the digit $1$ , favored by $h_{1}$ , while for small $\kappa$ the steering dominates and the digit $6$ emerges. The critical shift occurs between the third and fourth panels, at around $\kappa\approx 0.0015$ .

Diffusion Model Sampling

Diffusion model sampling studies how to generate high-quality samples efficiently by designing reverse processes, solvers, and inference-time controls.

See all 82 papers

Error-Corrected Inference-Time Scaling for Imperfect Diffusion Models
Oct 1, 2026Zuokai Wen, Louis Grenioux, Weinan E +1Markov Chain Monte Carlo Diffusion Model Sampling
School of Artificial Intelligence, Shanghai Jiao Tong University · Center for Computational Mathematics, Flatiron Institute · School of Mathematical Sciences, Peking University
General and Efficient Steering of Diffusion Models
Feb 11, 2026Qingsong Wang, Mikhail Belkin, Yusu WangDiffusion Model Guidance Activation Steering
Halıcıoğlu Data Science Institute, University of California San Diego, La Jolla, CA, USA
Correction Training for Diffusion Models
Oct 9, 2026Donghoon Ahn, Runqian Wang, Alexei A. Efros +2Diffusion-Based Image Generation Diffusion Models
University of California, Berkeley · Reve

Oct 1, 2026J/K move · Enter open · S save

Error-Corrected Inference-Time Scaling for Imperfect Diffusion Models

Zuokai Wen, Louis Grenioux, Weinan E, Jiequn Han

School of Artificial Intelligence, Shanghai Jiao Tong University · Center for Computational Mathematics, Flatiron Institute · School of Mathematical Sciences, Peking University

Markov Chain Monte Carlo Diffusion Model Sampling Inference-Time Scaling Diffusion Models

Open paper AskPDF arXiv

Inference-time scaling adapts pretrained diffusion models to new sampling tasks without additional training. Existing methods rely primarily on Monte Carlo sampling with more particles, yet are premised on the pretrained model being exact. In practice, data and training limitations make the model imperfect, and these methods inherit its error. More particles reduce Monte Carlo error but cannot remove the mismatch between the endpoint and the desired target or the error in tracking the prescribed probability path. We introduce the Energy-based Feynman-Kac Corrector (EBFKC), a framework for energy-based diffusion models that corrects these errors on the fly given a reference energy. We first derive Feynman-Kac dynamics that track a prescribed path exactly in the continuous-time population limit even when the model is imperfect, and approximate these dynamics using sequential Monte Carlo with variance-controlling guidance. To remove the endpoint mismatch, we use the pretrained energy as a surrogate along the diffusion path and progressively incorporate the discrepancy between the learned and target terminal energies. Experiments on Gaussian mixture models, particle systems, alanine dipeptide, and alanine tetrapeptide show that our method closely matches target distributions and molecular free-energy profiles under annealing and reward tilting, whereas standard inference-time scaling baselines retain substantial sampling errors.

Abstract

Diffusion Model Sampling

Explore similar work

Error-Corrected Inference-Time Scaling for Imperfect Diffusion Models

General and Efficient Steering of Diffusion Models

Correction Training for Diffusion Models