Variational Control for Guidance in Diffusion Models

Kushagra Pandey, Farrin Marouf Sofian, Felix Draxler, Theofanis Karaletsos, Stephan Mandt

TL;DR: We propose a plug and play method for guidance in diffusion models using optimal control

Abstract: Diffusion models exhibit excellent sample quality, but existing guidance methods often require additional model training or are limited to specific tasks. We revisit guidance in diffusion models from the perspective of variational inference and control, introducing \emph{Diffusion Trajectory Matching (DTM)} that enables guiding pretrained diffusion trajectories to satisfy a terminal cost. DTM unifies a broad class of guidance methods and enables novel instantiations. We introduce a new method within this framework that achieves state-of-the-art results on several linear, non-linear, and blind inverse problems without requiring additional model training or specificity to pixel or latent space diffusion models. Our code will be available at https://github.com/czi-ai/oc-guidance.

Lay Summary: Diffusion models power most state-of-the-art breakthroughs in text-to-image and video synthesis. However, in some applications, the user may be interested in generating samples that adhere to some external constraints. For instance, the user may be interested in generating samples that borrow style from a reference image. Additionally, since training a new model for each such conditional task can be cumbersome, the goal is to leverage existing large-scale pretrained diffusion models as powerful priors for such downstream tasks. In this work, we propose a method to achieve this goal using ideas from optimal control and Bayesian inference.

Link To Code: https://github.com/czi-ai/oc-guidance

Primary Area: Deep Learning->Generative Models and Autoencoders

Keywords: Diffusion Models, Inverse Problems

Submission Number: 2281