💬 NLP

Routing by Analogy: kNN-Augmented Expert Assignment for Mixture-of-Experts

Dieser Beitrag stellt kNN-MoE vor, ein durch Retrieval erweitertes Routing-Framework, das Mixture-of-Experts-Modelle dynamisch verbessert, indem es ein Gedächtnis optimierter vergangener Routing-Entscheidungen und vertrauensbasierte Ausweichmechanismen nutzt, um Verteilungsverschiebungen effektiver zu bewältigen als statische Router.

Ursprüngliche Autoren: Boxuan Lyu, Soichiro Murakami, Hidetaka Kamigaito, Peinan Zhang

Veröffentlicht 2026-05-26

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Boxuan Lyu, Soichiro Murakami, Hidetaka Kamigaito, Peinan Zhang

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich eine riesige, hochintelligente Bibliothek von Experten vor (ein „Mixture-of-Experts"-Modell). Wenn Sie dieser Bibliothek eine Frage stellen, fungiert ein Router wie ein Bibliothekar. Seine Aufgabe besteht darin, Ihre Frage zu betrachten und schnell zu entscheiden, welcher spezifische Experte in der Bibliothek am besten geeignet ist, sie zu beantworten.

Normalerweise wird dieser Bibliothekar einmal trainiert und dann eingefroren. Er verlässt sich auf sein Gedächtnis dessen, was er während des Trainings gelernt hat. Doch hier liegt das Problem: Wenn Sie eine seltsame, schwierige oder brandneue Art von Frage stellen, die der Bibliothekar noch nie gesehen hat, könnte er falsch raten und Sie zum falschen Experten schicken. Das ist so, als würde ein Bibliothekar eine medizinische Frage an einen Geschichtsexperten weiterleiten, weil die Wörter oberflächlich ähnlich aussehen.

Die Arbeit stellt kNN-MoE vor, ein kluges Upgrade, das diesem Bibliothekar eine „Spickzettel"-Hilfe basierend auf vergangenen Erfolgen gibt. So funktioniert es, aufgeteilt in einfache Konzepte:

1. Der „Spickzettel" (Speicherkonstruktion)

Bevor das System jemals eine echte Frage beantwortet, nehmen die Forscher einen Satz von Übungsfragen (einen Referenzdatensatz) und führen sie durch die Bibliothek.

Das Experiment: Für jedes einzelne Wort in diesen Übungsfragen fragen sie: „Wenn wir die Entscheidung des Bibliothekars jetzt magisch ändern könnten, welcher Experte hätte die absolut beste Antwort gegeben?"
Das Ergebnis: Sie finden den „perfekten" Experten für jeden spezifischen Moment und notieren ihn. Sie speichern diese Paare: „Dieser spezifische Frageeingang" + „Die perfekte Expertenwahl."
Die Analogie: Stellen Sie sich vor, der Bibliothekar lernt für eine Abschlussprüfung. Anstatt nur die Regeln auswendig zu lernen, erstellt er einen riesigen Karteikartenstapel. Auf der einen Seite steht eine knifflige Frage; auf der anderen der genaue Experte, der sie in der Vergangenheit perfekt gelöst hat.

2. Die „intelligente Nachschlageaktion" (Inferenz)

Wenn nun ein echter Benutzer eine Frage stellt, tut das System gleichzeitig zwei Dinge:

Der eingefrorene Bibliothekar: Der ursprüngliche, eingefrorene Router gibt seine beste Schätzung basierend auf seinem Training ab.
Der Spickzettel: Das System schlägt die Frage des Benutzers im Karteikartenstapel nach, um die ähnlichsten vergangenen Fragen zu finden.

3. Die „Vertrauensabstimmung" (Adaptives Mischen)

Dies ist der wichtigste Teil. Das System vertraut dem Spickzettel nicht blind. Es prüft, wie ähnlich die vergangenen Fälle der aktuellen Frage sind.

Hohe Zuversicht: Wenn die vergangenen Fälle fast identisch mit der aktuellen Frage sind, sagt das System: „Der Bibliothekar ist bei dieser kniffligen Frage wahrscheinlich unsicher, aber unser Spickzettel sagt, dass der perfekte Experte X ist." Es mischt dann die Schätzung des Bibliothekars mit dem Rat des Spickzettels und lehnt sich stark auf den Spickzettel.
Niedrige Zuversicht: Wenn die aktuelle Frage völlig einzigartig ist und nichts im Spickzettel gut passt, sagt das System: „Der Spickzettel ist hier nutzlos; wir fügen vielleicht nur Rauschen hinzu." Es ignoriert die Nachschlageaktion und vertraut der ursprünglichen, eingefrorenen Entscheidung des Bibliothekars.

Warum das wichtig ist

Die Arbeit behauptet, dass diese Methode ein „Sweet Spot" zwischen zwei Extremen ist:

Nichts tun (Zero-Shot): Einfach den eingefrorenen Bibliothekar verwenden. Das ist schnell, kann aber bei schwierigen oder neuen Fragen versagen.
Neu trainieren (Supervised Fine-Tuning): Dem Bibliothekar von Grund auf neue Regeln beibringen. Das funktioniert gut, ist aber unglaublich langsam, teuer und erfordert, den gesamten Prozess für jede neue Aufgabe erneut durchzuführen.

kNN-MoE erzielt den Leistungsboost des Neutrainierens ohne die hohen Kosten. Es ist so, als würde man dem Bibliothekar ein dynamisches, durchsuchbares Gedächtnis seiner besten vergangenen Züge geben, anstatt ihn zurück in die Schule zu schicken.

Wichtige Erkenntnisse aus der Arbeit

Es funktioniert bei schwierigen Fragen: Das System hilft am meisten, wenn der ursprüngliche Bibliothekar verwirrt ist (hohe „Perplexity"). Wenn der Bibliothekar bereits zuversichtlich ist, tritt das System zurück, um nichts zu verderben.
Weniger ist mehr: Überraschenderweise funktionierte das Betrachten von nur einem vergangenen Beispiel (der einzelne nächste Treffer) besser als das Mitteln vieler Beispiele. Die Arbeit legt nahe, dass für das Experten-Routing zu viele „Meinungen" aus der Vergangenheit das Signal tatsächlich verwässern.
Geschwindigkeit vs. Genauigkeit: Es ist viel schneller, diesen „Spickzettel" (offline) zu erstellen, als das gesamte Modell neu zu trainieren. Während der eigentlichen Beantwortungsphase fügt es eine winzige Verzögerung hinzu (etwa 3–4 % langsamer), verbessert aber die Genauigkeit bei schwierigen Aufgaben wie medizinischen Prüfungen und Programmieren erheblich.

Der Haken (Einschränkungen)

Die Arbeit stellt fest, dass dieses System einen „Referenzsatz" gelabelter Daten benötigt, um den Spickzettel zu erstellen. Wenn Sie sich in einer Situation befinden, in der Sie absolut keine ähnlichen vergangenen Beispiele haben, von denen Sie lernen können, kann diese Methode nicht helfen. Sie beruht auf der Annahme, dass „was vorher funktioniert hat" ein guter Leitfaden für „was jetzt funktionieren wird" ist.

Kurz gesagt ist kNN-MoE eine Möglichkeit, KI-Experten schlauer zu machen, indem man ihnen erlaubt, einen Blick auf ihre eigene Geschichte perfekter Entscheidungen zu werfen, aber nur, wenn es sicher ist, dies zu tun.

Technische Zusammenfassung: kNN-MoE

Problemstellung

Mixture-of-Experts (MoE)-Architekturen skalieren Large Language Models (LLMs) effizient, indem sie einen parametrischen „Router" verwenden, um Tokens an eine sparse Teilmenge von Experten zu verteilen. Standardrouter sind jedoch leichtgewichtige Klassifikatoren, die einmal trainiert und dann eingefroren werden. Obwohl sie bei Daten innerhalb der Verteilung effektiv sind, fehlt ihnen während der Inferenz Flexibilität bei Verteilungsverschiebungen (z. B. spezialisierte Domänen oder Aufgaben außerhalb der Verteilung). In solchen Szenarien muss der Router von festen Priors extrapolieren, was häufig zu einer suboptimalen Expertenauswahl und hoher Perplexität führt. Bestehende Methoden zur Anpassung zur Laufzeit, wie C3PO, versuchen, Routing-Logits online zu feinabstimmen, verursachen jedoch erheblichen Rechenaufwand (ca. das 5-fache des Baselines). Andere Ansätze erfordern permanente Parameteraktualisierungen durch supervised fine-tuning (SFT), was rechenintensiv und für die Inferenz weniger effizient ist.

Methodik: kNN-MoE

Die Autoren schlagen kNN-MoE vor, ein retrieval-augmentiertes Routing-Framework, das Entscheidungen zur Expertenzuweisung verfeinert, ohne die Modellparameter zu ändern. Die Methode arbeitet in zwei distincten Phasen:

1. Speicheraufbau (Offline)

Das System baut einen nicht-parametrischen Speicher $\mathcal{M}$ für jede MoE-Schicht auf, bestehend aus Schlüssel-Wert-Paaren $(k_t, v_t)$ :

Schlüssel ( $k_t$ ): Der Router-Eingabe (versteckter Zustand) $x_t$ aus einem Referenzdatensatz $D_{ref}$ .
Wert ( $v_t$ ): Eine „lokal optimale" Expertenzuweisung, die durch Optimierung der Routing-Logits gewonnen wird, um die Wahrscheinlichkeit des Ground-Truth-Tokens $y_t$ $y_{t}$ zu maximieren.
- Anstatt die Ausgabe des eingefrorenen Routers zu speichern, löst die Methode für jeden Token in $D_{ref}$ ein Optimierungsproblem: Finden eines Logit-Vektors $r^*_t$ , der die negative Log-Likelihood des Ziel-Tokens unter Berücksichtigung der eingefrorenen Netzwerkparameter minimiert.
- Diese Optimierung erfolgt über wenige Schritte des Gradientenabstiegs (typischerweise $S=1$ ), beginnend von den ursprünglichen parametrischen Logits.
- Die resultierenden optimierten Logits $r^{(S)}_t$ dienen als Wert $v_t$ im Speicher.

2. Vertrauensbewusste adaptive Mischung (Inferenz)

Während der Inferenz wird für eine gegebene Router-Eingabe $x$ :

Retrieval: Das System ruft die $K$ nächsten Nachbarn $\mathcal{N}(x)$ aus dem Speicher $\mathcal{M}$ basierend auf der Ähnlichkeit zwischen $x$ und den gespeicherten Schlüsseln ab.
Nicht-parametrischer Vorschlag: Ein speicherbasierter Logit-Vorschlag $r_{mem}(x)$ wird berechnet, indem die abgerufenen optimalen Logits aggregiert und mit ihrer Ähnlichkeit zur Eingabe gewichtet werden.
Vertrauensgesteuerte Mischung: Die endgültigen Routing-Logits sind eine lineare Interpolation zwischen den ursprünglichen parametrischen Router-Logits $r(x)$ $r (x)$ und dem speicherbasierten Vorschlag $r_{mem}(x)$ $r_{m e m} (x)$ :
$r_{final}(x) = (1 - \lambda(x))r(x) + \lambda(x)r_{mem}(x)$
Der Mischkoeffizient $\lambda(x)$ $λ (x)$ ist definiert als die durchschnittliche Ähnlichkeit der abgerufenen Nachbarn. Dies fungiert als Vertrauensscore:
- Hohe Ähnlichkeit ( $\lambda \approx 1$ ): Das Modell vertraut der abgerufenen, lokal optimalen Zuweisung.
- Niedrige Ähnlichkeit ( $\lambda \approx 0$ ): Das Modell weicht auf den ursprünglichen eingefrorenen Router aus und vermeidet Rauschen durch irrelevante Fälle.
Ausgabe: Die endgültigen Logits werden durch einen Top-K-Softmax geleitet, um die sparse Expertenzuweisung zu generieren, wobei die ursprünglichen Sparsitätsbeschränkungen des Modells erhalten bleiben.

Hauptbeiträge

Retrieval-Augmentiertes Routing: Einführung eines Frameworks, das offline berechnete, lokal optimale Expertenzuweisungen aus einem Speicher ähnlicher vergangener Fälle wiederverwendet, um Routing-Entscheidungen zu verbessern.
Nicht-parametrische Anpassung: Im Gegensatz zu SFT oder Online-Anpassungsmethoden aktualisiert kNN-MoE keine Modellparameter. Es verlagert den Rechenaufwand in eine Offline-Phase des Speicheraufbaus und ermöglicht so eine effiziente Inferenz.
Vertrauensgesteuerter Fallback: Die Methode nutzt die durchschnittliche Ähnlichkeit abgerufener Nachbarn als dynamischen Mischkoeffizienten und gewährleistet Robustheit, indem sie zum eingefrorenen Router zurückkehrt, wenn keine relevanten historischen Fälle gefunden werden.
Effizienz: Der Ansatz erfordert nur einen einzigen Gradientenabstiegsschritt pro Token während des Speicheraufbaus und vermeidet die hohen Rechenkosten der Online-Optimierung oder des vollständigen Modell-Fine-Tunings.

Experimentelle Ergebnisse

Die Autoren evaluierten kNN-MoE an drei MoE-Modellen (OLMoE, GPT-OSS und Qwen3) über sechs Benchmarks hinweg (GPQA, SuperGPQA, MMLU, USMLE, MedMCQA und MBPP).

Leistung vs. Baselines: kNN-MoE übertraf konsistent die Zero-Shot-Baseline und die „SFT (Nur Router)"-Baseline über alle Modelle und Datensätze hinweg.
Wettbewerbsfähigkeit mit SFT: Die Methode erreichte eine Leistung, die mit vollständigem Supervised Fine-Tuning (SFT) vergleichbar war, und übertraf in mehreren datenarmen Settings (z. B. GPQA, USMLE) sogar SFT, das bei kleinen Referenzsets gelegentlich unter Overfitting litt.
Stabilität: Im Gegensatz zum 5-Shot-Prompting, das bei bestimmten Modellen (GPT-OSS, OLMoE) aufgrund von Eingabelängenproblemen die Leistung verschlechterte, zeigte kNN-MoE stabile Verbesserungen.
Ablationsstudien:
- Nachbarn ( $K$ ): Entgegen typischen kNN-LM-Ergebnissen erreichte die Leistung ihr Maximum bei $K=1$ . Die Aggregation weiterer Nachbarn ( $K>1$ ) führte zu Rauschen, was darauf hindeutet, dass wirklich vorteilhafte Routing-Signale selten und spezifisch sind.
- Optimierungsschritte ( $S$ ): Eine Erhöhung der Gradientenabstiegsschritte über $S=1$ hinaus erbrachte vernachlässigbare Genauigkeitsgewinne bei gleichzeitiger Erhöhung der Aufbaukosten.
- Größe des Referenzsets: Die Leistung verbesserte sich monoton mit der Größe des Referenzsets $|D_{ref}|$ , wobei ab 0,5k–1k Beispielen abnehmende Grenzerträge beobachtet wurden.
Effizienz: Der Speicheraufbau benötigte für 1k Beispiele 0,46 Stunden (vs. ~1,34 h für SFT). Die IVFPQ-Variante (Inverted File + Product Quantization) reduzierte den Inferenz-Latenzaufwand auf ~3,67 % und den VRAM-Verbrauchszuwachs auf ~9,43 % im Vergleich zur Baseline.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass kNN-MoE die Lücke zwischen statischem parametrischem Routing und dynamischen Inferenzbedürfnissen schließt. Ihre primäre Bedeutung liegt darin, einen skalierbaren Weg zur Verbesserung der MoE-Anpassungsfähigkeit ohne den Rechenaufwand des Fine-Tunings zu bieten.

Die Autoren betonen, dass die Methode besonders effektiv ist, wenn der parametrische Router unsicher ist (korreliert mit hoher Perplexität), da der Retrieval-Mechanismus ein korrigierendes Signal aus ähnlichen, hochwertigen historischen Fällen liefert. Obwohl die Methode die Verfügbarkeit eines gelabelten Referenzsets mit einer Verteilung benötigt, die den Testdaten ähnelt, bietet sie eine praktische, nicht-parametrische Alternative zu teuren Anpassungstechniken und erzielt SFT-konkurrenzfähige Ergebnisse mit deutlich geringerem Vorbereitungszeitraum und ohne Parameteraktualisierungen.