Elastic Diffusion Transformer

1. Problem Statement

Diffusion Transformers (DiTs) have become the state-of-the-art architecture for generative modeling across images, videos, and 3D assets. However, they suffer from high computational costs due to large model sizes and the requirement for dozens of denoising steps.

• Limitations of Existing Methods: Traditional acceleration techniques like pruning and distillation rely on fixed, static architectures. They apply a uniform reduction in capacity (e.g., removing specific layers or reducing parameters) regardless of the input sample or the denoising step.
• The Core Issue: This static approach ignores the sample-dependent sparsity inherent in the generation process. Different inputs (e.g., simple vs. complex scenes) and different denoising timesteps require varying levels of computational effort. Static methods often lead to a suboptimal trade-off, either sacrificing generation quality for speed or failing to achieve significant acceleration.

2. Methodology: Elastic Diffusion Transformer (E-DiT)

E-DiT is an adaptive acceleration framework designed to dynamically allocate computational resources based on the specific needs of each sample and timestep. It introduces three complementary mechanisms:

A. Router Architecture

Each DiT block is equipped with a lightweight router conditioned on the input latent representation and the current denoising timestep. The router outputs two predictions:

1. Skip Probability ($p_t^i$): A scalar indicating whether the current block can be skipped entirely.
2. Width Distribution ($q_t^i$): A probability distribution over predefined MLP width reduction ratios (e.g., 1/4, 1/2, 3/4, 1) for blocks that are not skipped.

B. Adaptive Block Skipping

• Mechanism: If the router's skip probability falls below a threshold $\tau$ (e.g., 0.5), the block is bypassed, and the input passes directly to the next block.
• Training: To handle the non-differentiable nature of skipping, the authors use the Straight-Through Estimator (STE). They also apply a gating loss to encourage the average skip rate to match a target budget.
• Inference: Blocks with $p_t^i < \tau$ are completely skipped, eliminating their computation.

C. Adaptive MLP Width Reduction

• Mechanism: For blocks that are not skipped, the router dynamically selects an optimal MLP width. Instead of using the full hidden dimension ($H$), the model activates only a fraction of the channels (e.g., 50% or 75%) based on the predicted width ratio.
• Implementation: During training, this is achieved via masking; during inference, it is implemented via explicit matrix slicing to avoid computing on deactivated channels.

D. Block-Level Feature Caching (Training-Free)

• Motivation: Some blocks have skip probabilities close to the threshold (the "borderline" region). Skipping them directly might degrade quality, but computing them is redundant.
• Mechanism: E-DiT leverages temporal redundancy. If a block's probability falls within a borderline range $[\tau, \tau + \delta]$, the model reuses the residual update ($\Delta$) from a previous timestep where the block was fully computed.
• Constraint: Cached features are reused up to a maximum limit $K$ to prevent error accumulation. This mechanism requires no additional training.

3. Key Contributions

1. Sample-Adaptive Framework: E-DiT is the first framework to jointly optimize block skipping and MLP width reduction dynamically based on input content and timestep, rather than using a static architecture.
2. Lightweight Router Design: The introduction of a minimal router per block that predicts both skip decisions and width ratios without significantly increasing model parameters.
3. Training-Free Caching: A novel inference-time mechanism that uses router predictions as a principled criterion for feature reuse, further accelerating inference without retraining.
4. General Applicability: The method is designed to be backbone-agnostic, successfully applied to diverse models (Qwen-Image, FLUX) and modalities (2D images, 3D assets).

4. Experimental Results

The authors evaluated E-DiT on Qwen-Image, FLUX (2D image generation), and Hunyuan3D-3.0 (3D asset generation).

• Performance vs. Speed:
  • Qwen-Image: Achieved a ~2x speedup (latency reduced from 2431ms to 1283ms for the "turbo" version) with negligible loss in quality (GenEval score dropped only slightly from 0.870 to 0.853).
  • FLUX: Reduced latency from 715ms to 374ms while maintaining competitive scores on DPG-Bench and GenEval.
  • Hunyuan3D-3.0: Achieved nearly 2x speedup (5012ms $\to$ 2587ms) with comparable ULIP-I and Uni3D-I scores.
• Comparison with Baselines: E-DiT outperformed static pruning methods (like PPCL, TinyFusion) and other dynamic methods (like DyDiT) in terms of the quality-efficiency trade-off.
• Ablation Studies:
  • Adaptive vs. Static: Static models with random pruning showed significantly worse quality, proving that adaptive computation is crucial.
  • Initialization: Starting training with full capacity (no skipping) and gradually learning to prune yielded better stability and final performance than random initialization.
  • Caching: The block-wise caching mechanism provided additional latency reduction without quality degradation, whereas directly skipping "borderline" blocks hurt quality.

5. Significance

• Efficiency without Compromise: E-DiT demonstrates that high-fidelity generative models can be made significantly more efficient by exploiting the inherent sparsity of the generation process, rather than forcing a smaller, static model.
• Scalability: The framework is compatible with large-scale foundation models (e.g., 20B+ parameters) and diverse modalities, making it highly relevant for deploying generative AI on resource-constrained devices.
• New Paradigm: It shifts the paradigm from "static compression" to "dynamic elastic computation," suggesting that future generative models should adapt their computational depth and width in real-time based on input complexity.

Code Availability: The authors have released the code at https://github.com/wangjiangshan0725/Elastic-DiT.