arXiv:2505.07477cs.LGcs.CV2025-05TPAMI

通过只保留一步的反向传播,让扩散模型采样提速90%。

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

  • 只保留生成过程中的一步计算图,实现梯度快速传播。
  • 相比全量反向传播,计算成本降低约90%且性能相当。
  • 适用于控制生成内容或微调模型参数,通用性强。

扩散模型在建模大规模数据分布方面取得了显著成功。然而,许多下游任务需要基于可微指标引导生成内容,通常需在生成过程中进行反向传播,这导致计算开销巨大——生成过程常需数十到上百次递归网络调用,带来高内存占用和显著时间消耗。本文从并行去噪角度提出新思路:全程反向传播并非必需。只需保留生成过程中的单步计算图,即可为梯度传播提供捷径。由此提出的快捷扩散优化方法(SDO)具有通用性、高性能和计算轻量的特点,可优化扩散采样中所有参数类型。我们在多个真实任务中验证了SDO的有效性,包括通过优化潜在变量控制生成内容、以及微调网络参数对齐扩散模型。相比全量反向传播,本方法将计算成本降低约90%,同时保持优异性能。代码已开源:https://github.com/deng-ai-lab/SDO。

原文摘要 · Abstract (English)

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.

扩散模型反向传播加速采样优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。