Feynman-Kac-Korrektoren in der Diffusion: Tempern, Lenkung und Produkt von Experten

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

March 4, 2025

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

cs.AI

Zusammenfassung

Während score-basierte generative Modelle das Modell der Wahl in verschiedenen Domänen sind, gibt es nur begrenzte Werkzeuge, um das Verhalten zur Inferenzzeit auf prinzipielle Weise zu steuern, z. B. für die Komposition mehrerer vortrainierter Modelle. Bestehende Methoden zur klassifikatorfreien Steuerung verwenden eine einfache Heuristik, um bedingte und unbedingte Scores zu mischen und so näherungsweise aus bedingten Verteilungen zu sampeln. Solche Methoden approximieren jedoch nicht die Zwischenverteilungen, was zusätzliche „Korrektor“-Schritte erforderlich macht. In dieser Arbeit stellen wir eine effiziente und prinzipielle Methode vor, um aus einer Folge von getemperten, geometrisch gemittelten oder Produktverteilungen, die aus vortrainierten score-basierten Modellen abgeleitet wurden, zu sampeln. Wir leiten ein gewichtetes Simulationsschema ab, das wir Feynman-Kac-Korrektoren (FKCs) nennen, basierend auf der berühmten Feynman-Kac-Formel, indem wir sorgfältig die Terme in den entsprechenden partiellen Differentialgleichungen (PDEs) berücksichtigen. Um diese PDEs zu simulieren, schlagen wir Sequential-Monte-Carlo (SMC)-Resampling-Algorithmen vor, die die Skalierung zur Inferenzzeit nutzen, um die Sampling-Qualität zu verbessern. Wir demonstrieren den Nutzen unserer Methoden empirisch, indem wir amortisiertes Sampling durch Temperaturtemperierung zur Inferenzzeit vorschlagen, die Erzeugung von Molekülen mit mehreren Zielen unter Verwendung vortrainierter Modelle verbessern und die klassifikatorfreie Steuerung für die Text-zu-Bild-Generierung optimieren. Unser Code ist verfügbar unter https://github.com/martaskrt/fkc-diffusion.

English

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.