Fast constrained sampling in pre-trained diffusion models

TL;DR

This work proposes an algorithm that enables fast, high-quality generation under arbitrary constraints and shows that in denoising diffusion models, it can employ an approximation to Newton's optimization method that allows us to speed up inference and avoid the expensive backpropagation operations.

Abstract

Large denoising diffusion models, such as Stable Diffusion, have been trained on billions of image-caption pairs to perform text-conditioned image generation. As a byproduct of this training, these models have acquired general knowledge about image statistics, which can be useful for other inference tasks. However, when confronted with sampling an image under new constraints, e.g. generating the missing parts of an image, using large pre-trained text-to-image diffusion models is inefficient and often unreliable. Previous approaches either utilized backpropagation through the denoiser network, making them significantly slower and more memory-demanding than simple text-to-image generation, or only enforced the constraint locally, failing to capture critical long-range correlations in the sampled image. In this work, we propose an algorithm that enables fast, high-quality generation under arbitrary constraints. We show that in denoising diffusion models, we can employ an approximation to Newton's optimization method that allows us to speed up inference and avoid the expensive backpropagation operations. Our approach produces results that rival or surpass the state-of-the-art training-free inference methods while requiring a fraction of the time. We demonstrate the effectiveness of our algorithm under both linear (inpainting, super-resolution) and non-linear (style-guided generation) constraints. An implementation is provided at https://github.com/cvlab-stonybrook/fast-constrained-sampling.

Figures (24)

• Figure 1: When tasked with completing the missing half of an image, previous methods are slow and fail to capture the important long-range dependencies between pixels. The proposed algorithm generates a reasonable image at a fraction of the time.
• Figure 2: We showcase the difference between the proposed method to compute the update direction (${\bm{J}}$) and gradient descent (${\bm{J}}^T$). The heatmaps indicate where the input ${\bm{x}}_t$ would change when perturbing a single pixel in the output, denoted in green. The two directions are considerably different, with ours capturing better longer-range correlations and maintaining shapes. Even though we use finite differences, the direction computed from ${\bm{J}}$ is sharper in some regions, like the outlines.
• Figure 3: Comparison between our method and existing algorithms on free-form inpainting and $8\times$ super-resolution. We directly use the images and results from chung2023prompt since there is no code available to replicate their method.
• Figure 4: Qualitative evaluation of large area (box) inpainting on ImageNet. Our method achieves results closer to the fine-tuned inpainting model while requiring a fraction of the time to run per image compared to baselines.
• Figure 5: We guide the style of Stable Diffusion images with a CLIP (or OpenCLIP) model, using classifier-free guidance $w=5$. The images generated by our algorithm are closer to the reference style while maintaining faithfulness to the text prompt.