Routing Matters in MoE: Scaling Diffusion Transformers with Explicit Routing Guidance

Die Arbeit stellt ProMoE vor, ein Mixture-of-Experts-Framework für Diffusion-Transformer, das durch einen zweistufigen Router mit expliziter Leitlinie und einem kontrastiven Verlust die spezialisierte Zuweisung von Bild-Token-Experten verbessert und so den aktuellen Stand der Technik auf ImageNet übertrifft.

Das Wesentliche

Stell dir vor, du hast eine riesige, hochintelligente Kunstwerkstatt, die Bilder aus dem Nichts erschafft. Diese Werkstatt heißt Diffusion Transformer (DiT). Früher arbeitete dort ein einziger, extrem talentierter, aber müder Künstler, der alles selbst machen musste – vom Zeichnen eines Hundes bis zum Malen eines Himmels. Das war langsam und ineffizient, wenn die Werkstatt wachsen sollte.

Um das zu lösen, haben Forscher eine Idee aus der Welt der Sprach-KIs (wie Chatbots) übernommen: MoE (Mixture of Experts). Das ist wie ein Team aus vielen spezialisierten Künstlern. Statt dass einer alles macht, entscheidet ein "Kurator" (der Router), welcher Spezialist gerade an welchem Teil des Bildes arbeiten soll.

Das Problem:
Wenn man dieses MoE-System einfach auf Bild-KIs überträgt, funktioniert es nicht gut. Warum?

1. Text ist wie Perlenkette: Jeder Buchstabe oder jedes Wort hat eine klare, eigene Bedeutung. Ein "Hund" ist etwas ganz anderes als eine "Katze". Das ist leicht zu unterscheiden.
2. Bilder sind wie ein riesiges, verschwommenes Ölgemälde: Ein Bild besteht aus Millionen winziger Flecken (Pixel). Viele dieser Flecken sehen fast identisch aus (z. B. der blaue Himmel links und rechts). Außerdem gibt es im Bild zwei völlig verschiedene Arten von "Aufträgen":
   • Bedingte Aufträge: "Male einen Hund" (hier ist der Kontext wichtig).
   • Unbedingte Aufträge: "Male einfach nur etwas" (hier gibt es keinen Kontext).

Frühere Versuche haben alle diese Flecken einfach durcheinander gewürfelt und den Künstlern gegeben. Das Ergebnis? Die Spezialisten wurden verwirrt, lernten alle das Gleiche und konnten sich nicht auf ihre Stärken konzentrieren. Es war, als würde man einem Landschaftsmaler einen Porträtauftrag geben und einem Porträtmaler einen Landschaftsauftrag, ohne sie vorher zu fragen.

Die Lösung: ProMoE (ProMixture of Experts)
Die Autoren haben eine neue Methode namens ProMoE entwickelt. Sie ist wie ein sehr kluger Chef, der das Team neu organisiert. Hier ist, wie es funktioniert, mit einfachen Analogien:

1. Der Zwei-Schritt-Router (Der kluge Chef)

Statt alle Bilderflecken einfach zu verteilen, macht ProMoE zwei Dinge:

• Schritt 1: Die Funktion trennen (Conditional Routing)
  Der Chef fragt zuerst: "Ist dieser Fleck Teil des 'Hund'-Auftrags oder des 'leeren' Auftrags?"

  • Flecken ohne Kontext (leere Aufträge) gehen zu einer speziellen Gruppe von Künstlern, die nur für das "Grundgerüst" zuständig sind.
  • Flecken mit Kontext (z. B. "Hund") gehen in den nächsten Raum.
  • Analogie: Es ist wie in einer Küche. Man trennt zuerst die Leute, die nur den Tisch decken (Grundlage), von denen, die das eigentliche Gericht kochen (Inhalt).

• Schritt 2: Die Bedeutung verstehen (Prototypical Routing)
  Jetzt hat der Chef eine Liste von "Muster-Karten" (Prototypen). Jede Karte repräsentiert einen Spezialisten (z. B. "Experte für Fell", "Experte für Augen", "Experte für Gras").
  Der Chef vergleicht jeden Bild-Fleck mit diesen Karten. Wenn ein Fleck wie "Fell" aussieht, geht er zum Fell-Experten. Wenn er wie "Gras" aussieht, zum Gras-Experten.

  • Analogie: Statt zu raten, schaut der Chef genau hin und sagt: "Du siehst aus wie ein Hund, also geh zum Hund-Spezialisten. Du siehst aus wie ein Auto, also zum Auto-Spezialisten."

2. Der "Gegensatz-Lern-Trick" (Routing Contrastive Loss)

Damit die Experten wirklich gut werden, braucht es noch einen Trainings-Trick. Die Forscher haben eine neue Art zu lernen eingeführt.

