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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Skreta, Marta, Akhound-Sadegh, Tara, Ohanesian, Viktor, Bondesan, Roberto, Aspuru-Guzik, Alán, Doucet, Arnaud, Brekelmans, Rob, Tong, Alexander, Neklyudov, Kirill
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866918049396817920
author Skreta, Marta
Akhound-Sadegh, Tara
Ohanesian, Viktor
Bondesan, Roberto
Aspuru-Guzik, Alán
Doucet, Arnaud
Brekelmans, Rob
Tong, Alexander
Neklyudov, Kirill
author_facet Skreta, Marta
Akhound-Sadegh, Tara
Ohanesian, Viktor
Bondesan, Roberto
Aspuru-Guzik, Alán
Doucet, Arnaud
Brekelmans, Rob
Tong, Alexander
Neklyudov, Kirill
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2503_02819
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Feynman-Kac Correctors in Diffusion: Annealing, Guidance, and Product of Experts
Skreta, Marta
Akhound-Sadegh, Tara
Ohanesian, Viktor
Bondesan, Roberto
Aspuru-Guzik, Alán
Doucet, Arnaud
Brekelmans, Rob
Tong, Alexander
Neklyudov, Kirill
Machine Learning
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.
title Feynman-Kac Correctors in Diffusion: Annealing, Guidance, and Product of Experts
topic Machine Learning
url https://arxiv.org/abs/2503.02819