Un modèle, plusieurs budgets : Interfaces latentes élastiques pour les transformeurs de diffusion

Résumé

Les transformateurs de diffusion (DiT) atteignent une qualité générative élevée mais verrouillent les FLOPs à la résolution de l'image, limitant les compromis principés entre latence et qualité, et allouent le calcul uniformément sur les tokens spatiaux d'entrée, gaspillant ainsi l'allocation des ressources dans les régions non importantes. Nous présentons Elastic Latent Interface Transformer (ELIT), un mécanisme interchangeable et compatible DiT qui découple la taille de l'image d'entrée du calcul. Notre approche insère une interface latente, une séquence de tokens de longueur variable et apprenable sur laquelle des blocs transformeurs standards peuvent opérer. Des couches de lecture et d'écriture par attention croisée légères transfèrent l'information entre les tokens spatiaux et les latents, et hiérarchisent les régions importantes de l'entrée. En entraînant avec un abandon aléatoire des latents de queue, ELIT apprend à produire des représentations ordonnées par importance, les premiers latents capturant la structure globale tandis que les derniers contiennent l'information pour affiner les détails. Lors de l'inférence, le nombre de latents peut être ajusté dynamiquement pour correspondre aux contraintes de calcul. ELIT est délibérément minimal, n'ajoutant que deux couches d'attention croisée tout en laissant inchangés l'objectif de flux rectifié et l'architecture DiT. Sur divers ensembles de données et architectures (DiT, U-ViT, HDiT, MM-DiT), ELIT apporte des gains constants. Sur ImageNet-1K 512px, ELIT offre un gain moyen de 35,3 % et 39,6 % pour les scores FID et FDD. Page du projet : https://snap-research.github.io/elit/

English

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/

Un modèle, plusieurs budgets : Interfaces latentes élastiques pour les transformeurs de diffusion

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

Résumé

Support