Three Phases of Expert Routing: How Load Balance Evolves During Mixture-of-Experts Training

1. Problemstellung und Motivation

Hintergrund:
Mixture-of-Experts (MoE)-Architekturen skalieren die Kapazität von Sprachmodellen, indem sie jedes Token nur an eine Teilmenge spezialisierter Expert-Netzwerke weiterleiten. Die zentrale ingenieurtechnische Herausforderung hierbei ist das Load Balancing (Lastverteilung). Ohne Eingriffe konzentrieren sich Token oft auf wenige „hochwertige" Experten, während andere untätig bleiben (Routing Collapse).

Das bestehende Problem:
Der Standardansatz zur Lösung ist ein auxiliärer Balance-Verlust (auxiliary balance loss), der eine Lastkonzentration bestraft. Während diese Methoden in der Praxis effektiv sind, fehlt ein theoretisches Verständnis dafür, wie sich der Trade-off zwischen Qualität (Expertenauswahl) und Balance (Lastverteilung) während des Trainings dynamisch entwickelt. Bisherige Analysen konzentrieren sich meist auf den konvergierten Endzustand und übersehen die Trainingsdynamik.

Ziel:
Die Arbeit modelliert das MoE-Routing als Kongestionsspiel (Congestion Game) und verfolgt einen einzigen effektiven Parameter, den Kongessionskoeffizienten $\gamma_{eff}$, um die evolutionäre Dynamik des Load Balancings über den gesamten Trainingsverlauf hinweg zu quantifizieren.

---

2. Methodik und Theoretischer Rahmen

Die Arbeit verbindet Spieltheorie (Mean-Field Games, MFG) mit empirischer Analyse von Trainings-Checkpoints.

A. Das Kongestionsspiel-Modell

• Formulierung: Token werden als Spieler, Experten als Ressourcen betrachtet. Die Kostenfunktion eines Agents (Tokens) hängt von der Qualität des Experten ($q_i$) und der Last auf diesem Experten ($\mu_i$) ab:
  $$\ell(i, \mu) = -q_i + \gamma \mu_i$$
  Dabei ist $\gamma$ der Kongessionskoeffizient, der die Stärke des Trade-offs zwischen Qualität und Balance bestimmt.
• Gleichgewicht: Das Gleichgewicht dieses Spiels entspricht einem Softmax mit Temperatur, wobei die Temperatur dem Kehrwert des Regularisierungsterms entspricht.
• Erkenntnis: Für gut balancierte Modelle ist das MFG-Gleichgewicht theoretisch äquivalent zu einem temperatur-skalierten Softmax. Das Spiel bietet also keinen Vorteil bei der statischen Vorhersage der Lastverteilung, liefert aber eine Erklärung, warum Softmax funktioniert und was die Temperatur physikalisch bedeutet (nämlich den Kongessionskoeffizienten).

B. Effektiver Kongessionskoeffizient ($\gamma_{eff}$)

Da Modelle sowohl durch den expliziten Loss als auch durch die implizite Dynamik des Gradientenabstiegs lernen, wird ein effektiver Kongessionskoeffizient $\gamma_{eff}$ definiert.

• Schätzung: $\gamma_{eff}$ wird durch Minimierung des $L_1$-Fehlers zwischen der beobachteten Lastverteilung und der theoretischen Gleichgewichtsverteilung (basierend auf geschätzten Expertenqualitäten) an jedem Trainings-Checkpoint geschätzt.
• Zerlegung: $\gamma_{eff} = \gamma_{explicit} + \gamma_{implicit}$.
  • $\gamma_{explicit}$: Stammt vom auxilliären Loss ($\alpha \cdot M$).
  • $\gamma_{implicit}$: Repräsentiert das Balance-Verhalten, das der Optimierer intern durch die Gewichte gelernt hat.

C. Empirische Analyse

• Modelle: Analyse von zwei Open-Source-MoE-Modellen: OLMoE-1B-7B (20 Checkpoints, dichte Abtastung im kritischen Bereich) und OpenMoE-8B (6 Checkpoints).
• Daten: Nutzung von 50 Text-Batches pro Checkpoint mit Bootstrap-Resampling für Konfidenzintervalle.
• Erweiterung: Einführung eines Multi-Type MFG, das Token-Heterogenität durch Clustering berücksichtigt, um die Vorhersagegenauigkeit zu verbessern.

---

3. Schlüsselbeiträge und Ergebnisse

Die zentrale Entdeckung ist eine nicht-monotone, dreiphasige Trajektorie von $\gamma_{eff}$ während des Pre-Trainings, die in konvergierten Modellen unsichtbar bleibt.

Die drei Phasen der Expert-Routing-Dynamik:

