DLLG: Dynamic Logit-Level Gating of LLM Experts

Technische Zusammenfassung: DLLG: Dynamische Logit-Level-Gating von LLM-Experten

Problemstellung
Die Zunahme spezialisierter Large Language Models (LLMs), die für enge Domänen optimiert sind (z. B. mathematisches Denken, Codegenerierung), stellt eine Herausforderung dar: Wie lassen sich unabhängig trainierte „Black-Box“-Experten in ein einheitliches System integrieren, das deren komplementäre Stärken nutzt, ohne prohibitive Kosten für das erneute Training zu verursachen? Bestehende Ansätze leiden unter einer anhaltenden Spannung zwischen Adaptivität, Robustheit und Praktikabilität:

• Routing-basierte Methoden verlassen sich auf eine harte, frühe Expertenwahl. Dies führt zu einer vorzeitigen Festlegung; wenn der gewählte Experte bei einer spezifischen Teilaufgabe innerhalb einer einzelnen Antwort scheitert, ist dieser Fehler nicht mehr korrigierbar. Zudem erfordern viele Router eine umfangreiche, benchmark-spezifische Kalibrierung.
• Heuristische Ensembling-Verfahren (z. B. Perplexitätsgewichtung, Token-Voting) nutzen während der Inferenz Proxy-Signale. Diese Metriken sind oft empirisch und mit der tatsächlichen Aufgabenkorrektheit nicht synchronisiert, was zu suboptimalen Ausgaben führt.
• Parameter-Merging (z. B. Model Souping, Task Arithmetic) fusioniert Gewichte statisch. Dies opfert die Modularität und kann oft destruktive Interferenzen hervorrufen, wenn Experten mit gegensätzlichen Gradientenrichtungen kombiniert werden, was zu einer Leistungsverschlechterung führt.
• Mixture-of-Experts (MoE)-Architekturen erfordern typischerweise ein End-to-End-Joint-Training von Gates und Experten, was die Nutzung fertiger, vortrainierter Spezialisten einschränkt.

Methodik: DLLG
Die Autoren schlagen DLLG (Dynamic Logit-Level Gating) vor, ein Framework, das eine dynamische, Token-basierte Logit-Fusion mittels eines leichtgewichtigen Gating-Modells durchführt, das auf spärlicher, Antwort-basierter Überwachung trainiert wurde.

1. Architektur:

   • Gefrorene Experten: $K$ spezialisierte LLMs bleiben eingefroren. An jedem Dekodierungsschritt $t$ erzeugt jeder Experte $E_i$ gegeben einen Prompt $x$ und ein Präfix $y_{<t}$ einen Hidden State $h^{(i)}_t$ und die nächsten-Token-Logits $\ell^{(i)}_t$.
   • Gating-Modell: Ein leichtgewichtiges Modell ($g_\phi$) verarbeitet die konkatinierten Hidden States aller Experten. Es projiziert diese Repräsentationen in einen gemeinsamen Embedding-Raum (mittels Low-Rank-Projektion) und verarbeitet sie durch einen autoregressiven Stem (z. B. Qwen2.5-0.5B-Instruct mit LoRA-Adaptern).
   • Gewichtsvorhersage: Ein Prediction Head gibt Token-weise Fusionsgewichte $w_{t,i}$ über eine Sigmoid-Transformation aus. Diese Gewichte werden auf die Logits der Experten angewendet: $\ell_t = \sum_{i=1}^K w_{t,i} \ell^{(i)}_t$. Die Gewichte sind nicht dazu beschränkt, sich zu eins zu addieren, was flexible, weiche Beiträge ermöglicht.

2. Training (Antwort-basierte Überwachung):

   • Das Framework adressiert den Mangel an Token-basierten Annotationen durch die Verwendung von Antwort-basierten Korrektheits-Signalen ($s_i \in \{0, 1\}$), die aus automatischen Verifizierern abgeleitet werden (z. B. Unit-Tests für Code, Antwortprüfung für Mathematik).
   • Teacher Forcing: Während des Trainings generieren die Experten Hidden States, die auf das Ground-Truth-Präfix $y_{<t}$ konditioniert sind. Das Gating-Modell sagt die Gewichte $w_{t,i}$ für jede Token-Position voraus.
   • Zielsetzung: Das Modell wird darauf trainiert, den mittleren quadratischen Fehler (MSE) zwischen den vorhergesagten Gewichten und dem auf die gesamte Antwort übertragten Korrektheits-Signal zu minimieren: $\mathcal{L} = \frac{1}{T} \sum_{t=1}^T \sum_{i=1}^K (w_{t,i} - s_i)^2$. Dies lehrt das Gate, die Token-basierte Fusion mit dem Gesamterfolg des Pfades in Einklang zu bringen.

