cs.CVMar 12, 2026
SaveOne Model, Many Budgets: Elastic Latent Interfaces for Diffusion Transformers
Organizations: Rice University · Snap Inc.
Abstract
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 and 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 and patch size (★). 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 | IS | TF | 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 speedup on ImageNet-1K 256px and 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 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 while improving quality. Compared to DiT, ELIT-DiT achieves a 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.
| Baseline | ||||||
| –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 | 90.51 | 92.21 | 91.03 | 91.34 | 91.70 | 91.27 | ||
| ELIT-Qwen-Image: | ||||||||
| 100% Tokens | 4096 | 90.30 | 92.18 | 88.97 | 90.34 | 89.18 | 90.45 | |
| 50% Tokens | 2048 | 90.15 | 89.94 | 89.05 | 90.09 | 89.06 | 89.81 | |
| 25% Tokens | 1024 | 89.31 | 91.87 | 89.71 | 88.28 | 84.79 | 89.79 | |
| 12.5% Tokens | 512 | 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 | 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 , groups count , latent tokens per group , and hidden size . (right) Relationship between latent tokens per group and model FLOPs for a DiT-XL with 8 spatial blocks, 20 latent core blocks, and groups, varying input tokens and latent tokens per group . FLOPs are shown relative to 64 tokens per group.
| Baseline | Params | TFLOPs | ( ) | ( ) | 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) |
Table 5: Compute-quality tradeoff efficiency of baselines on ImageNet-1K 512px. 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 | 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 |
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 | IS | ||
| 11.60 | 86.68 | 188 | |
| 11.98 | 90.18 | 154 |
Table 7: Budget scheduling across noise levels. ImageNet 512px, ELIT-DiT-XL/2. : uses of tokens for high-noise steps ( ), otherwise.
| Method | 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 |
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 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.” |
Table 9: Prompts used to produce the showcased qualitative results for Qwen-Image [ 56 ] .