Prefix-Guided On-Policy Distillation: Mining Golden Trajectories from Rollouts

2026-06-20

AI Summary

This work addresses the inefficiency of traditional on-policy distillation in long-horizon mathematical reasoning, where uniform allocation of computational budget leads to suboptimal performance and introduces noisy supervision signals. The authors propose a prefix-guided distillation framework that leverages, for the first time, the top-k overlap between teacher and student models over fixed-length prefixes as a trajectory value metric. Only trajectories exhibiting high compatibility are extended to full sequences, enabling highly efficient allocation of computational resources. Evaluated on the AMC, AIME, and HMMT benchmarks, the method achieves up to a 4.80-point improvement in average accuracy while reducing training time by up to 2.46×, significantly enhancing training efficiency without compromising model performance.

Abstract

On-policy distillation (OPD) improves reasoning models by applying dense teacher supervision on student-sampled trajectories. However, scaling OPD to long-horizon mathematical reasoning exposes a reliability and efficiency problem: standard OPD assigns every sampled candidate the same long rollout budget, even though some trajectories may quickly become weakly aligned with the teacher and provide less useful supervision. Prior analyses suggest that successful OPD depends on local teacher-student compatibility, which can be measured by top-k overlap on student-visited prefixes. When this overlap is low, continuing to generate or train on long suffixes may waste computation and introduce noisy learning signal. To address this, we introduce Prefix-Guided On-Policy Distillation (PG-OPD), a simple rollout-allocation framework that uses fixed-length prefixes to estimate trajectory value before expensive long-horizon generation. PG-OPD first decodes every sampled candidate to the same prefix length, computes teacher-student top-k overlap within an early probe window of that prefix, and selectively continues high-overlap candidates to a fixed long length. Low-overlap candidates stop at the fixed prefix, avoiding unnecessary suffix generation. Across diverse teacher-student combinations on AMC, AIME, and HMMT benchmarks, PG-OPD improves average accuracy by up to 4.80 points while reducing training time by up to 2.46x. These results suggest that prefix-level compatibility provides a practical signal for directing OPD computation toward trajectories that remain learnable from the teacher.

Problem

• On-policy distillation
• long-horizon reasoning
• trajectory efficiency
• teacher-student alignment
• rollout allocation

Innovation

• on-policy distillation
• prefix-guided
• trajectory selection
• teacher-student compatibility
• rollout allocation