Neubewertung von On-Policy-Distillation: Empirische Fehlermodi und einfache Lösungsansätze

Zusammenfassung

On-Policy-Distillation (OPD) ist für das Post-Training großer Sprachmodelle (LLMs) attraktiv, da es Lehrerfeedback auf studentengenerierten Rollouts und nicht auf festen Lehrer-Traces auswertet. In Langzeitszenarien erweist sich jedoch die gängige Variante mit Stichprobentoken als instabil: Sie reduziert Distribution Matching auf ein Ein-Token-Signal und wird zunehmend unzuverlässig, sobald sich Rollouts von Präfixen entfernen, die der Lehrer häufig besucht. Wir untersuchen OPD neu aus der Perspektive des Schätzers und der Implementierung. Theoretisch ist Token-Level-OPD gegenüber Sequence-Level-Reverse-KL verzerrt, weist jedoch eine deutlich engere Varianzschranke im Worst Case auf; unsere Spielzeugstudie zeigt denselben Tradeoff empirisch, wobei eine stärkere Kopplung zukünftiger Belohnungen zu höherer Gradientenvarianz und weniger stabilem Lernen führt. Empirisch identifizieren wir drei Fehlermodi von Stichprobentoken-OPD: ein unausgeglichenes Ein-Token-Signal, unzuverlässige Lehrerführung auf studentengenerierten Präfixen sowie Verzerrungen durch Tokenizer- oder Sonder-Token-Fehlanpassung. Wir adressieren diese Probleme durch Teacher-Top-K-Local-Support-Matching, implementiert als truncated reverse-KL mit Top-p-Rollout-Sampling und Sonder-Token-Masking. Sowohl beim mathematischen Einzelaufgaben-Reasoning als auch im multi-task agentischen und mathematischen Training führt dieses Ziel im Vergleich zu Stichprobentoken-OPD zu stabilerer Optimierung und besserer Downstream-Leistung.

English

On-policy distillation (OPD) is appealing for large language model (LLM) post-training because it evaluates teacher feedback on student-generated rollouts rather than fixed teacher traces. In long-horizon settings, however, the common sampled-token variant is fragile: it reduces distribution matching to a one-token signal and becomes increasingly unreliable as rollouts drift away from prefixes the teacher commonly visits. We revisit OPD from the estimator and implementation sides. Theoretically, token-level OPD is biased relative to sequence-level reverse-KL, but it has a much tighter worst-case variance bound; our toy study shows the same tradeoff empirically, with stronger future-reward coupling producing higher gradient variance and less stable learning. Empirically, we identify three failure modes of sampled-token OPD: an imbalanced one-token signal, unreliable teacher guidance on student-generated prefixes, and distortions caused by tokenizer or special-token mismatch. We address these issues with teacher top-K local support matching, implemented as truncated reverse-KL with top-p rollout sampling and special-token masking. Across single-task math reasoning and multi-task agentic-plus-math training, this objective yields more stable optimization and better downstream performance than sampled-token OPD.