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

TL;DR

ELIT (Elastic Latent Interface Transformer) is a drop-in DiT modification from Snap Research that decouples compute from input resolution by inserting a variable-length sequence of latent tokens between two lightweight cross-attention layers (Read / Write). The transformer trunk operates on the latent sequence rather than the spatial grid, and tail-token dropping during training makes the latents importance-ordered — at inference, the number of retained latents is a user-facing FLOPs–quality knob on a single set of weights. Across DiT, U-ViT, HDiT, and MM-DiT on ImageNet-256/512 and on Qwen-Image 20B, ELIT reports up to 53% FID / 58% FDD improvement, ~4× faster convergence at 512px, and “Cheap CFG” guidance using a low-budget forward as the unconditional pass for ~33% guidance-cost reduction.

Key claims

A frozen-uniform-allocation DiT cannot reallocate compute from trivial regions to informative ones even when given extra tokens — motivating learned, content-adaptive token allocation rather than uniform spatial patching [§Synthetic experiment, Project page].
A latent interface of K learnable tokens, with one Read cross-attention pulling spatial → latent and one Write cross-attention pushing latent → spatial, suffices to redirect compute to harder regions; the standard transformer stack is unchanged and the rectified-flow objective is unchanged [Abstract].
Grouped cross-attention partitions the N spatial tokens into G non-overlapping groups, dropping cross-attention cost from O(NK) to O(NK/G) [Method].
Random tail-token dropping during training produces an importance-ordered latent sequence: early tokens encode global structure, later tokens encode fine detail — enabling graceful quality degradation when latents are truncated at inference [Method; “What does each latent token encode” visualization].
Read-attention maps learned without supervision correlate with per-patch flow-matching loss: latents preferentially attend to object boundaries, fine textures, and semantically complex regions, while easy background regions receive minimal attention [Project page, Read-attention visualization].
On ImageNet-256 (no CFG), ELIT-MultiBudget achieves FID 7.8 / FDD 203.7 on DiT-XL vs baseline 13.0 / 346.3 — a 40% FID / 41% FDD reduction at comparable training TFLOPs [Table].
On ImageNet-512 (no CFG), DiT-XL improves from FID 18.8 / FDD 339.2 to FID 10.1 / FDD 164.1 with ELIT-MB — 46% FID and 52% FDD reduction [Table].
ELIT reaches the same FDD as a vanilla DiT-XL at 512px in roughly 4× fewer training steps [Project page, training-efficiency claim].
ELIT’s gains grow with model size and resolution: FDD improvement is larger at 512px than 256px across DiT, U-ViT, HDiT [Table].
“Cheap CFG” (CCFG): using a low-token-count forward pass of the same model as the unconditional branch yields AutoGuidance-style guidance for free; combined with class dropping, it cuts guidance FLOPs by ~33% at improved quality [Project page, Cheap CFG section].
At a matched FLOPs budget, the token-count knob produces a better compute–quality trade-off than reducing the number of sampling steps [Project page, FLOPs budget comparison].
ELIT is orthogonal to training-free caching: composing with TeaCache yields speedup gains comparable to applying TeaCache on the baseline, stacking both sources of efficiency [Project page, TeaCache composition].
ELIT scales to Qwen-Image (20B MM-DiT): up to 63% sampling FLOPs reduction (~2.7× speedup); DPG-Bench average degrades gracefully from 90.45 (100% tokens, 0.69× FLOPs) to 88.02 (12.5% tokens, 0.37× FLOPs); the 100%-token ELIT variant nearly matches the original 4096-token Qwen-Image (90.45 vs 91.27) [Table, ELIT-Qwen-Image].

Method

ELIT inserts a single latent sequence of $K$ learnable tokens after patchification. A lightweight Read cross-attention layer treats the latent tokens as queries and the $N$ spatial tokens as keys/values, pulling region-importance-weighted information into the latents. The standard transformer trunk (self-attention + FFN blocks, with the rectified-flow objective unchanged) then operates only on the $K$ latents — decoupling compute from the spatial resolution that determined $N$ . A symmetric Write cross-attention with spatial queries / latent keys-values returns the updated information to the spatial grid for the final velocity prediction. Both cross-attentions use grouped attention (partition $N$ into $G$ groups, each attending independently), reducing the per-layer cross-attention cost from $O(NK)$ to $O(NK/G)$ .

