Are Full Rollouts Necessary for On-Policy Distillation?

TL;DR

Proposed horizon-control strategies—Progressive OPD and Truncated OPD—significantly improve long-horizon on-policy distillation efficiency, achieving 3× faster training and comparable performance with only 10% rollout length.

cs.CL 🔴 Advanced 2026-05-30 49 views
Yaocheng Zhang Jiajun Chai Yuqian Fu Songjun Tu Xiaohan Wang Wei Lin Guojun Yin Qichao Zhang Yuanheng Zhu Dongbin Zhao
deep learning knowledge distillation long-horizon reasoning efficiency reinforcement learning

Key Findings

Methodology

This work introduces two horizon-control strategies for OPD: Progressive OPD (POPD), which gradually increases the rollout horizon during training, and Truncated OPD (TOPD), which distills only partial, reliable segments of the rollout. The approach is grounded in theoretical analysis of noise accumulation in long sequences, emphasizing token-level alignment to mitigate future noise propagation. Experiments on mathematical reasoning datasets (AIME24, AMC23) demonstrate that POPD accelerates training by up to 3×, while TOPD achieves similar accuracy using only 10% of the full rollout length, drastically reducing computational resources. The core mechanism involves curriculum learning and selective distillation, focusing on reliable segments to optimize efficiency without sacrificing performance.

Key Results

  • POPD reduces training time by 3× compared to standard OPD, maintaining or improving accuracy on reasoning benchmarks. It effectively avoids noisy late-stage feedback during early training phases.
  • TOPD, with only 10% rollout length, matches the performance of full-sequence OPD on datasets like AIME24, achieving accuracy improvements from 49.3% to 51.3%, while saving over 80% of GPU memory and training time.
  • Analysis shows that local, partial rollouts contain sufficient reasoning signals, and models trained with TOPD can generalize reasoning patterns independently of the full trajectory, validated through prefix-continuation and reverse distillation experiments.

Significance

This research addresses the critical bottleneck of efficiency in long-horizon on-policy distillation, providing practical strategies that significantly reduce training costs while preserving or enhancing model performance. It opens new avenues for deploying large language models in resource-constrained environments and accelerates progress in complex reasoning tasks. The methods are compatible with existing frameworks and can be extended to multi-task and multi-modal scenarios, offering scalable solutions for industry applications.

Technical Contribution

The paper introduces a novel horizon control paradigm, combining curriculum-based expansion and fixed truncation, underpinned by a rigorous theoretical analysis of noise propagation. It innovates by applying token-level alignment to avoid the pitfalls of return-to-go signals in sequence-level methods. The algorithms are simple yet effective, with practical scheduling schemes that adaptively balance exploration and reliability. Extensive experiments validate the approach across reasoning and control tasks, establishing new benchmarks for efficiency and robustness in on-policy distillation.

Novelty

This is the first systematic study applying horizon control—via progressive expansion and truncation—to OPD for long-horizon tasks. Unlike prior work focusing solely on model architecture or reward design, this work emphasizes the importance of inference scope management. The combination of theoretical insights and empirical validation makes it a pioneering contribution, offering a fundamentally new perspective on improving distillation efficiency in sequence modeling.

Limitations

  • While effective in reasoning tasks, the truncated approach may underperform in scenarios requiring long-range dependencies, such as narrative generation or complex dialogue, where full context is crucial.
  • The curriculum schedule and truncation ratios require careful tuning; inappropriate settings can hinder learning or reduce robustness.
  • The current validation is limited to mathematical reasoning and control tasks; applicability to natural language understanding or multimodal tasks remains to be tested.

Future Work

Future directions include developing adaptive horizon scheduling algorithms that dynamically adjust based on task complexity and model confidence. Extending the framework to multi-modal tasks, multi-task learning, and larger models will be explored. Additionally, integrating reinforcement signals or auxiliary objectives could further enhance robustness and generalization, making the approach more versatile across diverse AI applications.

AI Executive Summary

