MMOE: Modernizing Diffusion Transformers with Efficient Expert Design

Technische Samenvatting: MMOE: Modernisering van Diffusion Transformers met Efficiënt Expert-ontwerp

Probleemstelling
Hoewel Large Language Models (LLMs) succesvol hebben geschaald door de totale parametercapaciteit te ontkoppelen van de geactiveerde berekening via sparse Mixture-of-Experts (MoE) en efficiëntiemechanismen (bijv. zero-computation experts, gate-residual routing), hebben AIGC Foundation Models (AFMs), specifjes diffusion-transformer backbones, grotendeels een ander traject gevolgd. De huidige inspanningen om MoE toe te passen op diffusion transformers (bijv. DiT-MoE, Race-DiT) richten zich primair op het vergroten van de totale expert-pool en sparsity-ratio's zonder de efficiëntiemechanismen te importeren die de schaling van LLM's praktisch maken. Deze aanpak resulteert vaak in modellen die moeilijk te trainen en te implementeren zijn op apparaten met beperkte middelen, waardoor ze er niet in slagen de generatiekwaliteit te balanceren tegenover de trainings- en implementatiekosten. Het artikel stelt de vraag: Kunnen de architecturale principes achter efficiënte LLM-schaling worden aangepast aan AFMs om de generatiekwaliteit te verbeteren zonder proportioneel de effectieve berekening of de trainingslast te verhogen?

Methodologie: ModernMOE (MMOE)
De auteurs stellen ModernMOE (MMOE) voor, een modernisering van SiT-stijl diffusion transformers die systematisch vier in LLM's bewezen efficiëntiemechanismen integreert in de diffusion backbone. Cruciaal is dat MMOE de externe interface (denoising objectief, conditionering, latente representatie en sampler) ongewijzigd laat, en alleen de interne block-niveau berekeningen aanpast.

De architectuur introduceert de volgende componenten:

1. MoE++ Expert Layer: In plaats van standaard gerouteerde MLPs, bevat de expert-pool:
   • Heavy Experts: Standaard MLPs.
   • Lightweight Zero-Computation Experts: "Copy" (identiteit), "Zero" (output nul) en "Constant" (geleerd vector-mengsel) experts. Deze stellen het model in staat om de volledige MLP-executie over te slaan voor tokens waar een zware update niet nodig is.
2. Gate-Residual Routing: De router incorporeert een residu van de beslissing van de vorige expert-laag, waardoor routing-tendensen over lagen heen kunnen worden overgenomen, wat de beslissingen stabiliseert.
3. Shared Experts: Een gedeelde dense expert is altijd actief om een gemeenschappelijk berekeningspad te bieden, naast geruteerde experts voor token-afhankelijke specialisatie.
4. Attention-Residual Aggregation: Het model hergebruikt informatie van eerder voltooide block-states. Vóór zowel de attention als de MLP sub-lagen aggregeert het model representaties van een lijst met voltooide block-states met behulp van een geleerde pseudo-query, wat kruis-diepte informatiehergebruik mogelijk maakt zonder nieuwe generatieve objectieven.

De forward pass houdt recurrente states bij (voltooide block-states, huidige partiële state en gate-residuen) om deze aggregatie en routing te faciliteren.

Belangrijkste Bijdragen

• Formulering van AFM Modernisering: Het artikel frameert AFM-schaling als een moderniseringsprobleem, waarbij systematisch wordt onderzocht of LLM MoE-ontwerpen de diffusion transformers kunnen verbeteren zonder simpelweg de parameter-aantallen of sparsity-ratio's op te blazen.
• Unified Architectuur: Het voorstel van MMOE, dat geruteerde experts, gedeelde experts, lightweight zero-computation experts, gate-residual routing en attention-residual aggregatie integreert in een enkele SiT-stijl backbone.
• Gecontroleerde Empirische Studie: De auteurs bieden een rigoureuze, gecontroleerde vergelijking van het moderniseringspad (Dense SiT $\to$ Standard MoE $\to$ Shared MoE $\to$ MoE++ $\to$ AttnRes-MoE $\to$ MMOE) onder een gematcht, toegankelijk trainingsbudget (enkele acht-GPU H100 node, batch size 256, 400k stappen).