1. Phase 1: Surge (Anstieg) – Schritte 5K bis 50K

   • Dynamik: $\gamma_{eff}$ steigt stark an von ca. 13,7 auf einen Peak von 36–39 (bei 30K–40K Schritten).
   • Interpretation: Der Router lernt aktiv, die Last zu balancieren. Die Routing-Entropie steigt, während die Qualitätsstreuung ($B_0$) der Experten noch hoch ist. Der Optimierer priorisiert hier stark die Balance.
   • Bedeutung: Dies ist der kritische Bereich, in dem das Gleichgewicht etabliert wird.

2. Phase 2: Stabilization (Stabilisierung) – Schritte 100K bis 400K

   • Dynamik: $\gamma_{eff}$ plateauiert bei ca. 24–28.
   • Interpretation: Der Router hat seinen „Betriebspunkt" für Balance gefunden. Unter dieser stabilen Bedingung spezialisieren sich die Experten weiter (die Qualitätsstreuung $B_0$ sinkt von 2,41 auf 2,25).
   • Entkopplung: Die Balance-Strategie des Routers bleibt stabil, während die Expertenlernen innerhalb dieses Rahmens fortschreitet.

3. Phase 3: Relaxation (Entspannung) – Schritte 400K bis 1,2M

   • Dynamik: $\gamma_{eff}$ fällt deutlich von ca. 27 auf 8,5.
   • Interpretation: Sobald die Expertenrollen feststehen, lockert der Router die Balance-Strikte. Er tauscht Balance gegen Qualität, da er sich nun selektiver verhalten kann, ohne dass das System kollabiert.
   • Ergebnis: Das Modell priorisiert in der späten Phase die Qualität der Expertenwahl über eine strikte Lastverteilung.

Weitere wichtige Ergebnisse:

• Implizite Dominanz: Am Ende des Trainings (Konvergenz) ist $\gamma_{eff}$ (8,5) etwa 13-mal größer als der explizite Loss-Term $\gamma_{explicit}$ (0,64). Das bedeutet, der Optimierer internalisiert den Balance-Aspekt durch die Gradientendynamik weit stärker, als der auxilliäre Loss allein vorschreibt. Der Loss dient nur als „Samen".
• Multi-Type MFG: Die Erweiterung auf Token-Cluster (Multi-Type) verbessert die Lastvorhersage um durchschnittlich 30% im Vergleich zum Single-Type-Modell, insbesondere in den frühen Schichten, wo Token-Repräsentationen differenzierter sind.
• Stabilität: Die dreiphasige Trajektorie ist spezifisch für das Pre-Training. Während des Fine-Tunings (Annealing) bleibt $\gamma_{eff}$ stabil, und die Phasen treten nicht auf.
• Statische Äquivalenz: Im konvergierten Zustand ist das MFG-Modell statistisch nicht besser als ein einfaches temperatur-skaliertes Softmax (L1-Fehler 0,199 vs. 0,200). Der Wert des Modells liegt also rein in der Dynamik und der Zerlegung, nicht in der statischen Vorhersage.

---

4. Signifikanz und Implikationen

• Neues Verständnis der MoE-Optimierung: Die Arbeit widerlegt die Annahme, dass Balance und Qualität statisch optimiert werden. Stattdessen zeigt sie einen fundamentalen Zielkonflikt: Frühes Training priorisiert Balance, spätes Training priorisiert Qualität.
• Diagnostisches Werkzeug: Der Parameter $\gamma_{eff}$ dient als neues Diagnoseinstrument. Ein vorzeitiger Abfall von $\gamma_{eff}$ in Phase 2 könnte auf einen bevorstehenden Experten-Kollaps hindeuten. Der Sicherheitsabstand zum kritischen Schwellenwert $\gamma_c$ kann als Alarmmechanismus dienen.
• Ursache von Balance: Die Erkenntnis, dass $\gamma_{implicit} \gg \gamma_{explicit}$, stellt die Rolle des auxilliären Losses in Frage. Er ist nicht die primäre Quelle der Balance, sondern ein Katalysator, der den Optimierer anleitet, Balance intern zu lernen.
• Grenzen: Das Modell gilt nur für $K > 1$ (Top-K Routing). Bei Top-1-Routing (Switch-Transformer ohne Top-K) bricht die Dichte-Annahme zusammen, und das MFG-Modell ist nicht anwendbar.

Fazit

Dieses Paper liefert den ersten theoretisch fundierten und empirisch validierten Blick auf die Dynamik des Load Balancings in MoE-Modellen. Es identifiziert eine charakteristische „Inverted-U"-Kurve des Kongessionskoeffizienten, die zeigt, wie MoE-Modelle von einer Phase der strikten Lastverteilung zu einer Phase der qualitätsbasierten Selektivität übergehen. Dies bietet neue Ansatzpunkte für adaptive Trainingspläne und robustere Monitoring-Verfahren für große Sprachmodelle.