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.
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) achieve high-quality generation but are costly due to iterative sampling. Dynamic-resolution sampling reduces early-stage cost by denoising at low resolution; however, uniformly upsampling all latent tokens at resolution transitions incurs redundant computation and may degrade fine-detail consistency. Existing partial upsampling strategies typically rely on local latent structure cues or single-step statistics, making it difficult to jointly capture token-text semantic relevance and token-wise representation dynamics across diffusion steps. We propose AViTS, an adaptive spatiotemporal token selection framework for dynamic-resolution DiTs. AViTS models spatial importance via latent-text attention and temporal importance via token-level feature variation across diffusion timesteps, and fuses them to enable spatiotemporal importance-aware selective upsampling: it prioritizes resolution refinement for critical tokens while deferring less important ones, thereby reducing redundant high-resolution computation and improving the quality-efficiency trade-off. AViTS achieves up to 6.34x on FLUX and nearly 9x FLOPs reduction on Qwen-Image-Edit and FLUX.1-Kontext-dev, orthogonal to distillation, quantization, and feature caching, and reaching 14.76x with distilled models. Code: https://github.com/QHR69/AViTS
Haoran Qin, Zhengan Yan, Shikang Zheng +9
Shanghai Jiao Tong University, China · Shandong University, China · Terminal Intelligent Computing Division, Alibaba Cloud, China +2
Image and video diffusion models allocate equal computation to every region, even when the intended scene calls for varying levels of detail. The spatial distribution of detail can often be anticipated before generation, indicating where computation can be reduced. We introduce Level-of-Token (LoT) Diffusion, a framework that turns this knowledge into an explicit multiresolution token layout (Level-of-Token layout) for adaptive and efficient generation. Tokens represent rectangular patches of varying sizes and shapes, allocating finer tokens where detail is needed and coarser tokens elsewhere. We adapt pretrained diffusion transformers to LoT layouts through a patch-wise asymmetric flow parametrization and embeddings for multiresolution tokens, preserving full-resolution flow prediction at every denoising step while processing only a reduced token sequence. LoT Diffusion enables layout-adaptive generation while preserving pretrained generative priors. We demonstrate LoT with layouts derived from semantic masks, bounding boxes, texture variance, and depth-of-field cues, as well as agentic plans. Across image and video generation, LoT offers favorable quality-efficiency tradeoffs, with significant speedups determined by the layout's token budget. Our project website is at https://georgenakayama.github.io/lotdiffusion/.