One Model, Many Budgets: Elastic Latent Interfaces for Diffusion Transformers

Moayed Haji-Ali; Willi Menapace; Ivan Skorokhodov; Dogyun Park; Anil Kag; Michael Vasilkovsky; Sergey Tulyakov; Vicente Ordonez; Aliaksandr Siarohin

Source-linked AI summary

One Model, Many Budgets: Elastic Latent Interfaces for Diffusion Transformers

Moayed Haji-Ali, Willi Menapace, Ivan Skorokhodov, Dogyun Park, Anil Kag, Michael Vasilkovsky, Sergey Tulyakov, Vicente Ordonez, Aliaksandr Siarohin

arXiv:2603.12245v1Mar 12, 2026cs.CV

TL;DR

DiTs tie computation to resolution and distribute it uniformly across spatial tokens, limiting principled latency-quality control. ELIT inserts a variable-length latent interface with lightweight Read and Write layers, allowing inference compute to be adjusted by changing latent count. Across architectures and settings, it consistently improves quality and enables more flexible compute-quality trade-offs, while large-scale from-scratch training benefits remain unverified.

Problem
DiTs lock per-step computation to image resolution and allocate it uniformly across spatial tokens, limiting latency-quality trade-offs and wasting computation on less important regions.
Method
ELIT adds a variable-length latent interface with lightweight Read and Write cross-attention layers that transfer information between spatial and latent tokens and make compute user-controlled.
Results
ELIT consistently improves generation quality and compute-quality trade-offs across varying architectures, with FDD improvements of 58.0% for DiT, 34.0% for U-ViT, and 37.4% for HDiT.
Takeaways & Limitations
A single DiT-compatible model can serve a spectrum of inference budgets while concentrating computation on more difficult input regions.
Takeaways & Limitations
Large-scale, from-scratch training benefits remain unverified, and CCFG tends to saturate images faster than CFG.

Abstract

from arXiv · show

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/

1. Introduction

DiTs deliver strong generation quality but use fixed, spatially uniform computation, limiting budget control and wasting computation on regions that need less processing. ELIT introduces a variable-length latent interface that reallocates computation and supports adjustable inference budgets while preserving the DiT training objective.

Motivation: DiTs tie per-step computation to input resolution and allocate computation uniformly across image regions.Increasing tokens helps on standard images, but zero-padding adds tokens without improving quality, indicating inefficient allocation.
Related limitations: Existing adaptive and masking-based methods either retain uniform token computation, have high complexity, or disable token dropping during inference.These limitations leave inference compute unchanged for masking-based approaches.
ELIT: ELIT adds Read and Write cross-attention layers around a variable-size latent interface to prioritize difficult regions and redistribute computation.Read transfers information from spatial tokens to latents, while Write broadcasts updated latent representations back.
ELIT: The number of latent tokens directly controls per-step compute, allowing one model to serve multiple latency-quality budgets without changing the training objective.ELIT retains the vanilla rectified-flow objective and supports drop-in integration with DiT-like architectures.
Results: ELIT redistributes computation non-uniformly across input tokens across multiple base architectures.The analysis covers DiT, U-ViT, and HDiT, with improvements reported on ImageNet-1K 512px FDD.
Results: 58.0% FDD improvement is reported for DiT, alongside 34.0% for U-ViT and 37.4% for HDiT on ImageNet-1K 512px.These figures compare ELIT with a fixed-grid model.
Results: ELIT supports inference-time compute-quality trade-offs and reduces inference cost by approximately 33% with autoguidance without affecting generation quality.Selecting the latent-token count provides variable test-time compute and remains compatible with training-free acceleration techniques.

2. Related Work

Prior adaptive generators seek variable inference budgets, but ELIT places adaptivity in a simple variable-length latent interface that integrates with DiTs. Its adaptive computation mechanism filters uninformative spatial regions instead of spending added computation uniformly on them.

Adaptive generators: Adaptive generators support multiple inference budgets through shared weights, varied patchification, or learnable routing, but differ in where adaptivity is implemented.These approaches share the goal of scaling compute at inference time with one model.
ELIT: ELIT uses a variable-length latent interface to integrate flexible compute allocation directly into DiT-like transformers.The interface is designed to preserve compatibility while enabling resizable computation.
Adaptive computation: ELIT-DiT filters zero-padded spatial regions through Read attention and matches the real-data baseline at equal FLOPs.Vanilla DiT instead directs attention among zeroed regions, so the added tokens increase cost without benefit.

3. Method

