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

TL;DR

This work tackles the high computational cost of backpropagating through diffusion sampling for downstream differentiable objectives. It introduces Shortcut Diffusion Optimization (SDO), a principled gradient shortcut based on a Picard-iteration fixed-point formulation that preserves gradients for only a single denoising step, while keeping the forward diffusion process unchanged. Theoretical guarantees and practical implementations show that one-step gradients can closely approximate or even outperform full backpropagation, yielding substantial speedups (up to ~90% lower computation) and memory reductions with broad applicability across latent optimization, reward alignment, and multi-framework diffusion systems. Empirically, SDO enables text-guided manipulation, style transfer, aesthetic enhancement, and adversarial perturbation tasks with improved or comparable performance to full backpropagation, and it supports end-to-end fine-tuning via differentiable rewards. The approach significantly broadens the practicality of gradient-based diffusion adjustments in real-world applications, while also highlighting areas for further work on non-differentiable metrics and reward robustness.

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.

Figures (16)

• Figure 1: Shortcut Diffusion Optimization (SDO) enables diverse applications through backpropagation in diffusion sampling. Parts (a)-(d) illustrate controlled generation by optimizing latent variables: (a) modifying an image based on text instructions using CLIP radford2021learning; (b) controlling texture according to a given style reference while preserving content; (c) enhancing image aesthetics guided by an aesthetic predictor; and (d) generating adversarial perturbations that bypass well-trained classifiers. Part (e) demonstrates SDO’s capability for end-to-end fine-tuning of parameters, allowing the diffusion model to maximize specific rewards, such as aligning with human preferences.
• Figure 2: Graphical models of DMs illustrating (a) sequential sampling song2020denoising and (b) parallel sampling using Picard iteration shih2024parallel. Picard iteration introduces skip dependencies between $\boldsymbol{x}_0$ and $\{\boldsymbol{x}_N,\cdots,\boldsymbol{x}_1\}$, which inspired the development of SDO.
• Figure 3: Plain backpropagation inefficiently differentiates the objective function through the entire forward sampling process. In contrast, SDO backpropagation requires differentiation only at the specific timestep associated with the variable being optimized (e.g., $\boldsymbol{x}_N$ at step $N$ or $\theta$ at step $i^\prime$), providing an efficient shortcut for the backward pass of the gradient.
• Figure 4: PyTorch implementation of SDO for optimizing $\boldsymbol{x}_N$. The function set_grad_enabled ensures that only the one-step computation graph required for backpropagation is retained. scheduler stands for solver scheduling, e.g. DDIM song2020denoising.
• Figure 5: PyTorch implementation of SDO for optimizing model parameters. The program samples the timesteps and retains the computational graph only for the sampled timesteps, ensuring lightweight backpropagation while enabling optimization across all timesteps.

Theorems & Definitions (2)

• proof
• proof