ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

Technische Zusammenfassung: ProbMoE: Differenzierbares probabilistisches Routing für Mixture-of-Experts

Problemstellung
Mixture-of-Experts (MoE)-Modelle skalieren effizient, indem sie pro Token nur eine kleine Teilmenge von Experten aktivieren. Das Training dieser Modelle bleibt jedoch eine Herausforderung, da das Standard-Routing auf einem diskreten, nicht differenzierbaren Top-$k$-Operator basiert. Herkömmliche Trainingsverfahren umgehen dies, indem sie die Abhängigkeit der Top-$k$-Auswahl von den Router-Logits ignorieren und die Gradienten nur durch die Softmax-Wahrscheinlichkeiten der ausgewählten Experten propagieren. Diese Approximation begrenzt das Lernsignal für alternative Experten-Teilmengen, was oft zu konzentrierten Routing-Verteilungen, schlechter Experten-Auslastung und instabilen Trainingsdynamiken führt. Bestehende Versuche, dies zu mildern, wie etwa dichte Backpropagation oder heuristische Gradientenschätzer, scheitern oft daran, die Verteilung über zulässige Experten-Teilmengen explizit zu modellieren oder bieten nur begrenzte Mechanismen zur Exploration alternativer Kombinationen während des Trainings.

Methodik
Die Autoren schlagen ProbMoE vor, ein Framework, das das Experten-Routing als ein Problem der probabilistischen Inferenz über kardinalitätsbeschränkte Experten-Teilmengen formuliert. Anstatt das Routing als deterministischen Operator zu behandeln, modelliert ProbMoE die Auswahl von Experten als eine Verteilung über diskrete Teilmengen.

1. Probabilistische Formulierung:

   • Exact-$k$ Routing: Das Framework definiert eine Verteilung über Teilmengen $S$, wobei $|S|=k$. Es bedingt unabhängige Bernoulli-Auswahrscheinlichkeiten (abgeleitet aus den Router-Logits via Sigmoid) auf der Bedingung, dass genau $k$ Experten ausgewählt werden. Die Normierungskonstante für diese Verteilung, $Z_k$, wird exakt mittels dynamischer Programmierung berechnet, wodurch die Formulierung handhabbar wird.
   • Dynamic-$k$ Routing: Das Framework generalisiert auf ein Bereich-beschränktes Setting, in dem die Teilmengen-Größe $|S|$ innerhalb von $[k_{min}, k_{max}]$ liegt. Die Verteilung ist auf diesen Bereich bedingt, was es dem Modell ermöglicht, sowohl die Identitäten der Experten als auch die Routing-Kardinalität für jeden Token gemeinsam zu inferieren.

2. Forward Pass:

   • Während des Trainings sampelt ProbMoE eine Experten-Teilmenge $S$ aus der definierten Verteilung (Exact-$k$ oder Bereich-beschränkt).
   • Die MoE-Schicht berechnet die Ausgabe unter Verwendung nur der gesampelten Experten, wodurch das spärliche Ausführungsmuster (Sparse Execution) von Standard-MoE-Modellen gewahrt bleibt.
   • Bei der Inferenz wird die Maximum-A-Posteriori (MAP)-Menge ausgewählt, um sicherzustellen, dass die Inferenzkosten mit konventionellem MoE-Routing übereinstimmen.

3. Backward Pass:

   • Um eine gradientenbasierte Optimierung zu ermöglichen, verwendet ProbMoE den SIMPLE Gradientenschätzer (Ahmed et al., 2023).
   • Anstatt durch die diskrete Stichprobe zu differenzieren, verwendet der Backward Pass die marginalen Wahrscheinlichkeiten jedes Experten, der unter der beschränkten Verteilung ausgewählt wird.
   • Die Routing-Gewichte werden unter Verwendung eines Straight-Through-Estimators konstruiert: $w = (\text{stopgrad}(z - m) + m) \odot \pi$, wobei $z$ die gesampelte Maske ist, $m$ der Vektor der exakten marginalen Wahrscheinlichkeiten und $\pi$ die Softmax-Routing-Wahrscheinlichkeiten sind.
   • Dies ermöglicht es, dass Gradienten durch die marginalen Wahrscheinlichkeiten fließen, was informative Updates an den Router liefert, die widerspiegeln, wie Änderungen der Auswahlwahrscheinlichkeiten die erwartete Verlustfunktion beeinflussen, während die spärliche Experten-Ausführung beibehalten wird.

Wesentliche Beiträge

1. Probabilistische Inferenz-Formulierung: Die Arbeit stellt das MoE-Routing als probabilistische Inferenz über kardinalitätsbeschränkte Experten-Teilmengen dar und bewegt sich damit über heuristische Gradientenapproximationen hinaus.
2. Handhabbare marginalbasierte Optimierung: Durch die Nutzung des SIMPLE-Schätzers kombinieren die Autoren stochastisches Teilmengen-Sampling mit exakten marginalbasierten Router-Gradienten, was ein differenzierbares Training ohne dichte Berechnung auf der Experten-Seite ermöglicht.
3. Vereinheitlichtes Dynamic-$k$ Framework: Die Methode lässt sich natürlich auf dynamisches Routing erweitern, bei dem das Modell sowohl die Identitäten der Experten als auch die Anzahl der aktiven Experten pro Token innerhalb eines vordefinierten Bereichs inferiert, ohne sich auf heuristische Gating-Regeln zu verlassen.
4. Empirische Verbesserungen:
   • Exact-$k$: ProbMoE erzielt eine starke Performance über Benchmarks (GSM8K, Law, Translation, MBPP, MMLU) auf OLMoE- und Qwen-Backbones und übertrifft oft kompetitive Baselines wie DenseMixer und konventionelles Top-$k$.
   • Dynamic-$k$: ProbMoE Dynamic-$k$ erreicht eine vergleichbare Performance wie Exact-$k$, aktiviert dabei jedoch im Durchschnitt weniger Experten, was eine effektive adaptive Allokation demonstriert.
5. Routing-Diversität: Analysen zeigen, dass ProbMoE im Vergleich zu Baselines breitere Routing-Verteilungen und eine ausgewogenere Experten-Auslastung induziert, wodurch der Experten-Kollaps gemildert wird.

Ergebnisse und Bedeutung
Die Arbeit zeigt, dass ProbMoE die Experten-Auslastung und die Routing-Diversität verbessert. In Ablationsstudien wird gezeigt, dass die Kombination aus stochastischem Sampling und marginalbasierter Optimierung der primäre Treiber für die Leistungssteigerungen ist, wobei sie Konfigurationen, die nur eine dieser Komponenten nutzen, übertrifft.

Die Autoren behaupten, dass ProbMoE eine fundierte Methode zur Expertenauswahl bietet, welche die Trainingsstabilität und die rechnerische Effizienz verbessert. Durch die probabilistische Modellierung des Routings ermöglicht das Framework die Exploration alternativer Experten-Kombinationen während des Trainings, was zu besser kalibrierten Routing-Verteilungen führt. Die Dynamic-$k$-Variante zeigt zudem, dass eine adaptive Kardinalität gelernt werden kann, um sich an die Token-Komplexität anzupassen, wodurch die effektiven Routing-Kosten reduziert werden, ohne die Aufgaben-Performance zu beeinträchtigen. Die Arbeit legt nahe, dass zukünftige Richtungen in der Erweiterung von ProbMoE auf das Pre-Training und systemseitige Optimierungen liegen, um die Sparsamkeit in dynamischen Settings voll auszuschöpfen.