💻 computer science

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

The paper introduces ELIT, a minimal, drop-in mechanism for Diffusion Transformers that decouples image resolution from computation by using a learnable, variable-length latent interface with cross-attention layers, enabling dynamic latency-quality trade-offs and prioritizing important image regions to significantly improve FID and FDD scores across various architectures.

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

Published 2026-03-13

📖 4 min read☕ Coffee break read

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are hiring a team of artists to paint a massive, intricate mural.

The Old Way (Standard Diffusion Transformers):
In the traditional approach, you hire a fixed number of artists, and you tell them: "No matter what part of the wall we are painting, every single square inch gets exactly the same amount of attention and time."

If you are painting a simple blue sky, those artists waste time staring at the empty blue. If you are painting a complex, detailed face, those same artists might not have enough time to get the eyes right because they spent too long on the sky. Furthermore, if you want the painting done faster, you can't just ask them to work smarter; you have to hire fewer artists, which often means the whole picture looks blurry or unfinished. The cost (time/money) is locked to the size of the wall, not the difficulty of the art.

The New Way (ELIT - Elastic Latent Interface Transformer):
This paper introduces a clever new system called ELIT. Think of it as giving the art director a "smart assistant" and a set of elastic sticky notes.

Here is how it works, broken down into simple steps:

1. The "Smart Assistant" (The Read Layer)

Before the main painting starts, a smart assistant looks at the wall. Instead of treating every square inch equally, the assistant asks: "Which parts are hard? Which parts are easy?"

The Sky? Easy.
The Face? Hard.
The Background? Medium.

The assistant then grabs a stack of sticky notes (these are the "latent tokens"). It writes down the most important details on the first few notes and the fine details on the later notes. Crucially, it only uses as many notes as it needs for the current task.

2. The "Elastic Brain" (The Latent Interface)

The main artists (the AI model) don't look at the wall directly anymore. They look at the sticky notes.

If the wall is simple, the assistant puts down only 5 notes. The artists work fast, focusing only on those 5 notes.
If the wall is complex, the assistant puts down 50 notes. The artists take their time, diving deep into the details.

This is the "Elastic" part. The system stretches or shrinks its own brain size based on how much work is actually required.

3. The "Feedback Loop" (The Write Layer)

Once the artists finish working on the sticky notes, they hand the notes back to the assistant. The assistant then paints the final details onto the wall based on those notes.

If the assistant only used 5 notes, the wall gets a quick, good-quality sketch.
If the assistant used 50 notes, the wall gets a photorealistic masterpiece.

Why is this a Big Deal?

1. You Control the Budget (The "Knob")
Imagine a dial on your computer that says "Quality vs. Speed."

Turn it to "Speed": The system uses fewer sticky notes. It ignores the tiny details and focuses on the big shapes. You get a great image in half the time.
Turn it to "Quality": The system uses all the sticky notes. It spends more time on the fine details.
The Magic: You can do this with one single model. You don't need to train a "fast version" and a "slow version." The same model just changes how many sticky notes it uses on the fly.

2. No More Wasted Effort
In the old way, the computer spent 50% of its energy painting a boring blue sky. With ELIT, the computer realizes the sky is easy, skips the extra work, and uses that saved energy to make the eyes of the portrait look perfect. It's like a chef who spends 10 minutes chopping a carrot but only 1 minute plating the garnish, instead of spending 10 minutes on both.

3. It Works Everywhere
The researchers tested this on different types of "painting styles" (different AI architectures) and for both pictures and videos. In every case, ELIT made the images sharper and the videos smoother, while using less computing power.

The Bottom Line

ELIT is like giving an AI a flexible budget. Instead of forcing the AI to spend the same amount of energy on a simple cloud as it does on a complex face, it learns to spend energy where it matters most.

Need a quick sketch? Use fewer notes.
Need a masterpiece? Use more notes.
Result: Better images, faster speeds, and a single model that can do it all.

It's the difference between hiring a team that works at a fixed, rigid pace and hiring a team that knows exactly how to allocate their time to get the best result for the money you have.

1. Problem Statement

Diffusion Transformers (DiTs) have achieved state-of-the-art generative quality, but they suffer from two critical limitations regarding computational efficiency:

Rigid Compute-Resolution Coupling: In standard DiTs, the computational cost (FLOPs) is a fixed function of the input image resolution. This prevents principled trade-offs between latency and quality during inference.
Uniform Compute Allocation: DiTs allocate computation uniformly across all spatial tokens (image patches), regardless of the visual complexity of the region. Experiments show that adding "dummy" tokens (zero-padded patches) increases FLOPs but fails to improve generation quality, indicating that DiTs waste resources on unimportant regions and cannot dynamically prioritize difficult areas.

