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.
Figures & tables
Figure 1: Left: By learning to re-allocate compute across spatial regions and timesteps, DC-DiT improves the FID–FLOPs Pareto frontier of DiT. Right: DC-DiT enables elastic inference, where a single trained checkpoint can be evaluated at multiple inference budgets while preserving visual quality.
Figure 2: Architecture of DC-DiT. The isotropic encoder aggregates local context across the input tokens. The chunking layer selects a subset of boundary tokens via a learned routing module, yielding a compressed sequence that is processed by the DiT blocks. The de-chunking layer restores the original resolution through spatial smoothing followed by plug-back.
Figure 3: Adaptive compute allocation learned by DC-DiT. Across diffusion timesteps, the router retains boundary tokens in spatially informative regions such as object structure, edges, and texture, while compressing more predictable background regions. The retained-token pattern changes over the denoising trajectory, allocating compute differently as images evolve from noisy structure to fine detail.
Scale
Model
Params (M)
Tail drop
TFLOPs/img ↓
FID-50K ↓
256×256, no classifier-free guidance
S
DiT-S/2
33
–
3.02
68.40
S
DC-DiT-S
34.7
0.2
3.06
58.43 (-14.6%)
S
DC-DiT-S
34.7
0.5
2.14 (-29.1%)
66.56
B
DiT-B/2
131
–
11.46
43.47
B
DC-DiT-B
137
0.2
11.94
33.91 (-22.0%)
Table 1: Main ImageNet results across model scale, guidance setting, and resolution. In Lite-CFG rows, the tail-drop value applies to the unconditional branch.
Figure 4: DCDiT is complementary to TeaCache and DyDiT, improving FID at comparable compute budgets.
Model
Tail drop
TFLOPs/img
Relation
Entity
Other
Attribute
Global
DPG-Bench
Z-Image
–
5603
92.72
89.95
90.02
91.08
93.06
85.97
DC-Z-Image
0
5901
93.15
92.38
89.47
90.72
91.03
86.82
DC-Z-Image
30%
4211
92.89
91.70
88.15
92.02
89.34
87.64
DC-Z-Image
60%
2570
92.03
90.79
92.91
89.23
91.64
86.32
Table 2: DPG-Bench preservation after upcycling Z-Image. DC-Z-Image preserves prompt-following quality across elastic inference budgets while reducing inference compute at higher tail-dropping fractions.
Figure 6: Tail-dropping robustness on DC-DiT-S. Multi-budget training keeps FID stable as ρ increases.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7: Extended qualitative elastic inference results. This figure expands the tail-dropping grid from Figure 1 , showing additional generations from the same checkpoint across user-selected compute budgets.
Figure 8: Extended adaptive compute visualization. This figure expands Figure 3 with additional examples of the router’s learned boundary predictions across diffusion timesteps.
Resolution
Scale
DiT TFLOPs/img
DiT steps
DC-DiT steps
256×256
S
3.02
400K
421,230
256×256
B
11.46
400K
408,071
256×256
L
40.21
400K
486,372
256×256
XL
59.00
400K
494,943
512×512
B
53.09
400K
432,097
512×512
XL
261.25
400K
495,365
Appendix
Table 4: Compute-matched training lengths for the ImageNet DC-DiT experiments. DC-DiT trains longer than 400K steps when multi-budget tail dropping reduces the average cost per post-warmup step.
Configuration
Enc. blocks
Transformer blocks
Dec. blocks
Hidden width
Heads
Enc./Dec. bottleneck
DC-DiT-S
2
12
2
384
6
96
DC-DiT-B
2
12
2
768
12
192
DC-DiT-L
2
24
2
1024
16
256
DC-DiT-XL
2
28
2
1152
16
288
Appendix
Table 5: Basic architecture shapes for the DC-DiT scale configurations.
Hyperparameter
Value
Dataset
ImageNet, 1000 classes
Latent resolution
256×256
Diffusion schedule
Linear, 1000 training steps
Sampler
DDPM, 250 sampling steps
Global batch size
256
Optimizer
AdamW
Appendix
Table 6: Main hyperparameters for ImageNet DC-DiT training.