PaperlayerView exact revision on arXiv

Source-linked AI summary

Feynman-Kac Correctors in Diffusion: Annealing, Guidance, and Product of Experts

Marta Skreta, Tara Akhound-Sadegh, Viktor Ohanesian, Roberto Bondesan, Alán Aspuru-Guzik, Arnaud Doucet, Rob Brekelmans, Alexander Tong, Kirill Neklyudov

arXiv:2503.02819v2cs.LG

TL;DR

Diffusion models lack principled tools for controlling inference-time distributions, especially when combining pretrained models or using classifier-free guidance. The paper introduces Feynman-Kac Correctors with weighted SDEs and SMC resampling, demonstrating temperature annealing, improved multi-objective molecule generation, and improved image guidance.

  • Problem

    Existing inference-time score modifications, including classifier-free guidance, do not generally preserve control over intermediate marginal distributions.

  • Method

    Feynman-Kac Correctors derive weighted SDEs from relevant PDE terms and simulate them with Sequential Monte Carlo resampling for target distributions from pretrained diffusion models.

  • Results

    The framework demonstrates inference-time temperature annealing, improved multi-objective molecule generation, and improved classifier-free guidance for image generation.

  • Takeaways & Limitations

    FKCs provide fine control over diffusion sample distributions for compositional generation, classifier-free guidance, and annealed amortized samplers.

  • Takeaways & Limitations

    The framework assumes access to an exact score model.

Abstract

from arXiv · show

