Technische Samenvatting: Sticky Routing voor Geheugenefficiënte MoE-inferentie

1. Probleemstelling: De Expert-Swapping Bottleneck

Mixture-of-Experts (MoE) modellen zijn ontworpen om Large Language Models (LLMs) te schalen door slechts een ijle subset van experts per token te activeren. Theoretisch maakt dit massale parameteraantallen mogelijk met een lage per-token berekening, wat ze aantrekkelijk maakt voor edge-deployment waar slechts een paar experts tegelijkertijd in het snelle geheugen (VRAM/SRAM) hoeven te verblijven.

De paper identificeert echter een kritieke foutmodus in de praktijk: temporele routing-inconsistentie. Standaard MoE-routers maken onafhankelijke beslissingen voor elk token op basis van enkel de hidden representatie van dat specifieke token. Gevolgen daarvan activeren opeenvolgende tokens frequent geheel andere experts. Op hardware met beperkt geheugen vereist elke wisseling het verwijderen van gecachte expert-gewichten en het laden van nieuwe gewichten uit traag opslagmedium (CPU RAM/Flash). Omdat de PCIe/NVMe-bandbreedte orden van grootte lager is dan de GPU-geheugenbandbreedte, domineren deze frequente "expert swaps" de inferentie-latentie, waardoor MoE-modellen vaak trager zijn dan vergelijkbare dense modellen.

Bestaande oplossingen worden als onvoldoende beschouwd:

• Systeem-niveau benaderingen (bijv. caching-heuristieken, prefetchers) zijn reactief; ze benutten de aanwezige lokaliteit, maar kunnen intrinsiek inconsistente routers niet corrigeren.
• Post-hoc benaderingen (bijv. alleen de router finetunen) falen omdat de expert-representaties gevormd zijn door een trainingsregime waarin lokale consistentie nooit werd gestimuleerd. De router kan beslissingen niet gemakkelijk heruitlijnen in een representatieruimte die daar niet voor ontworft is.
• Architecturale herontwerpen (bijv. Oracle-MoE) vereisen vaak significante structurele wijzigingen aan het model of de trainingspipeline.

2. Methodologie: StickyMoE

De auteurs stellen StickyMoE voor, een benadering tijdens de training die de temporele routing-lokaliteit behandelt als een expliciete doelstelling. De kerninnovatie is een differentieerbare routing-consistentieverlies ($L_{cons}$) dat abrupte veranderingen in expert-toewijzing tussen aangrenzende tokens bestraft.

De Routing Consistentie Loss

De verliesfunctie minimaliseert de $\ell_2$-afstand tussen de softmax gate-waarschijnlijkheidsvectoren ($g_t$) van opeenvolgende tokens:
$$L_{cons} = \frac{1}{T-1} \sum_{t=2}^{T} \| g_t - g_{t-1} \|_2^2$$
Waarbij $g_t \in \mathbb{R}^N$ de output van de router is voor token $t$.

Dit verlies wordt geïntegreerd in de totale trainingsdoelstelling naast standaard cross-entropy ($L_{CE}$) en load-balancing ($L_{bal}$) verliezen:
$$L = L_{CE} + \lambda L_{cons} + \mu L_{bal}$$

Kernmechanismen

• Co-adaptatie: In tegen tegenstelling tot post-hoc methoden traint StickyMoE zowel de router als de expert-representaties simultaan. Het minimaliseren van $L_{cons}$ werkt de routergewichten ($W_r$) bij en propageert via de hidden states ($h_t$) terug, wat het model stimuleert om representaties te leren die van nature consistent zijn over semantisch coherente segmenten.
• Soft-Hard Variant: Om "long-range drift" (waarbij kleine stap-voor-stap veranderingen accumuleren tot een verre switch) aan te pakken, introduceren de auteurs een anker-constraint. De sequentie wordt verdeeld in vensters; het eerste token van een venster stelt een anker vast, en daaropvolgende tokens worden bestraft voor afwijkingen van dit anker, waarbij de straf lineair toeneemt met de afstand tot het begin van het venster.
• Compatibiliteit: De methode vereist geen architecturale wijzigingen en voegt verwaarloosbare computationele overhead toe. Het is compatibel met elke standaard top-$k$ MoE en bestaande load-balancing technieken.

