Correttori di Feynman-Kac nella diffusione: Ricottura, Guida e Prodotto di Esperti

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

March 4, 2025

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

cs.AI

Abstract

Sebbene i modelli generativi basati su punteggi siano il modello di riferimento in diversi domini, gli strumenti disponibili per controllare il comportamento durante l'inferenza in modo principiato sono limitati, ad esempio per comporre più modelli pre-addestrati. I metodi esistenti di guida senza classificatore utilizzano un'euristica semplice per mescolare punteggi condizionali e non condizionali, campionando approssimativamente da distribuzioni condizionali. Tuttavia, tali metodi non approssimano le distribuzioni intermedie, rendendo necessari ulteriori passaggi di "correzione". In questo lavoro, forniamo un metodo efficiente e principiato per campionare da una sequenza di distribuzioni ricotte, mediate geometricamente o prodotte, derivate da modelli basati su punteggi pre-addestrati. Deriviamo uno schema di simulazione ponderato che chiamiamo Feynman-Kac Correctors (FKC), basato sulla celebre formula di Feynman-Kac, tenendo conto con attenzione dei termini nelle appropriate equazioni differenziali alle derivate parziali (PDE). Per simulare queste PDE, proponiamo algoritmi di ricampionamento Sequential Monte Carlo (SMC) che sfruttano il ridimensionamento durante l'inferenza per migliorare la qualità del campionamento. Dimostriamo empiricamente l'utilità dei nostri metodi proponendo un campionamento ammortizzato tramite ricottura della temperatura durante l'inferenza, migliorando la generazione di molecole multi-obiettivo utilizzando modelli pre-addestrati e migliorando la guida senza classificatore per la generazione di immagini da testo. Il nostro codice è disponibile all'indirizzo 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.