Real-time video generation, interactive world models, and on-device image synthesis are almost a reality. Diffusion Transformers (DiT) achieve excellent sample quality, shifting the bottleneck to the throughput of the models themselves. Natural images are inherently sparse, as large portions of an image are often background, contribute little to the scene content, and are easy to denoise. DiTs fail to exploit this sparsity, as they process image tokens statically and uniformly. We propose MoDDiT, a simple architecture modification that learns to route tokens through both MLP and self-attention operations. By routing tokens through both blocks independently, MoDDiT allows for dynamic compute allocation in MLPs that varies across tokens, images, and denoising timesteps. Using MoDDiT to train image generation models from scratch significantly reduces computational costs of training and inference, processing as few as 12.5% of the tokens. By increasing sparsity and model depth, MoDDiT reduces costs and improves generation performance over unrouted baselines. The learned allocation is interpretable, revealing how computation varies with visual content and noise level. We also use MoDDiT to train video generation models, where allocation correlates with motion, and to distill a dense bidirectional video model into a sparse autoregressive student.
Figures & tables
Figure 1: Image generation and learned compute allocation in MoDDiT for a red panda example. (a) Generated images along the denoising trajectory, from left to right. (b) Self-attention allocation. (c) MLP allocation. Patch brightness indicates the proportion of blocks applied to each token at the given timestep; the strips below the allocation panels show register-token allocation.
Figure 2: Token routing in MoDDiT. Top: self-attention uses fixed top- K routing per sample, selecting two of four tokens at the illustrated 50% capacity. Bottom: MLP routing shares the training budget across samples, selecting three green tokens and one blue token. Colors denote samples; striped overlays show additive updates. Funnels gather selected tokens and scatter their updates back to the original positions. The router gates update ( × ) before residual addition ( + ); unselected tokens retain their original values.
Figure 3: ImageNet performance—we report FID-50K, sFID-50K, and IS for samples generated with varying numbers of diffusion timesteps. Precision and recall evaluation is in section A.2 .
Figure 4: MoDDiT compute allocation increases at lower noise levels, and shows a preference towards classes with higher clean latent variance.
Model
Parameters
Denoiser
PSNR (dB) ↑
(M)
(TFLOPs) ↓
+0.5s
+1.0s
+2.0s
ST-DiT
651.2
7,746.3
24.8
22.6
21.0
ST-MoDDiT (12.5%)
942.7
4,886.4
25.4
23.0
21.2
Δ (ST-MoDDiT − ST-DiT)
+291.5
− 2,859.8
+ 0.68
+ 0.38
+ 0.21
95% CI
[0.486, 0.871]
[0.229, 0.515]
[0.075, 0.344]
Table 1: 1X video prediction at 200k EMA checkpoints on 400 paired validation clips from 249 segments. The final two rows give sparse-minus-dense deltas and their pointwise 95% paired segment-bootstrap confidence intervals for PSNR gains. Denoiser TFLOPs cover the first clip’s full rollout, including both guidance branches ( section B.1 ). Deltas are calculated before rounding.
Figure 5: Compute allocation is strongly correlated with video motion—we show mean routing percentage for binned values of patch difference (clipped at the 99 th percentile). Patch difference is computed as the magnitude of the difference between generated patches in adjacent frames, and the model is trained on the 1X World Modeling Dataset, with a capacity of 12.5%.
Figure 6: Results of distilling a MoDDiT model on the VBench benchmark, as well as qualitative results compared to the dense baseline ( section C.2 ).
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7: ImageNet performance.
Model
TFLOPs/image
FID ↓
sFID ↓
IS ↑
Precision ↑
Recall ↑
DiT-L/2
165.380
5.069
4.714
208.861
0.834
0.488
MoDDiT (50%)
170.316
4.390
4.649
233.120
0.834
0.505
MoDDiT (25%)
129.888
4.795
4.790
212.914
0.820
0.510
MoDDiT (12.5%)
109.783
5.081
4.889
202.856
0.809
0.516
Appendix
Table 2: Evaluation at 512 DDPM steps with CFG scale 1.5. Quality metrics use 50,000 images; compute estimates average 1,024 images and includes both CFG branches.
Figure 8: Allocation is consistently correlated with mean latent variance of each class, and the correlation becomes smaller for smaller capacities.
Figure 9: Earlier routing sweep: FID and sFID versus counted denoiser arithmetic across sampling lengths and nominal capacities, using 10,000 images per point at 64, 128, 256, 512 denoising steps with classifier-free guidance scale of 4.0 .
Figure 10: Routing selection and capacity ablations. Shared whole-block routing, separate fixed-capacity routing, separate adaptive routing and random selection use 25% MoDDiT models trained for 400k updates from a common initialization.
Metric
Self-Forcing
Funnel50
Funnel20
Total
82.07
−4.04
−4.39
Quality
84.89
−3.73
−3.82
Semantic
70.82
−5.29
−6.68
Multiple objects
70.93
−18.05
−21.71
Spatial relationship
80.03
−8.72
−15.26
Object class
83.80
−9.79
−11.27
Appendix
Table 3: VBench scores for dense Self-Forcing and percentage-point changes for sparse EMA models relative to it (positive is better). Category scores are raw values multiplied by 100; aggregates use official normalization and weighting. Deltas are computed before rounding; categories are ordered by the Funnel20 point reduction, largest first. The final row gives absolute transformer TFLOPs for the matched clip described below (lower is better).
Figure 11: Matched MovieGen prompt 0043, seed 0, for dense Self-Forcing, Funnel50 and Funnel20. Columns show frames 0, 26, 53 and 79 at 16 fps.
Diffusion Transformers use attention-based networks to generate and transform images, video, audio, and other structured data through iterative denoising.
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 (DiTs) deliver remarkable image and video generation quality but incur high computational cost, limiting scalability and on-device deployment. We introduce CoReDiT, a structured token pruning framework for DiTs across vision tasks. CoReDiT uses a linear-time spatial coherence score to estimate local redundancy in the latent token lattice and skips high coherence (redundant) tokens in self-attention. To maintain a dense representation and avoid visual discontinuities, we reconstruct skipped attention outputs via coherence-guided aggregation of spatially neighboring retained tokens. We further introduce a progressive, block-adaptive pruning schedule that increases pruning gradually and allocates larger budgets to blocks and denoising steps with higher redundancy. Across state-of-the-art diffusion backbones including PixArt-α and MagicDrive-V2, CoReDiT achieves up to 55% self-attention FLOPs reduction and inference speedups of 1.33x on cloud GPUs and 1.72x on mobile NPUs, while maintaining high visual quality. Notably, CoReDiT also increases on-device memory head-room, enabling higher-resolution generation.
Diffusion Transformers rely on static patchify tokenization, assigning the same token budget to smooth backgrounds, detailed object regions, noisy early timesteps, and late-stage refinements. We introduce the Dynamic Chunking Diffusion Transformer (DC-DiT), which replaces fixed patchification with a learned encoder-router-decoder scaffold that adaptively compresses the 2D input into a shorter token sequence through a chunking mechanism learned end-to-end with diffusion training. DC-DiT allocates fewer tokens to predictable regions and noisy timesteps, and more tokens to detailed regions and later refinement stages, yielding meaningful spatial segmentations and timestep-adaptive compression schedules without supervision. Furthermore, the router provides an importance ordering over retained tokens, enabling elastic inference: a single checkpoint can be evaluated at flexible compute budgets with a smooth quality-compute tradeoff. Additionally, DC-DiT can be upcycled from pretrained DiT checkpoints and is also compatible with orthogonal dynamic computation approaches. On class-conditional ImageNet generation, DC-DiT reduces inference FLOPs by up to 36.8% and improves FID by up to 37.8% over DiT baselines, yielding a stronger quality--compute Pareto frontier across model scales, resolutions, and guidance settings. More broadly, these results suggest that adaptive tokenization is a general mechanism for making visual generation both more efficient and more flexible at inference time.