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

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.

Diffusion models are powerful tools for generating data like images, molecules, or text, but it is generally difficult to control their generation process. This paper introduces a method called Feynman-Kac Correctors (FKC), which allows for precise control over what a diffusion model generates without retraining it. FKC works by adjusting the way samples are drawn from the model, based on the Sequential Monte Carlo framework and, in particular, the Feynman-Kac formula. This enables a principled approach to sampling from combined target distributions, like mixtures or products of multiple pretrained models, or temperature-annealed target distributions. We show that FKC improves sampling in three settings: 1. classifier-free guidance, which is widely used in text-to-image generation, 2. generating molecules that satisfy multiple objectives (binding to two proteins simultaneously) and 3. sampling from physical systems at different temperatures using a model trained at a single temperature. Unlike traditional methods, FKC allows for flexible and efficient sampling with little added computation. This opens up new possibilities for applications in AI, drug discovery, and scientific simulations.