Revisiter la distillation en politique : modes d'échec empiriques et correctifs simples

Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes

March 26, 2026

Auteurs: Yuqian Fu, Haohuan Huang, Kaiwen Jiang, Yuanheng Zhu, Dongbin Zhao

cs.AI

Résumé

La distillation sur politique (OPD) est attractive pour le post-entraînement des grands modèles de langage (LLM) car elle évalue les retours du modèle enseignant sur des séquences générées par l'élève plutôt que sur des traces fixes de l'enseignant. Cependant, dans des contextes à long horizon, la variante courante par token échantillonné est fragile : elle réduit l'appariement des distributions à un signal à un seul token et devient de plus en plus peu fiable à mesure que les séquences s'éloignent des préfixes que l'enseignant visite couramment. Nous revisitons l'OPD sous les angles de l'estimateur et de l'implémentation. Théoriquement, l'OPD au niveau token est biaisée par rapport à la divergence KL inverse au niveau séquence, mais elle possède une borne de variance bien plus serrée dans le pire des cas ; notre étude jouet montre empiriquement le même compromis, un couplage plus fort des récompenses futures produisant une variance de gradient plus élevée et un apprentissage moins stable. Empiriquement, nous identifions trois modes de défaillance de l'OPD par token échantillonné : un signal à un token déséquilibré, des guidages de l'enseignant peu fiables sur les préfixes générés par l'élève, et des distorsions causées par des incompatibilités de tokeniseur ou de tokens spéciaux. Nous abordons ces problèmes par un appariement local du support top-K de l'enseignant, implémenté comme une KL inverse tronquée avec un échantillonnage de séquences par top-p et un masquage des tokens spéciaux. Que ce soit pour un raisonnement mathématique sur tâche unique ou un entraînement multi-tâches agentique et mathématique, cet objectif permet une optimisation plus stable et de meilleures performances en aval que l'OPD par token échantillonné.

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.