PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling Reasoning Efficiently via Relaxed On-Policy Distillation

Jongwoo Ko, Sara Abdali, Young Jin Kim, Tianyi Chen, Pashmina Cameron

arXiv:2603.11137v1cs.LGcs.CL

TL;DR

Capacity-constrained models need alternatives to direct reward optimization, while on-policy distillation remains vulnerable to instability and negative transfer. The paper interprets distillation as policy optimization through a teacher–student log-likelihood reward and introduces REOPOLD, which regulates those signals. REOPOLD improves sample efficiency and test-time scaling across mathematical, visual, and agentic reasoning tasks.

  • Problem

    Capacity-constrained models struggle with direct reward optimization, while on-policy distillation remains prone to instability and negative transfer.

  • Method

    REOPOLD treats the teacher–student log-likelihood ratio as a fixed reward and stabilizes distillation with reward clipping, token-level dynamic sampling, and unified exploration-to-refinement training.

  • Results

    REOPOLD achieves superior sample efficiency and performance across mathematical, visual, and agentic reasoning, including up to 3.32× inference efficiency in visual reasoning.

  • Takeaways & Limitations

    Relaxing strict imitation enables a more stable distillation framework that supports efficient training and test-time scaling for smaller reasoning models.

  • Takeaways & Limitations

    Rapid entropy collapse creates a trade-off between exploration and teacher–student alignment, and higher sampling temperatures can exacerbate reward variance.

Abstract

from arXiv · show

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.

1. Introduction

Small language models struggle with direct reward optimization, motivating alternative reasoning-transfer mechanisms. The paper frames on-policy distillation as policy optimization and proposes REOPOLD to regulate teacher signals for more stable, efficient learning.

  • Limited representational capacity makes direct reward optimization ineffective for small language models, motivating alternative mechanisms for transferring reasoning abilities.
  • On-policy distillation trains students on their own trajectories under teacher guidance, but its methodological development lags behind modern reinforcement learning.
  • Vanilla on-policy distillation exhibits instability, negative transfer, rapid entropy collapse, and practical restrictions on teacher–student model-size differences.
  • REOPOLD interprets the teacher–student log-likelihood ratio as a fixed reward and uses stop-gradient to connect distillation with policy-gradient optimization.
  • REOPOLD filters harmful signals through temperate and selective teacher guidance, reward clipping, token-level dynamic sampling, and multi-stage training.
  • REOPOLD achieves superior training sample efficiency and state-of-the-art performance across mathematical, visual, and tool-use reasoning tasks.

2. Background and Related Work

Reasoning research uses policy optimization with sampled trajectories and estimated advantages, while on-policy distillation minimizes reverse KL on student-generated trajectories. Prior work reports that distillation can outperform RL with substantially lower computational cost.

  • Policy optimization samples response groups from a previous policy and updates parameters using token-level importance ratios and estimated advantages.
  • Reasoning RL methods commonly use scalable final-answer correctness rewards, alongside refinements involving normalization, clipping, KL regularization, and adaptive sampling.
  • Figure 3 compares vanilla RKL with stop-gradient RKL through training loss, gradient norm, and validation performance.
  • On-policy distillation trains on student-sampled trajectories and minimizes reverse KL, thereby addressing exposure bias associated with static teacher outputs.
  • 10× reduction in GPU hours is reported for on-policy distillation while surpassing RL in reasoning performance.

3. Analysis of On-Policy Distillation

The analysis establishes on-policy distillation as an on-policy policy-gradient method with teacher–student log-likelihood ratios serving as token rewards. It then identifies finite-sample variance, heavy-tailed negative rewards, near-zero rewards, and entropy collapse as key optimization challenges.

  • 3.1. Theoretical Equivalence to RL and Strong Baseline: On-policy distillation is formally equivalent to on-policy policy optimization when the teacher-defined term is treated as a token reward.The reward-term gradient vanishes in expectation under standard regularity conditions.
  • 3.1. Theoretical Equivalence to RL and Strong Baseline: Stop-gradient acts as a control variate that suppresses high-variance gradient noise and produces more stable training dynamics.The paper links reduced gradient norms with variance reduction and better validation performance.
  • 3.2. Optimization Challenges in On-policy Distillation: Heavy-tailed negative rewards arise when sampled tokens receive negligible teacher probability, producing extreme log-likelihood ratios that destabilize optimization.These rewards can dominate gradient estimates and drive destructive updates that move the student away from its original distribution.
  • 3.2. Optimization Challenges in On-policy Distillation: Near-zero rewards dilute the learning signal because most aligned student-teacher tokens provide negligible advantages while consuming computation.At the token level, this reduces effective sample size and worsens sample efficiency.
  • 3.2. Optimization Challenges in On-policy Distillation: Entropy decreases rapidly during training, causing premature convergence while higher sampling temperatures worsen student-teacher reward variance.The resulting trade-off is between exploring diverse reasoning paths and maintaining alignment with the teacher.

4. REOPOLD: Relaxed On-Policy Distillation for Compact Reasoning Models

