BlockDance: Reuse Structurally Similar Spatio-Temporal Features to Accelerate Diffusion Transformers

TL;DR

BlockDance tackles slow inference in Diffusion Transformers by caching structurally similar STSS features after their structure stabilizes and reusing them to cut redundant computation. It introduces BlockDance, a training-free acceleration scheme, and BlockDance-Ada, a lightweight RL-based method that dynamically decides per-sample caching policy to balance speed and quality. Across image and video tasks on models such as PixArt-$\alpha$, DiT-XL/2, and Open-Sora, BlockDance achieves 25%–50% speedups with fidelity preserved, while BlockDance-Ada further improves quality at similar acceleration. The approach is plug-and-play and shows strong potential for real-time diffusion applications, with ablations supporting the design choices and a limitation identified for very low step counts.

Abstract

Diffusion models have demonstrated impressive generation capabilities, particularly with recent advancements leveraging transformer architectures to improve both visual and artistic quality. However, Diffusion Transformers (DiTs) continue to encounter challenges related to low inference speed, primarily due to the iterative denoising process. To address this issue, we propose BlockDance, a training-free approach that explores feature similarities at adjacent time steps to accelerate DiTs. Unlike previous feature-reuse methods that lack tailored reuse strategies for features at different scales, BlockDance prioritizes the identification of the most structurally similar features, referred to as Structurally Similar Spatio-Temporal (STSS) features. These features are primarily located within the structure-focused blocks of the transformer during the later stages of denoising. BlockDance caches and reuses these highly similar features to mitigate redundant computation, thereby accelerating DiTs while maximizing consistency with the generated results of the original model. Furthermore, considering the diversity of generated content and the varying distributions of redundant features, we introduce BlockDance-Ada, a lightweight decision-making network tailored for instance-specific acceleration. BlockDance-Ada dynamically allocates resources and provides superior content quality. Both BlockDance and BlockDance-Ada have proven effective across various generation tasks and models, achieving accelerations between 25% and 50% while maintaining generation quality.

Figures (14)

• Figure 1: BlockDance accelerates DiT models DiT-XL/2, PixArt-$\alpha$ and Open-Sora by 37.4%, 25.4% and 34.8% respectively, while maintaining fidelity and high consistency with the original image.
• Figure 2: Feature similarity and redundancy in DiTs. (a) In the denoising process, the outputs of DiT blocks exhibit high similarity in adjacent steps, particularly in the gray shadow-masked region where the structure is stabilized. (b) This high similarity is mainly concentrated in the shallow and middle blocks within the transformer, i.e. between 0 and 20 blocks, which focus on low-level structures. Thus, redundant computation related to highly similar structural features in the denoising process can be saved by reusing them to accelerate DiTs inference while maintaining quality.
• Figure 3: Visualization of PixArt-$\bm{\alpha}$ blocks. Feature maps with blue borders exhibit high similarity across adjacent steps.
• Figure 4: An overview of BlockDance. The reuse step generates $\mathbf{z}_{t-1}$ based on reusing the structural features from the cache step, saving the computation of the first $i$ blocks to accelerate inference.
• Figure 5: An overview of BlockDance-Ada. Given the intermediate latent and prompt embedding, the reuse decision network learns the structural complexity of each sample and derives the corresponding reuse policy. These policies determine whether each subsequent step in DiTs is a cache step or a reuse step. The reward function balances the trade-off between image quality and speed.