Technischer Überblick: MIXSD – Gemischte kontextuelle Selbst-Distillation für Wissensinjektion

Problemstellung

Supervised Fine-Tuning (SFT) ist der Standardansatz zum Einbringen neuen Wissens in Large Language Models (LLMs). Dieser Prozess führt jedoch häufig zu katastrophalem Vergessen, bei dem das Modell in vorab trainierten Fähigkeiten wie Schlussfolgern, Befehlsbefolgung und Leistung im allgemeinen Domänenbereich nachlässt.

Die Autoren argumentieren, dass dieses Vergessen aus einer Verteilungsinkongruenz zwischen den Überwachungszielen und der nativen autoregressiven Verteilung des Modells resultiert. In Standard-Pipelines werden Zielsequenzen von Menschen oder externen Systemen verfasst und enthalten oft Formulierungen, Formatierungen oder Schlussfolgerungsstrukturen, die in den Wahrscheinlichkeits-Schwänzen der Verteilung des Basismodells liegen. Das Erzwingen einer tokenweisen Imitation dieser unwahrscheinlichen Tokensequenzen führt zu Parameter-Updates entlang sensibler Richtungen und stört damit zuvor erlerntes Verhalten.

Methodik: MIXSD

Um dies zu adressieren, schlägt die Arbeit MIXSD (Mixed Contextual Self-Distillation) vor, eine einfache, externe-Lehrer-freie Methode für verteilungsausgerichtete Wissensinjektion. Anstatt auf festen, extern verfassten Zielen zu trainieren, konstruiert MIXSD die Überwachung dynamisch unter Verwendung des Basismodells selbst.

Kernmechanismus

Zu jedem Dekodierschritt $t$ wird, gegeben ein gemeinsames autoregressives Präfix, der Überwachungs-Token aus einer Mischung zweier bedingter Verteilungen, definiert durch das Basismodell $p_{\theta^*}$, gesampelt:

1. Experten-Bedingung: $p_{\theta^*}(\cdot \mid x^+_i, y_{mix, <t})$, wobei die Eingabe $x^+_i$ die injizierte Tatsache (Ground Truth) im Kontext enthält. Dies erzeugt Tokens, die die korrekte Tatsache in der eigenen Oberflächenform des Modells ausdrücken.
2. Naive Bedingung: $p_{\theta^*}(\cdot \mid x_i, y_{mix, <t})$, wobei die Eingabe $x_i$ nur den ursprünglichen Prompt ohne das injizierte Wissen enthält. Dies spiegelt das ursprüngliche Prior des Modells wider.

Das Überwachungsziel $y^{mix}_{i,t}$ wird wie folgt gesampelt:
$$y^{mix}_{i,t} = \begin{cases} \tilde{y}^+_{i,t} & \text{mit Wahrscheinlichkeit } 1 - \lambda \\ \tilde{y}^-_{i,t} & \text{mit Wahrscheinlichkeit } \lambda \end{cases}$$
wobei $\lambda \in [0, 1]$ eine Mischrate ist. Das Modell wird dann unter Verwendung des Standard-Negative-Log-Likelihood (NLL)-Verlusts auf diesen gemischten Zielen aktualisiert.

Wichtige Design-Entscheidungen

• Kein Externer Lehrer: Im Gegensatz zur traditionellen Wissensdistillation benötigt MIXSD kein stärkeres Lehrer-Modell und vermeidet so die Rechenkosten für die Generierung mehrerer Rollouts pro Abfrage (wie sie bei On-Policy-Selbst-Distillation-Baselines zu sehen sind).
• Verteilungsausrichtung: Durch das Einmischen naiver Tokens bleiben die Überwachungsziele der nativen Verteilung des Basismodells erheblich näher, was den NLL der Trainingsziele unter dem Basismodell reduziert.
• Hyperparameter $\lambda$: Steuert den Trade-off. Ein niedrigeres $\lambda$ betont das Expertensignal (Memorisierung), während ein höheres $\lambda$ das Modell an sein Prior bindet (Erhaltung).

Experimentelles Setup

Die Autoren evaluierten MIXSD über mehrere Modellskalen hinweg (Qwen3-1.7B, 4B und 8B) und zwei synthetische Korpora, die entwickelt wurden, um spezifische Wissensarten zu isolieren:

