You Only Look One Step: Accelerating Backpropagation in Diffusion Sampling with Gradient Shortcuts

TL;DR

This paper introduces Shortcut Diffusion Optimization (SDO), a method that significantly accelerates diffusion model backpropagation by using gradient shortcuts, reducing computational costs by about 90% while maintaining performance.

Contribution

The paper proposes a novel gradient shortcut technique for diffusion models that drastically reduces backpropagation complexity during generation.

Findings

SDO reduces computational costs by approximately 90%.

SDO maintains or improves performance compared to full backpropagation.

The method is applicable to various downstream tasks involving diffusion models.

Abstract

Diffusion models (DMs) have recently demonstrated remarkable success in modeling large-scale data distributions. However, many downstream tasks require guiding the generated content based on specific differentiable metrics, typically necessitating backpropagation during the generation process. This approach is computationally expensive, as generating with DMs often demands tens to hundreds of recursive network calls, resulting in high memory usage and significant time consumption. In this paper, we propose a more efficient alternative that approaches the problem from the perspective of parallel denoising. We show that full backpropagation throughout the entire generation process is unnecessary. The downstream metrics can be optimized by retaining the computational graph of only one step during generation, thus providing a shortcut for gradient propagation. The resulting method, which we call Shortcut Diffusion Optimization (SDO), is generic, high-performance, and computationally lightweight, capable of optimizing all parameter types in diffusion sampling. We demonstrate the effectiveness of SDO on several real-world tasks, including controlling generation by optimizing latent and aligning the DMs by fine-tuning network parameters. Compared to full backpropagation, our approach reduces computational costs by $\sim 90\%$ while maintaining superior performance. Code is available at https://github.com/deng-ai-lab/SDO.