Fast and Effective On-policy Distillation from Reasoning Prefixes

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Zhang, Dongxu, Yang, Zhichao, Janghorbani, Sepehr, Han, Jun, Ressler II, Andrew, Qian, Qian, Lyng, Gregory D., Batra, Sanjit Singh, Tillman, Robert E.
Format: Preprint
Publié: 2026
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866914334119034880
author Zhang, Dongxu
Yang, Zhichao
Janghorbani, Sepehr
Han, Jun
Ressler II, Andrew
Qian, Qian
Lyng, Gregory D.
Batra, Sanjit Singh
Tillman, Robert E.
author_facet Zhang, Dongxu
Yang, Zhichao
Janghorbani, Sepehr
Han, Jun
Ressler II, Andrew
Qian, Qian
Lyng, Gregory D.
Batra, Sanjit Singh
Tillman, Robert E.
contents On-policy distillation (OPD), which samples trajectories from the student model and supervises them with a teacher at the token level, avoids relying solely on verifiable terminal rewards and can yield better generalization than off-policy distillation. However, OPD requires expensive on-the-fly sampling of the student policy during training, which substantially increases training cost, especially for long responses. Our initial analysis shows that, during OPD, training signals are often concentrated in the prefix of each output, and that even a short teacher-generated prefix can significantly help the student produce the correct answer. Motivated by these observations, we propose a simple yet effective modification of OPD: we apply the distillation objective only to prefixes of student-generated outputs and terminate each sampling early during distillation. Experiments on a suite of AI-for-Math and out-of-domain benchmarks show that on-policy prefix distillation matches the performance of full OPD while reducing training FLOP by 2x-47x.
format Preprint
id arxiv_https___arxiv_org_abs_2602_15260
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Fast and Effective On-policy Distillation from Reasoning Prefixes
Zhang, Dongxu
Yang, Zhichao
Janghorbani, Sepehr
Han, Jun
Ressler II, Andrew
Qian, Qian
Lyng, Gregory D.
Batra, Sanjit Singh
Tillman, Robert E.
Machine Learning
Artificial Intelligence
On-policy distillation (OPD), which samples trajectories from the student model and supervises them with a teacher at the token level, avoids relying solely on verifiable terminal rewards and can yield better generalization than off-policy distillation. However, OPD requires expensive on-the-fly sampling of the student policy during training, which substantially increases training cost, especially for long responses. Our initial analysis shows that, during OPD, training signals are often concentrated in the prefix of each output, and that even a short teacher-generated prefix can significantly help the student produce the correct answer. Motivated by these observations, we propose a simple yet effective modification of OPD: we apply the distillation objective only to prefixes of student-generated outputs and terminate each sampling early during distillation. Experiments on a suite of AI-for-Math and out-of-domain benchmarks show that on-policy prefix distillation matches the performance of full OPD while reducing training FLOP by 2x-47x.
title Fast and Effective On-policy Distillation from Reasoning Prefixes
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2602.15260