Technische Zusammenfassung: Nicht jede Diskrepanz ist lernbar: Token-Lehrbarkeit in der On-Policy-Distillation

Problemstellung

On-Policy-Distillation (OPD) trainiert ein studentisches Sprachmodell anhand seiner eigenen Rollouts unter Verwendung einer Token-Level-Überwachung durch eine Lehrer-Policy. Während neuere selektive OPD-Methoden versuchen, die Effizienz zu steigern, indem sie „saliente" Token priorisieren – wie etwa solche mit hoher Entropie oder hoher Lehrer-Schüler-Diskrepanz (KL-Divergenz) – basieren diese Ansätze auf der Annahme, dass eine hohe Diskrepanz einem hohen Lernwert entspricht.

Die Arbeit identifiziert einen kritischen Fehler in dieser Annahme: Die rohe KL-Diskrepanz ist ein grober Proxy für das Lernen. Sie vermischt zwei qualitativ unterschiedliche Szenarien:

1. Lernbare Diskrepanz: Der Lehrer weist Korrektur-Wahrscheinlichkeitsmasse auf Token zu, die bereits innerhalb der aktuellen top-$K$-Vorhersageunterstützung des Schülers liegen.
2. Inkompatible Diskrepanz: Der Lehrer weist Masse auf Token zu, die weit außerhalb der aktuellen Unterstützung des Schülers liegen. Obwohl dies eine große KL-Lücke erzeugt, kann der Schüler diese Korrektur nicht leicht aufnehmen, was zu schwachen oder sogar negativen lokalen Lerneffekten führt.

Das Kernproblem besteht darin, zu unterscheiden, welche Token-Level-Lehrersignale tatsächlich lernbar sind (d. h. lokale Verbesserungen bewirken können) und welche lediglich salient oder inkompatibel sind.

Methodik

1. Diagnose mit festem Kontext

Um den Lernwert spezifischer Token von Rollout-Rauschen und nachgelagerten Kontextverschiebungen zu isolieren, führen die Autoren eine Diagnose mit festem Kontext ein.

• Protokoll: Vom Schüler generierte Präfixe werden gesammelt und als Kontextbank eingefroren. Sowohl das initiale (vor dem Update) als auch das trainierte Schülermodell werden gegen dieselbe Lehrer-Verteilung auf diesen festen Kontexten neu bewertet.
• Metrik: Der Fixed-Context Token Gain ($G^{\text{fix}}_{i,t}$) ist definiert als die Reduktion der KL-Divergenz zwischen Lehrer und Schüler nach dem Training, gemessen am selben Präfix.
• Ziel: Dies ermöglicht den Autoren, vor-Update-Signale (wie die rohe KL-Divergenz) mit tatsächlicher lokaler Verbesserung zu korrelieren, wobei die Stichprobenvarianz kontrolliert wird.

2. Formalisierung der Token-Lehrbarkeit

Mittels der Diagnose zerlegen die Autoren die Diskrepanz in zwei Komponenten basierend auf der lokalen Unterstützungs-Kompatibilität:

• Lokale Unterstützung ($S^S_t(K)$): Die Menge der top-$K$-Token des Schülers.
• Kompatibilitätsmasse ($C_t$): Die gesamte Wahrscheinlichkeitsmasse, die der Lehrer der lokalen Unterstützung $S^S_t(K)$ des Schülers zuweist.
• Zerlegung:
  • Lernbare Diskrepanz ($D^L_t$): Hohe Diskrepanz, bei der die Korrektur des Lehrers mit der Unterstützung des Schülers übereinstimmt (Hohe $C_t$).
  • Inkompatible Diskrepanz ($D^I_t$): Hohe Diskrepanz, bei der die Korrektur des Lehrers außerhalb der Unterstützung des Schülers liegt (Niedrige $C_t$).

Der Lehrbarkeits-Score ist definiert als das Produkt aus normalisierter Diskrepanz und Kompatibilität:
$$s^{\text{teach}}_t = \tilde{D}_t \cdot \tilde{C}_t$$
wobei $\tilde{D}_t$ die normalisierte Diskrepanz und $\tilde{C}_t$ die normalisierte Kompatibilitätsmasse ist.

3. Lehrbarkeitsbewusste OPD (TA-OPD)

Basierend auf der Erkenntnis, dass $D^L_t$ den Fixed-Context-Gain besser vorhersagt als die rohe KL-Divergenz, schlagen die Autoren TA-OPD vor, eine leichte Methode zur Auswahl von Token-Positionen:

• Auswahl: Für ein gegebenes Budget (z. B. 5 % oder 10 % der Token) wählt TA-OPD die top-$N$ Positionen mit den höchsten Lehrbarkeits-Scores ($s^{\text{teach}}_t$) aus.
• Implementierung: Es werden nur die Token-Wahrscheinlichkeiten von Lehrer und Schüler benötigt (insbesondere die top-$K$-Log-Probabilities), die während der Standard-OPD verfügbar sind. Es werden keine Reward-Modelle, Verifizierer oder zusätzliche Labels benötigt.
• Kontrast: Im Gegensatz zu früheren Methoden, die auf Basis von Entropie oder roher Divergenz auswählen (TIP), filtert TA-OPD explizit nach „unterstützungsausgerichteten" Korrekturen.

