Корректоры Фейнмана-Каца в диффузии: отжиг, управление и произведение экспертов

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

March 4, 2025

Авторы: Marta Skreta, Tara Akhound-Sadegh, Viktor Ohanesian, Roberto Bondesan, Alán Aspuru-Guzik, Arnaud Doucet, Rob Brekelmans, Alexander Tong, Kirill Neklyudov

cs.AI

Аннотация

Хотя модели генерации на основе оценок являются предпочтительным выбором в различных областях, существует ограниченное количество инструментов для контролируемого управления поведением на этапе вывода в принципиальном ключе, например, для композиции нескольких предобученных моделей. Существующие методы без использования классификаторов применяют простую эвристику для смешивания условных и безусловных оценок с целью приближенной выборки из условных распределений. Однако такие методы не аппроксимируют промежуточные распределения, что требует дополнительных шагов "коррекции". В данной работе мы предлагаем эффективный и принципиальный метод для выборки из последовательности аннелированных, геометрически усредненных или произведенных распределений, полученных из предобученных моделей на основе оценок. Мы выводим взвешенную схему моделирования, которую называем Корректорами Фейнмана-Каца (FKC), основываясь на знаменитой формуле Фейнмана-Каца, тщательно учитывая члены соответствующих дифференциальных уравнений в частных производных (ДУЧП). Для моделирования этих ДУЧП мы предлагаем алгоритмы повторной выборки на основе метода последовательного Монте-Карло (SMC), которые используют масштабирование на этапе вывода для улучшения качества выборки. Мы эмпирически демонстрируем полезность наших методов, предлагая амортизированную выборку через аннелирование температуры на этапе вывода, улучшая генерацию молекул с несколькими целями с использованием предобученных моделей и повышая качество генерации изображений из текста без использования классификаторов. Наш код доступен по адресу 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.