Emergent Mind

Accelerating Parallel Sampling of Diffusion Models

(2402.09970)
Published Feb 15, 2024 in cs.LG and stat.ML

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~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.

Residual convergence comparison across various timesteps in computational analysis.

Overview

  • Diffusion models have become prominent in generative AI for producing high-quality images but are held back by the computational intensity of their sampling process.

  • A novel approach introduces the use of fixed-point iteration and parallel computing to accelerate the sampling process without sacrificing image quality.

  • The study develops Triangular Anderson Acceleration (TAA) to enhance efficiency by leveraging prior iterations' information in solving the system of nonlinear equations.

  • Empirical validation shows the proposed ParaTAA algorithm significantly outperforms conventional methods in speed and quality, indicating substantial implications for generative AI.

Accelerating the Sampling Process in Diffusion Models through Parallelization

Introduction

Diffusion models have recently taken center stage in the field of generative AI due to their superior capability in producing high-quality images. However, the extensive computational demands of these models, particularly in their sampling process, have posed significant challenges. The conventional autoregressive nature of this process makes it time-consuming, which restricts the practical deployment and iterative exploration of generated outputs.

Accelerating Sampling via Parallelization

In addressing the inefficiencies of the sampling process, recent efforts have specifically targeted the acceleration of this procedure. This work introduces a novel approach that reframes the sampling process as solving a system of triangular nonlinear equations using fixed-point iteration. This reformulation is pivotal as it allows the application of parallel computing techniques to significantly reduce the computational time without compromising the generated image quality.

Triangular Anderson Acceleration (TAA)

A standout contribution of this study is the development of what is termed as Triangular Anderson Acceleration (TAA). This technique stems from the classical Anderson Acceleration method, modified to cater specifically to the triangular structure of the nonlinear equations underlying the sampling process. By leveraging information from previous iterations more effectively, TAA substantially enhances the efficiency of fixed-point iteration, demonstrating a remarkable reduction in the number of steps required to converge to a high-quality image.

Practical Enhancements

The study further explores two practical enhancements: early stopping and initialization from existing trajectories. Early stopping capitalizes on the observation that high-quality images often materialize before the convergence criteria are met, enabling a significant reduction in sampling steps. On the other hand, initializing the sampling process with a trajectory from a similar input condition has shown to facilitate even faster convergence, presenting a valuable strategy for iterative image generation tasks.

Empirical Validation

The efficacy of the proposed approach is empirically validated across multiple scenarios, including sampling with the DiT and Stable Diffusion models. The findings reveal that the ParaTAA algorithm, incorporating both TAA and the practical enhancements, substantially outperforms both the naive fixed-point iteration and its variant with an optimized order of nonlinear equations. This is quantitatively supported by metrics such as FID, Inception Score, and CLIP Score, along with qualitative examples illustrating the accelerated convergence towards high-quality images.

Implications and Future Directions

This research not only presents a significant advancement in the practical application of diffusion models but also opens up new avenues for exploration in parallel computing techniques within the realm of generative AI. The foundational concept of accelerating sampling through parallelization holds potential for adaptation and further innovation, particularly in extending its application to other domains such as video generation and beyond.

In conclusion, the proposed parallel sampling strategy signifies a pivotal step towards overcoming the computational barriers associated with diffusion models, thereby paving the way for more efficient and versatile generative AI applications. The remarkable reduction in sampling time, coupled with the preservation of image quality, underscores the foundational advancements that parallel computation and algorithmic refinements can deliver in the ongoing evolution of generative modeling techniques.

Create an account to read this summary for free:

Newsletter

Get summaries of trending comp sci papers delivered straight to your inbox:

Unsubscribe anytime.