Desvendando a destilação on-policy: onde ajuda, onde prejudica e por quê

Resumo

A destilação on-policy oferece supervisão densa por token para treinar modelos de raciocínio; no entanto, ainda não está claro sob quais condições esse sinal é benéfico e sob quais é prejudicial. Qual modelo professor deve ser utilizado e, no caso de auto-destilação, qual contexto específico deve servir como sinal de supervisão? A escolha ótima varia de um token para o próximo? Atualmente, abordar essas questões geralmente requer execuções de treinamento custosas, cujas métricas de desempenho agregadas obscurecem a dinâmica no nível de tokens individuais. Introduzimos uma estrutura de diagnóstico livre de treinamento que opera na mais alta resolução: por token, por questão e por professor. Derivamos um gradiente ideal por nó, definido como a atualização de parâmetros que maximiza o aumento da probabilidade de sucesso do estudante. Em seguida, desenvolvemos um algoritmo escalável de rollout direcionado para estimar esse gradiente de forma eficiente, mesmo para longas cadeias de pensamentos intermediários. A pontuação de alinhamento de gradiente, definida como a similaridade de cosseno entre esse gradiente ideal e qualquer gradiente de destilação dado, quantifica o quanto uma configuração específica se aproxima do sinal ideal. Em uma gama de configurações de auto-destilação e modelos professores externos, observamos que a orientação por destilação exibe um alinhamento substancialmente maior com o ideal em rollouts incorretos do que em corretos, onde o estudante já tem um bom desempenho e o sinal do professor tende a se tornar ruidoso. Além disso, constatamos que o contexto ótimo de destilação depende conjuntamente da capacidade do modelo estudante e da tarefa alvo, e que não emerge uma única configuração universalmente eficaz. Essas descobertas motivam o uso de análises de diagnóstico por tarefa e por token para destilação.

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.