Technisches Resümee: DOT-MoE: Differenzierbarer Optimaler Transport für die MoEfikation

1. Problemstellung

Die schnelle Skalierung von Large Language Models (LLMs) hat einen erheblichen Engpass in der Inferenz-Effizienz geschaffen. Während Mixture-of-Experts (MoE)-Architekturen die Modellgröße von den Inferenzkosten entkoppeln, indem sie pro Token nur eine Teilmenge von Parametern aktivieren, ist das Training von MoEs von Grund auf bekanntlich instabil, datenhungrig und rechenintensiv.

Folglich hat sich die MoEfikation – die Umwandlung vortrainierter dichter Modelle in sparse MoEs – als eine praktikable Alternative etabliert. Bestehende Methoden zur Umwandlung dichter Feed-Forward-Networks (FFNs) in sparse Experten verlassen sich jedoch stark auf heuristische Strategien (z. B. zufälliges Splitting, gewichtsbasierte Clusterung oder aktivierungsbasierte Clusterung). Diese Ansätze leiden unter drei primären Einschränkungen:

1. Trennung der Zuständigkeiten (Separation of Concerns): Sie behandeln die Neuronenzuweisung (Partitionierung des FFN) und das Router-Training typischerweise als getrennte, sequentielle Stadien, was verhindert, dass beide ko-adaptieren.
2. Proxy-Optimierung: Sie optimieren basierend auf intermediären Statistiken (Gewichte oder Aktivierungen) statt auf dem tatsächlichen Beitrags der Neuronen zum Output, was oft zu einer suboptimalen Rekonstruktion des Verhaltens des dichten Modells führt.
3. Mangelnde Differenzierbarkeit: Die diskrete Natur der Neuronenzuweisung verhindert einen End-to-End-Gradientenfluss, was es schwierig macht, die strukturelle Dekomposition und die Routing-Policy gemeinsam zu optimieren.

2. Methodik: DOT-MoE

Die Autoren schlagen DOT-MoE vor, ein Framework, das die Dekomposition dichter Schichten in Experten als ein differenzierbares Optimales Transportproblem (DOT) formuliert. Dieser Ansatz ermöglicht das End-to-End-Lernen sowohl der Experten-Dekomposition als auch des Routing-Mechanismus.

2.1. Formulierung des Optimalen Transports

Die zentrale Erkennt-nis besteht darin, die Neuronenzuweisung als ein balanciertes Transportproblem zu betrachten, bei dem "Masse" (Neuronen) von einer Quellverteilung (die Neuronen des dichten FFN) zu einer Zielverteilung (den Experten) unter strengen Kapazitätsbeschränkungen transportiert werden muss.

• Quelle: $dffn$ Neuronen, jedes mit der Einheitmasse 1.
• Ziel: $E$ Experten, von denen jeder genau $s$ Neuronen benötigt ($s \cdot E = dffn$).
• Kostenfunktion: Anstatt einer festen Kostenmatrix verwendet die Methode eine lernbare Affinitätsmatrix $A$, wobei $A_{i,e}$ die Affinität darstellt, Neuron $i$ dem Experten $e$ zuzuweisen. Das Ziel ist es, die Gesamtaffinität zu maximieren und gleichzeitig die Randbedingungen zu erfüllen.

2.2. Differenzierbare Relaxation via Sinkhorn-Knopp

Um das diskrete Zuweisungsproblem differenzierbar zu machen, nutzt DOT-MoE eine entropische Regularisierung.

• Das Problem des optimalen Transports wird mittels des Sinkhorn-Knopp-Algorithmus relaxiert, welcher iterativ die Zeilen und Spalten der Transportplan-Matrix normalisiert.
• Dies liefert eine Soft-Assignment-Matrix $M_{soft} \in [0, 1]^{dffn \times E}$, welche die Kapazitätsbeschränkungen in Erwartung erfüllt.
• Wenn der Temperaturparameter $\tau \to 0$ geht, nähert sich die Lösung dem diskreten Optimum an.

2.3. Gemeinsame Optimierung mit Straight-Through Estimators (STE)

Eine kritische Herausforderung ist, dass der Einsatz diskrete Zuweisungen erfordert, und das Routing eine diskrete Top-k-Selektion beinhaltet, was beides nicht differenzierbar ist. DOT-MoE adresset dies mittels Straight-Through Estimators (STE):

• Forward Pass: Verwendet harte Entscheidungen (Greedy Rounding von $M_{soft}$ für die Zuweisung; Top-k-Selektion für das Routing).
• Backward Pass: Leitet Gradienten durch die Soft-Counterparts ($M_{soft}$ und die Router-Logits $P$) weiter.
• Dies erlaubt es, dass die Router-Gewichte ($W_r$) und die Affinitätsmatrix der Zuweisung ($A$) während des Trainings ko-adaptieren.

2.4. Trainingsziel

Das Framework wird über eine Alignment-Phase trainiert, in der die Gewichte des dichten Modells eingefroren sind. Das Modell minimiert eine Kombination aus Verlustfunktionen:

