Source-linked AI summary
DC-DiT: Adaptive Compute and Elastic Inference for Visual Generation via Dynamic Chunking
Akash Haridas, Utkarsh Saxena, Parsa Ashrafi Fashi, Mehdi Rezagholizadeh, Vikram Appia, Emad Barsoum
TL;DR
Diffusion Transformers use static tokenization even though image detail and denoising stages vary in computational needs. DC-DiT learns adaptive tokenization and multi-budget routing, reducing inference FLOPs by up to 36.8% and improving FID by up to 37.8% over DiT baselines on ImageNet.
Problem
Fixed patchification gives smooth and detailed regions, as well as noisy and refined timesteps, the same token budget despite differing computational needs.
Method
DC-DiT uses an encoder–router–decoder scaffold with supervised-free adaptive tokenization and multi-budget training for elastic inference.
Results
Up to 36.8% lower inference FLOPs and up to 37.8% better FID are reported than fixed-patch DiT baselines across ImageNet model scales, resolutions, and guidance settings.
Takeaways & Limitations
Adaptive tokenization serves as a practical primitive for efficient diffusion models with flexible inference budgets and compatibility with other dynamic computation methods.
Takeaways & Limitations
At 512px, vanilla DiT-S/2 and DiT-L/2 baselines repeatedly encountered loss divergence despite multiple restarts.
Abstract
from arXiv · showhide
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.
1 Introduction
DC-DiT replaces fixed patchification with learned adaptive tokenization that reallocates compute across spatial regions and denoising timesteps. Multi-budget routing enables elastic inference, while ImageNet results improve the quality–compute frontier over fixed-patch DiT baselines.
- Fixed patchification assigns identical token budgets to regions and timesteps despite differing spatial detail and coarse-to-fine diffusion trajectories.
- DC-DiT replaces fixed patchification with an encoder–router–decoder scaffold that learns to compress inputs into informative tokens during diffusion training.
- The router allocates fewer tokens to smooth regions and noisy steps, but denser allocations to boundaries, high-frequency details, and later refinement stages without supervision.
- Multi-budget training enables one checkpoint to operate at multiple token budgets, producing elastic inference and asymmetric compute for Lite-CFG.
- Up to 36.8% lower inference FLOPs and up to 37.8% better FID are reported versus fixed-patch DiT baselines across ImageNet scales, resolutions, and guidance settings.
- DC-DiT extends beyond class-conditional generation through lightweight Z-Image upcycling and composes with orthogonal methods including DyDiT and TeaCache.
2 Related Work
Related work reduces diffusion-transformer computation through adaptive backbones, latent representations, token routing, and content-adaptive tokenization. These methods vary token density, representation granularity, or computation according to spatial content, timestep, or network depth.
- Compute-adaptive diffusion transformers modify backbone computation, latent representations, or token routing to reduce redundancy.
- Examples adapt hidden width or prune spatial tokens across timesteps, encode regions at different downsampling rates, or learn layer- and timestep-dependent compression.
- Content-adaptive tokenizers vary token number, size, or granularity according to input content, producing compact representations for simple regions and denser ones for complex regions.
3 Method
DC-DiT replaces fixed patchification with an encoder–router–decoder scaffold that selects content-dependent boundary tokens, processes a shortened sequence, and reconstructs the full-resolution prediction. Its routing adapts token allocation spatially and supports elastic inference through ranked token dropping.
- Overall architecture: DC-DiT uses an encoder–router–decoder scaffold to convert the latent grid into a shorter sequence before DiT processing and restore the original resolution afterward.The encoder mixes local context, the router selects boundary tokens, and de-chunking combines smoothing with plug-back.
- Routing and chunking: The router assigns higher boundary probabilities to tokens that are difficult to predict locally, while predictable non-boundaries can be dropped and reconstructed from nearby representatives.Boundary probabilities are produced from locally mixed features and determine the retained token subset.
- Routing and chunking: Hard chunking thresholds boundary probabilities at 0.5 and uses a straight-through estimator so discrete selection remains end-to-end trainable.The forward pass uses a hard mask, while backward gradients treat the mask as continuous probabilities.
- Elastic inference: Elastic inference ranks retained tokens by router probability and drops the lowest-ranked tail to support user-specified compute budgets.Lite-CFG can apply more aggressive dropping to the unconditional branch than the conditional branch.
- De-chunking: De-chunking smooths retained representations using spatial distances and confidence, then assigns each original position the representation of its nearest retained boundary.Confident representatives remain close to their original features, whereas uncertain ones borrow context from nearby retained tokens.
- Training: Training combines the standard diffusion objective with a ratio regularizer targeting average compression, while multi-budget training exposes the model to several tail-dropping fractions.The ratio loss anchors natural routing to the target compression factor, and tail dropping adds an inference-time budget.
4 Experiments and Results
Experiments show that DC-DiT learns content- and timestep-adaptive token allocation, supports elastic and composable inference, and improves quality–compute tradeoffs across ImageNet settings and upcycled text-to-image generation.
- Main Results: DC-DiT improves the FID–FLOPs Pareto frontier across ImageNet model scales, guidance settings, and resolutions.At 512×512, DC-DiT-XL reduces compute by 36.8% with only a small FID tradeoff; under Lite-CFG, FID improves by up to 37.8%.
- Learned spatio-temporal compression: The router allocates fewer tokens to predictable backgrounds and noisy early steps, while retaining more tokens for detailed regions and later denoising stages.This learned spatio-temporal pattern emerges without supervision for segmentation, boundary detection, or timestep scheduling.
- Elastic inference: Multi-budget training enables elastic inference from one checkpoint, with tail dropping exposing cheaper operating points while preserving graceful quality degradation.Training randomizes the tail-dropping fraction, encouraging reconstruction-critical tokens to appear earlier in the router ranking.
- Training behavior: At 512px, DC-DiT-S and DC-DiT-L trained stably, whereas the corresponding vanilla DiT-S/2 and DiT-L/2 baselines repeatedly encountered loss divergence.The comparison involved multiple restarts for the vanilla baselines.
- Upcycling Z-Image: Lightweight upcycling of Z-Image produces DC-Z-Image, which preserves prompt-following quality across elastic inference budgets while reducing inference compute.The experiment tests dynamic chunking beyond class-conditional ImageNet models using distillation and flow-matching adaptation.
- Composability: DC-DiT remains compatible with TeaCache and DyDiT, achieving additional compute reductions while preserving or improving FID.The backbone is unchanged, so post-hoc dynamic execution methods can be applied inside DC-DiT.
5 Conclusion
DC-DiT replaces fixed patchification with supervised-free adaptive tokenization and reallocates computation across spatial regions and denoising timesteps. Across ImageNet settings, it improves the quality–compute Pareto frontier while supporting elastic inference and model composition.
- DC-DiT replaces fixed patchification with adaptive tokenization learned end-to-end through diffusion training.
- The router reallocates compute across spatial regions and denoising timesteps without explicit supervision.
- 36.8% lower inference FLOPs and 37.8% improved FID are reported across ImageNet settings.
- The same routing supports elastic inference and Lite-CFG, while the approach upcycles Z-Image and composes with DyDiT and TeaCache.
- The results position adaptive tokenization as a practical primitive for efficient diffusion models.
A.1 Additional visual results
The appendix extends the paper’s qualitative elastic-inference and adaptive-compute visualizations with additional samples. These examples illustrate learned spatial chunking and timestep-dependent token allocation.
- Figure 7 extends the qualitative elastic-inference grid with additional generations from the same checkpoint across user-selected compute budgets.
- Figure 8 extends the adaptive-compute visualization with samples illustrating learned spatial chunking and timestep-dependent token allocation.
A.2 FLOPs accounting
The FLOPs accounting reports end-to-end generation cost under each sampling protocol, including DC-DiT’s routing scaffold and variable-length packed computation. The accounting distinguishes unguided, standard classifier-free guidance, and Lite-CFG costs.
- TFLOPs/img measures the total floating-point operation count required to generate one image under the evaluation sampling protocol.Unguided sampling sums costs across 250 DDPM steps; guided protocols count conditional and unconditional evaluations.
- Lite-CFG counts conditional and unconditional branches separately, using conservative and aggressive tail-dropping budgets, respectively.
- The reported FLOPs include matrix multiplications, convolutions, normalizations, elementwise operations, embeddings, output projections, and encoder-router-decoder overhead.This reflects realized end-to-end DC-DiT cost at each tail-dropping fraction.
- Variable-length sequence packing uses each sample’s actual retained sequence length rather than padding all samples to the longest sequence.For fixed-patch DiT baselines, the retained length is constant across the batch and reduces to dense attention accounting.
A.3 Multi-budget training and compute matching
Multi-budget training exposes DC-DiT to several tail-dropping fractions after a no-dropping warmup, enabling one checkpoint to operate across token budgets. Training duration is compute-matched to the corresponding 400K-step DiT baseline.
- DC-DiT samples tail-dropping fractions from 0.0 through 0.6 after a 5K-step no-dropping warmup.The warmup lets the router converge toward the target compression ratio before aggressive compression reaches the DiT blocks.
- The multi-budget objective enables a single checkpoint to be evaluated at several tail-dropping fractions.
- DC-DiT training length is chosen to match the training compute of the corresponding 400K-step DiT baseline.The calculation accounts for warmup cost and average post-warmup FLOPs across sampled budgets.
- DC-DiT trains longer than 400K steps when multi-budget tail dropping lowers the average cost per post-warmup step.
A.4 Architecture details by scale
This section summarizes the architecture shapes used across DC-DiT scale configurations, including how bottleneck width is determined.
- The encoder/decoder bottleneck width is computed from the configured hidden width and dimension reduction factor.
- Table 5 reports the basic architecture shapes for the DC-DiT scale configurations used in the paper.
A.5 Hyperparameters
This section lists the main hyperparameters for ImageNet DC-DiT training and notes which settings are shared across model scales.
- Table 6 lists the main hyperparameters used for the ImageNet DC-DiT-B N = 4 experiments.
- Unless otherwise noted, the optimizer, diffusion, routing, and multi-budget settings are shared across ImageNet model scales.