ELIT introduces a variable-length latent interface between spatial inputs and the main transformer blocks, enabling flexible computation while preserving DiT-compatible processing. Read and Write cross-attention layers transfer information across domains, and prefix-based latent training supports adaptive inference budgets.

Elastic latent interface: ELIT inserts a variable-length latent interface where most transformer blocks operate, decoupling latent-domain computation from fixed spatial token mappings.The interface contains K latent tokens, whose inference-time length directly adjusts FLOPs.
Read and Write layers: Read cross-attention transfers information from spatial tokens into the latent interface, while Write cross-attention broadcasts updated latent representations back to spatial tokens.A short spatial head precedes Read, and a short spatial tail follows Write to produce the output velocity.
Training convergence: 3.3× and 4.0× training-convergence speedups are reported for ELIT-DiT on ImageNet-1K at 256px and 512px, respectively.The reported values compare convergence speed with DiT at the two resolutions.
Grouped cross-attention: Grouped cross-attention reduces Read and Write complexity from O(NK) to O(NK/G) by partitioning spatial and latent tokens into corresponding groups.Cross-attention operates only within corresponding groups rather than across all spatial-latent pairs.
Importance ordering: Randomly retaining a shared latent prefix during training makes earlier latents capture globally useful information while later latents refine details.At each iteration, the same sampled prefix length is kept across groups, and subsequent tail latents are dropped in Read, Write, and latent-core blocks.
Elastic computation: ELIT supports asymmetric classifier-free guidance by evaluating the main term with a larger latent budget and the guidance term with a smaller one.The multi-budget framework provides a weaker model version by varying the inference budget.

4. Experiments

ELIT is evaluated across diffusion backbones, resolutions, inference budgets, guidance strategies, and large-scale image generation. Results show broad quality improvements, flexible compute–quality trade-offs, and substantial speedups in selected settings.

Experimental setup: ELIT is evaluated on DiT, U-ViT, and HDiT using matched transformer blocks, rectified-flow training, and comparable training compute.The study also evaluates video generation and applies ELIT to a 20B MM-DiT-based Qwen-Image model.
Comparison to baselines: ELIT-MB achieves FID reductions of 40%, 14%, and 27% over DiT, U-ViT, and HDiT, respectively.The multi-budget variant delivers gains across all reported metrics.
Comparison to baselines: ELIT outperforms DiT across model sizes from DiT-S/4 to DiT-XL/2.The result is reported in the model-scaling evaluation.
Elastic inference capabilities: ELIT provides a superior compute–quality trade-off to reducing sampling steps by varying the number of retained latent tokens.The optimal combination of sampling steps and token count varies by FLOP target.
Elastic inference capabilities: Autoguidance achieves comparable performance to CFG with approximately 33% fewer FLOPs, while CCFG delivers the best results across metrics with the same speedup.The guidance comparison includes classifier-free guidance, autoguidance, and cheap classifier-free guidance.
Large-scale multi-budget model: On Qwen-Image, ELIT cuts sampling FLOPs by up to 63%, achieves approximately 2.7× speedup, and retains DPG-Bench scores from 90.45 to 88.02 across budgets.The ELIT model initially trails the original Qwen-Image by 0.82 average score points.

5. Discussion

The discussion positions ELIT as a lightweight extension that preserves the baseline architecture while supporting flexible compute allocation. Its FLOP structure and implementation choices are designed to retain compatibility with existing DiT components and acceleration methods.

Design and implementation: ELIT leaves each baseline architecture and training procedure unchanged apart from adding Read and Write operations.For XL-size models, Read is placed at block 4 and Write at block 24; main experiments use 16 spatial groups.
Compute analysis: The FLOP breakdown separates spatial blocks, latent blocks, and Read/Write layers as functions of input tokens, groups, latent tokens, and hidden size.The analysis also relates latent tokens per group to total model FLOPs for DiT-XL.
Training setup: The method uses group-wise latent processing, with synchronized compute and increased batch size to match training FLOPs.The video setup uses 64 groups and 2,048 spatial tokens.

D. Compute-quality Tradeoff Efficiency

ELIT is designed to preserve quality more efficiently when compute is constrained. Its variable latent-token budget produces favorable compute–quality behavior relative to alternative low-compute configurations and sampling-step reduction.

