Elastic Diffusion Transformer

TL;DR

The paper introduces Elastic Diffusion Transformer (E-DiT), an adaptive framework that dynamically skips computations in diffusion models, significantly improving efficiency with minimal quality loss across image and 3D asset generation.

Contribution

E-DiT is the first adaptive acceleration method for DiT that uses sample-dependent sparsity and lightweight routers to skip blocks and reduce computation during inference.

Findings

Achieves up to 2x speedup in diffusion model generation.

Maintains high generation quality with negligible loss.

Effective across 2D images and 3D assets.

Abstract

Diffusion Transformers (DiT) have demonstrated remarkable generative capabilities but remain highly computationally expensive. Previous acceleration methods, such as pruning and distillation, typically rely on a fixed computational capacity, leading to insufficient acceleration and degraded generation quality. To address this limitation, we propose \textbf{Elastic Diffusion Transformer (E-DiT)}, an adaptive acceleration framework for DiT that effectively improves efficiency while maintaining generation quality. Specifically, we observe that the generative process of DiT exhibits substantial sparsity (i.e., some computations can be skipped with minimal impact on quality), and this sparsity varies significantly across samples. Motivated by this observation, E-DiT equips each DiT block with a lightweight router that dynamically identifies sample-dependent sparsity from the input latent. Each router adaptively determines whether the corresponding block can be skipped. If the block is not skipped, the router then predicts the optimal MLP width reduction ratio within the block. During inference, we further introduce a block-level feature caching mechanism that leverages router predictions to eliminate redundant computations in a training-free manner. Extensive experiments across 2D image (Qwen-Image and FLUX) and 3D asset (Hunyuan3D-3.0) demonstrate the effectiveness of E-DiT, achieving up to $\sim$2$\times$ speedup with negligible loss in generation quality. Code will be available at https://github.com/wangjiangshan0725/Elastic-DiT.