In the realm of large language models and complex reasoning tasks, traditional on-policy distillation (OPD) relies on generating full-length rollouts to align student and teacher policies. While effective, this approach incurs high computational costs and suffers from noise accumulation in late-stage feedback, especially when the student policy diverges from the teacher. These issues hinder scalability and efficiency, limiting practical deployment in resource-constrained environments.

To address these challenges, this study proposes two horizon-control strategies: Progressive OPD (POPD) and Truncated OPD (TOPD). POPD employs a curriculum learning approach, starting with short rollouts and gradually increasing the horizon, allowing the student to first master reliable, short sequences before tackling longer reasoning paths. TOPD, on the other hand, distills only fixed-length, reliable segments of the rollout, significantly reducing resource consumption while maintaining performance.

Experimental results on mathematical reasoning benchmarks such as AIME24 and AMC23 demonstrate that POPD accelerates training by up to three times, achieving comparable or better accuracy than standard OPD. Meanwhile, TOPD, with only 10% of the original rollout length, attains similar accuracy levels, drastically reducing GPU memory and training time by over 80%. These findings confirm that controlling the inference horizon effectively mitigates noise propagation and computational burden.

The core insight is that local, partial rollouts contain sufficient reasoning signals for effective distillation, and models trained with these strategies can learn horizon-independent, generalizable reasoning patterns. Theoretical analysis supports the empirical results, revealing that limiting the rollout horizon prevents the accumulation of unreliable future signals, thus enabling more robust and efficient training.

This work advances the field by offering practical, easy-to-implement methods that significantly improve the scalability of on-policy distillation. Its implications extend to real-world applications such as question-answering, code generation, and autonomous decision-making, where resource efficiency and robustness are paramount. Future research will explore adaptive horizon scheduling, multi-task scenarios, and broader applicability across NLP and multimodal tasks, paving the way for more capable and efficient AI systems.

Deep Dive

Plain Language Accessible to non-experts

想象你在学习一项新技能,比如弹钢琴。刚开始,你不会弹整首复杂的曲子,而是先练习简单的片段。你会反复练习这些短小的部分,直到熟练,然后逐渐把它们连接起来,弹出完整的曲子。这样做可以让你更快掌握技巧,也不会被复杂的部分搞晕。类似的,训练一个智能模型也是这样:不要一开始就让它学习整段长的推理,而是先让它专注于短的推理段,逐步扩展到更长的推理路径。这样,模型可以在有限的时间和计算资源内学到核心的推理能力,就像你学琴一样,从基础开始,逐步变得更厉害。

ELI14 Explained like you're 14

想象你在玩一个超级难的游戏,每次都要玩完整一整关,花费时间又累,又容易出错。其实,如果你只练习游戏的前几部分,熟悉了再挑战更难的关卡,不但省时间,还能学得更快。这就像训练模型一样,传统的方法让它一次性学习整个长的推理过程,但这样很慢,还容易出错。新方法告诉我们:只学一部分内容,等熟悉了再逐步增加难度,效果反而更好。这样,模型可以在短时间内掌握关键技能,变得更聪明、更快。这就像你逐步成为游戏高手一样,先打好基础,再挑战更高难度,最终变得无敌!

Abstract

On-policy distillation (OPD) provides dense teacher feedback along student-generated rollouts rather than fixed teacher traces and has emerged as a promising post-training paradigm. However, standard OPD typically generates full rollouts during training, which is computationally expensive and may expose the student to unreliable teacher feedback at late rollout positions, especially during early training. We identify the rollout horizon as a key bottleneck in OPD that substantially impacts training efficiency. Unlike Reinforcement Learning with Verifiable Rewards (RLVR), OPD does not require a final answer reward to provide learning signals. Therefore, full rollouts may not always be necessary for OPD. Motivated by this insight, we propose two simple horizon-control strategies: Progressive OPD (POPD), which gradually expands the rollout horizon during training, and Truncated OPD (TOPD), which permanently performs distillation on reliable truncated rollouts. Experiments on mathematical reasoning show that POPD improves the training efficiency of OPD by up to 3$\times$, while TOPD matches OPD performance using only 10\% of the rollout horizon, leading to substantial wall-clock and memory reductions. These results demonstrate that controlling the rollout horizon offers a simple and practical path to more efficient OPD.

cs.CL