Diffusion transformers (DiTs) achieve high generative quality but lock FLOPs to image resolution, limiting principled latency-quality trade-offs, and allocate computation uniformly across input spatial tokens, wasting resource allocation to unimportant regions. We introduce Elastic Latent Interface Transformer (ELIT), a drop-in, DiT-compatible mechanism that decouples input image size from compute. Our approach inserts a latent interface, a learnable variable-length token sequence on which standard transformer blocks can operate. Lightweight Read and Write cross-attention layers move information between spatial tokens and latents and prioritize important input regions. By training with random dropping of tail latents, ELIT learns to produce importance-ordered representations with earlier latents capturing global structure while later ones contain information to refine details. At inference, the number of latents can be dynamically adjusted to match compute constraints. ELIT is deliberately minimal, adding two cross-attention layers while leaving the rectified flow objective and the DiT stack unchanged. Across datasets and architectures (DiT, U-ViT, HDiT, MM-DiT), ELIT delivers consistent gains. On ImageNet-1K 512px, ELIT delivers an average gain of 35.3% and 39.6% in FID and FDD scores. Project page: https://snap-research.github.io/elit/
Figures & tables
Figure 2: Adaptive computation. We test whether DiT and ELIT-DiT can reallocate compute across image regions by training on synthetic inputs formed by zero-padding real images, artificially increasing the token count (◆). We compare its performance to baselines trained on real data using patch size 2×2 and patch size 1×1 (★). Vanilla DiT does not improve: attention in zeroed regions targets other zeroed regions (see “DiT Attention”), so extra tokens raise cost without benefits. In contrast, ELIT-DiT uses the Read layer to pull informative spatial tokens into the latent interface (see “Read Attention”), effectively filtering out zeroed areas (see “ELIT-DiT Attention”). Consequently, it leverages the larger token budget and matches the real-data baseline at equal FLOPs.
Figure 3: Architecture of ELIT. We extend a DiT-like generator with a variable-length set of latent tokens—the latent interface —and lightweight Read/Write cross-attention layers. A short spatial DiT head processes patchified inputs; Read pulls information into the latent domain where core blocks operate. Write broadcasts updated latents back to spatial tokens, and a small spatial tail produce output. Spatial tokens and latents are partitioned into corresponding groups, with cross-attention operate only within groups. During training, we randomly drop tail latents, yielding an importance-ordered interface. At inference, the number of latents serves as a user-controlled compute knob.
Model
Params
ImageNet 256 × 256
ImageNet 512 × 512
TF
FID50K↓
FDD50K↓
IS ↑
TF
FID50K↓
FDD50K↓
IS ↑
@256
–G
+G
–G
+G
–G
+G
@512
–G
+G
–G
+G
–G
+G
DiT-XL
675M
182
13.0
5.7
346.3
232.9
66.2
115.3
806
18.8
9.5
339.2
233.6
53.0
86.4
└ ELIT
698M
188
8.2
3.8
200.2
124.5
93.0
160.1
831
11.1
4.9
175.6
106.1
80.0
134.1
└ MB
698M
190
7.8 -40%
3.8 -33%
203.7 -41%
128.6 -45%
99.0 +50%
167.6 +45%
804
10.1 -46%
4.5 -53%
164.1 -52%
98.2 -58%
88.8 +68%
147.0 +70%
UViT-XL
707M
196
8.3
3.8
220.2
138.0
84.4
145.1
861
11.6
5.3
202.7
125.9
72.5
117.2
Table 1: Comparative performance on ImageNet-1K at 256px and 512px resolutions. We evaluate FID ↓ , FDD ↓ , and IS ↑ without (–G) and with 0.25 CFG (+G). TFLOPs (TF) indicate single training iteration TFLOPs. Superscripts show percentage of improvement of ELIT MultiBudget (MB) relative to the baseline.
Figure 4: Training convergence. ELIT-DiT significantly accelerates convergence, achieving 3.3× speedup on ImageNet-1K 256px and 4.0× on ImageNet-1K 512px.
Figure 5: Guidance strategies. ELIT enables autoguidance out of the box by providing a well-aligned weaker model that runs at ≈35% of the cost for the unconditional path. When paired with classifier-free guidance (CFG), denoted as cheap CFG (CCFG), it reduces overall generation cost by ≈33% while improving quality. Compared to DiT, ELIT-DiT achieves a 19% better best FID.
Figure 6: Qualitative assessment of ELIT on ImageNet-1K 512px. We compare DiT against ELIT-DiT and ablate over different guidance settings and number of latent tokens. ELIT improves structural details while allowing per-step selection of inference budget through token dropping and giving access to autoguidance (AG) and cheap classifier free guidance (CFG) for improved sampling quality and cost. See results in Appendix Appendix L .
Figure 7: ELIT improves over DiT across model sizes.
Figure 8
Baseline
FDD10K↓
FID10K↓
FVD10K↓
–G
+G
–G
+G
–G
+G
DiT-XL
371.5
309.1
14.0
11.3
135.9
100.5
└ ELIT
277.4
222.0
13.3
10.7
116.5
90.5
Table 2: Comparative performance on Kinetics-700 256px. We report metrics without (–G) and with 0.25 CFG (+G).
Model
Tokens
FLOPs
Entity
Relation
Attribute
Other
Global
Avg.
Qwen-Image
1×
90.51
92.21
91.03
91.34
91.70
91.27
ELIT-Qwen-Image:
└ 100% Tokens
4096
0.688×
90.30
92.18
88.97
90.34
89.18
90.45
└ 50% Tokens
2048
0.494×
90.15
89.94
89.05
90.09
89.06
89.81
└ 25% Tokens
1024
0.409×
89.31
91.87
89.71
88.28
84.79
89.79
└ 12.5% Tokens
512
0.369×
91.20
90.35
88.77
89.94
79.84
88.02
Table 3: Qwen-Image evaluation on DPG-Bench. Results are reported with different inference budgets.
(a) Group Size Ablation
Group Size
Groups
FID10K↓
FDD10K↓
IS ↑
ImageNet-1K 256px (16 × 16 tokens)
1 × 1
256
29.94
638.8
38.39
2 × 2
64
25.48
546.8
45.66
4 × 4
16
26.53
531.8
45.95
8 × 8
4
27.73
552.5
44.64
Table 4: Ablations. (a) Read/Write group size. (b) Blocks allocation to head-latent core-tail blocks. (c) Strategies for achieving variable budget inference. (d) Batching strategy for multibudget training.
Appendix figures & tables18 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 10: (left) FLOPs for spatial blocks, latent blocks, and Read/Write layers as a function of input tokens N , groups count G , latent tokens per group J , and hidden size d . (right) Relationship between latent tokens per group and model FLOPs for a DiT-XL with 8 spatial blocks, 20 latent core blocks, and N/64 groups, varying input tokens N and latent tokens per group J~ . FLOPs are shown relative to 64 tokens per group.
Baseline
Params
TFLOPs
FID50K↓ ( ρ↓ )
FDD50K↓ ( ρ↓ )
IS ↑ ( ρ↓ )
–G
+G
–G
+G
–G
+G
DiT
675M
806
18.8 (1.00)
9.5 (1.00)
339.2 (1.00)
233.6 (1.00)
53.0 (1.00)
86.4 (1.00)
└ Patch size 2x4
675M
377
22.5 (0.56)
12.3 (0.61)
434.0 (0.60)
317.9 (0.74)
45.7 (0.54)
73.8 (0.55)
HDiT
1.4B
776
13.0 (1.00)
6.0 (1.00)
260.3 (1.00)
170.5 (1.00)
69.4 (1.00)
114.2 (1.00)
└ Smaller backbone
703M
392
22.2 (0.85)
11.5 (0.96)
435.2 (0.83)
315.4 (0.93)
48.8 (0.71)
80.0 (0.71)
ELIT-DiT
698M
831
11.1 (1.00)
4.9 (1.00)
175.6 (1.00)
106.1 (1.00)
80.0 (1.00)
134.1 (1.00)
Appendix
Table 5: Compute-quality tradeoff efficiency of baselines on ImageNet-1K 512px. ρ=(MetricRatio)/(FLOPsRatio) indicates the model degradation with respect to change in FLOPs between the low- and high-compute variants.
Figure 11: Lowering inference budget by using fewer latent tokens per group yields correlated reductions in forward time and FLOPs.
Baseline
FID10K↓
FDD10K↓
IS ↑
ELIT-DIT
26.53
531.8
45.95
Qformer Read
30.49
589.9
41.10
Self-Attn Read
28.38
602.5
40.12
Self-Attn Read/Write
29.46
631.1
38.49
↑ Read Capacity
27.45
540.7
45.40
↑ Write Capacity
25.23
516.9
47.59
Appendix
Table 6: Architectural ablations on DiT-B/2. Using cross-attn in Read/Write is superior to alternatives. Increasing the model capacity is only beneficial in Write and FFN.
Method
FID10K↓
IS ↑
IterFLOPs
100%_100%
11.60
86.68
188
50%_100%
11.98
90.18
154
Appendix
Table 7: Budget scheduling across noise levels. ImageNet 512px, ELIT-DiT-XL/2. 50%_100% : uses 50% of tokens for high-noise steps ( t<0.5 ), 100% otherwise.
Method
FID10K↓
FDD10K↓
IS ↑
Indep. (100% tok.)
13.60
205.23
80.90
Joint (100% tok.)
12.00
189.50
90.29
Indep. (50% tok.)
14.14
222.43
77.99
Joint (50% tok.)
12.95
203.58
85.18
Indep. (25% tok.)
15.36
247.77
74.04
Joint (25% tok.)
14.21
228.08
79.60
Appendix
Table 8: Joint multi-budget vs. independently trained single-budget models. When tested on ImageNet 512px, ELIT-DiT-XL/2, joint multi-budget models consistently outperform single budget models.
Figure 12: Read attention masks averaged over noise levels. Early latent tokens attend to broad, semantically important image regions, while later tokens exhibit sparser attention focusing on fine-grained details.
Figure 13: ELIT vs. token merging methods on ImageNet 512px (DiT-XL/2). Training-free methods (ToMe, SDTM) are bounded by DiT quality, while ELIT surpasses it even at 25% tokens.
Figure 14: When tested with CFG 0.25, ELIT provides better quality-compute tradeoff than reducing the number of sampling steps.
Figure 15: HSL saturation comparison between CFG and CCFG, across guidance scales on ImageNet 512px (DiT-XL/2). CCFG exhibits slightly higher saturation than CFG, attributed to its autoguidance component.
Figure 16: Qualitative results produced by ELIT-DiT on ImageNet-1K 512px with CCFG 4.0 for varying number of tokens in the latent interface. As the tokens and model FLOPs are reduced, the model preserves structure, while varying image details, producing gradual image changes. FLOPs are expressed relative to the model variant where no latent tokens are dropped.
Figure 17: Qualitative comparison of classifier-free guidance (CFG), autoguidance [ 30 ] (AG), cheap classifier-free guidance (CCFG) with different weights, when applied to ELIT-DiT trained on the ImageNet-1K 512px dataset. AG produces the most varied samples, generating results with similar structure across guidance weights, as opposed to CFG and CCFG which favor object-centric generations. Both AG and CCFG produce better generations of complex concepts such as human faces.
Figure 18: Qualitative comparison of ELIT against baselines on the ImageNet-1K 512px dataset. Results are produced using CFG with weight 4.0 for all methods.
Figure 19: Qualitative results produced by ELIT-Qwen-Image for varying number of tokens in the latent interface. As the number of tokens is decreased and model FLOPs are reduced, our method can preserve structural details, while prioritizing changes in image details, preserving perceptual quality. Reported FLOPs are expressed relative to the original Qwen-Image and account for both the sampling FLOPs reductions brought by CCFG and the reduction in the number of tokens in the latent interface.
Figure 20: Uncurated Qualitative samples comparing DiT with ELIT-DiT using CFG and CCFG on ImageNet-1k 512px. Results are produced using CFG with weight 4.0 for all methods.
Figure 21: Uncurated Qualitative samples comparing DiT with ELIT-DiT using CFG and CCFG on ImageNet-1k 512px. Results are produced using CFG with weight 4.0 for all methods.
Figure 22: Uncurated Qualitative samples comparing DiT with ELIT-DiT using CFG and CCFG on ImageNet-1k 512px. Results are produced using CFG with weight 4.0 for all methods.
Figure
Prompt
Figure 1
“The image portrays a woman with dark skin wearing a gold headpiece adorned with a blue jewel. Her gaze is directed towards something off-camera, giving her a focused expression. The background appears to be blurred, drawing attention to her face and headpiece.”
Figure 1
“The image features actor Liev Schreiber in a snowy scene from a movie or TV show. He is dressed in black tactical gear, including a vest with “ARCTIC OCEAN” written on it, and a helmet with goggles. The setting appears to be a bustling city street filled with people and vehicles, all covered in snow.”
Figure 1
“The image features a woman walking down a city street at night. She is wearing a black leather jacket, a white crop top, and a short black skirt. The street is illuminated by neon signs and streetlights, creating a vibrant atmosphere. There are other people visible in the background, but they are not the main focus of the image.”
Figure 19
“The image portrays a man with long black hair and red eyes, wearing a black hooded cloak. He has a red gem on his forehead and holds a red orb-like object in his hand. The background features a circular pattern with red and black colors.”
Figure 19
“The image features a large, white robot-like creature with wings standing on a desert landscape. The creature has sharp claws and appears to be looking down at something. Its body structure resembles a fusion of humanoid and bird-like characteristics. The background consists of a clear blue sky and rocky terrain.”
Figure 19
“The image features a man wearing a blue knit cap, looking upwards with a serious expression. The background is dark blue, creating a contrast with the man’s face and hat.”
Appendix
Table 9: Prompts used to produce the showcased qualitative results for Qwen-Image [ 56 ] .
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.
Diffusion Transformers (DiTs) are a dominant backbone for high-fidelity text-to-image generation due to strong scalability and alignment at high resolutions. However, quadratic self-attention over dense spatial tokens leads to high inference latency and limits deployment. We observe that denoising is spatially non-uniform with respect to aesthetic descriptors in the prompt. Regions associated with aesthetic tokens receive concentrated cross-attention and show larger temporal variation, while low-affinity regions evolve smoothly with redundant computation. Based on this insight, we propose AccelAes, a training-free framework that accelerates DiTs through aesthetics-aware spatio-temporal reduction while improving perceptual aesthetics. AccelAes builds AesMask, a one-shot aesthetic focus mask derived from prompt semantics and cross-attention signals. When localized computation is feasible, SkipSparse reallocates computation and guidance to masked regions. We further reduce temporal redundancy using a lightweight step-level prediction cache that periodically replaces full Transformer evaluations. Experiments on representative DiT families show consistent acceleration and improved aesthetics-oriented quality. On Lumina-Next, AccelAes achieves a 2.11× speedup and improves ImageReward by +11.9% over the dense baseline. Code is available at https://github.com/xuanhuayin/AccelAes.
Xuanhua Yin, Chuanzhi Xu, Haoxian Zhou +2
School of Computer Science, The University of Sydney, NSW 2006, Australia
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.