While score-based generative models are the model of choice across diverse domains, there are limited tools available for controlling inference-time behavior in a principled manner, e.g. for composing multiple pretrained models. Existing classifier-free guidance methods use a simple heuristic to mix conditional and unconditional scores to approximately sample from conditional distributions. However, such methods do not approximate the intermediate distributions, necessitating additional `corrector' steps. In this work, we provide an efficient and principled method for sampling from a sequence of annealed, geometric-averaged, or product distributions derived from pretrained score-based models. We derive a weighted simulation scheme which we call Feynman-Kac Correctors (FKCs) based on the celebrated Feynman-Kac formula by carefully accounting for terms in the appropriate partial differential equations (PDEs). To simulate these PDEs, we propose Sequential Monte Carlo (SMC) resampling algorithms that leverage inference-time scaling to improve sampling quality. We empirically demonstrate the utility of our methods by proposing amortized sampling via inference-time temperature annealing, improving multi-objective molecule generation using pretrained models, and improving classifier-free guidance for text-to-image generation. Our code is available at https://github.com/martaskrt/fkc-diffusion.

1. Introduction

Diffusion models are widely used, but inference-time score modifications for composition and guidance can lose control over intermediate marginal distributions. The paper introduces Feynman-Kac Correctors and weighted SDE/SMC methods to sample prescribed annealed, geometric-averaged, or product distributions from pretrained models.

  • Motivation: Diffusion models generate samples by learning scores of noise-corrupted distributions and reversing the corruption process.They are used across image generation, natural language, and protein simulation.
  • Motivation: Inference-time score combinations support concept composition and classifier-free guidance but lose control over generated samples’ marginal distributions.Additional corrector methods are therefore used to better match prescribed intermediate distributions.
  • Contribution: The paper introduces Feynman-Kac Correctors for efficient and principled sampling from annealed, geometric-averaged, or product distributions derived from pretrained diffusion models.The approach accounts for additional terms arising when manipulating generated-sample distributions.
  • Contribution: The proposed framework constructs weighted SDEs and uses Sequential Monte Carlo resampling to correct simulated batches toward intermediate target distributions.The correction terms cover heuristic schemes including classifier-free guidance.
  • Applications: For unnormalized densities, FKC enables sampling across temperatures without retraining and supports high-temperature learning followed by lower-temperature inference.The paper also reports enhanced multi-property molecule generation and classifier-free guidance for image generation.

2. Background

Diffusion generation can be described by a reverse-time denoising SDE whose marginals evolve according to associated PDEs. Feynman-Kac formulations add time-dependent reweighting, enabling weighted simulation schemes for target marginals beyond the standard forward-noising path.

  • Diffusion Models: Diffusion models train by gradually corrupting data with a noising SDE and generate samples by simulating the corresponding process in reverse time.The reverse process starts from an approximately standard normal distribution and uses learned score functions.
  • Diffusion Models: The forward noising SDE uses drift, time-varying noise scale, and a Wiener process to move the data density toward standard normality.The drift is usually linear, with f_τ(x_τ) = α_τx_τ.
  • PDE Formulations: Continuity, diffusion, and reweighting equations describe density evolution caused by transport, Brownian motion, and time-dependent sample weights.The standard SDE combines flow and diffusion terms through its Fokker-Planck PDE.
  • Feynman-Kac PDEs: The Feynman-Kac PDE is simulated by first evolving samples with an SDE and then reweighting them, producing a density of weighted samples distinct from the ordinary SDE density.Self-normalized importance sampling estimates test-function expectations and converges as the particle count K approaches infinity.
  • Simulation Transformations: Flow and diffusion terms can be translated into reweighting factors when an exact score function is available.These transformations motivate the weighted SDE constructions used to simulate marginals other than the forward noising distributions.

3. Modifying Diffusion Inference using Feynman-Kac Correctors

The paper modifies diffusion inference by deriving Feynman-Kac-corrected weighted SDEs for annealed, geometric-averaged, and product target distributions. These corrections support principled sampling with pretrained models, including CFG and product-of-experts generation.

  • Target distributions: Pretrained diffusion models can be combined or modified at inference time to target distributions formed from their marginals.The paper considers models trained on different datasets or with different conditioning, then constructs modified targets from them.
  • Motivation: Heuristic score simulation generally fails to sample the prescribed intermediate marginals, except in special cases.This motivates adding Feynman-Kac correction terms rather than relying on the target score alone.
  • Feynman-Kac construction: The authors derive weighted SDEs by converting the target-distribution PDE into a Feynman-Kac PDE and collecting residual terms into weights.Conversion rules organize the transformations for annealed and product densities, while resampling schemes simulate the resulting weighted SDEs.
  • Computational considerations: For linear drifts, the common heuristic can yield Feynman-Kac weights estimated without additional computational overhead.The paper focuses on these schemes in its examples because constant divergence terms cancel or can be ignored.
  • Classifier-Free Guidance: Classifier-free guidance is treated as geometric averaging, and FKC supplies consistent trajectory reweighting for this target.The resulting framework includes a continuous family of weighted SDEs for geometric-average distributions and combinations with product-of-experts targets.
  • Annealing and Product of Experts: Annealed targets use densities proportional to q_t(x)^β, while product-of-experts targets favor samples supported by multiple models and improve multi-condition molecule generation.The paper derives weighted SDEs for annealing and products, and reports improved molecule generation when multiple conditions must be satisfied simultaneously.

4. Resampling Methods

The paper uses particle reweighting and Sequential Monte Carlo to simulate Feynman-Kac PDEs while controlling weight variance and preserving sample diversity. It also interprets reweighting as a Markov jump process that enables alternative simulation algorithms.

  • Sequential Monte Carlo: Full-trajectory importance weights may accumulate error and have high variance, motivating intermediate resampling of particle batches.The simplest approach simulates independent weighted trajectories and applies self-normalized importance sampling only at the end.
  • Sequential Monte Carlo: Resampling over an active interval and setting weights to zero outside it improves sample quality while preserving diversity.Within the interval, particles are resampled at each step using systematic sampling proportional to exponentiated weight increments.
  • Sequential Monte Carlo: Systematic resampling preserves all particles when weights are uniform.This makes it an attractive selection mechanism for small discretization steps and relatively low-variance weights.
  • Jump Process Interpretation: Feynman-Kac reweighting can be reformulated as a Markov jump process with a state-dependent rate and transition kernel.Jump rates govern event frequency, while the transition kernel selects the next state; the forward equation captures probability inflow and outflow.
  • Jump Process Interpretation: In the jump formulation, events trigger preferentially from low-excess-weight states and transition toward high-excess-weight states.This provides an efficient population-improvement mechanism and supports alternative algorithms such as birth-death exponential clocks.

5. Empirical Study

Across image generation, Boltzmann sampling, and molecule design, FKC resampling improves agreement with target distributions or downstream quality, while introducing computational or diversity trade-offs.

  • 5.1. Image Generation with EDM2: FKC resampling improves both ImageReward and CLIP score over vanilla CFG, with benefits observed even at K = 2 and best performance at K = 8.The comparison uses CLIP and ImageReward averaged over 10,000 images across noise levels and sampling steps.
  • 5.2. Samplers from the Boltzmann Density: Target score + FKC performs best on a highly multimodal tractable distribution, while tempered noise tends to drop modes.FKC outperforms SDE-only simulation in both target score and tempered noise settings.
  • 5.2. Samplers from the Boltzmann Density: Both tempered noise+FKC and target score SDE+FKC outperform DEM directly trained at lower temperatures for TS ∈[2.0, 0.8].Tempered noise+FKC performs best at higher temperatures, whereas target score SDE+FKC performs best at lower temperatures; the starting temperature is TL = 2.
  • 5.3. Multi-Target Structure-Based Drug Design: Target noise SDE with β > 0.5 yields better average docking scores than DualDiff and TargetDiff, with FKC improving those scores further.The evaluation covers 100 protein pairs and averages results across tasks.
  • 5.3. Multi-Target Structure-Based Drug Design: FKC achieves the lowest maximum docking score and the highest fraction of molecules better than known binders, indicating stronger simultaneous binding across protein pairs.Lower maximum docking scores are better; the reference molecules’ mean docking score is −7.915±2.841.
  • 5.3. Multi-Target Structure-Based Drug Design: FKC resampling improves generation of small-molecule SMILES satisfying multiple functional properties, with a slight sacrifice in diversity and uniqueness.The additional latent-diffusion analysis reports this multi-property improvement, while the diversity trade-off is observed in the molecule-design experiments.

6. Related Work

Related work connects FKC to predictor-corrector methods, energy-based composition, annealing, Jarzynski-style identities, and recent SMC applications, while distinguishing its guidance-weight correction.

  • SMC for diffusion models: Prior diffusion-model SMC methods address reward-guided generation, conditional generation, inverse problems, and discrete diffusion settings.These works establish SMC as a broadly used tool for diffusion sampling.
  • Energy-based and Boltzmann sampling: Energy-based approaches use learned energies or energy-based scores with MCMC or SMC to sample annealed or product densities.The cited methods apply energy parameterizations to compositional generation and Boltzmann-density sampling.
  • Density-path methods: Feynman-Kac provides an elegant proof framework for Jarzynski equality variants, alongside transport-based methods that learn additional terms to follow prescribed density paths.These connections place FKC within broader work on density-path evolution and nonequilibrium sampling.
  • Predictor-corrector methods: Predictor-corrector methods add Langevin steps to better match intermediate marginals, but related work also uses Metropolis corrections for annealed or product targets.Standard CFG has been interpreted as targeting a different guidance or geometric-mixture weight.
  • FKC distinction: FKC resampling correctors are tailored to the original guidance weight β rather than the alternative β′ = 1/2(1 + β) identified for standard CFG SDE simulation.This distinguishes the correction target used by FKC from the guidance mixture analyzed in prior work.

7. Conclusion

Feynman-Kac Correctors provide weighted diffusion tools for controlling sample distributions, including annealed, compositional, and reward-informed targets. The framework supports intermediate-time weighting and SMC resampling for estimating expectations and normalization constants.

  • Conclusion: FKCs provide fine control over diffusion sample distributions for compositional modeling, geometric averaging, and temperature annealing.The framework combines specialist models and supports annealing from higher-temperature amortized samplers to lower-temperature targets.
  • Conclusion: The framework extends to reward models and time-dependent annealing schedules when the required log-density terms can be estimated.The annealing schedule is denoted βt, and weighting terms may use estimated log densities.
  • Feynman-Kac derivation: The derivation uses a backward Kolmogorov equation and a Feynman-Kac PDE to obtain an expectation estimator for weighted SDE simulations.The proof combines PDE evolution for unnormalized densities with a backward equation and Ito calculus.
  • Expectation estimation: Self-normalized importance sampling is consistent as K →∞, with normalization constants independent of the state variable.The same reasoning supports estimating intermediate-time expectations and normalization constants.
  • SMC and jump-process simulation: Reweighting can be simulated through jump processes whose adjoint generators match the reweighting generator, with thresholding choices affecting resampling events.The construction uses positive and negative parts of the weights and assumes normalized jump kernels.

C. Proofs for Table 1

The proofs establish PDE conversion rules for annealed, scaled, reweighted, and product distributions. These rules identify corresponding simulation forms for combining or transforming diffusion and continuity equations.

  • Annealing: The annealed continuity and diffusion propositions derive PDEs for marginals obtained by raising a reference density to a temperature exponent.The appendix separately treats ordinary and scaled annealed continuity and diffusion equations.
  • Proof strategy: Each derivation differentiates the target density and rearranges the resulting terms into a PDE that can be simulated.The displayed intermediate calculations include score-gradient and Laplacian terms.
  • Annealing: The annealed reweighting proposition derives a reweighting PDE and an equivalent simulation form for the annealed density.The time-dependent annealing proposition extends the target family to pt,βt(x) ∝ qt(x)βt.
  • Products: For products of densities, the appendix derives PDEs for products formed from continuity equations, diffusion equations, and reweighting equations.The product constructions treat both shared and distinct underlying evolutions.

D. Proofs of Propositions

The propositions translate annealing, product-of-experts, classifier-free guidance, reward tilting, and target-score products into weighted SDEs. Their proofs match PDE terms while selecting diffusion and drift coefficients to obtain usable simulation schemes.

  • Annealing and guidance: The annealed SDE proposition gives an SDE whose marginals follow the annealed target family.This supplies the SDE-level counterpart to the appendix’s annealed PDE rules.
  • Product of Experts: The Product of Experts proposition constructs an SDE family for marginals proportional to products or geometric averages of component marginals.The construction imposes a parameter constraint on the admissible coefficient a.
  • Annealing and guidance: Classifier-free guidance is represented as an SDE family whose marginals are annealed versions of the guidance-related target distribution.The proof first applies annealing and then matches diffusion coefficients to avoid additional Laplacian terms in the weights.
  • Reward tilting: Reward tilting samples from pt(x) ∝ qt(x) exp(βt r(x)) by incorporating the reward gradient into the drift and deriving the associated weighted SDE.The construction explicitly adds a∇r(x) as an additional drift term.
  • Scope and connections: The appendix also states that FKCs can be applied to any diffusion-based sampler and relates their reweighting terms to Fisher-Rao gradient flows.These statements extend the method beyond the DEM sampler used for the LJ-13 setting.

F.1. Sampling Metrics

The evaluation uses complementary distributional metrics, with task-specific choices reflecting geometry and the properties each metric can capture. Energy-based filtering and metric definitions impose explicit evaluation boundaries.

  • Evaluation protocol: The study evaluates generated samples with multiple metrics capturing different aspects of distribution quality.Samples with energy > 100 are filtered before computing metrics, although the filter affects only non-resampled metrics.
  • LJ-13 metrics: For LJ-13, Distance-W2 compares flattened pairwise-distance histograms and measures global fidelity while respecting SE(3) equivariance.A 10,000-sample set produces 700,000 pairwise distances for the 13-particle system.
  • LJ-13 metrics: Energy-W1 and Energy-W2 measure deviation between reference and generated energy distributions but cannot determine whether modes have been dropped.A small Energy Wasserstein distance can still occur when a mode with similar energy is missed.
  • Evaluation protocol: The LJ-13 filtering excludes roughly 2-3% of samples for Target Score and Tempered Noise SDE baselines without FKC, as well as lower-temperature DEM.The authors note that this exclusion helps those baselines.
  • Distributional metrics: MMD measures local distribution fit, total variation measures grid-discretized density fit, and W1/W2 measure Euclidean fit in the 40-Gaussian-mixture setting.Euclidean Wasserstein distances are not considered suitable for LJ-13 because its particles are SE(3) equivariant.
  • Benchmark setting: The 40-Gaussian-mixture setting has 40 randomly initialized two-dimensional modes with equal standard deviation and permits exact density and score calculation.It is used as an experimental setting without modeling error.

F.2.1. ADDITIONAL RESULTS

The tractable GMM experiment anneals from temperature TL = 3 to TS = 1/3 and finds Target Score sampling with systematic resampling performs best across more metrics.

  • Target Score sampling with a = 0 and systematic resampling performs best across more metrics in the GMM example.The experiment uses 10k samples, 1000 integration steps, and dt = 0.001.
  • The annealing experiment starts at temperature TL = 3 and targets TS = 1/3.
  • Because the GMM is represented by |π|TL Gaussians for integer TL, the experiments are restricted to small integer starting temperatures.The 40 Gaussians setting uses TL = 3, with β = TS/TL.

F.3. LJ-13 Sampling Task

The LJ-13 task evaluates temperature-annealed inference from a model trained at TL = 2.0, comparing FKC samplers with direct lower-temperature training and other sampling choices. FKC performs well across temperatures, while trade-offs arise among fidelity, diversity, computation, and latent-space applicability.

  • Experimental setup: Lennard-Jones-13 models use an EGNN architecture, geometric noise schedules, score clipping, and 1000 integration steps during sampling.Models are trained for 166 epochs on 4 NVIDIA A100 80GB GPUs.
  • Limitations: DEM score estimates become noisier at high times because estimator variance grows with time and the energy-gradient magnitude.The objective is easier to train on smooth energies.
  • Results: At intermediate temperatures, DEM at target temperature 1.0 outperforms FKC with β = 2.0 on Energy-W1 and Energy-W2, but FKC retains better global fidelity.
  • Results: FKC performs well across target temperatures on interatomic-distance and energy-Wasserstein metrics, although Energy-W1 worsens at low temperatures for all methods.The comparison excludes roughly 2–3% of samples with unacceptably bad energy from two SDEs.
  • Ablations: Larger FKC batch sizes improve scores, with performance plateauing after batch size approximately 8.The reported sweep includes batch size 2.
  • Limitations: For SDXL latent diffusion on GenEval, FKC shows no significant aggregate increase over vanilla CFG across 1000 prompts.The paper reports some qualitative improvements but finds the gain inconsistent in this setting.

F.6. Multi-Target Structure-Based Drug Design

The multi-target drug-design experiments use Product of Experts guidance to generate molecules conditioned on protein pockets and functional properties. FKC generally improves molecular quality and target performance, but resampling trades diversity for quality and requires equal-dimensional samples.

  • Structure-based drug design: For dual-target generation, molecules are sampled with an SE(3)-equivariant graph neural network and guided using the Product of Experts scheme.
  • Implementation constraints: PoE weight computation requires equal sample dimensionality, motivating fixed-size molecule bins with the same number of atoms within each batch.
  • Structure-based drug design: Increasing β from 0.5 to 2.0 increases the product of docking scores for the protein pair, with larger delta increases at smaller β values.
  • Structure-based drug design: FKC improves dual-target molecule performance at β = 1.5 and β = 2.0, with a larger improvement at β = 2.0.The improvement comes at a cost to molecular diversity and uniqueness.
  • Structure-based drug design: Using tmax = 0.6 with β = 2.0 provides a practical trade-off between molecule quality and maintaining diversity.Resampling only occurs when τ <= tmax.
  • Evaluation: The study reports docking, validity, uniqueness, diversity, and quality metrics for generated molecules across protein-pair and multi-property tasks.Quality requires unique and valid molecules with QED ≥0.6 and SA ≤4.0.
  • SMILES generation: For functional-property generation, FKC further increases average binding fitness and tends to improve uniqueness, validity, and drug-like quality while slightly reducing diversity.The experiments use LDMol, a latent diffusion model conditioned on natural-language molecule-property descriptions.
  • SMILES generation: FKC produces molecules with better docking scores in the text-prompt generation experiments.The paper compares these scores with known binders to ATP1A1 and CPT2.
Loading 2503.02819v2…