ReNIO : Répondération de l'importance des trajectoires négatives pour la distillation on-policy des LLM

Résumé

La distillation en ligne (on-policy distillation, OPD) améliore le raisonnement des modèles de langue de grande taille (LLM) en entraînant un modèle étudiant sur ses propres sorties générées, mais l'OPD standard traite toutes les sorties générées par l'étudiant (SGO) de manière égale, sans tenir compte de leur caractère informatif. Nous observons une asymétrie constante dans des expériences de filtrage contrôlé : dans l'OPD comme dans l'auto-distillation en ligne (OPSD), l'entraînement uniquement sur les SGO incorrectes surpasse l'entraînement uniquement sur les SGO correctes. Notre analyse approfondie suggère que les modèles entraînés uniquement sur les SGO correctes ont tendance à générer des traces de raisonnement plus courtes et à montrer un comportement de réflexion plus faible, tandis que les SGO incorrectes préservent mieux le raisonnement exploratoire près de la frontière des capacités du modèle. Pour exploiter ce signal sans nécessiter de déploiements complets contenant la réponse, nous introduisons ReNIO (Reweighting Negative trajectory Importance for On-policy distillation des LLM), qui réévalue l'importance des trajectoires négatives pour la distillation en ligne des LLM. En utilisant le rapport de probabilité étudiant-enseignant, ReNIO identifie les jetons pivots menant à des traces de raisonnement erronées et agrège leur information en un poids d'échantillon normalisé, attribuant intrinsèquement des poids plus importants aux trajectoires vraisemblablement négatives sans observer l'exactitude de la réponse finale. Comme ReNIO n'utilise que les probabilités de jetons conditionnées par un préfixe, il préserve l'avantage de l'entraînement par préfixe de l'OPD par rapport à l'apprentissage par renforcement sur déploiements complets. À la fois dans les tâches de raisonnement mathématique et de génération de code, ReNIO améliore l'OPD et l'OPSD, avec des gains relatifs représentatifs allant jusqu'à 8,90 % pour Qwen3-1,7B et 10,00 % pour R1-Distill-Qwen-7B sur les benchmarks de raisonnement mathématique. Dépôt du code : https://github.com/BDML-lab/ReNIO.

English

On-policy distillation (OPD) improves LLM reasoning by training a student model on its own generated outputs, but standard OPD treats all student-generated outputs (SGOs) equally regardless of their informativeness. We observe a consistent asymmetry in controlled filtering experiments: in both OPD and on-policy self distillation (OPSD), training only on incorrect SGOs outperforms training only on correct ones. Our further analysis suggests that models trained on correct-only SGOs tend to generate shorter reasoning traces and show weaker reflection behavior, while incorrect SGOs better preserve exploratory reasoning near the model's capability boundary. To exploit this signal without requiring full answer-containing rollouts, we introduce ReNIO, which Reweights Negative trajectory Importance for LLM On-policy distillation. By using the student-to-teacher probability ratio, ReNIO identifies pivotal tokens leading to wrong reasoning traces and aggregates their information into a normalized sample weight, inherently assigning larger weights to likely negative trajectories without observing the correctness of final-answer. Since Re-NIO only uses prefix-conditioned token probabilities, it preserves OPD's prefix training advantage over full-rollout reinforcement learning. Across both mathematical reasoning and code generation tasks, ReNIO improves both OPD and OPSD, with representative relative gains of up to 8.90% for Qwen3-1.7B and 10.00% for R1-Distill-Qwen-7B on mathematical reasoning benchmarks. Code repo: https://github.com/BDML-lab/ReNIO.