Technische Samenvatting: DOT-MoE: Differentiabele Optimale Transport voor MoEfication

1. Probleemstelling

De snelle schaling van Large Language Models (LLMs) heeft een aanzienlijke flessenhals gecreëerd in de efficiëntie van inferentie. Hoewel Mixture-of-Experts (MoE) architecturen de modelgrootte ontkoppelen van de inferentiekosten door slechts een subset van parameters per token te activeren, is het vanaf nul trainen van MoEs berucht instabiel, data-hongerig en computationeel duur.

Consequent is MoEfication—de conversie van gepre-trainde dichte (dense) modellen naar sparse MoEs—een levensvatbaar alternatief geworden. Bestaande methoden voor het converteren van dichte Feed-Forward Networks (FFNs) naar sparse experts vertrouwen echter sterk op heuristische strategieën (bijv. willekeurige splitsing, gewichtsgebaseerde clustering of activatiegebaseerde clustering). Deze benaderingen lijden onder drie primaire beperkingen:

1. Scheiding van Belangen: Ze behandelen neuron-toewijzing (het partitioneren van de FFN) en router-training doorgaans als afzonderlijke, sequentiële fasen, waardoor deze niet gezamenlijk kunnen co-adapteren.
2. Proxy-optimalisatie: Ze optimaliseren op basis van tussenliggende statistieken (gewichten of activaties) in plaats van de werkelijke output-bijdrage van neuronen, wat vaak leidt tot een suboptimale reconstructie van het gedrag van het dichte model.
3. Gebrek aan Differentiabiliteit: De discrete aard van neuron-toewijzing verhindert end-to-end gradiëntstroom, wat het moeilijk maakt om de structurele decompositie en het routingbeleid gezamenlijk te optimaliseren.

2. Methodologie: DOT-MoE

De auteurs stellen DOT-MoE voor, een framework dat de decompositie van dichte lagen in experts formuleert als een Differentiabel Optimaal Transport (DOT) probleem. Deze benadering maakt end-to-end leren van zowel de expert-decompositie als het routingmechanisme mogelijk.

2.1. Formule voor Optimaal Transport

Het kerninzicht is om neuron-toewijzing te zien als een gebalanceerd transportprobleem waarbij "massa" (neuronen) moet worden getransporteerd van een bronverdeling (de dichte FFN-neuronen) naar een doelverdeling (de experts) onder strikte capaciteitsbeperkingen.

• Bron: $dffn$ neuronen, elk met een eenheidsmassa.
• Doel: $E$ experts, die elk precies $s$ neuronen vereisen ($s \cdot E = dffn$).
• Kostenfunctie: In plaats van een vaste kostenmatrix gebruikt de methode een leerbare affiniteitsmatrix $A$, waarbij $A_{i,e}$ de affiniteit vertegenwoordigt van het toewijzen van neuron $i$ aan expert $e$. Het doel is om de totale affiniteit te maximaliseren terwijl de marginale beperkingen worden nageleefd.

2.2. Differentiabele Relaxatie via Sinkhorn-Knopp

Om het discrete toewijzingsprobleem differentieerbaar te maken, gebruikt DOT-MoE Entropische Regularisatie.

• Het optimale transportprobleem wordt gerelaxeerd met behulp van het Sinkhorn-Knopp algoritme, dat iteratief de rijen en kolommen van de transportplanmatrix normaliseert.
• Dit levert een zachte toewijzingsmatrix $M_{soft} \in [0, 1]^{dffn \times E}$ op die de capaciteitsbeperkingen in verwachting naleeft.
• Naarmate de temperatuurparameter $\tau \to 0$ gaat, nadert de oplossing het discrete optimum.

2.3. Gezamenlijke Optimalisatie met Straight-Through Estimators (STE)

Een cruciale uitdaging is dat implementatie discrete toewijzingen vereist, en routing discrete top-k selectie omvat, wat beide niet-differentieerbaar zijn. DOT-MoE pakt dit aan met behulp van Straight-Through Estimators (STE):

• Forward Pass: Gebruikt harde beslissingen (greedy rounding van $M_{soft}$ voor toewijzing; top-k selectie voor routing).
• Backward Pass: Routeert gradiënten door de zachte tegenhangers ($M_{soft}$ en de router-logits $P$).
• Dit stelt de router-gewichten ($W_r$) en de toewijzingsaffiniteitsmatrix ($A$) in staat om tijdens de training te co-adapteren.

2.4. Trainingsdoelstelling

Het framework wordt getraind via een Alignment Phase waarbij de gewichten van het dichte model bevroren zijn. Het model minimaliseert een combinatie van verliezen:

1. Output Reconstructie: Gemeten via de KL-divergentie tussen de dichte teacher en de sparse MoE student outputs.
2. Taalmodellering: Cross-entropy verlies op de taalmodelleringsoefening.
3. Stabiliteit Auxiliairen: Router $z$-loss (om logit-explosie te voorkomen) en load-balancing loss (om expert-instorting te voorkomen).

3. Belangrijkste Bijdragen

1. Optimaal Transport Framework: Het paper introduceert een nieuwe formulering van dichte-naar-MoE conversie als een gebalanceerd transportprobleem, opgelost via differentiabele Sinkhorn-iteraties om strikte expert-capaciteitsbeperkingen af te dwingen.
2. Dual-Level Co-Adaptatie: In tegen tegenstelling tot eerdere methoden die de router eerst vastzetten voordat de training begint, optimaliseert DOT-MoE de neuron-naar-expert decompositie en de token-naar-expert routing gezamenlijk met complementaire STEs.
3. Output-bewuste Constructie: Door direct de bijdrage van het neuron aan de uiteindelijke output te modelleren (in plaats van te optimaliseren op basis van proxies zoals gewichten of tussenliggende activaties), bereikt de methode een hogere reconstructiefaillibiliteit.

4. Experimentele Resultaten

De auteurs hebben DOT-MoE geëvalueerd over drie modelfamilies (LLaMA-2, LLaMA-3, Qwen2.5) en zes benchmarks (ARC-Challenge, Winogrande, HellaSwag, PIQA, SciQ, BoolQ).

• Prestaties vs. Pruning: Op LLaMA-2 7B met een 50% parameterbudget behaalde DOT-MoE een perplexiteit van 7.99 op WikiText-2, wat aanzienlijk beter is dan de state-of-the-art gestructureerde pruning methode DISP-LLM (9.84).
• Prestaties vs. Bestaande MoEfication: In zero-shot settings presteerde DOT-MoE aanzienlijk beter dan methoden zoals CMoE en LLaMA-MoE-v2. Bijvoorbeeld, op LLaMA-2 7B behaalde DOT-MoE een gemiddelde nauwkeurigheid van 61.5% vergeleken met 44.5% van CMoE.
• Efficiëntie: DOT-MoE behoudt ongeveer 90% van de oorspronkelijke dichte modelprestaties terwijl de actieve parameters met 50% worden verminderd.
• Fine-Tuning: Met slechts 1.2B tokens aan continue fine-tuning herstelden DOT-MoE modellen prestaties die dicht bij de dichte tegenhangers lagen (bijv. 66.6% vs 76.6% op LLaMA-2 7B), waarmee ze baselines versloegen die vergelijkbare of grotere data-budgetten vereisten.
• Ablatie Studies:
  • Granulariteit: Prestaties verbeteren met toenemende expert-granulariteit tot een verzadigingspunt, waarna de winsten afnemen.
  • Inference Snelheid: Ondanks een hoger aantal experts blijft de inferentie-throughput stabiel bij gebruik van gefuseerde MoE kernels (bijv. vLLM), aangezien de totale gefuseerde tussenliggende dimensie constant blijft.
  • Sparsity: Modellen getraind bij hogere sparsity-niveaus (75%) generaliseren robuust naar nog extremere sparsity-regimes tijdens inferentie.

5. Betekenis en Claims

Het paper claimt dat DOT-MoE een significante vooruitgang vormt in de efficiëntie van LLM-inferentie door een principieel, differentieerbaar methode te bieden voor het converteren van dichte modellen naar sparse MoEs.

• Paradigmaverschuiving: Het beweegt weg van heuristische, statische partitionering naar een dynamische, leerbare decompositie die de functionele samenhang van experts respecteert.
• Praktische Bruikbaarheid: De methode maakt efficiënte post-hoc sparsificatie van pre-trainde modellen mogelijk zonder de noodzaak voor volledige hertraining vanaf nul, waarbij een hoge capaciteit behouden blijft terwijl de actieve parameters drastisch worden verminderd.
• Robuustheid: Het framework toont aan dat het behoud van FFN-getrouwheid vereist dat de bijdrage van het neuron aan de output expliciet wordt gemodelleerd, een principe dat bestaande proxy-gebaseerde methoden niet vatten.

De auteurs concluderen dat DOT-MoE een superieure initialisatie biedt voor sparse modellen, waardoor de prestatiekloof tussen dichte en sparse architecturen wordt overbrugd met minimale aanvullende trainingsdata.