Technisches Resümee: Hin zu robuster Internalisierung in der On-Policy-Distillation

1. Problemstellung

Jüngste Fortschritte in der On-Policy-Distillation (OPD) haben es Schüler-Modellen ermöglicht, privilegierten Kontext (z. B. System-Prompts, Aufgabenhinweise oder Spielzustandsbeschreibungen) in ihre Parameter zu internalisieren, sodass sie auch ohne diesen Kontext zur Inferenzzeit gut abschneiden. Dieses Paper identifiziert jedoch ein bisher ununtersuchtes Phänomen namens kontextinduzierte Degradation (context-induced degradation).

Entgegen der Intuition, dass internalisiertes Wissen den Kontext redundant machen sollte (eine idempotente Operation), beobachten die Autoren, dass die Wiedereinführung des privilegierten Kontexts in einen destillierten Schüler die Leistung oft verschlechtert und die Antwortlängen aufbläht, selbst bei Instanzen, die das Modell ohne Kontext korrekt löst. Dies geschieht, weil die Standard-OPD ausschließlich auf privilegierte Fidelität (das Abgleichen mit der kontextabhängigen Ausgabe des Lehrers) optimiert, aber keine Einschränkungen für das Verhalten des Schülers auferlegt, wenn der Kontext vorhanden ist. Folglich scheitert der Schüler an der Erreichung von Kontext-Entfernbarkeit (context removability): der Eigenschaft, dass die Ausgabe des Modells unabhängig davon stabil bleibt, ob der privilegierte Kontext bereitgestellt wird oder nicht.

Die Autoren kategorisieren die Interaktionen mit dem Kontext nach der Distillation in drei Regime:

• Regime A: Der Schüler performt gut ohne Kontext, aber die Wiedereinführung des Kontexts führt zu signifikanten Genauigkeitsverlusten und hohen „Schadensraten“ (Umkehrung korrekter Vorhersagen in falsche). Dies ist das am weitesten verbreitete Regime in ihren Experimenten.
• Regime B: Der Kontext fungiert weiterhin als nützliches Gerüst, was auf eine unvollständige Internalisierung hindeutet.
• Regime C: Der Schüler hat den Kontext erfolgreich internalisiert, und die Leistung ist sowohl mit als auch ohne Kontext stabil.

2. Methodik: No-Context Anchoring (NCA)

Um die kontextinduzierte Degradation zu adressen, schlagen die Autoren No-Context Anchoring (NCA) vor, einen leichtgewichtigen Konsistenz-Regularisierer. Der Kern der Idee besteht darin, vorzuschreiben, dass die kontextabhängige Ausgabe des Schülers ($q_c$) konsistent mit seiner kontextfreien Ausgabe ($q_x$) bleibt, wobei letztere als Anker dient.

Zielfunktion

Die Standard-OPD minimiert die KL-Divergenz zwischen der kontextfreien Ausgabe des Schülers und der kontextabhängigen Verteilung des Lehrers:
$$\mathcal{L}_{OPD} = \mathbb{E}_{x \sim D} [\text{KL}(q_x \parallel f(y \mid x, c))]$$

NCA erweitert dies um einen Konsistenzterm:
$$\mathcal{L} = \mathcal{L}_{OPD} + \beta \cdot \mathcal{L}_{NCA}$$
wobei $\mathcal{L}_{NCA}$ definiert ist als:
$$\mathcal{L}_{NCA} \triangleq \text{KL}(\text{sg}[g_\theta(\cdot \mid x)] \parallel g_\theta(\cdot \mid x, c))$$

Wichtige Designentscheidungen

1. Stop-Gradient Anchoring: Die kontextfreie Sicht $q_x$ wird durch einen Stop-Gradient-Operator (sg) geleitet. Dies stellt sicher, dass $q_x$ nur durch den primären OPD-Loss Gradienten erhält, während $q_c$ nur durch den NCA-Loss Gradienten erhält. Dies eliminiert die Gradienten-Konkurrenz zwischen den beiden Zielen.
2. Forward KL Divergenz: Die Autoren wählen $\text{KL}(q_x \parallel q_c)$ (Forward KL) anstelle von Reverse KL. Diese Wahl ist „modus-deckend“ (mode-covering) und verhindert, dass $q_c$ Modi von $q_x$ kollabieren lässt. Zudem fällt die Erwartung unter $q_x$ an, was die Wiederverwendung der für die OPD generierten On-Policy-Samples ermöglicht, ohne dass ein separates Rollout von $q_c$ erforderlich ist.
3. Recheneffizienz: NCA erfordert pro Trainingsschritt nur einen zusätzlichen Forward Pass (zur Berechnung von $q_c$). Die Logits für $q_x$ werden vom Berechnungsgraph getrennt, und die On-Policy-Sequenzen werden wiederverwendet, wodurch der Overhead vernachlässigbar bleibt.

3. Experimenteller Aufbau

Die Autoren evaluierten NCA über 12 Konfigurationen hinweg, die folgende Bereiche abdecken:

