Temporally Extended Mixture-of-Experts Models

Titel: Temporally Extended Mixture-of-Experts Models

Auteurs: Zeyu Shen en Peter Henderson (Princeton University)

1. Het Probleem: De "Churn" in Bestaande MoE-modellen

Moderne Large Language Models (LLMs) gebruiken vaak Mixture-of-Experts (MoE) lagen om de capaciteit te vergroten zonder de inferentiële rekentijd per token te verhogen. In deze architectuur wordt voor elk gegenereerde token een subgroep van experts (bijv. top-4 van 32) geselecteerd.

• De Uitdaging: Bestaande MoE-modellen wisselen de actieve set experts bij bijna elk token (een "switch rate" van >90-100%).
• Geheugenbeperkingen: Wanneer het totale aantal experts groter wordt dan het beschikbare GPU-geheugen, moeten experts worden uitgewisseld (offloaded) naar host-geheugen of schijf.
• Inefficiëntie: Door de constante wisselingen ("churn") worden optimalisaties zoals caching, prefetching en offloading ineffectief. Elke laadoperatie introduceert latentie, wat de doorvoer verlaagt en workflows onderbreekt. Bestaande oplossingen (zoals expert-pruning of heuristieken voor caching) lossen dit fundamentele probleem van gebrek aan temporele continuïteit niet op.

2. Methodologie: Het Options Framework en Option-Critic

De auteurs stellen dat dit probleem een directe analogie heeft met versterkende leer (Reinforcement Learning), specifiek het Options Framework.

• Concept: In plaats van bij elk token een nieuwe expert te kiezen, wordt een "optie" gedefinieerd als een set van toegestane experts (een binair masker). Deze optie blijft actief over meerdere tijdstappen.
• De Controller: Er wordt een lichtgewicht controller toegevoegd aan elke laag van het transformer-model. Deze controller leert:
  1. Wanneer de huidige set experts te behouden (niet switchen).
  2. Wanneer te switchen naar een nieuwe set (terminatie van de optie).
  3. Welke nieuwe set experts moet worden geladen.
• Formulering als s-MDP: Het probleem wordt gemodelleerd als een Semi-Markov Decision Process (s-MDP).
  • Acties: Het selecteren van een expert-masker (optie).
  • Kosten: Een expliciete "deliberation cost" ($\eta$) wordt toegevoegd aan de doelstelling. Dit straalt de kosten af van het switchen (latentie van het laden van weights).
• Architectuur:
  • De controller gebruikt een DeepSets-encoder om de huidige expert-set te vertalen naar een embedding.
  • Er zijn drie hoofdkoppen: een Termination Head (beslist of er geswitcht wordt), een Selection Head (kiest de nieuwe set) en Value Heads ($V_\Omega$ en $Q_\Omega$) voor het schatten van de waarde.
  • De training gebruikt de Option-Critic architectuur, waarbij zowel de intra-optie beleid (de LLM-parameters) als de terminatie- en selectiepolitiek worden geoptimaliseerd.
• Training: Het model wordt getraind via on-policy distillation. Een "student" (het MoE-model met controller) probeert de output van een "teacher" (het originele, bevroren MoE-model) na te bootsen, terwijl de controller de switch-rates minimaliseert via de deliberation cost.

3. Belangrijkste Bijdragen

1. Design Filosofie: Het introduceren van "Temporally Extended MoE" als een nieuwe benadering, waarbij expert-selectie wordt gezien als een tijdelijk uitgestrekte beslissing in plaats van een per-token gebeurtenis.
2. Formalisatie: Het formaliseren van dynamisch laden van experts als een s-MDP met deliberation costs, en het toepassen van de Option-Critic framework voor optimalisatie.
3. Efficiëntie zonder Pre-training: Het aantonen dat bestaande, voorgeïntegreerde modellen (zoals gpt-oss-20b) kunnen worden omgezet naar temporally extended MoE's met slechts lichte training (LoRA-adapters) en zonder volledige herpre-training.

4. Resultaten

De methode is getest op gpt-oss-20b (24 lagen, 32 experts per laag, top-4 routing).

• Switch Rate Reductie: De switch rate daalde drastisch van >50% (bijna elke token) naar <5% (en zelfs <1% bij hogere deliberation costs).
• Behoud van Prestaties: Ondanks de drastische reductie in switch-frequentie, behield het model tot 90% van de nauwkeurigheid van het basismodel op benchmarks zoals MATH, MMLU en MMMLU.
  • Voorbeeld (MATH, $\hat{k}=16$): Basismodel 71.5% vs. Auteurs' methode (met $\eta=0.02$) 64.0%.
• Trade-off: Er is een duidelijke afweging tussen de deliberation cost ($\eta$) en de prestaties. Een hogere kostenparameter leidt tot minder switches en lagere geheugeneisen, maar een iets lagere nauwkeurigheid.
• Vergelijking met Baselines: De methode overtreft aanzienlijk traditionele pruning-methoden (zoals frequentie-based, reconstructie-loss, Wanda en random selection), die vaak leiden tot catastrofale degradatie van de prestaties of repetitieve output.

5. Significantie en Toekomstperspectief

• Geheugenefficiëntie: Door de actieve expert-set voor langere periodes vast te houden, kan de hoeveelheid VRAM die nodig is voor inferentie drastisch worden verlaagd. Alleen de huidige set van $\hat{k}$ experts hoeft in het snelle geheugen te zitten, in plaats van alle $N$ experts.
• Continuïteit in Training: Het stelt "temporal chunking" voor tijdens training toe, waarbij onactieve experts kunnen worden uitgeladen tijdens de forward/backward pass van een chunk.
• Continuïtair Leren: Het biedt een natuurlijk pad voor het toevoegen van nieuwe experts aan een model zonder de per-token rekentijd te verhogen, wat essentieel is voor toekomstige modellen die groeien met de schijfruimte in plaats van GPU-geheugen.
• Principiële Weg: Het paper opent een weg voor het ontwerpen van MoE-architecturen die inherent "temporally extended" zijn, waardoor de kloof tussen enorme modelcapaciteit en beperkte hardware-resource wordt overbrugd.

Conclusie:
Dit werk toont aan dat het introduceren van temporele continuïteit in MoE-routing, gestuurd door een Reinforcement Learning-controller met expliciete switch-kosten, een krachtige en praktische oplossing biedt voor de geheugen- en latentieproblemen van toekomstige, zeer grote taalmodellen.