Elastic Diffusion Transformer

Titel: Elastic Diffusion Transformer (E-DiT)

Auteurs: Jiangshan Wang, Zeqiang Lai, et al. (Tsinghua University, Tencent Hunyuan, CUHK, HITSZ)

---

1. Het Probleem

Diffusiemodellen, en specifiek Diffusion Transformers (DiT), hebben opmerkelijke generatieve vaardigheden getoond voor afbeeldingen, video's en 3D-objecten. Echter, deze modellen zijn extreem rekenintensief en vereisen vaak tientallen denoisings-stappen met zware transformer-berekeningen, wat leidt tot hoge latentie en beperkte praktische inzetbaarheid.

Bestaande versnellingsmethoden, zoals pruning (het verwijderen van parameters) en distillatie, hebben twee belangrijke beperkingen:

1. Statische architectuur: Ze passen een vaste, kleinere modelstructuur toe op alle denoisings-stappen en alle invoer.
2. Gebrek aan flexibiliteit: Ze negeren dat de rekenbehoeften per sample variëren. Sommige samples zijn "makkelijk" en vereisen minder rekenkracht, terwijl andere complex zijn en volledige capaciteit nodig hebben. Een statische aanpak leidt vaak tot een suboptimale afweging tussen snelheid en kwaliteit (of te traag, of te veel kwaliteitsverlies).

2. Methodologie: Elastic Diffusion Transformer (E-DiT)

E-DiT is een adaptief versnellingsframework dat de sample-afhankelijke sparsiteit (de variatie in rekenbehoeften per sample) benut. In plaats van een statisch model, past E-DiT de computationele capaciteit dynamisch aan tijdens de inferentie.

Het framework bestaat uit drie complementaire componenten:

A. Router Architectuur

Elke transformer-blok ($B_i$) in het DiT-model wordt uitgerust met een lichtgewicht router ($R_i$). Deze router neemt de invoer-latent en de denoising-tijdstap als input en voert twee taken uit:

1. Adaptief Blokspringen (Block Skipping): De router voorspelt een kans ($p_i$) of een blok overgeslagen kan worden. Als $p_i$ onder een drempelwaarde ($\tau$) ligt, wordt het blok volledig genegeerd.
2. Adaptieve MLP-breedtevermindering: Voor blokken die niet worden overgeslagen, voorspelt de router de optimale verhouding voor de breedte van de Multi-Layer Perceptron (MLP). De standaard breedte (vaak 4x de input-dimensie) wordt dynamisch verkleind (bijv. naar 1/4, 1/2, 3/4 of 1) op basis van de complexiteit van het sample.

B. Trainingsstrategie

Het model wordt getraind met een gezamenlijk doel dat kwaliteit en efficiëntie combineert:

• Performance Loss: De standaard flow-matching loss om de generatiekwaliteit te behouden.
• Efficiency Loss: Een regularisatieterm die de gemiddelde kans op blokspringen en de gemiddelde MLP-breedte bestraft om te sturen naar een gewenst efficiëntie-budget.
• Initialisatie: Training start met een volledig geopend model (geen blokken overgeslagen, volle breedte) om stabiliteit te garanderen, waarna de router geleidelijk leert welke onderdelen overbodig zijn.

C. Block-wise Feature Caching (Tijdredundantie)

Tijdens inferentie wordt een extra versnellingstechniek toegepast voor blokken die net boven de spring-drempel liggen (de "borderline" zone).

• In plaats van deze blokken direct te springen (wat kwaliteit zou kosten), worden hun tussenliggende features (residuen) gecacheerd.
• Bij volgende tijdstappen kunnen deze gecachete features worden hergebruikt als de router aangeeft dat het blok weinig bijdraagt.
• Dit elimineert redundante berekeningen zonder extra training te vereisen.

3. Belangrijkste Bijdragen

1. Adaptieve Architectuur: E-DiT introduceert een framework dat computationele middelen dynamisch toewijst op basis van de moeilijkheidsgraad van het gegenereerde sample, in plaats van een statische verkleining.
2. Dubbele Adaptiviteit: Het combineert blokspringen (diepe structuur) met breedtevermindering (MLP-structuur) binnen hetzelfde systeem.
3. Training-vrije Caching: Een innovatieve caching-mechanisme dat gebruikmaakt van de router-voorspellingen om tijdsredundantie te exploiteren zonder de modelkwaliteit te schaden.
4. Generaliseerbaarheid: Het framework werkt op diverse basismodellen (Qwen-Image, FLUX, Hunyuan3D) en modaliteiten (2D afbeeldingen, 3D assets).

4. Resultaten

Experimenten zijn uitgevoerd op grote foundation modellen:

• Afbeeldingsgeneratie (Qwen-Image & FLUX):
  • E-DiT bereikt een snelheidswinst van ongeveer 2x (tot ~2x versnelling) vergeleken met de basismodellen.
  • De kwaliteit (gemeten via GenEval, DPG-Bench, T2I-CompBench) blijft verwaarloosbaar achter bij het originele model.
  • Bijvoorbeeld op Qwen-Image: Latentie daalt van 2431ms naar 1283ms (Turbo-modus) met een GenEval-score van 0.853 (vs 0.870 van het basismodel).
• 3D Asset Generatie (Hunyuan3D-3.0):
  • Bereikt een 2x versnelling (latentie daalt van ~5000ms naar ~2600ms) met behoud van geometrische details en kwaliteit (ULIP-I en Uni3D scores).
• Vergelijking met State-of-the-Art:
  • E-DiT presteert beter dan statische pruning-methoden (zoals TinyFusion, PPCL) en andere dynamische methoden (zoals DyDiT) qua snelheid-kwaliteit trade-off.

5. Betekenis en Impact

Het paper demonstreert dat de generatieve processen van diffusiemodellen inherent spars zijn en dat deze sparsiteit sterk varieert per sample.

• Efficiëntie: E-DiT maakt het mogelijk om zware generatieve modellen veel sneller en goedkoper in te zetten, wat cruciaal is voor real-time toepassingen en grootschalige deploy.
• Kwaliteit: Het bewijst dat dynamische aanpassing superieur is aan statische compressie, omdat het de rekenkracht richt op de complexe delen van een sample en tijd/energie bespaart op de eenvoudige delen.
• Toekomst: De aanpak biedt een nieuwe richting voor het ontwerpen van "elastic" neurale netwerken die zich aanpassen aan de input, niet alleen in diffusiemodellen maar mogelijk ook in andere generatieve taken.

Kortom, E-DiT lost het probleem van de hoge rekenkosten van DiT's op door een slimme, adaptieve router te introduceren die "slim springt" en "slim verkleint", resulterend in een snellere generatie zonder kwaliteitsverlies.