Démasquer la distillation on-policy : où elle aide, où elle nuit, et pourquoi

Résumé

La distillation sur politique offre une supervision dense, jeton par jeton, pour l’entraînement des modèles de raisonnement ; cependant, il reste flou dans quelles conditions ce signal est bénéfique et dans lesquelles il est nuisible. Quel modèle enseignant faut-il utiliser, et dans le cas de l’auto-distillation, quel contexte spécifique doit servir de signal de supervision ? Le choix optimal varie-t-il d’un jeton à l’autre ? Actuellement, répondre à ces questions nécessite généralement des cycles d’entraînement coûteux dont les métriques de performance agrégées masquent la dynamique au niveau des jetons individuels. Nous introduisons un cadre diagnostique sans entraînement qui opère à la plus haute résolution : par jeton, par question et par enseignant. Nous dérivons un gradient idéal par nœud, défini comme la mise à jour des paramètres qui maximise l’augmentation de la probabilité de succès de l’étudiant. Nous développons ensuite un algorithme de déroulement ciblé et passant à l’échelle pour estimer ce gradient efficacement, même pour de longues chaînes de pensées intermédiaires. Le score d’alignement du gradient, défini comme la similarité cosinus entre ce gradient idéal et tout gradient de distillation donné, quantifie dans quelle mesure une configuration particulière se rapproche du signal idéal. À travers un ensemble de configurations d’auto-distillation et de modèles enseignants externes, nous observons que le guidage par distillation présente un alignement sensiblement plus élevé avec l’idéal sur les déroulements incorrects que sur les déroulements corrects, où l’étudiant performe déjà bien et où le signal de l’enseignant tend à devenir bruité. De plus, nous constatons que le contexte de distillation optimal dépend conjointement de la capacité du modèle étudiant et de la tâche cible, et qu’aucune configuration universellement efficace n’émerge. Ces résultats motivent l’utilisation d’analyses diagnostiques par tâche et par jeton pour la distillation.

English

On-policy distillation offers dense, per-token supervision for training reasoning models; however, it remains unclear under which conditions this signal is beneficial and under which it is detrimental. Which teacher model should be used, and in the case of self-distillation, which specific context should serve as the supervisory signal? Does the optimal choice vary from one token to the next? At present, addressing these questions typically requires costly training runs whose aggregate performance metrics obscure the dynamics at the level of individual tokens. We introduce a training-free diagnostic framework that operates at the highest resolution: per token, per question, and per teacher. We derive an ideal per-node gradient defined as the parameter update that maximally increases the student's probability of success. We then develop a scalable targeted-rollout algorithm to estimate this gradient efficiently, even for long chains of intermediate thoughts. The gradient alignment score, defined as the cosine similarity between this ideal gradient and any given distillation gradient, quantifies the extent to which a particular configuration approximates the ideal signal. Across a range of self-distillation settings and external teacher models, we observe that distillation guidance exhibits substantially higher alignment with the ideal on incorrect rollouts than on correct ones, where the student already performs well and the teacher's signal tends to become noisy. Furthermore, we find that the optimal distillation context depends jointly on the student model's capacity and the target task, and that no single universally effective configuration emerges. These findings motivate the use of per-task, per-token diagnostic analyses for distillation.