Temporally Extended Mixture-of-Experts Models

1. Problemstellung

Moderne Large Language Models (LLMs) nutzen zunehmend Mixture-of-Experts (MoE)-Architekturen, um die Kapazität des Modells zu skalieren, ohne die Inferenzgeschwindigkeit proportional zu erhöhen. Dabei wird für jedes Token nur eine spärliche Teilmenge von Experten (Experten-Subnetzwerken) aktiviert.

Das zentrale Problem, das in diesem Paper adressiert wird, ist die hohe Switch-Rate (Wechselfrequenz) der Experten:

• Aktueller Zustand: In bestehenden MoE-Modellen (z. B. gpt-oss, Qwen) ändert sich der aktive Satz von Experten fast bei jedem generierten Token (Switch-Rate > 50–100 %).
• Folgen: Sobald die Gesamtzahl der Experten den verfügbaren GPU-Speicher übersteigt, müssen Gewichte in den Host-Speicher (RAM) oder auf die Festplatte ausgelagert werden (Offloading).
• Ineffizienz: Die ständige Notwendigkeit, bei jedem Token neue Experten zu laden, macht Optimierungen wie Prefetching oder Caching wirkungslos. Die Latenz durch das Laden von Gewichten dominiert die Inferenzzeit, und die Speichernutzung bleibt hoch, da keine temporale Kontinuität ausgenutzt werden kann.

2. Methodik: Temporally Extended MoE

Die Autoren schlagen vor, das Problem der Experten-Auswahl als Problem der temporalen Erweiterung im Rahmen des Options-Frameworks aus dem Reinforcement Learning (RL) zu betrachten.

Kernkonzept: Options-Framework

Statt bei jedem Token einen neuen Experten zu wählen, lernt ein Controller, wann ein aktueller Satz von Experten (ein „Option") beibehalten werden soll und wann ein Wechsel notwendig ist.

• Option: Ein binärer Maskenvektor $\omega$, der festlegt, welche $\hat{k}$ Experten in einer Schicht für eine Sequenz von Tokens aktiv sind.
• Terminierung: Ein Switch tritt nur auf, wenn eine Termination-Funktion $\beta(s)$ entscheidet, dass der aktuelle Option nicht mehr optimal ist.
• Deliberation Cost ($\eta$): Eine explizite Kostenkomponente im Lernziel, die den Wechsel bestraft. Dies zwingt den Agenten, nur dann zu wechseln, wenn der erwartete Qualitätsgewinn die Latenzkosten des Ladens überwiegt.

Architektur und Training

• Controller: Ein leichter Controller wird pro MoE-Schicht eingefügt. Er beobachtet den versteckten Zustand des LLM ($h_t$) und die aktuelle Option ($\omega_{t-1}$).
  • Termination Head: Berechnet die Wahrscheinlichkeit $\beta_t$, die Option zu beenden.
  • Selection Head: Wählt eine neue Experten-Maske aus, falls ein Wechsel stattfindet (verwendet Plackett-Luce-Verteilung).
• Training (Option-Critic):
  • Das Modell wird mit dem Option-Critic-Framework trainiert, das Policy-Gradient-Theoreme auf temporale Aktionen anwendet.
  • Reward-Design: Da das Ziel ist, ein vortrainiertes Modell in ein temporales MoE zu verwandeln, ohne die Leistung zu verlieren, wird On-Policy Self-Distillation verwendet. Der Reward ist die negative Reverse-KL-Divergenz zwischen dem Output des trainierten Modells (Student) und dem des originalen, eingefrorenen MoE-Modells (Lehrer).
  • Intra-Option Policy: Die Gewichte des LLM (Experten und Router) werden gleichzeitig mit dem Controller feinabgestimmt (z. B. via LoRA), um sich an die eingeschränkte Experten-Auswahl anzupassen.

3. Wichtige Beiträge

1. Paradigmenwechsel: Die Einführung des Konzepts „Temporally Extended MoE", das die Experten-Auswahl als sequentielle Entscheidungsfindung mit Kosten für den Wechsel behandelt.
2. Formalisierung: Die Abbildung des dynamischen Ladens von Experten auf ein Semi-Markov Decision Process (s-MDP), wobei Experten-Masken als Optionen und die Ladeverzögerung als Deliberation Cost modelliert werden.
3. Praktische Umsetzung: Entwicklung eines leichten Controllers, der in bestehende MoE-Architekturen integriert werden kann, ohne das gesamte Modell neu vorzutrainieren.
4. Nachweis der Machbarkeit: Demonstration, dass vortrainierte Modelle (gpt-oss-20b) mit minimalem Aufwand (LoRA + Controller) in temporale MoEs umgewandelt werden können, wobei die Switch-Rate drastisch gesenkt wird, ohne signifikante Genauigkeitsverluste.

4. Ergebnisse

Die Methode wurde am Modell gpt-oss-20b (32 Experten pro Schicht, Top-4 Routing) getestet und mit verschiedenen Pruning-Baselines verglichen.

• Reduktion der Switch-Rate:
  • Baseline (Standard MoE): Switch-Rate von ca. 50–100 % (Wechsel bei fast jedem Token).
  • Mit dem gelernten Controller: Die Switch-Rate sank auf unter 5 % (bei $\hat{k}=16$ und $\eta=0.02$) und sogar auf unter 1 % bei höheren Deliberation Costs.
• Genauigkeit (Accuracy):
  • Trotz der drastischen Reduktion der aktiven Experten pro Token behielt das Modell bis zu 90 % der Genauigkeit des Baseline-Modells auf Benchmarks wie MATH, MMLU und MMMLU.
  • Beispiel MATH: Baseline 71,5 %, Methode (Ours, $\eta=0.02$) 64,0 %.
  • Im Vergleich dazu brachen statische Pruning-Methoden (wie Frequency-based oder Wanda) bei gleicher Reduktion der Expertenanzahl massiv ein (z. B. auf 15–50 % Genauigkeit).
• Speichereffizienz:
  • Durch die temporale Kontinuität müssen nur $\hat{k}$ Experten pro Schicht im VRAM gehalten werden, anstatt alle $N$ Experten. Bei $\hat{k}=16$ (von 32) und $\hat{k}=8$ wurden VRAM-Anforderungen um 37 % bzw. 55 % reduziert.
• Verhalten: Die Visualisierung der Experten-Aktivierung zeigt eine starke zeitliche Kontinuität (lange Blöcke gleicher Experten), im Gegensatz zum chaotischen Wechsel bei Standard-Modellen.

5. Bedeutung und Ausblick

Dieses Paper bietet einen prinzipiellen Weg, um MoE-Modelle speichereffizienter zu machen und die Skalierung von Experten über den verfügbaren GPU-Speicher hinaus zu ermöglichen.

• Inferenz: Ermöglicht das Serving von riesigen MoE-Modellen auf Hardware mit begrenztem VRAM, da Experten seltener geladen werden müssen.
• Training: Eröffnet die Möglichkeit für „temporales Chunking" beim Training, bei dem nur die für einen Chunk aktiven Experten im Speicher gehalten werden müssen.
• Continual Learning: Neue Experten können hinzugefügt werden, ohne die Inferenzkosten pro Token zu erhöhen, da der Controller lernt, wann diese neuen Optionen genutzt werden sollen.

Die Arbeit zeigt, dass die Integration von RL-Konzepten (Options-Framework) in die Architektur von LLMs nicht nur die Effizienz steigert, sondern auch neue Wege für das Management von Modellkapazität und -kosten eröffnet.