[论文解读] Parallel Sampling of Diffusion Models

Andy Y. Shih, Suneel Belkhale+3 more|arXiv (Cornell University)|May 25, 2023

Advanced Neuroimaging Techniques and Applications被引用 9

一句话总结

ParaDiGMS 通过 Picard 迭代引入并行去噪以加速扩散模型采样，在机器人与图像生成任务中实现 2–4 倍的速度提升且无可测量的质量损失。

ABSTRACT

Diffusion models are powerful generative models but suffer from slow sampling, often taking 1000 sequential denoising steps for one sample. As a result, considerable efforts have been directed toward reducing the number of denoising steps, but these methods hurt sample quality. Instead of reducing the number of denoising steps (trading quality for speed), in this paper we explore an orthogonal approach: can we run the denoising steps in parallel (trading compute for speed)? In spite of the sequential nature of the denoising steps, we show that surprisingly it is possible to parallelize sampling via Picard iterations, by guessing the solution of future denoising steps and iteratively refining until convergence. With this insight, we present ParaDiGMS, a novel method to accelerate the sampling of pretrained diffusion models by denoising multiple steps in parallel. ParaDiGMS is the first diffusion sampling method that enables trading compute for speed and is even compatible with existing fast sampling techniques such as DDIM and DPMSolver. Using ParaDiGMS, we improve sampling speed by 2-4x across a range of robotics and image generation models, giving state-of-the-art sampling speeds of 0.2s on 100-step DiffusionPolicy and 14.6s on 1000-step StableDiffusion-v2 with no measurable degradation of task reward, FID score, or CLIP score.

研究动机与目标

• 通过以计算换取速度来推动更快的扩散模型采样，同时不牺牲样本质量。
• 开发一个通用的并行采样框架，兼容预训练的扩散模型与现有的快速采样器。
• 在机器人策略与图像生成模型上展示加速，同时保持如奖励、FID、CLIP 等指标。
• 展示与现有快速采样方法（DDIM、DPMSolver）以及无分类器引导的兼容性。

提出的方法

• 使用 Picard 迭代通过迭代细化完整去噪轨迹的猜测来求解反向时间扩散 ODE。
• 离散化 Picard 更新以实现跨时间步并行计算漂移项 s(x_i, t_i)，随后进行前缀和聚合。
• 引入滑动窗口小批量策略以适配 GPU 内存，并设定确保总变差距离（total-variation distance）保真度的停止准则。
• 在扩展到反向 SDE 时将固定的前向噪声实现吸收为确定性 ODE。
• 通过在 ParaDiGMS 中替换欧拉或更高阶求解器来证明兼容性，并与 ParaDDPM、ParaDDIM、ParaDPMSolver 集成。
• 采用基于容忍度的停止准则来约束与序列基线的分布距离。

实验结果

研究问题

• RQ1 在不牺牲样本质量的前提下，是否可以通过并行去噪步骤来加速扩散模型采样？
• RQ2 基于 Picard 迭代的并行采样在速度和质量方面与现有快速采样器（DDIM、DPMSolver）相比如何？
• RQ3 该方法是否与分类器自由引导和先前的加速兼容，是否能推广到图像和机器人策略任务？

主要发现

• ParaDiGMS 在机器人与图像生成模型上实现约 2–4 倍的加速，在任务奖励、FID 或 CLIP 分数方面没有可测量的降级。
• 在 DiffusionPolicy 的 100 步采样上，单样本所需时间从 0.74s 降至 0.2s，适用于 100 步的行动生成任务（示例给出）。
• 在 StableDiffusion-v2 的 1000 步采样上，单样本时间从 50.0s 降至 14.6s，且质量无损失。
• ParaDiGMS 与 DDPM、DDIM、DPMSolver 兼容，能够在任务中实现联合加速（ParaDDPM、ParaDDIM、ParaDPMSolver）。
• 在 StableDiffusion-v2 的潜在空间扩散中，使用多 GPU 设置后净加速超过 3 倍，加速取决于批量窗口大小与硬件。
• ParaDiGMS 在 COCO 标注上的 CLIP 分数等指标下仍能保持样本质量。