3. Belangrijkste Bijdragen

1. Routing Consistentie Loss: Een eenvoudige, differentieerbare trainingsdoelstelling die direct de expert-switching tussen aangrenzende tokens bestraft, waarbij slechts één hyperparameter ($\lambda$) nodig is.
2. Complementariteit met Load Balancing: Empirisch bewijs toont aan dat de consistentie-loss en de load-balancing loss complementair zijn. De consistentie-loss vormt wanneer experts worden gebruikt zonder te leiden tot "expert collapse" (de utilitatie-entropie blijft hoog, nabij het theoretische maximum).
3. Soft-Hard Variant: Een combinatie van per-stap zachte straffen en segment-niveau anker-constraints die sterke lokaliteit garandeert zonder extra kwaliteitskosten.
4. Uitgebreide Evaluatie: Gecontroleerde experimenten waarbij StickyMoE wordt vergeleken met vanilla MoE, post-hoc router fine-tuning (ReMoE) en architecturale baselines (Oracle-MoE) op kleine en middelgrote modellen.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op WikiText-2 met kleine (8.8M parameters) en middelgrote (22M parameters) MoE-modellen.

• Expert Switch Rate (SR): StickyMoE vermindert de expert switch rate met wel 59% (bijv. van 0.71 naar 0.29 op het middelgrote model).
• Cache Hit Rate (CHR): De reductie in switching vertaalt zich in een 3.92× reductie in cache misses (CHR steeg van ~0.54 naar ~0.88) op het middelgrote model met een 2-expert slot cache.
• Perplexity (Kwaliteit):
  • Bij gematigde $\lambda$-waarden (bijv. 0.05) verbetert StickyMoE de perplexity vergeleken met de baseline (bijv. -4.1% op het middelgrote model), wat suggereert dat de consistentie-loss fungeert als een gunstige regularisator.
  • Agressieve $\lambda$-waarden introduceren een trade-off, maar de Pareto-front van StickyMoE domineert andere methoden strikt.
• Falen van Post-Hoc Fine-Tuning: De ReMoE-baseline (het finetunen van enkel de router op een gepre-traind model) slaagde er niet in de switch rate betekenisvol te verlagen (<0.5% verandering), wat bevestigt dat lokaliteit niet achteraf kan worden toegevoegd zodra de representaties vaststaan.
• Layer-wise Analyse: De consistentie-loss is het meest effectief in diepere lagen (L1–L3). De input-laag (L0) blijft resistent tegen verbeteringen in lokaliteit, waarschijnlijk omdat token-embeddings een gebrek aan contextuele mixing vertonen.

5. Betekenis en Claims

De paper betoogt dat routing temporele lokaliteit een legitieme trainingsdoelstelling is die tijdens de pretraining moet worden geïnstilleerd in plaats van te worden behandeld als een systeem-niveau achteraf.

• Paradigmaverschuiving: Net zoals quantization-aware training de standaard is geworden, stellen de auteurs dat "locality-aware training" de standaard zou moeten zijn voor MoE-modellen bedoeld voor deployment op hardware met beperkt geheugen.
• Oplossing van de Grondoorzaak: StickyMoE pakt de grondoorzaak aan (inconsistente routing-beslissingen) in plaats van de symptomen (cache misses). Door de representaties en routing-beslissingen gezamenlijk te trainen, voorkomt het de mismatch in representatie die post-hoc methoden teistert.
• Praktische Bruikbaarheid: De methode is architectuur-agnostisch en vereist geen wijzigingen in de inferentie-pipeline. De getrainde router produceert van nature consistente gate-distributies, waardoor standaard caching-strategieën (LRU, LFU) aanzienlijk efficiënter kunnen functioneren zonder extra engineering.

De auteurs concluderen dat hoewel systeem-niveau optimalisaties en post-hoc aanpassingen nuttige aanvullingen blijven, ze niet kunnen dienen als vervanging voor een model dat nooit getraind is om cache-vriendelijk te zijn. StickyMoE biedt een pad naar het realiseren van de theoretische geheugenefficiëntie van MoE-architecturen op edge-apparaten.