Hauptbeiträge

1. Diagnose mit festem Kontext: Eine Methode zur Messung der KL-Reduktion im selben Kontext, die vor-Update-Token-Signale mit lokaler OPD-Verbesserung verknüpft, ohne durch Rollout-Rauschen zu verfälschen.
2. Formalisierung der Token-Lehrbarkeit: Der Nachweis, dass die rohe KL-Diskrepanz lernbare und inkompatible Signale vermischt. Die Arbeit beweist, dass unterstützungsausgerichtete Diskrepanz der zuverlässige Prädiktor für nützliches Lernen ist.
3. TA-OPD-Algorithmus: Eine praktische, leichte Auswahlmethode, die den OPD-Verlust nur auf Positionen mit hoher Lehrbarkeit anwendet und so die Effizienz ohne externe Tools verbessert.
4. Empirische Validierung: Umfassende Experimente zeigen, dass die Auswahl nach Lehrbarkeit bessere Downstream-Leistung erzielt als die Auswahl nach Salienz, selbst bei deutlich reduzierten Token-Budgets.

Experimentelle Ergebnisse

Die Autoren evaluierten TA-OPD über vier Lehrer-Schüler-Konfigurationen hinweg, die Modelle der Familien Qwen3 (4B, 8B, 14B) und Qwen2.5 (3B) umfassten, unter Verwendung von Benchmarks wie AIME, GPQA-Diamond, HumanEval und MATH-500.

• Leistung vs. Voll-OPD: In vielen Konfigurationen erreichte TA-OPD die Leistung der Voll-Token-OPD oder übertraf sie, wobei es nur 5 % der Token verwendete. Beispielsweise erzielte TA-OPD im Setting Qwen3-4B $\to$ 1.7B einen Durchschnittswert von 44,89 im Vergleich zu 42,37 bei der Voll-OPD, trotz eines 10%-Budgets (und übertraf die Voll-OPD in spezifischen Metriken).
• Vergleich mit Baselines: TA-OPD schnitt konsistent besser ab als entropiebasierte, divergenzbasierte und TIP-Baselines (Token Importance in OPD).
  • Im Cross-Backbone-Setting (DeepSeek-R1-Distill-Qwen-14B $\to$ Qwen2.5-3B) verbesserte TA-OPD den Durchschnittswert auf 30,62 und übertraf damit deutlich das Base-Modell (29,98) und die Voll-OPD (28,76).
• Budget-Sensitivität: Die Ergebnisse deuten darauf hin, dass die OPD-Leistung durch die Token-Qualität begrenzt ist und nicht rein durch das Budget. Eine Erhöhung des Token-Budgets über 10 % hinaus führte oft zu keinen weiteren Gewinnen und verschlechterte die Leistung manchmal sogar, was die Erkenntnis untermauert, dass nicht jede zusätzliche Überwachung vorteilhaft ist.
• Robustheit: Die positive „High-Support-Gain-Lücke" (der Leistungsunterschied zwischen Token mit hoher und niedriger Lehrbarkeit) blieb über verschiedene Kontexte (Mathematik, GSM8K-COT), Lehrerskalen (8B, 14B) und Unterstützungsgrößen ($K=8, 16, 32$) hinweg konsistent.

Bedeutung und Behauptungen

Die Arbeit stellt die selektive On-Policy-Distillation nicht als Problem der Auswahl „salienter" oder „unsicherer" Token neu dar, sondern als Problem der Identifizierung lokal lernbarer Lehrersignale.

• Kernbehauptung: Die Gewinne aus OPD hängen von der Kompatibilität zwischen der Korrektur des Lehrers und der aktuellen Vorhersageunterstützung des Schülers ab. Eine hohe KL-Divergenz allein ist unzureichend; der Lehrer muss Masse auf Token verteilen, die der Schüler bereits als plausibel betrachtet.
• Praktische Auswirkungen: TA-OPD bietet einen Weg, die Rechenkosten der Distillation drastisch zu senken (durch Filterung inkompatibler Token), während die Modellleistung erhalten bleibt oder verbessert wird. Dies wird ohne den Overhead von Reward-Modellen oder Verifizierern erreicht, was die Methode für Standard-Trainingspipelines anwendbar macht.
• Einschränkungen: Die Autoren weisen darauf hin, dass ihre Analyse sich auf mathematisch geprägtes Reasoning und Modelle der Qwen-Familie konzentriert. Obwohl die Diagnose modellagnostisch ist, sind breitere Validierungen für Dialoge, Code und nicht-Qwen-Backbones erforderlich, um die Generalisierbarkeit zu bestätigen. Zudem bezieht sich das „Token-Budget" auf die Dichte der Überwachung und nicht notwendigerweise auf eine Beschleunigung der Wandzeit, da die Methode Positionen auswählt, anstatt Transformer-Berechnungen zu beschneiden.

Zusammenfassend legt die Arbeit nahe, dass eine effektive Distillation eine kompatibilitätsbewusste Überwachungsstrategie erfordert, bei der der Schüler am besten von Lehrersignalen lernt, die sowohl korrigierend als auch lokal erreichbar sind.