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

TL;DR

This work introduces Feynman-Kac Correctors (FKCs) to enable principled inference-time control of diffusion models by sampling from annealed, geometric-average, and product distributions derived from pretrained scores. It derives weighted SDEs, explicit conversion rules, and SMC-based resampling to closely track intermediate target distributions, addressing limitations of heuristic CFG approaches. The framework accommodates annealing, product-of-experts, and reward-tilted densities, and demonstrates practical gains in tasks including multi-objective molecule design and text-to-image generation. Empirical results show improved image-generation quality and molecule docking performance with FKCs, along with actionable guidance on when to use target-score versus tempered-noise variants and how to scale resampling. Overall, FKCs provide a versatile toolkit for modular, inference-time customization of diffusion models with broad applicability to chemistry, imagery, and beyond.

Abstract

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.

Figures (14)

• Figure 1: Feynman-Kac Corrector Inference for annealed $\textcolor{plotorange}{p_{t,\beta}(x)} \propto \textcolor{plotgreen}{q_t(x)}^{\textcolor{plotorange}{\beta=10}}$ and product $\textcolor{plotorange}{p_{t}(x)} \propto \textcolor{plotgreen}{q_t^1(x)}\textcolor{plotpurple}{q_t^2(x)}$ densities.
• Figure 2: Samples from Mixture of 40 Gaussians.
• Figure 3: Samples from EDM2+CFG (top), EDM2+FKC (bottom).
• Figure 3: LJ-13 sampling task with various SDEs, with performance measured by mean $\pm$ standard deviation over 3 seeds. The starting temperature is $T_L=2$, annealed to target temperatures $T_S=0.8$ and $T_S=1.5$. The DEM samples are generated with a model trained at those corresponding target temperatures.
• Figure 4: 2-Wasserstein between energy distributions of MCMC samples from the annealed target distribution and our methods at different temperatures. Note the training temperature $T_L=2$.

Theorems & Definitions (43)

• Proposition 3.1: Classifier-Free Guidance + FKC
• Proposition 3.2: Annealed SDE + FKC
• Proposition 3.3: Product of Experts + FKC
• Proposition 3.4: Reward-tilted Target + FKC
• Proposition 4.1
• Proposition 1.1
• proof
• Lemma 2.1: Adjoint Generators
• proof
• Lemma 2.2: Jump Process Generators