Scaling Reasoning Efficiently via Relaxed On-Policy Distillation
REOPOLD enhances reasoning model training via reward clipping and token-level dynamic sampling, achieving 6.7× sample efficiency and 3.32× inference speedup.
Key Findings
Methodology
This work reformulates on-policy distillation as a form of policy optimization, interpreting the teacher-student log-likelihood ratio as a token reward. Based on this, REOPOLD employs reward clipping, entropy-guided token sampling, and multi-stage training to stabilize optimization. Specifically, it uses mixture-based reward clipping to mitigate heavy-tailed negative rewards, token-level entropy thresholds to focus on high-information tokens, and a phased training approach balancing exploration and refinement. These components collectively improve training stability and inference scalability, demonstrated across mathematical, visual, and tool-use reasoning benchmarks.
Key Results
- REOPOLD achieves a 6.7× reduction in training samples on the AIME-25 benchmark, significantly outperforming baseline methods. In visual reasoning tasks, a 7B student model matches the performance of a 32B teacher while accelerating inference by 3.32×. The approach reduces gradient variance, avoids negative transfer, and enhances convergence stability, outperforming standard RL and distillation techniques in accuracy and efficiency.
- Across multiple reasoning datasets, REOPOLD improves accuracy by 5-10 percentage points over SFT and vanilla RKL. Reward clipping suppresses extreme negative rewards, while entropy-based sampling emphasizes high-information tokens, resulting in more robust and sample-efficient training. Ablation studies confirm the importance of each component, with combined methods yielding the best results.
- The experiments demonstrate that the proposed framework not only boosts sample efficiency but also enhances model robustness and scalability, enabling smaller models to approach larger teacher performance with fewer samples and faster inference, thus opening new avenues for deploying lightweight yet capable reasoning models.
Significance
This research addresses the core challenge of training stable, efficient reasoning models with limited capacity. By integrating reinforcement learning principles into distillation, it overcomes instability issues like gradient explosion and negative transfer. The approach significantly reduces training costs and accelerates inference, making high-performance reasoning accessible for resource-constrained environments. Its success in mathematical, visual, and tool reasoning tasks indicates broad applicability, promising to influence future model compression, deployment, and real-world AI systems. Theoretically, it bridges the gap between RL and distillation, providing a new perspective on model optimization under capacity constraints.
Technical Contribution
The paper introduces a novel formulation of on-policy distillation as a policy gradient method with fixed rewards, stabilized via stop-gradient operations. It proposes reward clipping based on mixture distributions to prevent gradient explosion, and token-level entropy-guided sampling to focus learning on uncertain regions. The multi-stage training strategy balances exploration and refinement, ensuring convergence stability. These innovations collectively improve sample efficiency and inference scalability, surpassing existing RL and distillation methods, and enabling smaller models to effectively learn complex reasoning skills.
Novelty
This is the first work to systematically connect on-policy distillation with reinforcement learning optimization techniques, particularly reward clipping and token-level adaptive sampling. Unlike prior methods that treat distillation as mere imitation, this approach explicitly regulates learning signals through reward manipulation, leading to stable, efficient training. The combination of theoretical insights and practical algorithms marks a significant advancement in reasoning model training, setting a new standard for capacity-constrained AI systems.
Limitations
- The method relies on preset thresholds for reward clipping and entropy, which may require careful tuning across different tasks. In highly complex or noisy environments, these parameters might limit generalization.
- Multi-stage training introduces additional complexity and hyperparameter sensitivity, potentially hindering straightforward deployment in diverse scenarios. Further automation of parameter selection is needed.
- While effective on tested benchmarks, the approach's performance on extremely large models or real-world noisy data remains to be validated. Future work should explore adaptive mechanisms for broader robustness.
Future Work
Future directions include developing adaptive parameter tuning strategies for reward clipping and entropy thresholds, enabling more robust generalization. Extending the framework to multi-modal reasoning and real-world noisy data will broaden its applicability. Additionally, integrating meta-learning techniques could automate hyperparameter selection, further improving training stability and efficiency. Exploring unsupervised or semi-supervised variants may also reduce reliance on high-quality teacher models, making the approach more scalable and accessible.
AI Executive Summary
Training reasoning models that are both accurate and efficient remains a significant challenge in AI research. Existing methods like supervised fine-tuning and standard knowledge distillation often struggle with instability, especially when models are constrained in capacity. Reinforcement learning approaches have shown promise but are computationally expensive and prone to gradient instability. This paper introduces REOPOLD, a novel framework that bridges the gap between distillation and reinforcement learning, leveraging reward clipping, token-level entropy-guided sampling, and multi-stage training to stabilize and accelerate learning.
The core idea is to interpret the teacher-student log-likelihood ratio as a token reward within a policy optimization paradigm. By employing mixture-based reward clipping, the method prevents extreme negative rewards that cause gradient explosion. Simultaneously, entropy-guided sampling focuses training on high-information tokens, improving sample efficiency. The multi-stage training process balances exploration and refinement, avoiding premature convergence and entropy collapse.
Experimental results across mathematical, visual, and tool reasoning benchmarks demonstrate that REOPOLD significantly outperforms traditional distillation and RL approaches. It achieves a 6.7× reduction in training samples on the challenging AIME-25 benchmark and enables a 7B model to match a 32B teacher's performance in visual reasoning, with a 3.32× inference speedup. These advances suggest a new pathway for deploying lightweight models capable of complex reasoning tasks.
Overall, REOPOLD offers a robust, scalable solution to the longstanding instability issues in training capacity-constrained reasoning models. Its innovative integration of RL principles into distillation not only enhances efficiency but also opens new avenues for future research in AI model compression, multi-modal reasoning, and real-world deployment. While some hyperparameter tuning remains necessary, the framework's flexibility and effectiveness mark a significant step forward in the field.
Deep Dive
Abstract
On-policy distillation is pivotal for transferring reasoning capabilities to capacity-constrained models, yet remains prone to instability and negative transfer. We show that on-policy distillation can be interpreted, both theoretically and empirically, as a form of policy optimization, where the teacher-student log-likelihood ratio acts as a token reward. From this insight, we introduce REOPOLD (Relaxed On-Policy Distillation) a framework that stabilizes optimization by relaxing the strict imitation constraints of standard on-policy distillation. Specifically, REOPOLD temperately and selectively leverages rewards from the teacher through mixture-based reward clipping, entropy-based token-level dynamic sampling, and a unified exploration-to-refinement training strategy. Empirically, REOPOLD surpasses its baselines with superior sample efficiency during training and enhanced test-time scaling at inference, across mathematical, visual, and agentic tool-use reasoning tasks. Specifically, REOPOLD outperforms recent RL approaches achieving 6.7~12x greater sample efficiency and enables a 7B student to match a 32B teacher in visual reasoning with a ~3.32x inference speedup.