[论文解读] Accelerating Parallel Sampling of Diffusion Models

Zhiwei Tang, Jiasheng Tang+3 more|arXiv (Cornell University)|Feb 15, 2024

Advanced Mathematical Modeling in Engineering被引用 4

一句话总结

本文提出 ParaTAA，一种无需训练的、适用于扩散模型的通用并行采样算法，通过固定点迭代求解三角非线性方程组，将自回归采样过程重新表述为并行计算。通过利用三角 Anderson 加速和自适应窗口机制，ParaTAA 将推理步数减少 4 至 14 倍，实现与串行采样几乎相同的图像质量——例如，仅用 7 步即可匹配 Stable Diffusion 的 100 步 DDIM 结果。

ABSTRACT

Diffusion models have emerged as state-of-the-art generative models for image generation. However, sampling from diffusion models is usually time-consuming due to the inherent autoregressive nature of their sampling process. In this work, we propose a novel approach that accelerates the sampling of diffusion models by parallelizing the autoregressive process. Specifically, we reformulate the sampling process as solving a system of triangular nonlinear equations through fixed-point iteration. With this innovative formulation, we explore several systematic techniques to further reduce the iteration steps required by the solving process. Applying these techniques, we introduce ParaTAA, a universal and training-free parallel sampling algorithm that can leverage extra computational and memory resources to increase the sampling speed. Our experiments demonstrate that ParaTAA can decrease the inference steps required by common sequential sampling algorithms such as DDIM and DDPM by a factor of 4$\sim$14 times. Notably, when applying ParaTAA with 100 steps DDIM for Stable Diffusion, a widely-used text-to-image diffusion model, it can produce the same images as the sequential sampling in only 7 inference steps. The code is available at https://github.com/TZW1998/ParaTAA-Diffusion.

研究动机与目标

• 解决自回归扩散采样过程推理速度缓慢的问题，尽管其生成图像质量高，但限制了实时部署。
• 开发一种方法，将本质上串行的采样过程并行化，且无需模型微调或蒸馏。
• 在大幅减少推理步数的同时，保持与 DDIM 和 DDPM 等串行方法相当的图像质量和多样性。
• 有效利用额外的计算和内存资源，在不损害保真度的前提下加速采样过程。
• 为扩散模型中并行采样的收敛性分析与优化提供理论基础框架。

提出的方法

• 将扩散采样过程重新表述为通过固定点迭代（FP）求解三角非线性方程组，从而实现中间步骤的并行计算。
• 引入三角 Anderson 加速（TAA），通过在三角结构中利用历史迭代值，加速固定点迭代的收敛。
• 应用自适应窗口机制，动态调节收敛速度与计算成本之间的权衡，自适应调整所用历史迭代值的数量。
• 利用先前采样运行的轨迹进行初始化，以加速收敛，显著减少相似提示下的推理步数。
• 通过修改非线性系统中的系数结构，将该方法与现有串行采样器（如 DDIM、DDPM）集成，无需重新训练。
• 通过使用 16 位精度和高效的 GPU 内存访问优化实现，降低实际运行时间，且不改变模型权重。

(a) DDIM 100 steps

实验结果

研究问题

• RQ1 能否在保持图像质量的前提下，有效并行化扩散模型中的自回归采样过程？
• RQ2 如何系统性地优化固定点迭代，以加速求解扩散采样非线性方程组的收敛过程？
• RQ3 窗口大小和初始化策略对并行采样中收敛速度与推理步数减少的影响如何？
• RQ4 在不进行模型微调的前提下，可多大程度上利用额外计算资源以减少推理步数？
• RQ5 所提方法是否能在实现数量级加速的同时，达到与串行采样几乎完全一致的结果？

主要发现

• ParaTAA 将 DDIM 和 DDPM 等串行采样方法所需的推理步数减少 4 至 14 倍。
• 对于使用 100 步 DDIM 的 Stable Diffusion，ParaTAA 仅用 7 步即可达到等效图像质量，实现 14 倍加速。
• 使用窗口大小为 10 时，ParaTAA 将 Stable Diffusion 上 DDIM 的推理步数从 100 步减少至 25 步，实现 4 倍加速。
• 将窗口大小从 10 倍增至 20 时，步数仅从 25 步减少至 21 步，表明在某一窗口大小后收益递减。
• 使用相似提示的已有轨迹进行初始化，可将收敛时间减少最多 3 倍，高质量图像可在仅 3 步内生成。
• 该方法在保持与串行采样相当的图像保真度和多样性的同时，具备收敛性和解唯一性的理论保证。

(b) DDPM 100 steps