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

2026-03-12Computer Vision and Pattern Recognition

Computer Vision and Pattern Recognition
AI summary

The authors propose ELIT, a small change to existing diffusion transformers used for generating images. ELIT adds a special set of tokens that summarize important image parts and lets the model focus computing power on those areas instead of treating all parts equally. This method also allows adjusting how much computing to use based on available resources without changing the image size. Their experiments show ELIT improves image generation quality consistently across different models and datasets.

Diffusion transformersImage generationLatent tokensCross-attentionCompute efficiencyFID scoreImportance weightingImageNetTransformer blocks
Authors
Moayed Haji-Ali, Willi Menapace, Ivan Skorokhodov, Dogyun Park, Anil Kag, Michael Vasilkovsky, Sergey Tulyakov, Vicente Ordonez, Aliaksandr Siarohin
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 $35.3\%$ and $39.6\%$ in FID and FDD scores. Project page: https://snap-research.github.io/elit/