Diffusion models guide generation faster with learned guidance schedules
Learned End-to-End Guidance Schedules for Diffusion Models
Machine Learning
Summary
Diffusion models create images or data step-by-step, and guiding this process helps meet specific goals but makes it slow and costly. The authors found a way to learn how to adjust the guidance automatically during generation to get good results faster. They train this adjustment with a shortcut to save time, so the model needs fewer steps to create images or solve problems. This method works well for tasks like fixing images, creating faces, and solving physics equations from data.
What this means in practice
- •For image processing teams: Generate or repair images with better quality and faster completion times by using learned time-dependent guidance schedules.
- •For scientific computing teams: Solve forward and inverse partial differential equation problems faster using guided diffusion with learned guidance schedules.
Authors
Aneesh Barthakur, Mathias Niepert, Luiz F. O. Chamon
Abstract
Diffusion models are a powerful generative paradigm used across multimedia and scientific applications. Guided diffusion methods impose requirements on the generation by adding the gradient of a differentiable loss (the guidance function) as a drift term during inference. The weight of this drift (the guidance scale) is critical for the trade-off between data quality and requirement satisfaction. To achieve both of these goals, guided diffusion must resort to small guidance scales and lengthy sampling, incurring high computational costs. This work proposes learned end-to-end guidance schedules (LEEGS) to achieve these objectives with fewer sampling steps. LEEGS trains a time-dependent schedule by minimizing the guidance function over a small set of examples using stochastic gradient descent. Backpropagating through guided sampling is computationally expensive, so LEEGS uses an approximation of the gradient that cuts training time by a factor of 4. We evaluate LEEGS on diverse guidance tasks, including (a) image inpainting, (b) noisy image inverse problems, (c) face-ID-guided generation, and (d) forward and inverse PDE problems, outperforming baselines at equal budget (50 or 100 NFEs), or matching constant guidance with only 10% of the steps.