• Das Prinzip: "Ähnliche Dinge gehören zusammen, verschiedene Dinge gehören getrennt."
• Wenn zwei Bild-Flecken sich sehr ähnlich sind (z. B. beide sind braunes Fell), müssen sie vom gleichen Experten bearbeitet werden.
• Wenn sie unterschiedlich sind, müssen sie von anderen Experten bearbeitet werden.
• Analogie: Stell dir vor, du bist ein Lehrer. Du sagst deinen Schülern: "Wenn ihr zwei Fragen habt, die fast gleich klingen, arbeitet zusammen! Aber wenn die Fragen total unterschiedlich sind, arbeitet getrennt, damit ihr nicht durcheinandergeraten." Das sorgt dafür, dass jeder Schüler (Experte) ein echtes Spezialist wird.

Das Ergebnis

Mit dieser neuen Methode (ProMoE) passiert etwas Magisches:

• Die KI wird schneller und effizienter, weil nicht jeder Spezialist alles machen muss.
• Die Bilder werden besser und detaillierter, weil die Experten sich auf ihre Stärken konzentrieren können.
• Es braucht sogar weniger Rechenleistung als die alten Methoden, um bessere Ergebnisse zu erzielen.

Zusammenfassung:
Die Forscher haben erkannt, dass Bilder anders funktionieren als Text. Statt die Spezialisten blind arbeiten zu lassen, haben sie ihnen einen klaren Plan gegeben: Erst trennen, was wichtig ist, dann genau zuordnen, was gemacht werden muss, und sie dabei trainieren, sich auf ihre spezifischen Aufgaben zu konzentrieren. Das Ergebnis ist eine KI, die Bilder nicht nur schneller, sondern auch kreativer und präziser malt.

Technische Zusammenfassung

1. Problemstellung

Mixture-of-Experts (MoE) hat sich als erfolgreiches Paradigma zur Skalierung von Large Language Models (LLMs) etabliert, da es die Modellkapazität erhöht, ohne die Rechenkosten proportional zu steigern. Bei der Anwendung auf Diffusion Transformers (DiTs) für die Bildgenerierung waren die bisherigen Ergebnisse jedoch enttäuschend. Bestehende Ansätze wie DiT-MoE oder EC-DiT zeigen nur marginale Verbesserungen oder sogar schlechtere Leistungen im Vergleich zu dichten Modellen.

Die Autoren identifizieren zwei fundamentale Unterschiede zwischen sprachlichen und visuellen Token als Ursache für dieses Versagen:

1. Hohe räumliche Redundanz: Im Gegensatz zu semantisch dichten Text-Token sind visuelle Token (Bild-Patches) kontinuierlich, räumlich gekoppelt und stark redundant. Dies führt dazu, dass Experten oft homogene Merkmale lernen, anstatt sich zu spezialisieren.
2. Funktionale Heterogenität: Diffusionsmodelle nutzen häufig Classifier-Free Guidance (CFG), was zwei funktionell unterschiedliche Token-Typen erzeugt: bedingte (conditional) und unbedingte (unconditional) Token. Herkömmliche MoE-Router behandeln diese jedoch einheitlich, was die notwendige funktionale Trennung und Spezialisierung der Experten behindert.

Das Ergebnis ist eine mangelnde Experten-Diversität und Intra-Experten-Kohärenz, was die Effektivität von MoE in visuellen Modellen einschränkt.

2. Methodik: ProMoE Framework

Um diese Herausforderungen zu lösen, stellen die Autoren ProMoE vor, ein MoE-Framework, das einen Zwei-Schritt-Router mit expliziter Routing-Guidance verwendet. Das Ziel ist die Förderung von Experten-Spezialisierung durch zwei Kriterien: Intra-Experten-Kohärenz (ein Experte verarbeitet konsistent ähnliche Muster) und Inter-Experten-Diversität (unterschiedliche Experten übernehmen unterschiedliche Aufgaben).

Der Router arbeitet in zwei Schritten:

A. Conditional Routing (Funktionale Trennung)

Der erste Schritt adressiert die funktionale Heterogenität durch eine harte Aufteilung der Token basierend auf ihrer Rolle:

• Unconditional Image Tokens: Token, die aus Bild-Patches unter Null-Bedingung (z. B. leere Labels) stammen, werden deterministisch einem dedizierten Satz von unconditional Experten zugeführt.
• Conditional Image Tokens: Token unter spezifischen Bedingungen werden an den zweiten Schritt weitergeleitet, wo sie unter den standardmäßigen routed Experten verteilt werden.
  Dies erzwingt eine funktionale Segregation und ermöglicht es den Experten, sich auf ihre jeweilige Rolle zu spezialisieren.

