Opnieuw Bekijken van On-Policy Destillatie: Empirische Faalwijzen en Eenvoudige Oplossingen

Samenvatting

On-policy distillatie (OPD) is aantrekkelijk voor post-training van grote taalmodellen (LLM's) omdat het teacher-feedback evalueert op student-gegenereerde rollouts in plaats van vaste teacher-traces. In langetermijnsettings is de veelgebruikte variant met gesampelde tokens echter fragiel: het reduceert distributie-matching tot een één-token signaal en wordt steeds onbetrouwbaarder naarmate rollouts verder afdrijven van prefixes die de teacher vaak bezoekt. Wij herbezien OPD vanuit de estimator- en implementatiekant. Theoretisch gezien is token-level OPD bevooroordeeld ten opzichte van sequence-level reverse-KL, maar het heeft een veel strengere bovengrens voor variantie in het ergste geval; onze kleinschalige studie toont dezelfde afweging empirisch aan, waarbij sterkere koppeling van toekomstige beloning hogere gradientvariantie en instabieler leren veroorzaakt. Empirisch identificeren wij drie faalwijzen van gesampelde-token OPD: een ongebalanceerd één-token signaal, onbetrouwbare teacher-begeleiding op student-gegenereerde prefixes, en vervormingen veroorzaakt door tokenizer- of speciale-token-mismatch. Wij pakken deze problemen aan met teacher top-K lokale ondersteuningsmatching, geïmplementeerd als afgeknotte reverse-KL met top-p rollout sampling en maskering van speciale tokens. Voor zowel enkelvoudige wiskundige redeneertaken als multi-task agentische-plus-wiskunde training levert deze doelstelling stabielere optimalisatie en betere downstreamprestaties op dan gesampelde-token OPD.

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.