1. Output-Rekonstruktion: Gemessen via KL-Divergenz zwischen dem dichten Teacher- und dem sparse MoE-Student-Output.
2. Sprachmodellierung: Cross-Entropy-Loss bezüglich der Sprachmodellierungs-Zielsetzung.
3. Stabilitäts-Auxiliaren: Router $z$-loss (um Logit-Explosion zu verhindern) und Load-Balancing-Loss (um Experten-Kollaps zu verhindern).

3. Kernbeiträge

1. Optimaler Transport-Framework: Die Arbeit führt eine neuartige Formulierung der Dense-zu-MoE-Konvertierung als balanciertes Transportproblem ein, das mittels differenzierbarer Sinkhorn-Iterationen gelöst wird, um strikte Experten-Kapazitätsbeschränkungen durchzusetzen.
2. Dual-Level Ko-Adaptation: Im Gegensatz zu bisherigen Methoden, die Zuweisungen vor dem Training der Router festlegen, optimiert DOT-MoE die Neuron-zu-Experten-Dekomposition und das Token-zu-Experten-Routing gemeinsam unter Verwendung komplementärer STEs.
3. Output-bewusste Konstruktion: Durch die direkte Modellierung des Beitrags eines Neurons zum finalen Output (anstatt Proxies wie Gewichte oder intermediäre Aktivierungen zu optimieren) erreicht die Methode eine höhere Rekonstruktions-Fidelität.

4. Experimentelle Ergebnisse

Die Autoren evaluierten DOT-MoE über drei Modellfamilien (LLaMA-2, LLaMA-3, Qwen2.5) und sechs Benchmarks (ARC-Challenge, Winogrande, HellaSwag, PIQA, SciQ, BoolQ).

• Performance vs. Pruning: Bei LLaMA-2 7B mit einem Parameterbudget von 50 % erreichte DOT-MoE eine Perplexität von 7,99 auf WikiText-2 und übertraf damit die State-of-the-Art strukturierte Pruning-Methode DISP-LLM (9,84) signifikant.
• Performance vs. bestehende MoEfikation: In Zero-Shot-Settings übertraf DOT-MoE Methoden wie CMoE und LLaMA-MoE-v2 deutlich. Beispielsweise erreichte DOT-MoE auf LLaMA-2 7B eine durchschnittliche Genauigkeit von 61,5 % im Vergleich zu 44,5 % bei CMoE.
• Effizienz: DOT-MoE behält etwa 90 % der ursprünglichen Leistung des dichten Modells bei, während es die aktiven Parameter um 50 % reduziert.
• Fine-Tuning: Mit nur 1,2 Mrd. Tokens kontinuierlichem Fine-Tuning konnten DOT-MoE-Modelle eine Leistung nahe an den dichten Gegenstücken wiederherstellen (z. B. 66,6 % vs. 76,6 % auf LLaMA-2 7B) und übertrafen Baselines, die ähnliche oder größere Datenbudgets benötigten.
• Ablationsstudien:
  • Granularität: Die Performance verbessert sich mit zunehmender Experten-Granularität bis zu einem Sättigungspunkt, wonach der Gewinn abnimmt.
  • Inferenzgeschwindigkeit: Trotz höherer Expertenanzahl bleibt der Inferenz-Durchsatz stabil, wenn fusionierte MoE-Kernel (z. B. vLLM) verwendet werden, da die gesamte fusionierte intermediäre Dimension konstant bleibt.
  • Sparsity: Modelle, die bei höheren Sparsity-Leveln trainiert wurden, generalisieren robust auf noch extremere Sparsity-Regime zur Inferenzzeit.

5. Bedeutung und Ansprüche

Das Paper behauptet, dass DOT-MoE einen bedeutenden Fortschritt für die Effizienz der LLM-Inferenz darstellt, indem es eine prinzipielle, differenzierbare Methode bereitstellt, um dichte Modelle in sparse MoEs zu konvertieren.

• Paradigmenwechsel: Es bewegt sich weg von heuristischen, statischen Partitionierungen hin zu einer dynamischen, lernbaren Dekomposition, welche die funktionale Kohärenz der Experten respektiert.
• Praktikabilität: Die Methode ermöglicht eine effiziente Post-hoc-Sparsifizierung vortrainierter Modelle ohne die Notwendigkeit eines vollständigen Trainings von Grund auf, wobei eine hohe Kapazität beibehalten wird, während die aktiven Parameter drastisch reduziert werden.
• Robustheit: Das Framework zeigt, dass die Bewahrung der FFN-Fidelität die explizite Modellierung des Beitrags eines Neurons zum Output erfordert – ein Prinzip, das bestehende Proxy-basierte Methoden nicht erfassen können.

Die Autoren kommen zu dem Schluss, dass DOT-MoE eine überlegene Initialisierung für sparse Modelle bietet und die Performance-Lücke zwischen dichten und sparsen Architekturen mit minimalem zusätzlichem Trainingsdatenaufwand schließt.