【论文精读】SDO: 用梯度捷径加速扩散采样中的反向传播

Hongkun Dou, Zeyu Li, Xingyu Jiang, Hongjue Li, Lijun Yang, Wen Yao, and Yue Deng2024年5月8日大约 7 分钟

【论文精读】SDO: 用梯度捷径加速扩散采样中的反向传播

摘要

扩散模型 (DMs) 在下游任务中反向传播计算成本高昂。本文提出捷径扩散优化 (SDO)，通过仅保留一步计算图优化目标函数，显著降低约90%计算成本，同时保持或超越完整反向传播性能。SDO适用于潜变量优化、模型微调等任务，兼具通用性、高性能和轻量级特点。

研究背景与动机

扩散模型 (DMs) 在图像生成、视频生成等领域取得了巨大成功。然而，许多下游任务（如提升图像质量、实现个性化风格）常常需要对预训练模型进行定制。这通常被构建为一个优化问题，通过反向传播来解决。但由于DMs固有的迭代特性，完整的反向传播会导致巨大的计算量和内存消耗。

扩散模型中的梯度流。普通反向传播需要通过整个链计算梯度（红色箭头），而SDO通过捷径（紫色箭头）直接优化参数。

本文对这种完全反向传播的必要性提出了质疑。受到近期并行去噪研究（特别是 Picard迭代）的启发，作者提出捷径扩散优化 (SDO)。其核心思想在于，仅保留单步计算图进行优化就足以有效地传递梯度，从而大幅降低计算负担。实验表明，SDO在多种任务上表现优异，计算成本降低约90%，效果媲美甚至超越完全反向传播。

方法与技术创新

普通反向传播的局限性

在扩散模型中，通过整个采样过程（N步）进行普通反向传播 (Plain Backpropagation, PBP) 来优化目标函数，会导致计算图过深，带来巨大的内存消耗 (memory consumption) 和时间开销，且可能存在梯度爆炸/消失 (gradient explosion/vanishing) 问题。

(a) 普通反向传播和 (b) SDO 的计算图比较。SDO通过仅保留当前步骤的计算图显著减少了计算开销。

捷径扩散优化 (SDO)

核心思想

SDO 的核心思想是通过梯度捷径 (gradient shortcut)，在采样链的优化过程中，仅保留关键一步的梯度信息，而对其余步骤的梯度进行阻断。这种方法显著降低了反向传播的计算复杂度和内存需求，同时保持优化效果。

传统方法：完整反向传播需要通过整个采样链计算梯度，导致内存占用和计算量随采样步数线性增长。
SDO 方法：仅在关键一步（如最后一步）保留梯度，其余步骤通过 torch.no_grad() 或类似机制阻断梯度流。

实现细节

采样阶段：
- 使用标准扩散采样算法（如 DDIM 或 DPM-Solver）生成采样链。
- 所有采样步骤均正常执行，确保生成质量与普通采样一致。
梯度计算阶段：
- 在反向传播时，仅保留关键一步的计算图。
- 其余步骤的梯度通过 detach() 或 torch.no_grad() 阻断，避免冗余计算。
优化阶段：
- 使用关键一步的梯度更新优化目标（如潜变量或模型参数）。
- 通过减少梯度计算的深度，显著降低内存占用和计算时间。

伪代码示例

以下是 SDO 的简化实现伪代码：

params = {'params': model.parameters(), 'lr': lr}
optimizer = torch.optim.Adam([params])
for _ in range(epochs):
   optimizer.zero_grad()
   x_t = x_N
   backprop_step = random.randint(0, len(scheduler.timesteps) -1 ) # Ensure valid index
   for i, t in enumerate(scheduler.timesteps):
      is_grad = (i == backprop_step)
      with torch.set_grad_enabled(is_grad):
         noise_pred = model(x_t, t)
         x_t = scheduler.step(noise_pred, t, x_t) # .sample for some schedulers
         if is_grad:
            output = torch.clamp(x_t, -1, 1)
            loss = J(output)
            loss.backward()
   optimizer.step()

理论支持

SDO 的有效性基于 Picard 迭代的理论支持。Picard 迭代是一种用于求解固定点问题的高效方法，能够在并行采样的基础上提供梯度近似。理论分析表明，SDO 的梯度近似误差满足以下条件：

$$|\nabla_\theta L_{SDO} - \nabla_\theta L_{full}| \leq C$$

其中 $C$ 是一个有界常数，表明 SDO 的梯度近似在合理假设下是可靠的。

优势总结

计算效率：将反向传播的复杂度从 $O(N)$ 降低到 $O(1)$。
内存友好：显著减少内存占用，支持更大规模的模型和任务。
兼容性强：适用于各种扩散采样器（如 DDIM、DPM-Solver）和优化目标（如潜变量、模型参数）。
稳定性高：缓解梯度爆炸和梯度消失问题，优化过程更稳定。

实验与结果分析

质量评估

实验对比了 SDO 与现有方法在多个指标上的表现，包括 LPIPS（越低越好）、CLIP（越高越好）和 ID loss（越低越好）。结果表明，SDO 在所有指标上均优于其他方法，尤其在 CLIP 指标上表现突出。

方法	LPIPS ↓	CLIP ↑	ID loss ↓
DiffusionCLIP [77]	0.175	29.93	0.901
FlowGrad [38]	0.142	31.30	0.797
AdjointDPM [26]	0.188	28.14	0.908
DOODL [27]	0.183	28.76	0.914
SDO (Ours)	0.134	32.18	0.790