Resultaten
Geëvalueerd op class-conditional ImageNet-256 generatie:

• Convergentie en Kwaliteit: MMOE bereikt een lagere Fréchet Inception Distance (FID) bij elk geregistreerd checkpoint vergeleken met de dense SiT en tussenliggende sparse-expert baselines. Het convergeert sneller per trainingsstap en bereikt een FID van 3.75 bij 400k stappen, vergeleken met 5.20 voor de dense baseline.
• Efficiëntie: Terwijl standaard MoE-varianten de wall-clock trainingstijd verhoogden door routing-overhead, bereikte MMOE de beste kwaliteit-kostenbalans onder de sparse varianten. Het verminderde de trainingstijd van de duurste variant (AMOE) van 120 uur naar 67 uur terwijl het een superieure FID behield.
• Schaalbaarheid: De verbeteringen houden stand bij verschillende model-schalen (S/2, B/2, L/2) en resoluties (ImageNet-512), wat aangeeft dat de voordelen niet beperkt zijn tot de XL/2 setting.
• Routing Analyse: Analyse van het geconvergeerde model onthult:
  • Stabiele Specialisatie: Sterke block-wise expert specialisatie (hoge load Gini coëfficiënt) in plaats van uniforme benutting.
  • Diepte-afhankelijke Lightweight Gebruik: Lightweight experts (copy/zero/constant) worden intensief gebruikt in vroege blocks (44.6% van de routes) en nemen af in latere blocks (23.4%), wat de hypothese ondersteunt dat vroege denoising stappen profiteren van goedkope transformaties.
  • Lage Switch Rates: De top-2 expert set verandert voor slechts ~1.4% tot 2.7% van de tokens tussen aangrenzende denoising stappen, wat wijst op stabiele routing-beslissingen.
• Geheugenbesparing: Het gebruik van lightweight experts vermindert de forward activation memory met ~20% en de backward memory met ~32% per block vergeleken met attention-residual MoE zonder lightweight experts.

Betekenis en Claims
Het artikel claimt dat AFMs het "gebalanceerde schaalpad" van LLM's kunnen volgen door bewezen efficiëntie-ontwerpen te importeren in plaats van enkel de totale parameters en sparsity-ratio's te vergroten. De betekenis van MMOE ligt in het aantonen dat:

1. Kwaliteit-Kosten Balans: Hoogwaardige generatie kan worden bereikt met betere kwaliteit-kosten trade-offs door sparse routing te combineren met goedkope expert-paden en representatiehergebruik, in plaats van enkel te vertrouwen op grotere expert-pools.
2. Haalbaarheid op Toegankelijke Budgets: Deze verbeteringen zijn haalbaar binnen een single-machine budget (8x H100), wat geavanceerd diffusion transformer onderzoek toegankelijker maakt dan voorgaande werken die massale multi-node clusters en miljoenen iteraties vereisen.
3. Mechanisme Transfer: Specifieke LLM efficiëntiemechanismen (zero-computation experts, gate-residuals, attention-residuals) worden effectief overgedragen aan diffusion transformers, wat de convergentiesnelheid en specialisatie-stabiliteit verbetert.

De auteurs blijven bescheiden over hun claims en merken op dat zij niet beweren de state-of-the-art FID te hebben vergeleken met methoden die op veel grotere schaal zijn getraind (bijv. 7M iteraties), maar dat hun bijdrage de gecontroleerde moderniseringsstudie en de demonstratie van superieure prestaties onder een beperkt, single-node budget is. Zij erkennen beperkingen, waaronder de focus op class-conditional generatie en de noodzaak voor verdere studie naar gedistribueerde trainingsefficiëntie, die momenteel communication-bound is.