ReNIO: Reponderação da Importância de Trajetórias Negativas para Destilação On-Policy de LLMs

Resumo

A destilação on-policy (OPD) melhora o raciocínio de LLMs ao treinar um modelo aluno em suas próprias saídas geradas, mas a OPD padrão trata todas as saídas geradas pelo aluno (SGAs) igualmente, independentemente de sua informatividade. Observamos uma assimetria consistente em experimentos controlados de filtragem: tanto na OPD quanto na autodestilação on-policy (OPSD), treinar apenas nas SGAs incorretas supera treinar apenas nas corretas. Nossa análise adicional sugere que modelos treinados apenas em SGAs corretas tendem a gerar traços de raciocínio mais curtos e exibir comportamento de reflexão mais fraco, enquanto as SGAs incorretas preservam melhor o raciocínio exploratório próximo ao limite de capacidade do modelo. Para explorar esse sinal sem exigir rollouts completos contendo respostas, introduzimos o ReNIO, que Reatribui Peso à Importância de Trajetórias Negativas para Destilação On-policy de LLMs. Ao usar a razão de probabilidades aluno-professor, o ReNIO identifica tokens pivô que levam a traços de raciocínio errados e agrega suas informações em um peso amostral normalizado, atribuindo inerentemente pesos maiores a trajetórias provavelmente negativas sem observar a correção da resposta final. Como o ReNIO utiliza apenas probabilidades de tokens condicionadas ao prefixo, ele preserva a vantagem do treinamento com prefixo da OPD sobre o aprendizado por reforço com rollout completo. Tanto em tarefas de raciocínio matemático quanto de geração de código, o ReNIO melhora tanto a OPD quanto a OPSD, com ganhos relativos representativos de até 8,90% para Qwen3-1.7B e 10,00% para R1-Distill-Qwen-7B em benchmarks de raciocínio matemático. Repositório de código: 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.