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

Probleemstelling

Diffusion Transformers (DiTs) hebben de staat-of-the-art bereikt in beeld- en videogeneratie, maar ze lijden aan twee fundamentele beperkingen:

1. Rigiditeit in compute-toewijzing: De rekentijd (FLOPs) is strikt gekoppeld aan de beeldresolutie. Een DiT voert per stap een vast aantal berekeningen uit, ongeacht of de invoer complex of eenvoudig is. Dit maakt het moeilijk om een principieel afweging te maken tussen latentie (snelheid) en kwaliteit.
2. Uniforme resource-allocatie: DiTs verdelen rekenkracht gelijkmatig over alle ruimtelijke tokens (beeldpatches). Experimenten tonen aan dat DiTs extra rekenkracht niet effectief benutten om de kwaliteit te verbeteren; zelfs als men de token-aantal kunstmatig verhoogt (bijv. door nul-patchen toe te voegen), verspillen ze resources aan onbelangrijke gebieden in plaats van zich te concentreren op complexe regio's.

Bestaande oplossingen, zoals adaptieve generators of token-dropping, hebben vaak te maken met hoge complexiteit, vereisen extra trainingsdoelen (auxiliary losses), of werken niet tijdens de inferentie (inference) omdat informatie verloren gaat.

Methodologie: Elastic Latent Interface Transformer (ELIT)

De auteurs stellen ELIT voor, een "drop-in" mechanisme dat compatibel is met bestaande DiT-architecturen en de invoerbeeldgrootte ontkoppelt van de benodigde rekentijd.

Kerncomponenten:

1. Latent Interface: In plaats van dat alle transformer-blokken direct op de ruimtelijke tokens werken, wordt een variabele reeks van latent tokens (het "latent interface") geïntroduceerd. Dit is een compacte representatie waarop de meeste zware transformer-blokken opereren.
2. Read en Write Layers: Twee lichtgewicht cross-attention lagen verbinden de ruimtelijke domein (ruimtelijke tokens) met het latente domein:
   • Read: Haalt informatie uit de ruimtelijke tokens en projecteert deze naar de latent interface. Deze laag prioriteert complexe of "moeilijke" regio's van het beeld.
   • Write: Broadcastt de bijgewerkte latent states terug naar de ruimtelijke tokens om details te verfijnen.
3. Groepering (Grouped Cross-Attention): Om de kosten van Read/Write te verlagen, worden tokens opgesplitst in groepen. Cross-attention vindt alleen plaats binnen deze groepen, wat de complexiteit verlaagt van $O(NK)$ naar $O(NK/G)$.
4. Trainingsstrategie (Tail Dropping): Tijdens het trainen worden willekeurig de "staart" (tail) van de latent tokens per groep weggegooid. Hierdoor leert het model om informatie te ordenen op basis van belangrijkheid: de eerste tokens bevatten globale structuur, terwijl latere tokens fijne details toevoegen.
5. Inferentie: Het aantal latent tokens is een door de gebruiker instelbare parameter ("knop"). Door minder tokens te gebruiken, wordt de rekentijd direct verlaagd zonder de modelgewichten opnieuw te hoeven trainen.

Geavanceerde Inferentie-optimalisatie:
De auteurs introduceren CCFG (Cheap Classifier-Free Guidance). Omdat ELIT een model biedt dat natief werkt op verschillende budgetten, kan het een "zwakker" versie van zichzelf (met minder tokens) gebruiken als de conditionele guidance-term in plaats van een apart getraind zwak model. Dit reduceert de inferentiekosten met ongeveer 33% zonder kwaliteitsverlies.

Belangrijkste Bijdragen

• Decoupling van Resolutie en Compute: ELIT maakt het mogelijk om de rekentijd dynamisch aan te passen tijdens de inferentie, ongeacht de invoerresolutie.
• Niet-uniforme Compute-toewijzing: Het model leert automatisch resources te concentreren op visueel complexe gebieden, wat leidt tot efficiëntere generatie.
• Drop-in Compatibiliteit: Het vereist geen wijzigingen in het trainingsdoel (Rectified Flow) of de basis-DiT-stack. Het werkt naadloos met DiT, U-ViT, HDiT en MM-DiT.
• Autoguidance uit de doos: Door de variabele budgetten kan het model autoguidance technieken toepassen zonder extra training of handmatige modelcorrupties.

Resultaten

De prestaties zijn getest op ImageNet-1K (256px en 512px) en Kinetics-700 (video):

• Kwaliteitsverbetering: Op ImageNet-1K 512px levert ELIT een gemiddelde verbetering van 35,3% in FID en 39,6% in FDD (Fréchet Distance op DINOv2 features) ten opzichte van baselines, zelfs bij vergelijkbare trainingskosten.
• Multi-Budget Flexibiliteit: Een enkel ELIT-model kan een continuüm van latentie-kwaliteit trade-offs bedienen. Bij het verminderen van het aantal tokens behoudt het model de structuur beter dan het simpelweg verminderen van het aantal sampling-stappen.
• Schaalbaarheid: De voordelen zijn groter bij hogere resoluties en grotere modellen (bijv. DiT-XL), wat suggereert dat dynamische compute-toewijzing essentieel is voor schaalbare generatieve modellen.
• Grote Modellen (Qwen-Image): Toepassing op een 20B-parameter MM-DiT (Qwen-Image) toonde aan dat ELIT stabiel trainbaar is en de sampling-FLOPs met tot 63% kan reduceren (een snelheidswinst van ~2,7x) met minimale kwaliteitsverlies.
• Convergentie: ELIT-modellen convergeren sneller dan standaard DiT-modellen (tot 4x sneller op 512px).

Betekenis en Impact

Dit paper biedt een fundamentele verschuiving in hoe diffusion modellen compute toewijzen. In plaats van dat kwaliteit strikt afhankelijk is van een vast rekenbudget, introduceert ELIT een elasticiteit die het mogelijk maakt om een model te gebruiken dat zich aanpast aan de beschikbare hardware of latency-eisen.

De methode lost het probleem op van "verspilde" rekenkracht op eenvoudige beeldregio's en biedt een praktische route naar snellere, goedkopere generatieve AI zonder in te leveren op de uiteindelijke beeldkwaliteit. Het concept van een "latent interface" als schakel tussen invoer en verwerking kan een nieuwe standaard worden voor efficiënte transformer-architecturen in de toekomst.