Elastic Diffusion Transformer

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.

Diffusion-based AI models can generate highly realistic images and 3D content, but they are often slow and expensive to run because they perform many complex calculations during generation. Existing methods for speeding up these models usually use a fixed amount of computation for every input, which can either limit the speed improvement or reduce the quality of the generated results. In this work, we introduce Elastic Diffusion Transformer (E-DiT), a method that allows these models to adapt their computation dynamically for each generated sample. We found that many parts of the generation process are not always necessary, and the amount of useful computation can vary significantly depending on the input. Based on this observation, E-DiT uses lightweight decision modules to determine which computations can be skipped and how much computation is needed at different stages of the model. Furthermore, we also introduce a feature caching strategy that avoids repeating unnecessary calculations during generation without requiring additional training. Experiments on advanced image and 3D generation models show that E-DiT can make generation up to twice as fast while maintaining nearly the same visual quality. Our work could help make high-quality generative AI systems more efficient and practical for real-world applications.