• Domänen: System-Prompt QA (MedMCQA, Safety-Datensätze) und Textspiele (Sokoban, FrozenLake).
• Modelle: Verschiedene Llama-3 und Qwen Modelle (von 1,7B bis 8B Parametern), einschließlich sowohl Selbst-Distillation als auch Cross-Model-Distillation-Szenarien.
• Metriken: Genauigkeit/Win-Rate mit und ohne Kontext, Schadensrate (Wahrscheinlichkeit, dass Kontext eine korrekte Antwort in eine falsche dreht), Both-Correct/Won-Raten sowie relative Änderungen der Antwortlänge.

4. Kernergebnisse

Über die 12 Settings hinweg zeigte NCA folgende Effekte:

• Milderung der kontextinduzierten Degradation: NCA reduzierte die kontextinduzierte Schädigung in 11 von 12 Settings. In Regime-A-Settings (wo die Degradation schwerwiegend war) reduzierte NCA den Genauigkeitsabfall bei der Wiedereinführung des Kontexts signifikant.
• Verbesserte Stabilität: NCA steigerte konsistent die „BothCorrect“- oder „BothWon“-Raten, was darauf hindeutet, dass das Modell unter beiden Ansichten ein korrektes Verhalten beibehält.
• Eliminierung der Antwortlängen-Inflation: Die Standard-OPD führte oft dazu, dass das Modell deutlich längere Antworten generierte, wenn der Kontext wiedereingeführt wurde (bis zu +39,5 % Längensteigerung). NCA eliminierte diese Inflation in den meisten Settings effektiv und reduzierte die Längenänderung oft auf nahezu Null.
• Verbesserte No-Context-Performance: In 7 von 12 Settings verbesserte NCA das primäre Trainingsziel (No-Context-Genauigkeit) um 0,6 bis 3,6 Prozentpunkte. Die Autoren führen dies darauf zurück, dass die Konsistenzbeschränkung als Regularisierer wirkt, der verhindert, dass Parameter-Updates widersprüchliche Verhaltensweisen über die Ansichten hinweg anpassen.
• Grenzfälle: In einem Regime-B-Setting (FrozenLake, Cross-Model) erhöhte NCA die Schadensrate leicht, was darauf hindeutet, dass in Fällen, in denen der Kontext noch einzigartige, nicht-internalisierte Informationen trägt, das Erzwingen einer strikten Entfernbarkeit die kontextabhängige Performance zugunsten der Stabilität opfern kann.

5. Mechanistische Analyse

Das Paper liefert eine mechanistische Fallstudie zur Sokoban-Aufgabe, um zu erklären, wie NCA die Kontext-Entfernbarkeit erreicht:

• Repräsentations-Ebene: Eine Analyse der Hidden States zeigte, dass die Standard-OPD zwar die Divergenz zwischen Kontext- und No-Context-Ansichten reduzierte, eine merkliche Lücke jedoch in den mittleren und späten Layern bestehen blieb. NCA hingegen hielt eine nahezu perfekte Kosinus-Ähnlichkeit ($\ge 0,997$) über alle Layer aufrecht, was bestätigt, dass die Kontext-Entfernbarkeit auf der Repräsentations-Ebene erreicht wird.
• Gradienten-Ausrichtung: In der Standard-OPD waren die Gradienten für das Distillationsziel und die kontextabhängige Sicht nahezu orthogonal. NCA führt eine neue Gradientenrichtung ein, die aktiv die Kräfte ausgleicht, die den Kontext für die Distillation nutzen wollen und die Invarianz gegenüber diesem fordern.
• Parameter-Updates: Sowohl OPD als auch NCA konzentrieren Parameter-Updates in flachen Layern (insbesondere $k\_proj$ und $q\_proj$). Die spezifischen Updates, die durch NCA getrieben werden, sind jedoch von denen der OPD verschieden, was darauf hindeutet, dass die Kontext-Entfernbarkeit primär durch Modifikationen an den Wissensrepräsentationen der flachen Layer erreicht wird, die die OPD allein nicht produziert.

6. Bedeutung und Ansprüche

Die Autoren behaupten, dass ihre Arbeit das Evaluationsparadigma für die Privileged-Context-Distillation verschiebt. Sie argumentieren, dass robuste Internalisierung nicht nur das Abgleichen des Verhaltens des Lehrers (privilegierte Fidelität) erfordert, sondern auch die Sicherstellung der Kontext-Entfernbarkeit.

• Neuheit: Dies ist die erste Studie, die das Phänomen identifiziert und adressiert, bei dem die Wiedereinführung des internalisierten Kontexts die Modellleistung verschlechtert.
• Praktikabilität: Die vorgeschlagene Methode (NCA) ist einfach, benötigt nur einen einzigen Forward Pass und einen Standard-Konsistenz-Regularisierer, löst aber effektiv das Degradationsproblem, ohne die primären Distillationsziele zu gefährden.
• Generalisierbarkeit: Die Methode verbessert die Performance über verschiedene Modellfamilien, Aufgaben und Distillations-Settings (Self und Cross-Model) hinweg.

Das Paper schließt mit dem Appell an die Community, distillierte Modelle nicht nur nach ihrer No-Context-Performance zu bewerten, sondern auch nach ihrer Robustheit gegenüber der Wiedereinführung des Kontexts, um sicherzustellen, dass internalisiertes Wissen nicht zu einem Nachteil wird, wenn der Kontext versehentlich oder absichtlich bereitgestellt wird.