Communication-Efficient Diffusion Denoising Parallelization via Reuse-then-Predict Mechanism

TL;DR

ParaStep addresses diffusion-model inference latency by exploiting the similarity between adjacent denoising steps through a reuse-then-predict mechanism that enables step-wise parallelism with lightweight communication. It introduces a distributed sampling framework that reduces inter-device data transfer while preserving generation quality, achieving substantial speedups across image, video, and audio modalities. A single-device variant (BatchStep) and TeaCache integration offer additional flexibility and quality gains, especially for non-compute-intensive models. Comprehensive experiments and a detailed communication analysis demonstrate strong scalability and practical applicability on commodity hardware, making diffusion-model deployment more accessible in bandwidth-constrained environments.

Abstract

Diffusion models have emerged as a powerful class of generative models across various modalities, including image, video, and audio synthesis. However, their deployment is often limited by significant inference latency, primarily due to the inherently sequential nature of the denoising process. While existing parallelization strategies attempt to accelerate inference by distributing computation across multiple devices, they typically incur high communication overhead, hindering deployment on commercial hardware. To address this challenge, we propose \textbf{ParaStep}, a novel parallelization method based on a reuse-then-predict mechanism that parallelizes diffusion inference by exploiting similarity between adjacent denoising steps. Unlike prior approaches that rely on layer-wise or stage-wise communication, ParaStep employs lightweight, step-wise communication, substantially reducing overhead. ParaStep achieves end-to-end speedups of up to \textbf{3.88}$\times$ on SVD, \textbf{2.43}$\times$ on CogVideoX-2b, and \textbf{6.56}$\times$ on AudioLDM2-large, while maintaining generation quality. These results highlight ParaStep as a scalable and communication-efficient solution for accelerating diffusion inference, particularly in bandwidth-constrained environments.

Figures (14)

• Figure 1: Relative MAE (as defined in Equation \ref{['eq:rel_mae']}) between adjacent denoising steps $t$ and $t+1$ for the noisy sample $\mathbf{x}_t$ and the predicted noise $\boldsymbol{\epsilon}_t$. Here, 0% on the x-axis indicates the first step of the denoising process, and 100% indicates the last.
• Figure 2: Dotted lines show the difference in noisy sample $\mathbf{x}_t$ between the original model and direct reuse process (stride = 2), while solid lines compare the original model with our reuse-then-predict process (degree = 2). Results for predicted noise $\boldsymbol{\epsilon}_t$ follow the same pattern. Reuse-then-predict process results in smaller deviations from the original model compared to direct reuse process.
• Figure 3: Illustration of the computation process of a diffusion model. (a) The original computation process. (b) Reusing noise $\epsilon_{T}$ from the previous timestep to skip the computation of noise prediction in the current timestep $T-1$. (c) ParaStep: adjacent-step noise prediction is distributed across GPUs using reuse-then-predict, enabling parallel denoising with minimal communication. Since the computational cost of scheduler operations is negligible, the noise predictor computations on GPU0, GPU1, and GPU2 are fully parallelized.
• Figure 3: Generation latency and FAD on AudioLDM2-large. $p$ is the degree of parallelism, $s$ is the cycle length in BatchStep.
• Figure 4: Utilizing batching effect of audio diffusion models, parallel computation of noise predictor can be transformed batching execution on single device.