Training adds two design choices on top of standard rectified flow. (1) Tail latents are randomly dropped each step — the model never sees a fixed $K$ , and is forced to push the highest-priority information into the earliest tokens. This yields the importance-ordered property visualized on the project page: early tokens attend broadly (global layout), late tokens attend narrowly (texture detail). (2) The ELIT-MultiBudget variant trains across a distribution of budget values, so a single set of weights serves the full latency–quality curve. At inference, the user picks how many latents to retain; truncating the sequence is the only knob.

For guidance, ELIT exploits the budget knob to produce a “weaker” model from the same weights: running a forward with a small latent budget plays the role of AutoGuidance’s separate weaker network. Composed with class-condition dropping, this yields Cheap CFG with no additional model and ~33% guidance-FLOPs reduction.

Results

ImageNet-256, no CFG, DiT-XL: baseline FID 13.0 / FDD 346.3 → ELIT-MB FID 7.8 / FDD 203.7 at 190 vs 182 training TFLOPs (40% / 41% reductions) [Table].
ImageNet-256, no CFG, U-ViT-XL: 8.3 / 220.2 → 7.1 / 203.2 (ELIT-MB) at 204 vs 196 TFLOPs (14% / 8%) [Table].
ImageNet-256, no CFG, HDiT-XL: 12.8 / 361.6 → 9.4 / 271.8 (ELIT-MB) at 191 vs 182 TFLOPs (27% / 25%) [Table].
ImageNet-512, no CFG, DiT-XL: 18.8 / 339.2 → 10.1 / 164.1 (ELIT-MB) at 804 vs 806 TFLOPs (46% / 52%) — gains are larger at higher resolution [Table].
ImageNet-512, no CFG, U-ViT-XL: 11.6 / 202.7 → 7.7 / 135.8 (34% / 33%) [Table].
Training efficiency at 512px: ~4× fewer steps to reach the baseline’s FDD [Project page].
Qwen-Image (20B MM-DiT), DPG-Bench: full-budget ELIT (4096 tokens) 90.45 vs baseline 91.27; ELIT-50% (2048 tokens) 89.81 at 0.49× FLOPs; ELIT-12.5% (512 tokens) 88.02 at 0.37× FLOPs (~2.7× faster) [Table].
CCFG on ELIT-Qwen-Image: ~33% FLOPs savings vs standard CFG at comparable quality [Project page].
TeaCache composition: stacks with ELIT for additional inference speedup [Project page].

Why it’s interesting

ELIT introduces a fourth, largely orthogonal lever to the diffusion-efficiency landscape on the wiki. (1) DDiT: Dynamic Patch Scheduling for Efficient Diffusion Transformers also adapts compute per region — but by re-quantizing the spatial grid (per-step coarser patches when the latent’s third-order acceleration is low) rather than by routing into a variable-length latent bottleneck; both report stacking with TeaCache, which suggests they could compose with each other and with caching for ≥3 multiplicative levers. (2) Unlike RAEv2 / Self-Flow / L2P in Diffusion training efficiency, ELIT does not change the supervision (no DINO alignment, no encoder swap, no synthetic-data transfer) — it changes the architecture’s compute graph, so it is in principle composable with every recipe filed under that concept. (3) The “Cheap CFG” trick is a notable new entry under Classifier-Free Guidance variants: instead of training-time CFG (SoFlow) or training-noise augmentation (Ostris), it uses a budget-knob ablation of the same model as the AutoGuidance weak branch — a guidance recipe that exists only because the model is elastic at inference. For Luma’s video DiTs the obvious composition is ELIT (compute-per-pixel adaptive) × DDiT (patch-size adaptive) × FastGen/SoFlow (step-count adaptive) × FBC/TeaCache (step-caching) on a single backbone.