Nucleus-Image: Sparse MoE for Image Generation

Probleemstelling

De huidige staat van de kunst voor tekst-naar-afbeelding generatie (text-to-image) wordt gedomineerd door modellen met enorme parameterbudgetten. Hoewel deze modellen hoge kwaliteit leveren, zijn ze extreem rekenintensief en duur in inferentie. Er is een duidelijke behoefte aan modellen die een vergelijkbare of superieure kwaliteit bereiken met aanzienlijk lagere actieve parameters per doorloop (forward pass), zonder dat er post-training optimalisaties (zoals Reinforcement Learning of Direct Preference Optimization) nodig zijn. Bestaande Mixture-of-Experts (MoE) architecturen voor diffusiemodellen kampen vaak met stabiliteitsproblemen, inefficiënt routeren en hoge inferentiekosten door het meenemen van teksttokens in de volledige transformer-backbone.

Methodologie

Nucleus-Image introduceert een geavanceerde Sparse Mixture-of-Experts (MoE) Diffusion Transformer architectuur die is geoptimaliseerd voor zowel trainingsstabiliteit als inferentie-efficiëntie.

1. Architectuur en MoE-Design:

• Totale capaciteit: Het model heeft in totaal 17 miljard parameters, verdeeld over 32 transformerlagen.
• Actieve parameters: Dankzij Expert-Choice Routing worden er per doorloop slechts ongeveer 2 miljard parameters geactiveerd.
• Expert-Choice Routing: In tegenstelling tot traditionele token-choice routing (waar tokens experts kiezen), kiezen hier de experts hun top-k tokens. Dit garandeert een perfecte belastingverdeling zonder de noodzaak van zware auxiliaire load-balancing losses.
• Decoupled Routing: Een cruciale innovatie is het ontkoppelen van routeringsbeslissingen en expert-berekening. De router kijkt naar de niet-gemoduleerde tokenrepresentatie (gecombineerd met de tijdstap), terwijl de experts de volledig gemoduleerde representatie verwerken. Dit voorkomt dat de adaptieve modulatie (die varieert per tijdstap in diffusiemodellen) de routeringslogica overneemt, wat vaak leidt tot instabiliteit.
• Tekst-Optimalisatie: Teksttokens worden volledig uitgesloten van de MoE-lagen en de MLP's. Ze fungeren alleen als Key-Value bijdragers in de Joint Attention mechanismen. Dit elimineert routeringskosten voor tekst en maakt het mogelijk om tekst KV-cache te gebruiken over alle denoising-stappen heen.
• Gedeelde Expert: Elke MoE-laag bevat een gedeelde expert die alle tokens verwerkt, wat een basisberekeningspad biedt en redundante leerprocessen voorkomt.

2. Data-pijplijn:

• Het team heeft een corpus van 700 miljoen unieke afbeeldingen samengesteld uit webbronnen, resulterend in 1,5 miljard tekst-afbeelding paren.
• De pipeline omvat multi-stap filtering, deduplicatie, esthetische scoring (tiering A1-A5) en curatie van bijschriften op verschillende niveaus van detail (kort, medium, gedetailleerd).
• Er is speciale aandacht voor synthetische data voor tekstweergave om de leesbaarheid van tekst in gegenereerde afbeeldingen te verbeteren.

3. Training en Optimisatie:

• Curriculum: Training verloopt progressief in resolutie (256 → 512 → 1024 pixels) met een dynamische aanpassing van de "capacity factor" (hoeveel tokens per expert).
• Optimizer: Gebruik van de Muon optimizer (met Newton-Schulz orthogonalisatie) gecombineerd met een Warmup-Stable-Merge (WSM) leerplan.
• Geen EMA: In plaats van een dure Exponential Moving Average (EMA) shadow copy tijdens het trainen, worden checkpoints aan het einde van de training samengevoegd (merged) met gewichten gebaseerd op een inverse wortel-verdeling. Dit bespaart geheugen en behoudt flexibiliteit.
• Verliesfuncties: Rectified Flow loss als basis, aangevuld met Wavelet loss (voor hoge frequenties in late fasen) en specifieke regularisaties voor de router (Z-loss en Orthogonale loss) om expert-diversiteit te waarborgen.

4. Systeem- en Kernel-optimalisatie:

• Gebruik van aangepaste Triton kernels voor gefuseerde operaties (zoals SwiGLU activatie en LayerNorm) om geheugenbandbreedte-bottlenecks te verminderen.
• Expert Parallelism: De 64 experts zijn verspreid over 64 GPU's, wat een perfecte load-balancing mogelijk maakt met minimale communicatie-overhead dankzij geoptimaliseerde all-to-all communicatie.

Belangrijkste Resultaten

Nucleus-Image bereikt state-of-the-art prestaties op drie belangrijke benchmarks, terwijl het slechts 2B actieve parameters gebruikt:

1. GenEval: Behaalt een score van 0.87, wat gelijkstaat aan Qwen-Image en beter is dan SD3.5 Large (0.71) en FLUX.1 Dev (0.66). Het scoort bijzonder hoog op ruimtelijke posities (0.85).
2. DPG-Bench: Bereikt de hoogste totaalscore van 88.79, net voorbij Qwen-Image (88.32) en Seedream 3.0 (88.27). Het excelleert in het volgen van complexe prompts met meerdere eisen.
3. OneIG-Bench: Behaalt een score van 0.522, wat het plaatst boven Imagen4 (0.515) en Recraft V3 (0.502), met name sterk in stijl en tekstweergave.

Conclusie van de resultaten: Het model bereikt prestaties die vergelijkbaar zijn met of beter zijn dan modellen met aanzienlijk grotere actieve parameterbudgetten, maar tegen een fractie van de inferentiekosten.

Bijdragen en Betekenis

De belangrijkste bijdragen van dit paper zijn:

• Eerste Open Source MoE Diffusiemodel: Nucleus-Image is het eerste volledig open-source MoE diffusiemodel van deze kwaliteit, waarbij de weights, training code en dataset vrij beschikbaar zijn.
• Efficiëntie zonder Kwaliteitsverlies: Het bewijst dat sparse MoE-schaal een zeer effectieve route is voor hoogwaardige afbeeldinggeneratie, waarbij de "Pareto-grens" tussen kwaliteit en efficiëntie wordt verlegd.
• Stabiliteitsinnovaties: De introductie van decoupled routing lost een specifiek stabiliteitsprobleem op bij MoE diffusiemodellen (waarbij de tijdstap-modulatie de routering verstoort), en de WSM-schedule elimineert de noodzaak voor EMA-shadow copies.
• Inferentie-efficiëntie: Door teksttokens uit de MoE-lagen te houden en joint attention te gebruiken, wordt de inferentielatentie drastisch verlaagd zonder post-training optimalisatie (zoals RLHF of DPO).
• Schaalbaarheid: Het paper demonstreert dat het mogelijk is om modellen te trainen die concurreren met gesloten, proprietary modellen (zoals die van Google of OpenAI) door slimme architecturale keuzes en data-curatie in plaats van brute rekenkracht.

Kortom, Nucleus-Image toont aan dat door slimme architecturale innovaties (MoE, decoupled routing, WSM) en grondige data-engineering, hoogwaardige tekst-naar-afbeelding generatie toegankelijker en rendabeler kan worden gemaakt voor de bredere gemeenschap.