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

TL;DR

ELIT introduces a flexible, learnable latent interface in diffusion transformers, enabling dynamic adjustment of computation based on resource constraints while improving image generation quality.

Contribution

ELIT provides a minimal, compatible mechanism to decouple input size from compute, allowing dynamic resource allocation and improved performance in diffusion transformers.

Findings

Consistent performance gains across datasets and architectures.

Significant improvements in FID and FDD scores on ImageNet-1K.

Effective importance ordering of representations through training.

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/