Compute-quality tradeoff: ELIT-DiT uses fewer latent tokens for the low-compute variant, while DiT and HDiT use alternative architectural reductions.The comparison trains low- and high-compute variants for each baseline.
Evaluation metric: The degradation metric ρ divides metric degradation by the corresponding FLOP reduction between low- and high-compute variants.Lower ρ indicates less degradation per unit of reduced compute.
Results: ELIT outperforms baselines at similar training compute and consistently achieves lower ρ under compute constraints.The paper attributes this efficiency to the latent interface focusing on important input information.
Inference budget: Fewer latent tokens per group reduce forward time and FLOPs during inference.This provides a direct budget-control mechanism through latent-token retention.
Architectural ablations: A single cross-attention Read layer outperforms Q-Former-style Read and full self-attention, while extra Write or FFN capacity improves results at additional FLOPs.Stacking two Read cross-attention layers provides no measurable gain.

F. Compatibility with Distillation Methods.

The paper examines ELIT alongside distillation, pruning-oriented efficiency, and noise-level budget scheduling. These experiments indicate compatibility with orthogonal compression methods and suggest that adaptive token budgets can reduce compute in selected sampling regimes.

Distillation compatibility: Grafting ELIT MLPs with expansion ratio r = 3 causes 12.6% FID and 8.9% IS degradation.These degradations are compared with the original paper’s reported 17.2% FID and 9.4% IS degradation.
Distillation compatibility: ELIT remains compatible with orthogonal efficiency methods such as network pruning and distillation.The grafting experiment provides the supporting compatibility result.
Budget scheduling: Using 50% of tokens at high-noise steps and 100% otherwise maintains comparable performance at 154 versus 188 TFLOPs per iteration.The experiment uses ELIT-DiT-XL/2 on ImageNet 512px.

H. Joint vs. Independent Budget Training.

Joint multi-budget training consistently outperforms independently trained single-budget models on ImageNet 512px. ELIT’s latent representations are importance-ordered, with early tokens covering broad structure and later tokens refining details.

Joint multi-budget ELIT-DiT-XL/2 models consistently outperform independently trained single-budget models on ImageNet 512px.
ELIT surpasses the base DiT even when using only 25% of the tokens on ImageNet 512px.
Early latent tokens attend to broad, semantically important image regions, while later tokens focus more sparsely on fine-grained details.

J. Comparison with Token Merging Methods.

ELIT provides a more favorable quality–compute tradeoff than training-free token-merging methods on ImageNet 512px.

ELIT improves over the base DiT with only 25% of the tokens, achieving FID10K =14.2.
Training-free ToMe and SDTM trade compute for quality less favorably than ELIT.

K. Compute Analysis of ELIT

ELIT’s latent-token budget controls computation in practice, with latency generally tracking FLOPs. This produces a quality–compute tradeoff that is more favorable than reducing sampling steps.

When core blocks dominate computation, the Read and Write cross-attention costs are minimal relative to the model cost.
ELIT achieves a more favorable quality–compute tradeoff than reducing the number of sampling steps.
Forward time drops monotonically as the latent-token count decreases and closely follows the corresponding FLOPs reduction.At higher budgets, fixed overheads such as I/O and kernel launches weaken the correlation slightly.

L. Additional Results

Additional experiments show that varying ELIT’s latent-token budget preserves image structure while changing less prominent details, and that guidance variants provide further compute-quality tradeoffs.

Decreasing the latent-token count preserves image structure while changing less noticeable details.
CCFG exhibits slightly higher HSL saturation than CFG across guidance scales, attributed to its autoguidance component.
ELIT-Qwen-Image operates at 69% of the original model’s FLOPs with CCFG and reaches 35% in its cheapest shown configuration.
Qualitative ImageNet-1K 512px comparisons include baseline DiT and ELIT-DiT using CFG and CCFG on randomly selected classes and samples.

M. Failed Experiments

Spatial token dropping did not provide satisfactory flexible inference because dropped regions lose information irrecoverably. Latent-token reduction instead preserves structure while varying details and enables gradual quality–compute changes.

Spatial token masking: Spatial-domain token dropping produced unsatisfactory inference results because information in dropped regions cannot be recovered.
Latent-token reduction: Reducing latent tokens and FLOPs preserved image structure while allowing image details to vary gradually.The qualitative results express FLOPs relative to the variant with no latent tokens dropped.
Guidance comparisons: Autoguidance generated the most varied samples, while autoguidance and cheap classifier-free guidance improved complex concepts such as human faces.Classifier-free guidance and cheap classifier-free guidance favored more object-centric generations.
Baseline comparison: Qualitative comparisons evaluated ELIT against baselines on ImageNet-1K 512px using classifier-free guidance with weight 4.0.
Latent-token reduction: With fewer latent tokens, ELIT-Qwen-Image preserved structural details while prioritizing changes in image details and perceptual quality.Reported FLOPs account for both cheap classifier-free guidance sampling reductions and latent-interface token reduction.