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

2025-05-12

AI Summary

Diffusion models incur prohibitive computational overhead and memory consumption when optimizing downstream differentiable metrics via full backward propagation throughout the denoising process. This work observes that full backpropagation is unnecessary: a single-step gradient shortcut suffices for efficient optimization of both latent variables and network parameters. To this end, we propose Shortcut Diffusion Optimization (SDO), a lightweight framework grounded in a parallel denoising perspective, integrating gradient truncation, computational graph simplification, and single-step differentiable sampling. We provide the first theoretical guarantee that single-step gradients ensure convergence and preserve optimization quality. Experiments across multiple real-world tasks demonstrate that SDO reduces computational cost by approximately 90% while maintaining—often surpassing—the performance of full-backpropagation baselines. The method is broadly applicable, highly efficient, and practically deployable.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Learning to Search
• Machine Learning: Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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.

Problem

• Reducing computational cost in diffusion model backpropagation
• Optimizing downstream metrics with single-step gradient shortcuts
• Maintaining performance while minimizing memory and time usage

Innovation

• Parallel denoising for efficient gradient computation
• Single-step gradient shortcut reduces backpropagation cost
• Lightweight optimization maintains performance with 90% cost reduction