Technische Samenvatting: UMoE (Unlocking Every Expert in Domain-Specific Training)

1. Probleemstelling

Mixture-of-Experts (MoE) modellen zijn een standaardarchitectuur geworden voor frontier Large Language Models (LLMs) vanwege hun vermogen om capaciteit te schalen zonder dat de inferentiekosten proportioneel toenemen. Er ontstaat echter een cruciale inefficiëntie tijdens domeinspecifieke post-training (bijv. Supervised Fine-Tuning of SFT).

Pre-trained MoE-modellen erven een expert-pool die gevormd is door mixed-domain pre-training. Onder een specifiek doel-domein is deze pool heterogeen: een aanzienlijk deel van de experts vertoont een lage "domain-aligned saliency", wat betekent dat ze weinig bijdragen aan de doeltaak ondanks dat ze een niet-verwaarloosbare hoeveelheid routing ontvangen. Standaard SFT werkt deze geërfde experts in plaats aan, maar heralloceert geen expert-slots. Gevolgen daarvan blijft een deel van de vaste modelcapaciteit slecht afgestemd op het doel-domein, wat de effectiviteit van het fine-tuning proces beperkt.

De kernvraag die wordt behandeld is: Kan de expert-pool worden georganiseerd in een meer domein-afgestemde initialisatie vóór SFT, zodanig dat meer expert-slots nuttig worden voor het leren van doel-domein capaciteiten, zonder de architectuur of het inferentiebudget van het model te wijzigen?

2. Methodologie: UMoE

De auteurs stellen UMoE (Unlocking MoE Capacity) voor, een budget-preserverende pipeline die de expert-pool reorganiseert vóór standaard SFT. De methode werkt in drie afzonderlijke stappen en behoudt het oorspronkelijke aantal experts ($E$), het aantal actieve experts ($K$), het aantal parameters en de inferentiekosten.

Stap 1: Domain-Aligned Expert Pruning

Het model identificeert en verwijdert eerst de minst relevante experts voor het doel-domein.

• Kalibratie: Een subset van de doel-domein data ($C \subset D$) wordt gebruikt om experts te scoren.
• Saliency Metric: De auteurs maken gebruik van REAP (Lasby et al., 2026) saliency, die de activatiestärke van een expert koppelt aan het gewicht van de router-gate. Voor een expert $e$ op laag $l$ is de saliency $S(l, e)$ het conditionele gemiddelde van de gate-gewogen output-norm over tokens waar de expert is geselecteerd.
• Actie: De onderste 50% van de experts (per laag) op basis van deze saliency-score wordt gepruned.

Stap 2: Budget-Preserving Expansion

Om de oorspronkelijke modelgrootte te herstellen en het vaste inferentiebudget te handhaven, worden de geprunede slots geregen uit de resterende "survivor" experts.

• Initialisatie: Elk vrijgekomen slot wordt geïnitialiseerd door de gewichten (MLP en router-rij) van een behouden "parent" expert te kopiëren via een round-robin procedure.
• Perturbatie: Om symmetrie te doorbreken en specialisatie te stimuleren, worden zowel de parent als de geregen expert onafhankelijk geperturbeerd. Voor een parametertensor $W$ worden de geperturbeerde gewichten $W'$ berekend als:
  $$W' = W + \sigma s(W) \epsilon$$
  waarbij $\epsilon \sim \mathcal{N}(0, I)$, $s(W)$ de empirische standaarddeviatie van de parametergroep is, en $\sigma = 0.05$.
• Ontwerpprincipe: Cruciaal is dat de perturbatie symmetrisch wordt toegepast op zowel de parent als de geregen expert. Ablatie-studies tonen aan dat een eenzijdige perturbatie (alleen op de geregen expert) leidt tot routing-asymmetrie en prestatieverlies. Zowel de router-rij als de MLP-gewichten worden geperturbeerd om complementaire routing en functionele asymmetrie te bieden.

Stap 3: Domain SFT

Het gereorganiseerde model, nu geïnitialiseerd met een domein-afgestemde expert-pool, ondergaat standaard full-parameter Supervised Fine-Tuning op de doel-domein data. Er is geen per-domein hyperparameter tuning vereist; hetzelfde recept wordt toegepast op alle domeinen.

3. Belangrijkste Bijdragen