• KGFACT: Ein Korpus faktischen Wissens, abgeleitet aus einem simulierten Weltgraphen mit neuen Entitäten, um faktisches Abrufen und Retrieval zu testen.
• KGFUNC: Ein Datensatz für den Erwerb arithmetischer Funktionen, der vom Modell verlangt, Regeln aus Operationen auf Ziffernebene abzuleiten.
• Benchmarks: Die Evaluation umfasste zurückgehaltene Benchmarks im allgemeinen Domänenbereich (AIME, MATH, GSM8K, HumanEval, MMLU) zur Messung des Vergessens, neben der Trainingsgenauigkeit innerhalb der Domäne.

Wichtige Ergebnisse

1. Verbesserter Trade-off zwischen Memorisierung und Erhaltung: Über alle Datensätze und Modellskalen hinweg schnitt MIXSD konsistent besser ab als Standard-SFT und On-Policy-Selbst-Distillation (OPSD)-Baselines.
   • Standard-SFT erreichte eine nahezu perfekte Trainingsgenauigkeit, behielt jedoch nur 1% der zurückgehaltenen Fähigkeiten des Basismodells bei.
   • MIXSD (mit $\lambda=0.3$ oder $0.5$) behielt eine nahezu perfekte Trainingsgenauigkeit bei, während es bis zu 100% der zurückgehaltenen Fähigkeiten des Basismodells bewahrte (z. B. bei Qwen3-4B behielt MIXSD ~77% der allgemeinen Fähigkeiten bei, verglichen mit ~43% für SFT).
2. Verteilungsausrichtung: Die MIXSD-Überwachungsziele wiesen unter dem Basismodell einen signifikant niedrigeren per-Token-NLL auf als SFT-Ziele. Hoch-NLL-Tokens (die eine geringe Wahrscheinlichkeit unter dem Basismodell anzeigen) machten 27–42% der SFT-Ziele aus, aber nur 4–8% für MIXSD bei $\lambda=0.3$.
3. Fehlermodus-Analyse: SFT verursachte qualitative Verschiebungen im Generierungsverhalten, was zu „Leckagen" (Generieren fiktiver Entitäten, die nichts mit dem Prompt zu tun haben) und „Kollaps" (kurze, vorlagenartige Antworten ohne Schlussfolgerung) führte. MIXSD bewahrte weitgehend das ursprüngliche Fehlerprofil des Basismodells; die Mehrheit der Fehler waren „echte" Schlussfolgerungsversuche und keine pathologischen Ausfälle.
4. Richtung des Parameter-Updates: Die Arbeit führt eine Metrik namens Fisher-gewichtete Parameter-Verschiebung ein. Die Ergebnisse zeigen, dass die Richtung der Parameter-Updates (Ausrichtung mit sensiblen Fisher-Richtungen) stark mit Vergessen korreliert ($r \approx 0.56-0.82$), wohingegen die Größe der Updates dies nicht tut ($r \approx 0.02-0.34$). MIXSD reduziert die Ausrichtung mit diesen sensiblen Richtungen im Vergleich zu SFT.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass die Ausrichtung der Überwachung mit der nativen Generierungsverteilung des Modells ein einfaches und effektives Prinzip für die Wissensinjektion ist, das katastrophales Vergessen mildert.

• Theoretische Einsicht: Die Arbeit legt nahe, dass katastrophales Vergessen weniger durch die Größe der Parameteränderungen getrieben wird, sondern vielmehr durch Updates entlang sensibler Richtungen, die durch verteilungsinkongruente Überwachung induziert werden.
• Praktischer Nutzen: MIXSD bietet einen „Drop-in"-Ersatz für Standard-SFT-Ziele, der keinen externen Lehrer erfordert und nur minimalen zusätzlichen Rechenaufwand benötigt (nur einmalige Vorverarbeitungskosten zur Generierung gemischter Rollouts).
• Generalisierbarkeit: Die Vorteile wurden über verschiedene Modellfamilien (Qwen und Llama) und Settings hinweg beobachtet, einschließlich Wissensedition (Überschreiben bestehender Fakten) und dem Erwerb neuer arithmetischer Funktionen.

Die Autoren schließen, dass SFT Modelle zwingt, unwahrscheinliche externe Trajektorien zu imitieren, während MIXSD das faktische Lernsignal bewahrt, gleichzeitig die Überwachungs-Trajektorie jedoch innerhalb der hochwahrscheinlichen Regionen des Modells hält und so vorab trainierte Fähigkeiten stabilisiert.