확산에서의 파인만-카크 보정자: 어닐링, 가이던스, 그리고 전문가 곱

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

초록

스코어 기반 생성 모델은 다양한 도메인에서 선호되는 모델이지만, 추론 시점의 행동을 원칙적으로 제어하기 위한 도구는 제한적입니다. 예를 들어, 여러 사전 학습된 모델을 조합하는 경우가 그러합니다. 기존의 분류자 없는 지도 방법은 조건부 및 무조건부 스코어를 단순한 휴리스틱으로 혼합하여 조건부 분포에서 근사적으로 샘플링합니다. 그러나 이러한 방법은 중간 분포를 근사하지 않기 때문에 추가적인 '수정자' 단계가 필요합니다. 본 연구에서는 사전 학습된 스코어 기반 모델에서 유도된 어닐링, 기하 평균, 또는 곱 분포 시퀀스에서 샘플링하기 위한 효율적이고 원칙적인 방법을 제시합니다. 우리는 적절한 편미분 방정식(PDEs)의 항을 신중하게 고려하여 유명한 파인만-카츠 공식을 기반으로 한 가중 시뮬레이션 기법인 파인만-카츠 수정자(FKCs)를 도출했습니다. 이러한 PDEs를 시뮬레이션하기 위해, 우리는 샘플링 품질을 향상시키기 위해 추론 시점 스케일링을 활용하는 순차적 몬테카를로(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.