1. Expert Pool Reorganisatie: Een nieuwe methode om de expert-compositie van een pre-trained MoE af te stemmen op een specifiek doel-domein vóór fine-tuning, waarbij "redundante" capaciteit wordt omgezet in "domein" capaciteit.
2. Budget Preservatie: De methode behoudt strikt de oorspronkelijke architectuur, het aantal parameters en de inferentiekosten, waardoor het een drop-in vervanging is voor standaard SFT-pipelines.
3. Symmetrische Perturbatie Strategie: De identificatie dat symmetrische perturbatie van zowel de parent als de geregen experts cruciaal is voor het behoud van gebalanceerde routing-dynamiek en het voorkomen van de ineenstorting van de bruikbaarheid van de geregen helft.
4. Empirische Validatie: Uitgebreide evaluatie over twee MoE-architecturen (Qwen3-30B-A3B en Qwen3.5-35B-A3B), vijf diverse domeinen (wiskunde, code, wetenschap, tool-use, agentic coding) en 12 benchmarks.

4. Experimentele Resultaten

UMoE presteert consistent beter dan Direct SFT (het fine-tunen van het pre-trained model zonder reorganisatie) in alle geëvalueerde settings.

• Wiskunde Domein:
  • Op Qwen3-30B-A3B (OpenMathReasoning data), verbeterde UMoE de gemiddelde nauwkeurigheid over zes competitie-benchmarks met +3.40 punten, met een piekverbetering van +6.56 op AIME 2024.
  • Op een sterke in-house wiskunde corpus waar de Direct SFT baseline al de Qwen3-30B-A3B-Thinking model had overtroffen (82.81 vs. 81.06), verhoogde UMoE het gemiddelde verder naar 84.17 (+1.36 punten), wat de robuustheid aantoont zelfs tegen sterke baselines.
  • Op Qwen3.5-35B-A3B, verbeterde UMoE het gemiddelde met +1.67 punten.
• Cross-Domein Prestaties:
  • Code: +2.73 / +0.85 punten op LiveCodeBench v5/v6.
  • Wetenschap: +3.10 punten op GPQA-Diamond.
  • Tool-use: +1.20 / +1.40 punten op $\tau$-bench / $\tau^2$-bench.
  • Agentic Coding: Een significante verbetering van +6.0 punten op SWE-bench Verified (pass@1 steeg van 16.2% naar 22.2%).
• Data Schaling: Experimenten waarbij de SFT data grootte werd gevarieerd (van 0.5B tot 4.1B tokens) lieten zien dat de prestatiewinsten van UMoE standhouden en positief blijven over alle dataschalen.

5. Mechanisme Analyse

Het paper biedt drie belangrijke inzichten in waarom UMoE werkt:

1. Redundantie in Direct SFT: Analyse van het Direct SFT model laat zien dat het een aanzienlijk deel van de gerouteerde expert-compute (ongeveer 42%) toewijst aan een subset van experts met een lage saliency. Post-hoc pruning van deze subset in Direct SFT modellen resulteert in een verwaarloosbare prestatiedaling (−0.12 punten), wat bevestigt dat deze capaciteit functioneel redundant is.
2. Consequentiële Capaciteit in UMoE: In contrast hiermee veroorzaakt het prunen van het getrainde UMoE model naar 64 experts een significante prestatiedaling (−5.00 punten). Dit geeft aan dat UMoE er succesvol in is geslaagd de voorheen redundante helft van de expert-pool om te zetten in capaciteit die consequent is voor downstream prestaties.
3. Optimalisatie Dynamiek: UMoE begint met een iets hogere training loss dan Direct SFT vanwege de reorganisatie, maar convergeert naar een lagere uiteindelijke training loss (0.396 vs. 0.404), wat duidt op een efficiënter optimalisatiepad.
4. Foutcorrectie: Case studies (bijv. AIME 2026 Probleem 25) laten zien dat UMoE specifieke redeneerfouten corrigeert (bijv. breukreductie in de Angle Bisector Theorem) die Direct SFT voortplant, wat leidt tot correcte eindantwoorden.

6. Betekenis en Claims

Het paper claimt dat UMoE een eenvoudige, budget-preserverende en universeel toepasbare aanpak biedt voor domeinadaptatie van MoE-modellen. De betekenis ligt in:

• Efficiëntie: Het ontsluit latente capaciteit binnen bestaande modellen zonder extra parameters of inferentiekosten te vereisen.
• Generaliseerbaarheid: De methode werkt over verschillende model-architecturen en diverse domeinen (redeneren, coderen, agentic taken) met behulp van een enkel bevroren recept.
• Robuustheid: Het biedt winst, zelfs wanneer de baseline SFT al sterk is en wanneer de trainingsdata opschaalt.

De auteurs concluderen dat het reorganiseren van geërfde expert-capaciteit modellen in staat stelt om hun vaste compute-budget beter af te stemmen op de specifieke vereisten van een doel-domein, waardoor "dode" experts worden getransformeerd in actieve bijdragers aan domeinspecifieke prestaties.