B. Prototypical Routing (Semantische Zuweisung)

Der zweite Schritt weist die bedingten Token basierend auf ihrem semantischen Inhalt zu.

• Anstelle eines linearen Layers zur Berechnung von Affinitätsscores verwendet ProMoE lernbare Prototypen ($P$), die jedem Experten zugeordnet sind.
• Die Zuweisung erfolgt über die Cosine-Ähnlichkeit zwischen den Token-Embeddings und den Prototypen im latenten Raum.
• Dies ermöglicht eine flexible, semantisch getriebene Zuweisung, die natürlicher für die visuelle Redundanz ist als herkömmliche Token-Choice-Methoden.

C. Routing Contrastive Loss (RCL)

Um die semantische Guidance im prototypischen Routing zu verstärken, wird ein neuer Routing Contrastive Loss eingeführt:

• Ziel: Semantisch ähnliche Token sollen zum selben Experten geleitet werden (Intra-Experten-Kohärenz), während Token mit unterschiedlicher Semantik zu verschiedenen Experten gedrückt werden (Inter-Experten-Diversität).
• Mechanismus: Der Loss zieht die Prototypen an die Schwerpunkte ihrer zugewiesenen Token-Cluster heran und drückt sie von den Schwerpunkten anderer Cluster weg.
• Vorteil: Im Gegensatz zu klassischem Load-Balancing, das nur die Anzahl der Token reguliert, nutzt RCL die semantische Struktur des Raums und wirkt als effektiver Regularisierer für die Experten-Spezialisierung.

3. Schlüsselbeiträge

1. Analyse der Token-Unterschiede: Die Arbeit liefert eine fundierte Analyse, warum MoE in DiTs anders funktioniert als in LLMs (Redundanz vs. semantische Dichte).
2. Zwei-Schritt-Router: Einführung eines hybriden Routing-Mechanismus, der zuerst nach Funktion (bedingte/unbedingte) und dann nach Semantik (Prototypen) trennt.
3. Routing Contrastive Loss: Ein neuer Verlustterm, der die Experten-Spezialisierung explizit fördert, ohne manuelle Labels zu benötigen und robuster als herkömmliche Clustering-Methoden ist.
4. Skalierbarkeit: Das Framework ist so gestaltet, dass es mit verschiedenen Modellgrößen (S bis XL) und Trainingszielen (Rectified Flow und DDPM) skaliert.

4. Ergebnisse

Die Autoren evaluieren ProMoE umfassend auf dem ImageNet-Benchmark (256x256) unter Verwendung von Rectified Flow und DDPM.

• Leistungsvorteil: ProMoE übertrifft sowohl dichte DiT-Modelle als auch den aktuellen State-of-the-Art (SOTA) MoE-Methoden (wie DiffMoE, EC-DiT) deutlich.
  • Beispiel (Rectified Flow, CFG=1.5): ProMoE-L reduziert den FID um 29,4 % im Vergleich zum dichten DiT-L und um 20,5 % im Vergleich zu DiffMoE-L.
• Parameter-Effizienz: ProMoE erreicht diese Ergebnisse mit weniger aktivierten Parametern als die dichten Baseline-Modelle und schlägt sogar Modelle mit 1,7-fach mehr Gesamtparametern.
• Konvergenz: ProMoE zeigt eine schnellere Konvergenz und erreicht bei weniger Trainingsschritten (500k) bessere Ergebnisse als dichte Modelle bei 1M oder 2M Schritten.
• Expertennutzung: Visualisierungen (t-SNE) zeigen, dass ProMoE eine klare Experten-Spezialisierung erreicht, während Baseline-Modelle (wie DiT-MoE) eine homogene Verteilung ohne echte Diversität aufweisen.
• Text-to-Image: Die Methode generalisiert erfolgreich auf Text-zu-Bild-Aufgaben (GenEval-Benchmark) und übertrifft dort ebenfalls dichte und Token-Choice-MoE-Baselines.

5. Bedeutung und Fazit

Die Arbeit demonstriert, dass der Erfolg von MoE in Diffusionsmodellen nicht allein von der Architektur, sondern maßgeblich von der Qualität des Routing-Mechanismus abhängt. Durch die explizite Berücksichtigung der funktionalen Rolle (CFG) und der semantischen Inhalte der Token sowie die Einführung eines kontrastiven Lernansatzes für das Routing, kann die inhärente Redundanz visueller Daten überwunden werden.

ProMoE bietet einen robusten Weg, Diffusion Transformers effizient zu skalieren, ohne die Rechenkosten drastisch zu erhöhen, und setzt einen neuen Standard für die Anwendung von MoE in der visuellen Generierung. Der Code ist öffentlich verfügbar, was die Reproduzierbarkeit und weitere Forschung fördert.