Technische Samenvatting: kNN-MoE

Probleemstelling

Mixture-of-Experts (MoE)-architecturen schalen Large Language Models (LLM's) efficiënt door het gebruik van een parametrische "router" om tokens te distribueren naar een schaarse subset van experts. Echter, standaard routers zijn lichtgewicht classifiers die één keer worden getraind en vervolgens worden bevroren. Hoewel ze effectief zijn op data uit de verdeling, missen ze flexibiliteit tijdens inferentie wanneer ze geconfronteerd worden met verschuivingen in de verdeling (bijvoorbeeld gespecialiseerde domeinen of taken buiten de verdeling). In dergelijke scenario's moet de router extrapoleren vanuit vaste a priori-kennis, wat vaak leidt tot suboptimale expert-selectie en hoge perplexiteit. Bestaande methoden voor aanpassing tijdens de testtijd, zoals C3PO, proberen routing-logits online te fine-tunen, maar veroorzaken aanzienlijke rekenkundige overhead (ongeveer 5x de baseline). Andere benaderingen vereisen permanente parameterupdates via supervised fine-tuning (SFT), wat rekenkundig intensief is en minder efficiënt voor inferentie.

Methodologie: kNN-MoE

De auteurs stellen kNN-MoE voor, een retrieval-versterkt routingkader dat beslissingen over expert-toewijzing verfijnt zonder de parameters van het model te wijzigen. De methode werkt in twee distincte fasen:

1. Geheugenconstructie (Offline)

Het systeem bouwt een niet-parametrisch geheugen $\mathcal{M}$ voor elke MoE-laag, bestaande uit sleutel-waardeparen $(k_t, v_t)$:

• Sleutel ($k_t$): De routerinput (verborgen toestand) $x_t$ uit een referentiedataset $D_{ref}$.
• Waarde ($v_t$): Een "lokaal optimale" expert-toewijzing afgeleid door het optimaliseren van de routing-logits om de waarschijnlijkheid van de ground-truth token $y_t$ te maximaliseren.
  • In plaats van de output van de bevroren router op te slaan, lost de methode een optimalisatieprobleem op voor elke token in $D_{ref}$: het vinden van een logit-vector $r^*_t$ die de negatieve log-likelihood van de doeltoken minimaliseert, gegeven de bevroren netwerkparameters.
  • Deze optimalisatie wordt uitgevoerd via een paar stappen van gradient descent (typisch $S=1$), beginnend bij de oorspronkelijke parametrische logits.
  • De resulterende geoptimaliseerde logits $r^{(S)}_t$ dienen als de waarde $v_t$ in het geheugen.

2. Vertrouwensbewuste Adaptieve Menging (Inferentie)

Tijdens inferentie, voor een gegeven routerinput $x$:

• Retrieval: Het systeem haalt de $K$ dichtstbijzijnde buren $\mathcal{N}(x)$ op uit het geheugen $\mathcal{M}$ op basis van de gelijkenis tussen $x$ en de opgeslagen sleutels.
• Niet-parametrisch Voorstel: Een geheugen-gebaseerde logit-propositie $r_{mem}(x)$ wordt berekend door de opgehaalde optimale logits te aggregeren, gewogen naar hun gelijkenis met de input.
• Vertrouwens-gedreven Menging: De uiteindelijke routing-logits zijn een lineaire interpolatie tussen de oorspronkelijke parametrische routerlogits $r(x)$ en de geheugen-gebaseerde propositie $r_{mem}(x)$:
  $$r_{final}(x) = (1 - \lambda(x))r(x) + \lambda(x)r_{mem}(x)$$
  De mengcoëfficiënt $\lambda(x)$ wordt gedefinieerd als de gemiddelde gelijkenis van de opgehaalde buren. Dit fungeert als een vertrouwensscore:
  • Hoge Gelijkenis ($\lambda \approx 1$): Het model vertrouwt de opgehaalde, lokaal optimale toewijzing.
  • Lage Gelijkenis ($\lambda \approx 0$): Het model valt terug op de oorspronkelijke bevroren router, om ruis van irrelevante gevallen te vermijden.
• Output: De uiteindelijke logits worden door een Top-K Softmax gevoerd om de schaarse expert-toewijzing te genereren, waarbij de oorspronkelijke schaarsheidsbeperkingen van het model worden gehandhaafd.

Belangrijkste Bijdragen

1. Retrieval-Versterkte Routing: Introduceert een kader dat offline berekende, lokaal optimale expert-toewijzingen hergebruikt uit een geheugen van vergelijkbare eerdere gevallen om routingbeslissingen te verbeteren.
2. Niet-parametrische Aanpassing: In tegenstelling tot SFT of methoden voor online aanpassing, update kNN-MoE geen modelparameters. Het verplaatst de rekenkundige last naar een offline fase voor geheugenconstructie, waardoor efficiënte inferentie mogelijk wordt.
3. Vertrouwens-gedreven Terugval: De methode gebruikt de gemiddelde gelijkenis van opgehaalde buren als een dynamische mengcoëfficiënt, wat robuustheid garandeert door terug te vallen op de bevroren router wanneer geen relevante historische gevallen worden gevonden.
4. Efficiëntie: De benadering vereist slechts één stap gradient descent per token tijdens de geheugenconstructie en vermijdt de zware rekenkundige kosten van online optimalisatie of volledige model fine-tuning.

Experimentele Resultaten

De auteurs hebben kNN-MoE geëvalueerd op drie MoE-modellen (OLMoE, GPT-OSS en Qwen3) over zes benchmarks (GPQA, SuperGPQA, MMLU, USMLE, MedMCQA en MBPP).

• Prestaties versus Baselines: kNN-MoE presteerde consequent beter dan de Zero-shot-baseline en de "SFT (Router Only)"-baseline over alle modellen en datasets.
• Concurrentie met SFT: De methode bereikte prestaties die concurrerend waren met volledige Supervised Fine-Tuning (SFT) en, in verschillende data-schaarse instellingen (bijvoorbeeld GPQA, USMLE), presteerde beter dan SFT, dat soms last had van overfitting op kleine referentiestellen.
• Stabiliteit: In tegenstelling tot 5-shot prompting, wat de prestaties op bepaalde modellen (GPT-OSS, OLMoE) verslechterde door problemen met inputlengte, vertoonde kNN-MoE stabiele verbeteringen.
• Ablatiestudies:
  • Buren ($K$): In tegenstelling tot typische bevindingen bij kNN-LM, piekte de prestatie bij $K=1$. Het aggregeren van meer buren ($K>1$) introduceerde ruis, wat suggereert dat echt gunstige routing-signalen zeldzaam en specifiek zijn.
  • Optimalisatiestappen ($S$): Het verhogen van het aantal gradient descent-stappen boven $S=1$ leverde verwaarloosbare nauwkeurigheidswinst op terwijl de constructiekosten toenamen.
  • Grootte Referentiestel: De prestaties verbeterden monotoon met de grootte van het referentiestel $|D_{ref}|$, met afnemende meeropbrengst waargenomen boven 0,5k–1k voorbeelden.
• Efficiëntie: Geheugenconstructie duurde 0,46 uur voor 1k voorbeelden (tegenover ~1,34 uur voor SFT). De IVFPQ-variant (Inverted File + Product Quantization) verlaagde de inferentie-latentie-overhead tot ~3,67% en de groei in VRAM-gebruik tot ~9,43% ten opzichte van de baseline.

Betekenis en Claims

Het artikel beweert dat kNN-MoE de kloof overbrugt tussen statische parametrische routing en dynamische inferentiebehoeften. De primaire betekenis ligt in het bieden van een schaalbaar pad om MoE-adaptiviteit te verbeteren zonder de rekenkundige last van fine-tuning.

De auteurs benadrukken dat de methode bijzonder effectief is wanneer de parametrische router onzeker is (gecorreleerd met hoge perplexiteit), aangezien het retrieval-mechanisme een corrigerend signaal biedt vanuit vergelijkbare, hoogwaardige historische gevallen. Hoewel de methode afhankelijk is van de beschikbaarheid van een gelabeld referentiestel met een verdeling die lijkt op de testdata, biedt het een praktische, niet-parametrisch alternatief voor dure aanpassingstechnieken, waarbij SFT-concurrerende resultaten worden bereikt met aanzienlijk minder voorbereidingstijd en zonder parameterupdates.