Diffusion models demonstrate outstanding performance in image generation, but their multi-step inference mechanism requires immense computational cost. Previous works accelerate inference by leveraging layer or token cache techniques to reduce computational cost. However, these methods fail to achieve superior acceleration performance in few-step diffusion transformer models due to inefficient feature caching strategies, manually designed sparsity allocation, and the practice of retaining complete forward computations in several steps in these token cache methods. To tackle these challenges, we propose a differentiable layer-wise sparsity optimization framework for diffusion transformer models, leveraging token caching to reduce token computation costs and enhance acceleration. Our method optimizes layer-wise sparsity allocation in an end-to-end manner through a learnable network combined with a dynamic programming solver. Additionally, our proposed two-stage training strategy eliminates the need for full-step processing in existing methods, further improving efficiency. We conducted extensive experiments on a range of diffusion-transformer models, including DiT-XL/2, PixArt-α, FLUX, and Wan2.1. Across these architectures, our method consistently improves efficiency without degrading sample quality. For example, on PixArt-α with 20 sampling steps, we reduce computational cost by 54% while achieving generation metrics that surpass those of the original model, substantially outperforming prior approaches. These results demonstrate that our method delivers large efficiency gains while often improving generation quality.
Figures & tables
Figure 1: DiffSparse uses a learnable sparsity-cost predictor and dynamic programming to learn per-layer sparsity under target ratio R . We generate binary masks from the chosen sparsity maps and candidate masks. A token selector reuses features from previous diffusion steps to skip unimportant tokens and speed sampling. To enable gradient flow through the binary masks, we apply Straight-Through Estimation (STE) and train our model using full-step sampling targets with LPIPS loss.
Method
MACs (T) ↓
Speedup ↑
FID-30k ↓
CLIP ↑
PixArt- α ( Chen et al., 2024b )
2.86
1.00 ×
28.20
0.163
50% steps
1.43
1.74 ×
37.57
0.158
FORA (N=2) ( Selvaraju et al., 2024 )
1.43
1.64 ×
29.67
0.164
DeepCache ( N=2 ) ( Ma et al., 2024 )
1.48
1.61 ×
29.61
0.163
DiCache ( Bu et al., 2025 )
1.63
1.77 ×
28.19
0.164
ToCa ( Zou et al., 2025 )
1.64
1.75 ×
28.35
0.164
Table 1: Results of text-to-image generation on MS-COCO2017 with PixArt- α and 20 DPM++ steps.
Method
MACs (T) ↓
Speedup ↑
FID ↓
sFID ↓
Precision ↑
Recall ↑
DDIM-50 steps
11.44
1.00 ×
2.26
4.29
0.80
0.60
DDIM-40 steps
9.14
1.24 ×
2.39
4.28
0.80
0.59
DDIM-25 steps
5.73
1.96 ×
3.01
4.60
0.79
0.58
DDIM-20 steps
4.58
2.42 ×
3.48
4.64
0.79
0.56
FORA
4.13
2.12 ×
3.88
6.74
0.79
0.56
ToCa
4.97
2.09 ×
3.05
4.70
0.79
0.57
Table 2: Results of class-conditional generation with DiT-XL/2 and 50 DDIM steps on ImageNet.
Figure 2: Comparison of our method with the baseline (PixArt- α with DPM-Solver++ using 20 steps) and existing methods under different acceleration rates.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Method
MACs (T) ↓
Image Reward ↑
FLUX.1-schnell
13.247
1.064
75% Steps
9.936
1.063
ToCa
7.313
1.063
DiffSparse
7.316
1.184
Appendix
Table 9: Comparison in text-to-image generation for FLUX.1-schnell on PartiPrompts.
Variant
FID ↓
CLIP ↑
DiffSparse
26.91
0.164
−s1 (self-attention influence)
27.11
0.164
−s2 (cross-attention focus)
27.48
0.163
−s3 (cache-reuse frequency)
27.23
0.164
−B (spatial bonus)
27.05
0.164
Appendix
Table 10: Ablations of the attention-based score in text-to-image generation for PixArt- α .
Figure 3: Visualization of predicted layer sparsity of PixArt- α with 20 steps. In the figure, the x-axis denotes different network layers, the y-axis denotes sampling time steps, and the color gradient from blue to yellow indicates increasing sparsity.
Figure 4: Comparison of our method with the baseline (PixArt- α with DPM-Solver++ using 20 steps) under different acceleration rates.
Figure 5: Comparison between our DiffSparse, and ToCa with the baseline (PixArt- α with DPM-Solver++ using 20 steps under 512 × 512 resolution).
Diffusion Transformers (DiTs) have achieved state-of-the-art performance in image and video generation, but their success comes at the cost of heavy computation. This inefficiency is largely due to the fixed tokenization process, which uses constant-sized patches throughout the entire denoising phase, regardless of the content's complexity. We propose dynamic tokenization, an efficient test-time strategy that varies patch sizes based on content complexity and the denoising timestep. Our key insight is that early timesteps only require coarser patches to model global structure, while later iterations demand finer (smaller-sized) patches to refine local details. During inference, our method dynamically reallocates patch sizes across denoising steps for image and video generation and substantially reduces cost while preserving perceptual generation quality. Extensive experiments demonstrate the effectiveness of our approach: it achieves up to 3.52× and 3.2× speedup on FLUX-1.Dev and Wan 2.1, respectively, without compromising the generation quality and prompt adherence.
Diffusion Transformers (DiT) achieve strong performance in image generation but incur substantial inference costs. While prior work has reduced this cost via quantization and distillation, semi-structured sparsity, which can nearly halve FLOPs, remains underexplored. A key reason is that most existing approaches focus on weight sparsification, and pruning 50% of the weights can remove critical model capacity and degrade generation quality. Our study, however, shows that DiT activations are intrinsically sparse and significantly more robust to N:M semi-structured sparsification than weights. Motivated by this observation, we advocate a paradigm shift from weight sparsification to activation sparsification. We propose RT-Lynx, which applies N:M sparsification to activations and incorporates error-compensation techniques to mitigate accuracy loss. We further implement highly optimized CUDA kernels tailored to this setting, achieving up to a 1.55x speedup on average in linear layers. Extensive experiments across multiple diffusion models demonstrate that our method preserves the generation quality of the original models while substantially accelerating inference.
Diffusion Transformers have established a new state-of-the-art in image synthesis, but the high computational cost of iterative sampling severely hampers their practical deployment. While existing acceleration methods often focus on the temporal domain, they overlook the substantial spatial redundancy inherent in the generative process, where global structures emerge long before fine-grained details are formed. The uniform computational treatment of all spatial regions represents a critical inefficiency. In this paper, we introduce Just-in-Time (JiT), a novel training-free framework that addresses this challenge by acceleration in the spatial domain. JiT formulates a spatially approximated generative ordinary differential equation (ODE) that drives the full latent state evolution based on computations from a dynamically selected, sparse subset of anchor tokens. To ensure seamless transitions as new tokens are incorporated to expand the dimensions of the latent state, we propose a deterministic micro-flow, a simple and effective finite-time ODE that maintains both structural coherence and statistical correctness. Extensive experiments on the state-of-the-art FLUX.1-dev model demonstrate that JiT achieves up to a 7x speedup with nearly lossless performance, significantly outperforming existing acceleration methods and establishing a new and superior trade-off between inference speed and generation fidelity.
Wenhao Sun, Ji Li, Zhaoqiang Liu
University of Electronic Science and Technology of China · Capital Normal University