Organizations: School of Computer Science, Wuhan University · School of Artificial Intelligence, Wuhan University · School of Data Science, Fudan University
Diffusion Probabilistic Models (DPMs) have achieved great success in image generation but suffer from high inference latency due to their iterative denoising nature. Motivated by the evolving feature dynamics across the denoising trajectory, we propose a novel framework to optimize the computational graph of pre-trained DPMs on a per-timestep basis. By learning timestep-specific masks, our method dynamically determines which blocks to execute or bypass through feature reuse at each inference stage. Unlike global optimization methods that incur prohibitive memory costs via full-chain backpropagation, our method optimizes masks for each timestep independently, ensuring a memory-efficient training process. To guide this process, we introduce a timestep-aware loss scaling mechanism that prioritizes feature fidelity during sensitive denoising phases, complemented by a knowledge-guided mask rectification strategy to prune redundant spatial-temporal dependencies. Our approach is architecture-agnostic and demonstrates significant efficiency gains across a broad spectrum of models, including DDPM, LDM, DiT, and PixArt. Experimental results show that by treating the denoising process as a sequence of optimized computational paths, our method achieves a superior balance between sampling speed and generative quality. Our code will be released.
Figures & tables
Figure 1: Qualitative demonstration of our method on representative diffusion architectures. In these instances, our method achieves substantial speedups—specifically 1.63 × , 2.75 × , and 1.48 × —while preserving high-fidelity generation quality with high visual consistency.
Figure 2: Illustration of our method on a UNet-based model (left) and a Transformer-based model (right). During the denoising phase, upon receiving the input xt , the model checks the mask mt values corresponding to each block. If the mask value is 1, it performs computations and updates the cached features for that block; if the mask value is 0, it skips the computation and uses the existing cached features.
Figure 3: Illustration of the end-to-end mask optimization framework. The original DPM is frozen, serving as a teacher to provide reference features xt . For each timestep, the learnable mask mt is updated by the total loss Lt . The Detach operation ensures that gradients are restricted to the current timestep, enabling memory-efficient per-step optimization.
Figure 4: Illustration of mask rectification. Within the same timestep, if the subsequent block does not require its computation result as input, and the same block in the next timestep does not need to reuse its features, then the computation of this block can be safely skipped.
Figure 5: Acceleration for PixArt-Sigma-XL on MS-COCO 1024 × 1024. Our method not only achieves a 1.43 × speedup, but also maintains excellent image generation quality and text-to-image semantic consistency.
CIFAR-10 32 × 32
Method
Extra Data
Training Time ↓
MACs ↓
Speed ↑
FID ↓
DDPM [ 13 ]
–
–
0.61T
1 ×
4.19
DDPM*
–
–
0.61T
1 ×
4.25
Diff-Pruning [ 7 ]
✓
–
0.34T
1.37 ×
5.29
CT [ 42 ] *
✓
–
–
1.62 ×
4.68
DeepCache [ 30 ]
✗
✗
0.35T
1.61 ×
4.70
Table 1: Unconditional generation quality using DDPM on CIFAR-10, LSUN-Bedroom, and LSUN-Churches. All the methods here adopt 100 DDIM steps, except for CT. CT means consistency model (with 55 denoising steps). * means the reproduced results. Bold indicates the best performance.
ImageNet 256 × 256
Method
NFE
Extra Data
Training Time ↓
MACs ↓
Speed ↑
IS ↑
FID ↓
sFID ↓
Precision ↑
Recall ↑
LDM-4 [ 34 ] *
250
–
–
25.0T
1 ×
206.45
3.42
5.14
82.83
53.13
Diff-Pruning [ 7 ]
250
✓
–
13.2T
1.51 ×
201.81
9.16
10.59
87.87
30.87
DeepCache [ 30 ]
250
✗
✗
9.1T
2.65 ×
202.79
3.44
5.11
82.65
53.81
Ours
250
✗
2.9h
8.3T
2.75 ×
205.76
3.51
5.00
82.74
52.81
DiT-XL/2 [ 33 ] *
100
–
–
11.9T
1 ×
242.8
2.16
4.45
80.35
60.34
Table 2: Class-conditional generation quality on ImageNet using LDM-4-G and DiT-XL/2. The baselines here, as well as our methods, employ the DDIM scheduler. * means the reproduced results. Bold indicates the best performance.
Method
NFE
Extra Data
Training Time ↓
MACs ↓
Speed ↑
IS ↑
FID ↓
CLIP Score ↑
PixArt-Sigma-XL [ 4 ] *
50
–
–
107.8T
1 ×
55.65
24.33
31.27
PixArt-Sigma-XL
38
✗
✗
81.9T
1.32 ×
50.88
24.89
30.02
DitFastAttn [ 51 ]
50
✗
0.3h
–
1.24 ×
52.74
23.94
31.18
Ours
50
✗
4.2h
71.9T
1.43 ×
55.52
22.97
31.45
Table 3: Prompt-conditional generation quality on MS-COCO 1024 × 1024 using PixArt-Sigma-XL. The baselines here, as well as our methods, employ the DPM-Solver scheduler. * means the reproduced results. Bold indicates the best performance.
Figure 6: The visualization results of the mask trained by our method. The color of each rectangle represents the number of times the corresponding block performs computations over 5 timesteps; the darker the color, the more computations.
Figure 7: The visualization results of the mask’s values show that most elements have converged to near 0 or 1.
CIFAR-10 32 × 32
Sampling Method
Speed ↑
FID ↓
Random
1.63 ×
4.66
Gumbel
1.29 ×
4.67
Table 4: Performance comparison with random sampling and Gumbel-Softmax sampling.
CIFAR-10 32 × 32
Rectification
Loss Scaling
Bi-modal Loss
Speed ↑
FID ↓
✓
✓
✓
1.63 ×
4.66
✗
✓
✓
1.49 ×
4.66
✓
✗
✓
1.62 ×
4.74
✓
✓
✗
1.63 ×
4.72
Table 5: Ablation study on timestep-aware loss scaling, knowledge-based model rectification, and the introduction of bi-modal loss.
Diffusion models have demonstrated remarkable success in image and video generation, yet their practical deployment remains hindered by the substantial computational overhead of multi-step iterative sampling. Among acceleration strategies, caching-based methods offer a training-free and effective solution by reusing or predicting features across timesteps. However, existing approaches rely on fixed or locally adaptive schedules without considering the global structure of the denoising trajectory, often leading to error accumulation and visual artifacts. To overcome this limitation, we propose DPCache, a novel training-free acceleration framework that formulates diffusion sampling acceleration as a global path planning problem. DPCache constructs a Path-Aware Cost Tensor from a small calibration set to quantify the path-dependent error of skipping timesteps conditioned on the preceding key timestep. Leveraging this tensor, DPCache employs dynamic programming to select an optimal sequence of key timesteps that minimizes the total path cost while preserving trajectory fidelity. During inference, the model performs full computations only at these key timesteps, while intermediate outputs are efficiently predicted using cached features. Extensive experiments on DiT, FLUX, and HunyuanVideo demonstrate that DPCache achieves strong acceleration with minimal quality loss, outperforming prior acceleration methods by +0.031 ImageReward at 4.87× speedup and even surpassing the full-step baseline by +0.028 ImageReward at 3.54× speedup on FLUX, validating the effectiveness of our path-aware global scheduling framework. Code is available at https://github.com/argsss/DPCache.
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 models achieve state-of-the-art generative performance but suffer from high computational costs during inference due to the repeated evaluation of a heavy neural network. In this work, we propose Dual-Rate Diffusion, a method to accelerate sampling by interleaving the execution of a heavy high-capacity context encoder and a light efficient denoising model. The context encoder is evaluated sparsely to extract high-dimensional features, which are effectively reused by the light denoising model at every step to refine the sample efficiently. This approach significantly accelerates inference without compromising sample quality. On ImageNet benchmarks, Dual-Rate Diffusion matches the performance of standard baselines while reducing computational cost by a factor of 2-4. Furthermore, we demonstrate that our method is compatible with distillation techniques, such as Moment Matching Distillation, enabling further efficiency gains in few-step generation.
Grigory Bartosh, David Ruhe, Emiel Hoogeboom +3
*Work done while interning at Google DeepMind Amsterdam, 1University of Amsterdam, 2Google DeepMind Amsterdam · 2Google DeepMind Amsterdam