REOPOLD stabilizes on-policy distillation by treating the teacher–student log-likelihood ratio as a reward and relaxing imitation through clipped, entropy-selective, multi-stage optimization. Its unified objective combines mixture-based reward clipping, high-entropy token sampling, and exploration-to-refinement training to preserve diversity while improving reasoning-signal efficiency.

  • Unified framework: REOPOLD modernizes on-policy distillation with a unified dynamic objective that adapts the learning signal across training stages.The framework combines reward clipping, token-level dynamic sampling, and multi-stage training.
  • Reward clipping: Mixture-based clipping replaces divergent negative rewards with a finite penalty floor while selectively targeting outliers and preserving RKL’s mode-seeking behavior.The mixture term converges to a finite constant, preventing the gradient explosion associated with standard RKL.
  • Entropy-guided sampling: High-entropy token selection filters low-information tokens whose rewards concentrate near zero, mitigating gradient dilution and accelerating convergence.The mask retains tokens above the batch’s top β-percentile entropy threshold and focuses learning where teacher–student divergence is largest.
  • Exploration-to-refinement training: The multi-stage strategy outperforms the baseline in both quality, measured by Avg@32, and diversity, measured by Pass@32.Figure 6 reports that the proposed strategy prevents entropy collapse during exploration and achieves superior final performance.
  • Exploration-to-refinement training: The exploration phase masks strongly negative rewards to preserve diverse teacher-aligned modes, while refinement reintroduces negative feedback on high-entropy tokens.The two phases respectively resemble SFT-style positive reinforcement and RL-style discrimination among uncertain reasoning paths.

5. Experimental Results

Across mathematical, visual, and agentic reasoning evaluations, REOPOLD improves training efficiency, robustness, scaling, and component-level performance over distillation and RL baselines. It also stabilizes larger-student training and achieves stronger visual test-time scaling.

  • Mathematical reasoning: REOPOLD matches ProRL in 600 steps versus 2000, yielding a > 6.7× sample-efficiency gain, and exceeds other baselines at 300 steps with > 12× efficiency.It also outperforms vanilla RKL in 150 steps and surpasses GRPO under identical conditions.
  • Mathematical reasoning: REOPOLD consistently outperforms SFT across mathematical benchmarks and remains effective regardless of the teacher model.Vanilla RKL shows negligible improvements with the SkyWork-OR1-7B teacher.
  • Mathematical reasoning: With a 7B student, REOPOLD prevents the sharp early performance drop and failed improvement beyond the base model observed for vanilla RKL.Figure 7 reports stable and consistent validation improvement for REOPOLD on AIME.
  • Visual reasoning: REOPOLD surpasses GRPO and RKL across visual reasoning and perception benchmarks for both 3B and 7B models.It also exceeds NoisyRollout and PAPO, with extended experiments confirming consistent superiority over vanilla RKL.
  • Visual reasoning: REOPOLD yields uniform accuracy improvements over RKL across six visual benchmarks, averaging 1.34% with a 7B teacher and 0.93% with a 32B teacher.Longer training improves both the baseline and REOPOLD, while REOPOLD retains its advantage.
  • Analysis: Stop-gradient, reward clipping, token-level dynamic sampling, and multi-stage training each improve performance, with the complete pipeline achieving the best overall results.A stricter dynamic-sampling threshold, β = 0.2, outperforms β = 0.5 by filtering low-entropy tokens and reducing gradient dilution.
  • Agentic reasoning with visual tool-use: REOPOLD outperforms vanilla RKL and GRPO on agentic visual tool-use benchmarks and shows better sample efficiency than both.It slightly trails GRPO on the Pixel test split but performs better across other benchmarks and can be applied without intricate reward engineering.

6. Conclusion

The paper frames on-policy distillation as policy optimization with inherited instability and introduces REOPOLD to relax strict imitation. Across mathematical, visual, and agentic reasoning, the framework reports greater stability and stronger performance than recent RL algorithms.

  • Conclusion: On-policy distillation is theoretically equivalent to policy optimization and therefore inherits fundamental optimization instabilities.The paper states that standard RL solutions do not effectively resolve these instabilities.
  • Conclusion: REOPOLD stabilizes training through mixture-based reward clipping, token-level dynamic sampling, and a unified exploration-to-refinement strategy.These mechanisms replace rigid imitation with more flexible teacher-signal use.
  • Conclusion: REOPOLD achieves superior performance across mathematical, visual, and agentic reasoning compared with recent RL algorithms.The conclusion identifies relaxing strict imitation as essential for scaling reasoning capabilities in compact models.

A.1. RKL

The RKL derivation requires regularity and support-coverage assumptions, then expresses its gradient in policy-gradient form. The resulting advantage-like signal is the teacher-related log-likelihood ratio weighted by importance sampling.

  • Assumptions: The derivation assumes continuous differentiability and bounded gradients sufficient to interchange gradients and expectations.These conditions invoke the Leibniz integral rule and the Dominated Convergence Theorem.
  • Assumptions: Absolute continuity requires πθ ≪ πθold, ensuring importance-sampling and log-likelihood ratios are finite almost everywhere.If πθold assigns zero probability, πθ must also assign zero probability.
  • Derivation: The RKL objective and gradient are derived using techniques analogous to PPO or GRPO.The derivation introduces the importance-sampling ratio ρt and teacher-related reward Rt.
  • Interpretation: Optimizing RKL is mathematically equivalent to maximizing a standard policy-gradient objective with Rt as the advantage and ρt as its weight.This connects on-policy distillation directly to policy optimization.

