[Paper Review] Accelerating Parallel Sampling of Diffusion Models

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

Advanced Mathematical Modeling in Engineering4 citations

TL;DR

This paper proposes ParaTAA, a training-free, universal parallel sampling algorithm for diffusion models that reformulates the autoregressive sampling process as solving triangular nonlinear equations via fixed-point iteration. By leveraging triangular Anderson acceleration and adaptive windowing, ParaTAA reduces inference steps by 4x to 14x, achieving near-identical image quality to sequential sampling—e.g., matching 100-step DDIM results in just 7 steps for Stable Diffusion.

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.

Motivation & Objective

• Address the slow inference speed of autoregressive diffusion sampling, which limits real-time deployment despite high-quality generation.
• Develop a method to parallelize the inherently sequential sampling process without requiring model fine-tuning or distillation.
• Preserve image quality and diversity while drastically reducing the number of inference steps compared to sequential methods like DDIM and DDPM.
• Enable efficient use of extra computational and memory resources to accelerate sampling without compromising fidelity.
• Provide a theoretically grounded framework for convergence analysis and optimization of parallel sampling in diffusion models.

Proposed method

• Reformulate the diffusion sampling process as solving a system of triangular nonlinear equations using fixed-point iteration (FP), enabling parallel computation of intermediate steps.
• Introduce triangular Anderson acceleration (TAA) to accelerate convergence of the fixed-point iteration by leveraging past iterates in a triangular structure.
• Apply adaptive windowing to control the trade-off between convergence speed and computational cost, dynamically adjusting the number of past iterates used.
• Use trajectory initialization from prior sampling runs to bootstrap convergence, significantly reducing steps for similar prompts.
• Integrate the method with existing sequential samplers (e.g., DDIM, DDPM) by modifying the coefficient structure in the nonlinear system without retraining.
• Optimize implementation with 16-bit precision and efficient GPU memory access to reduce wall-clock time without altering model weights.

(a) DDIM 100 steps

Experimental results

Research questions

• RQ1 Can the autoregressive sampling process in diffusion models be effectively parallelized while maintaining image quality?
• RQ2 How can fixed-point iteration be systematically optimized to accelerate convergence in solving the nonlinear system of equations for diffusion sampling?
• RQ3 What is the impact of window size and initialization strategy on the convergence speed and inference step reduction in parallel sampling?
• RQ4 To what extent can extra computational resources be leveraged to reduce inference steps without model fine-tuning?
• RQ5 Can the proposed method achieve near-identical results to sequential sampling while enabling orders-of-magnitude speedup?

Key findings

• ParaTAA reduces the number of inference steps required by sequential sampling methods like DDIM and DDPM by a factor of 4 to 14 times.
• For Stable Diffusion using 100-step DDIM, ParaTAA achieves equivalent image quality in only 7 inference steps, demonstrating a 14x speedup.
• Using a window size of 10, ParaTAA reduces inference steps from 100 to 25 for DDIM on Stable Diffusion, achieving a 4x speedup.
• Doubling the window size from 10 to 20 only marginally reduces steps from 25 to 21, indicating diminishing returns beyond a certain window size.
• Initializing sampling with a pre-existing trajectory from a similar prompt reduces convergence time by up to 3x, with high-quality images generated in as few as 3 steps.
• The method maintains image fidelity and diversity comparable to sequential sampling, with theoretical guarantees on convergence and solution uniqueness.

(b) DDPM 100 steps