Existing solutions like adaptive generators often spread compute uniformly or suffer from high complexity, while token-dropping methods usually disable dropping at inference to prevent information loss, failing to reduce inference costs.

2. Methodology: Elastic Latent Interface Transformer (ELIT)

The authors propose ELIT, a minimal, drop-in mechanism compatible with DiT architectures that decouples input image size from computational cost.

Core Architecture

ELIT introduces a Latent Interface, a variable-length set of learnable tokens, and two lightweight cross-attention layers:

Read Layer: A cross-attention module that pulls information from the spatial tokens (input image patches) into the latent interface. It prioritizes "hard" or informative regions, effectively filtering out redundant areas.
Write Layer: A symmetric cross-attention module that broadcasts the updated latent state back to the spatial tokens, allowing the model to refine predictions and retain fine details.
Latent Core: The majority of the transformer blocks operate on this variable-length latent sequence rather than the full spatial grid.

Training Strategy: Importance-Ordered Representation

To enable flexible inference budgets, ELIT is trained with a random tail-dropping strategy:

During training, a random number of latent tokens ( $\tilde{J}$ ) are kept per group, while the "tail" tokens are dropped.
This forces the model to learn an importance-ordered hierarchy: early latent tokens capture global structure and essential information, while later tokens refine details.
Consequently, any prefix of the latent tokens forms a valid, lower-compute representation.

Inference: Dynamic Budget Control

At inference, the user controls the compute budget by selecting the number of latent tokens ( $\tilde{J}$ ) to retain.

Variable Compute: Reducing $\tilde{J}$ directly reduces FLOPs and latency without retraining.
Cheap Classifier-Free Guidance (CCFG): The multi-budget nature of ELIT enables a novel guidance strategy. Instead of running a full unconditional model for guidance (which doubles cost), ELIT uses a "weaker" version of itself (with fewer tokens) as the guidance term. This reduces inference cost by ~33% while maintaining or improving quality.

3. Key Contributions

Decoupling Resolution and Compute: ELIT allows a single model to operate across a spectrum of latency-quality points by simply adjusting the number of latent tokens, breaking the fixed FLOPs-resolution dependency of DiTs.
Adaptive Compute Allocation: Through the Read/Write mechanism, the model learns to concentrate computation on complex image regions and ignore redundant ones, effectively redistributing resources.
Drop-in Compatibility: The method requires no changes to the rectified flow objective or the core DiT stack. It only adds two cross-attention layers and latent token sampling, making it applicable to DiT, U-ViT, HDiT, and MM-DiT backbones.
Efficient Guidance: The introduction of CCFG leverages the multi-budget capability to accelerate sampling without external weaker models or handcrafted corruptions.

4. Experimental Results

The authors evaluated ELIT on ImageNet-1K (256px and 512px) and Kinetics-700 (video) across multiple backbones (DiT, U-ViT, HDiT).

Performance Gains:
- On ImageNet-1K 512px, ELIT-DiT achieved a 35.3% improvement in FID and 39.6% improvement in FDD (Fréchet Distance on DINOv2 features) compared to the baseline DiT, despite similar training FLOPs.
- Improvements were even more pronounced at higher resolutions (512px), where pixel redundancy is higher.
- ELIT consistently outperformed baselines across U-ViT and HDiT architectures.
Compute-Quality Trade-off:
- ELIT offers a superior quality-compute curve compared to simply reducing sampling steps.
- At 25% of the tokens, ELIT still outperformed the baseline DiT, whereas training-free token merging methods were capped by the baseline's quality.
Large-Scale Application:
- Applied to Qwen-Image (a 20B parameter MM-DiT), ELIT enabled stable training and inference.
- It achieved a 2.7x speedup (reducing FLOPs by up to 63%) while maintaining strong performance on the DPG-Bench benchmark, with only a negligible drop in average scores (90.45 to 88.02).
Convergence: ELIT accelerated training convergence, achieving up to 4.0x speedup in reaching target metrics compared to standard DiTs.

5. Significance

ELIT represents a significant shift in generative model design by moving from rigid, uniform computation to elastic, adaptive computation.

Practicality: It solves the "one model, one budget" problem, allowing a single set of weights to serve diverse hardware constraints (e.g., mobile vs. server) without retraining.
Efficiency: By learning to ignore unimportant regions, it addresses the fundamental inefficiency of uniform attention in diffusion models.
Scalability: The method scales effectively to large models (20B+ parameters) and high resolutions, offering a path toward more sustainable and accessible high-fidelity generation.

In summary, ELIT provides a simple yet powerful architectural modification that transforms Diffusion Transformers into flexible, budget-aware engines, delivering higher quality at lower costs and enabling new inference strategies like CCFG.