A.2. Derivation of Clipping Threshold

The derivation bounds the teacher–student log-likelihood ratio using a convex mixture of teacher and student policies, motivating selective clipping rather than changing the global objective. Empirical comparisons examine clipping behavior and robustness relative to skew RKL.

  • Derivation: Jensen’s inequality upper-bounds the teacher–student log-ratio by the log-ratio induced by their convex mixture.The mixture uses weight λ for the student policy and 1−λ for the teacher policy.
  • Clipping design: The resulting mixture-based bound is used as a clipping threshold for the log-likelihood ratio.The approach applies the bound selectively instead of modifying the global divergence objective.
  • Comparison with ρ(θ) clipping: ρ(θ) clipping clips fewer than 0.2% of samples after the initial steps throughout training.This indicates limited deviation from the behavior policy in the examined setting.
  • Comparison with skew RKL: REOPOLD remains robust to λ while skew RKL drops significantly at λ = 0.7.The passage reports that REOPOLD’s lowest accuracy exceeds skew RKL’s peak performance.
  • Context: The surrounding literature frames reasoning models and policy optimization around structured inference, sample efficiency, stability, and reasoning depth.On-policy distillation is presented as aligning training trajectories with the student’s current policy to reduce distribution mismatch.

C. Detailed Experimental Setup

The experiments evaluate REOPOLD across mathematical, visual, and agentic visual tool-use reasoning with specified rollout, optimization, and batch configurations. Extended visual results compare accuracy–latency trade-offs against the RKL baseline.

  • Hyperparameters: The REOPOLD experiments use hyperparameter settings summarized in Table 8.The supplied passage identifies the table but does not enumerate its values.
  • Mathematical reasoning: Mathematical distillation samples 8 responses per prompt, permits 8192-token responses, and trains for 300 iterations.Training uses batch size 128, mini-batch size 32, four gradient updates per rollout step, and AdamW at 1 × 10^-5.
  • Visual reasoning: Visual distillation samples 12 responses per prompt, permits 2048-token responses, and trains for 60 iterations.The setup uses batch size 128, mini-batch size 64, two gradient updates per step, and AdamW.
  • Agentic visual tool-use: Agentic visual tool-use distillation samples 8 trajectories per prompt with an 8192-token limit and a maximum of 2 rounds.Evaluation uses the InfoVQA, TallyQA, and VStar datasets.
  • Test-time evaluation: On Geometry3K and MathVerse, REOPOLD’s accuracy–latency trade-off is compared with vanilla RKL as the sample budget K increases.The visualization reports Pass@K and Maj@K, with K up to 64 for Geometry3K and 16 for MathVerse.

D. Additional Analyses and Discussions

Additional analyses assess REOPOLD’s test-time scaling, training-time overhead, and full-vocabulary distillation alternatives. The results show consistent consensus advantages, moderate teacher cost, and no meaningful benefit from full-vocabulary or Top-5 approximations.

  • Extended test-time scaling: Both 3B and 7B REOPOLD models maintain a consistent Maj@K lead on Geometry3K and MathVerse.Maj@K measures consensus robustness, extending the reported advantage beyond Pass@K coverage.
  • Training time analysis: The teacher accounts for 8%–22% of each training step, while student generation remains the primary computational bottleneck.Student generation accounts for 77.3% in 8192-token math tasks, and the teacher accounts for 21.8% in 2048-token visual tasks.
  • Full-vocabulary distillation: Full-vocabulary GKD incurs out-of-memory issues with a 150K vocabulary, and Top-5 approximation lowers efficiency relative to sampled-token approaches.Applying full vocabulary or Top-5 approximation to REOPOLD also yields no meaningful improvement.

E. Qualitative Evaluation

The qualitative evaluation contrasts vanilla RKL with REOPOLD on visual grounding and mathematical self-correction examples. REOPOLD preserves chart perception in one example and corrects an initially incorrect geometric derivation in another.

  • Evaluation examples: The evaluated examples include a UK bread-export yes-or-no question and a geometry question asking for angle AFC.The supplied passages provide the prompts but do not state their answers.
  • Visual reasoning: On Hallusion Bench, vanilla RKL hallucinates a flat chart trend, whereas REOPOLD identifies the visual peak.The comparison presents this as a difference in visual grounding robustness.
  • Mathematical reasoning: On MathVerse, vanilla RKL becomes trapped in circular logic and hallucinates the final calculation.The cited example identifies failures at Step 5 and Step 8.
  • Mathematical reasoning: REOPOLD initially derives 96◦ but triggers self-correction and reaches the correct 84◦ solution.The example explicitly includes a re-evaluation step before the corrected answer.
Loading 2603.11137v1…