3. Inferenz:

   • DLLG arbeitet autoregressiv. In jedem Schritt sagt das Gating-Modell Fusionsgewichte basierend auf den aktuellen Hidden States der Experten voraus, die Logits werden kombiniert und das nächste Token wird ausgewählt. Dies ermöglicht eine rekonfigurierbare Expertennutzung; wenn ein Experte frühzeitig einen Fehler macht, kann das Gating-Modell das Gewicht in den folgenden Schritten dynamisch auf einen kompetenteren Experten verschieben.

Wesentliche Beiträge

1. Framework-Vorschlag: Einführung von DLLG, einem dynamischen Logit-Level-Ensembling-Framework, das die Token-basierte Generierung mit spärlicher Antwort-basierter Überwachung verbindet und so eine feingliedrige Expertennutzung ohne Token-basierte Annotation oder Online-Reinforcement-Learning ermöglicht.
2. Mechanismus-Analyse: Demonstration, dass der gelernte Gating-Mechanismus die Beiträge der Experten während des Rollouts dynamisch anpasst und die Gewichte basierend auf dem sich entwickelnden Kontext anpasst (z. B. Wechsel von mathematischer Argumentation zu Codegenerierung innerhalb einer einzigen Antwort).
3. Empirische Validierung: Evaluierung über verschiedene Reasoning- und Code-Benchmarks (GSM8K, MATH, HumanEval usw.) und Modellskalen (0.5B, 1.5B), die konsistente Gewinne gegenüber Routing-, heuristischen Ensembling- und Parameter-Merging-Baselines zeigt.

Ergebnisse

• Performance: Über beide Modellskalen (0.5B und 1.5B) hinweg erreicht DLLG konsistent die höchste durchschnittliche Leistung im Vergleich zu starken Baselines, einschließlich RouterDC, EmbedLLM, verschiedenen Logit-Ensembling-Methoden (GaC, UniTe) und Parameter-Merging-Techniken (Linear, SLERP, Task Arithmetic).
• Robustheit: In Settings mit geringer Kapazität (0.5B), in denen einzelne Experten schwach sind, übertrifft DLLG spröde Routing- und heuristische Methoden, was seine Fähigkeit demonstriert, frühe Fehler durch weiche Fusion abzumildern.
• Generalisierung: Die Methode zeigt eine starke Performance bei Out-of-Domain-Benchmarks (z. B. HumanEval, MBPP, BBH), was darauf hindeutet, dass die gelernten Fusionsregeln über die spezifische Trainingsverteilung hinaus generalisieren.
• Ablation: Experimente bestätigen, dass eine ausreichend expressive Gating-Architektur (die ein Hilfs-LLM und eine Low-Rank-Projektion beinhaltet) für eine effektive Token-basierte Fusion notwendig ist und einfache MLP-Heads übertrifft.

Bedeutung und Ansprüche
Das Paper behauptet, dass DLLG ein robustes und skalierbares Paradigma zur Integration spezialisierter Experten bietet. Durch das Erlernen der Logit-Level-Fusion aus schwacher Überwachung erreicht es eine MoE-ähnliche Adaptivität, ohne die Kosten eines erneuten Trainings der Experten oder das Risiko von Parameterinterferenzen. Die Autoren betonen, dass DLLG die Experten-Modularität bewahrt und eine feingliedrige, rekonfigurierbare Expertennutzung ermöglicht, wodurch die Einschränkungen der vorzeitigen Festlegung durch Routing und die statische Natur des Parameter-Merging überwunden werden. Die Arbeit legt nahe, dass das Erlernen der feingliedrigen Expertennutzung aus schwacher Überwachung ein praktischer Weg nach vorne ist, um zunehmend spezialisierte Sprachmodelle zu nutzen.