ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

Technische Samenvatting: ProbMoE: Differentiabele Probabilistische Routing voor Mixture-of-Experts

Probleemstelling
Mixture-of-Experts (MoE) modellen schalen efficiënt door slechts een kleine subset van experts per token te activeren. Het trainen van deze modellen blijft echter uitdagend omdat standaard routing vertrouwt op een discrete, niet-differentiabele top-$k$ operator. Conventionele trainingsprocedures omzeilen dit door de afhankelijkheid van de top-$k$ selectie van router-logits te negeren en gradiënten alleen te propageren via de softmax-waarschijnlijkheden van de geselecteerde experts. Deze benadering beperkt het leersignaal voor alternatieve expert-subsets, wat vaak leidt tot geconcentreerde routing-distributies, slechte expert-benutting en instabiele trainingsdynamiek. Bestaande pogingen om dit te verzachten, zoals dense backpropagation of heuristische gradiënt-schatters, falen vaak omdat ze de distributie over haalbare expert-subsets niet expliciet modelleren of slechts beperkte mechanismen bieden voor het verkennen van alternatieve combinaties tijdens de training.

Methodologie
De auteurs stellen ProbMoE voor, een framework dat expert-routing herformuleert als een probabilistisch inferentieprobleem over cardinaliteit-beperkte expert-subsets. In plaats van routing te behandelen als een deterministische operator, modelleert ProbMoE de selectie van experts als een distributie over discrete subsets.

1. Probabilistische Formulering:

   • Exact-$k$ Routing: Het framework definieert een distributie over subsets $S$ waar $|S|=k$. Het conditioneert onafhankelijke Bernoulli-selectiekansen (afgeleid van router-logits via sigmoid) op de beperking dat er exact $k$ experts worden geselecteerd. De normaliserende constante voor deze distributie, $Z_k$, wordt exact berekend met behulp van dynamisch programmeren, waardoor de formulering tractabel wordt.
   • Dynamic-$k$ Routing: Het framework generaliseert naar een range-beperkte setting waarbij de subsetgrootte $|S|$ binnen $[k_{min}, k_{max}]$ ligt. De distributie is geconditioneerd op deze range, waardoor het model zowel de identiteiten van de experts als de routing-cardinaliteit voor elke token gezamenlijk kan afleiden.

2. Forward Pass:

   • Tijdens de training samplet ProbMoE een expert-subset $S$ uit de gedefinieerde distributie (exact-$k$ of range-beperkt).
   • De MoE-laag berekent de output met behulp van alleen de gesamplede experts, waardoor het sparse executiepatroon van standaard MoE-modellen behouden blijft.
   • Tijdens de inferentie wordt de Maximum-A-Posteriori (MAP) set geselecteerd, zodat de kosten van de inferentie overeenkomen met conventionele MoE-routing.

3. Backward Pass:

   • Om gradiënt-gebaseerde optimalisatie mogelijk te maken, gebruikt ProbMoE de SIMPLE gradiënt-schatter (Ahmed et al., 2023).
   • In plaats van te differentiëren door de discrete sample, gebruikt de backward pass de marginale kansen van elke expert die geselecteerd wordt onder de beperkte distributie.
   • Routing-gewichten worden geconstrueerd met behulp van een straight-through estimator: $w = (\text{stopgrad}(z - m) + m) \odot \pi$, waarbij $z$ de gesamplede mask is, $m$ de vector van exacte marginale kansen is, en $\pi$ de softmax routing-kansen zijn.
   • Dit maakt het mogelijk om gradiënten door de marginale kansen te laten stromen, wat informatieve updates aan de router geeft die reflecteren hoe veranderingen in selectiekansen de verwachte loss beïnvloeden, terwijl de sparse expert-executie behouden blijft.

Belangrijkste Bijdragen

1. Probabilistische Inferentie Formulering: Het artikel presenteert de MoE-routing als probabilistische inferentie over cardinaliteit-beperkte expert-subsets, waarmee het verder gaat dan heuristische gradiënt-benaderingen.
2. Tractabele Marginal-Based Optimalisatie: Door gebruik te maken van de SIMPLE-estimator, koppelen de auteurs stochastische subset-sampling aan exacte marginal-gebaseerde router-gradiënten, wat differentiabele training mogelijk maakt zonder dense expert-zijde berekeningen.
3. Verenigd Dynamic-$k$ Framework: De methode breidt zich natuurlijk uit naar dynamische routing, waarbij het model zowel de expert-identiteiten als het aantal actieve experts per token binnen een vooraf gedefinieerde range afleidt, zonder te vertrouwen op heuristische gating-regels.
4. Empirische Verbeteringen:
   • Exact-$k$: ProbMoE bereikt sterke prestaties over benchmarks (GSM8K, Law, Translation, MBPP, MMLU) op OLMoE en Qwen backbones, en presteert vaak gelijkwaardig aan of beter dan competitieve baselines zoals DenseMixer en Conventionele Top-$k$.
   • Dynamic-$k$: ProbMoE Dynamic-$k$ bereikt vergelijkbare prestaties als Exact-$k$ terwijl het gemiddeld minder experts activeert, wat effectieve adaptieve allocatie demonstreert.
5. Routing Diversiteit: Analyse toont aan dat ProbMoE bredere routing-distributies en een meer gebalanceerde expert-benutting induceert vergeleken met baselines, wat expert-collapse mitigeert.

Resultaten en Betekenis
Het artikel demonstreert dat ProbMoE de expert-benutting en routing-diversiteit verbetert. In ablatie-studies wordt aangetoond dat de combinatie van stochastische sampling en marginal-gebaseerde optimalisatie de primaire drijfveer is voor prestatiewinsten, waarbij het beter presteert dan configuraties die slechts één van deze componenten gebruiken.

De auteurs beweren dat ProbMoE een principieel methode biedt voor expert-selectie die de trainingsstabiliteit en computationele efficiëntie verbetert. Door routing probabilistisch te modelleren, staat het framework de exploratie van alternatieve expert-combinaties tijdens de training toe, wat leidt tot beter gekalibreerde routing-distributies. De Dynamic-$k$ variant laat verder zien dat adaptieve cardinaliteit geleerd kan worden om aan te sluiten bij de token-niveau complexiteit, wat de effectieve routing-kosten vermindert zonder de taakprestaties op te offeren. Het werk suggereert dat toekomstige richtingen onder meer het uitbreiden van ProbMoE naar pre-training en systeem-niveau optimalisaties om de sparsity in dynamische settings volledig te exploiteren.