PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Imitation: Filtering On-Policy Distillation by Reasoning Progress

Chen Yang, Haiyuan Wan, Rengrong Xiong, Yize Chen, Danny H. K. Tsang

arXiv:2608.19408v1cs.AIcs.LG

TL;DR

OPD can reward teacher imitation even when a student’s reasoning advances differently, making teacher-derived supervision an unreliable proxy for progress. R2-OPD estimates progress independently, compares within-trajectory rankings, and masks conflicting distillation rewards; experiments report consistent improvements over standard OPD, while future work targets estimation cost, variance, and broader generalization.

  • Problem

    Teacher-derived distillation rewards do not necessarily reflect whether intermediate reasoning steps advance the solution.

  • Method

    R2-OPD merges sign-consistent reasoning spans, compares process-progress rankings with teacher–student divergence rankings, and masks supervision for conflicting spans.

  • Results

    R2-OPD consistently improves reasoning performance over standard OPD across evaluated settings, with gains varying across benchmarks, model families, datasets, and metrics.

  • Takeaways & Limitations

    Selective filtering preserves teacher guidance where it agrees with reasoning progress while reducing supervision that may discourage productive reasoning.

  • Takeaways & Limitations

    Future work must reduce the cost and variance of process-reward estimation and test generalization to larger models and diverse domains.

Abstract

from arXiv · show

On-policy distillation (OPD) has emerged as an effective framework for post-training language models by pairing student-generated trajectories with dense token-level supervision from a teacher. However, OPD implicitly assumes that teacher-derived rewards are an appropriate proxy for reasoning progress, and therefore treats all teacher feedback equally during policy optimization. While in practice, this assumption does not always hold. We observe that teacher-derived rewards often conflict with genuine reasoning progress, as reasoning steps with clear reasoning advancement may still receive lower distillation rewards, simply due to deviation from teacher's outputs. To address this mismatch, we propose Reasoning-Progress-Aware Reward Filtering for On-Policy Distillation (R2-OPD), which constructs two within-trajectory rankings of reasoning spans, one from teacher-derived rewards and the other from independently estimated progress reward. Distillation rewards are selectively suppressed whenever the two rankings disagree, reducing supervision that conflicts with reasoning progress while preserving effective teacher guidance. Our approach shows consistent improvement over standard OPD especially regarding reasoning performances.

Introduction

Reasoning-oriented OPD can mistake agreement with the teacher for reasoning progress, so R2-OPD filters teacher supervision using independently estimated progress and ranking conflicts.

  • Motivation: OPD rewards student trajectories for matching the teacher, but teacher similarity does not directly measure whether a reasoning span advances the solution.A productive span may depart from the teacher, while a teacher-like span may not make meaningful progress.
  • Motivation: Process rewards provide a teacher-independent surrogate for reasoning progress by estimating how intermediate states change solve probability.Raw process rewards and token-level divergence differ in scale and contain fine-grained noise, complicating direct combination.
  • Method: R2-OPD merges adjacent spans with sign-consistent process rewards, aggregates their divergence, and compares progress rankings with teacher-divergence rankings.Merging reduces dependence on noisy internal boundaries and produces more stable units for calibration.
  • Method: R2-OPD masks distillation rewards for the most inconsistent spans instead of adding process rewards as a separate optimization objective.Teacher guidance is preserved where it agrees with progress while potentially discouraging supervision is reduced.
  • Contributions: The paper reports consistent reasoning-performance improvements over standard OPD by detecting local progress–distillation conflicts and selectively masking unreliable supervision.The contributions also include sign-consistent merging, segment-level divergence averaging, and theoretical analyses of noise reduction.

Preliminaries

OPD trains on student-generated trajectories using reverse KL distillation, while process rewards estimate intermediate reasoning progress from on-policy continuations and solve probabilities.

  • On-Policy Distillation: Reverse-KL OPD minimizes token-level divergence between student and teacher distributions along student-generated trajectories.The full-vocabulary objective weights tokens by the student policy.
  • On-Policy Distillation: Practical OPD restricts vocabulary computations to a support set of H tokens because full-vocabulary KL evaluation is computationally prohibitive.The support may use the highest-probability student tokens, teacher tokens, or their union.
  • Process Supervision: Outcome feedback identifies completed-response correctness but provides limited information about which intermediate steps advance the solution or introduce errors.This motivates finer-grained process supervision for reasoning trajectories.
  • Process Reward Estimation: Process rewards are defined as incremental changes in solve probability across consecutive reasoning states partitioned into contiguous segments.Segments are identified using discourse-marker boundaries and a minimum sentence-count condition.
  • Process Reward Estimation: Solve probability at intermediate boundary states is estimated with on-policy Monte Carlo rollouts using a ground-truth answer and verifiable reward function.The initial state receives zero credit, while the terminal state reuses the actual response’s correctness.
  • Process Reward Estimation: R2-OPD computes the process reward for segment σm as the marginal solve-probability gain, PRm = Ŝm − Ŝm−1.Process-reward computation is skipped when segmentation or the required answer check fails, and such trajectories pass through unfiltered.

Reasoning-Progress-Aware Reward Filtering

R2-OPD stabilizes noisy process and distillation signals by merging consistent spans, compares their within-response rankings, and masks spans whose teacher signal conflicts with progress.

  • Motivation: A reasoning span can improve solve probability while receiving low distillation reward because its token distribution differs from the teacher.Terminal correctness cannot attribute success or failure to individual reasoning spans.
  • Noise Reduction via Sign-Consistent Merging: Fine-grained process rewards and token-level divergences are noisy, so direct comparisons over short spans can frequently produce unreliable conflicts.Noise arises from finite continuation sampling, lexical choices, teacher uncertainty, and top-k approximation.
  • Noise Reduction via Sign-Consistent Merging: Adjacent segments with same-sign nonzero process rewards are merged, while zero rewards remain in the current run and sign changes define new runs.This preserves major reasoning-direction changes while reducing dependence on noisy internal boundaries.
  • Noise Reduction via Sign-Consistent Merging: Summing process rewards over a merged run telescopes, leaving only endpoint solve-probability estimates and endpoint estimation errors.Internal boundary errors cancel because process rewards are consecutive differences.
  • Rank-Based Conflict Detection and Masking: R2-OPD averages token-level distillation losses within each merged segment to compare teacher feedback with segment-level process rewards without systematic dependence on segment length.The resulting average segment-level loss κj measures the segment’s teacher-discrepancy signal.
  • Rank-Based Conflict Detection and Masking: Segment averaging reduces loss variance as O(1/Lj) when token correlations decay with distance, trading temporal resolution for stability.Theoretical analysis links this stabilization to more reliable teacher–student discrepancy estimates over coherent progress intervals.
  • Rank-Based Conflict Detection and Masking: A local conflict occurs when a higher-progress segment receives higher average distillation loss, treating the ordering as a diagnostic rather than an assumption of teacher quality.Within-response rankings avoid requiring process rewards and divergences to share an absolute scale.
  • Rank-Based Conflict Detection and Masking: The algorithm ranks segments by process progress, scores order violations, and masks the highest-inconsistency eligible segments under a masking budget.Masked tokens contribute neither to the loss numerator nor denominator, while the rest of policy optimization remains unchanged.

Experiments

Experiments show that R2-OPD improves reasoning performance across model pairings, with gains depending on benchmark and filtering design. Ablations identify moderate masking and sign-consistent merging as important choices for reliable supervision.

  • Experimental Setup: The primary experiments use DeepSeek-R1-Distill-Qwen-1.5B as student and JustRL-1.5B as teacher, with transfer evaluated using Qwen3-1.7B and e3-1.7B.Evaluation covers AIME 2024, AIME 2025, and OlympiadBench using avg@4 and pass@4.
  • Main Results: R2-OPD achieves 35.06 avg@4 and 51.83 pass@4, exceeding standard OPD by 2.51 and 4.46 points, respectively.It also exceeds Uni-OPD by 4.28 avg@4 and 5.17 pass@4 points, although Uni-OPD is stronger on OlympiadBench.
  • Transfer Results: R2-OPD raises aggregate pass@4 from 45.70 to 48.19 on the Qwen3-1.7B/e3-1.7B pair and improves pass@4 on all three benchmarks.The effective gain varies across datasets and metrics.
  • Sensitivity to the Masking Ratio: A masking ratio of q = 30 performs best across all three benchmarks, reaching 32.50 on AIME 2024, 25.83 on AIME 2025, and 46.86 on OlympiadBench.Lower q retains more progress-conflicting supervision, while higher q substantially degrades AIME performance.
  • Effectiveness of Sign-Consistent Merging: Sign-consistent merging increases smoothed PR–KL rank agreement from approximately 0.20–0.27 to 0.55–0.73 and keeps it above the random baseline for nearly the entire trajectory.The analysis measures average consistent-pair proportions, with 0.5 as the random baseline.
  • Effectiveness of Sign-Consistent Merging: Merging improves avg@4 by 15.00 and 14.16 points on AIME 2024 and AIME 2025, respectively, while decreasing OlympiadBench by 0.96 points.The experiments use the primary DeepSeek-R1-Distill-Qwen-1.5B/JustRL-1.5B configuration with other settings held fixed.

Related works

Related work uses process rewards for finer-grained reasoning credit and studies the reliability of dense OPD supervision. These lines of work motivate more selective supervision because process and teacher signals can be noisy or misaligned.

  • Process Rewards for Reasoning Language Models: Process rewards evaluate intermediate reasoning states and provide finer-grained credit assignment than terminal outcome rewards.Later methods automatically construct process supervision from sampled continuations, outcome verification, or model-based judgments.
  • Reliability and Selectivity of OPD Signals: Prior OPD studies report that large student–teacher discrepancies and teacher preferences for plausible solutions can misdirect exploration.Aggregated token-level guidance can also rank correct trajectories below incorrect ones, motivating calibration against outcome rewards.

Conclusion

The work identifies a limitation of reasoning-oriented OPD and introduces R2-OPD to filter supervision using estimated reasoning progress. Experiments report improved aggregate reasoning performance over standard OPD, with especially strong gains on challenging AIME benchmarks and consistent pass@4 improvements, while broader generalization remains future work.

  • R2-OPD addresses the limitation that teacher agreement does not necessarily indicate whether an intermediate reasoning step advances the solution.
  • R2-OPD estimates progress from on-policy continuations, merges adjacent segments with sign-consistent progress, and compares progress rankings with teacher–student divergence rankings.
  • Rather than adding a separate reward objective, R2-OPD treats ranking conflicts as a reliability test and selectively masks supervision that may discourage productive reasoning.
  • Experiments show improved aggregate reasoning performance over standard OPD, with particularly strong gains on the more challenging AIME benchmarks and consistent improvements in pass@4.
  • Future work will reduce the cost and variance of process-reward estimation and examine generalization to larger models and other diverse domains.

A.1 Proof of Proposition 2

The proposition establishes that the segment-level estimator is unbiased for the locally constant segment-average distillation loss, with a variance bound that decreases inversely with segment length under bounded correlation parameters.

  • κj is an unbiased estimator of the locally constant segment-average distillation loss µj.
  • The variance analysis separates the unchanged constant contribution from diagonal and off-diagonal covariance terms over the contiguous token-index set.
  • The covariance bound uses contiguity to count token pairs by distance and bounds off-diagonal terms through their absolute values.
  • The geometric-series bound applies when 0 ≤ρ < 1, yielding a multiplicative factor σ2(1+ρ)/(1−ρ).
  • When σ2 and ρ are bounded independently of Lj, the variance upper bound scales as O(1/Lj).

B Training Details

The training-details section specifies that R2-OPD follows the standard OPD loop except for progress-aware masking, while separating optimization settings from process-reward estimation and filtering settings.

  • R2-OPD retains the standard OPD training loop except for its progress-aware masking stage.
  • The student generates responses on-policy, the teacher provides token-level distributional supervision on those contexts, and the student is updated with a reverse-KL objective.
  • Table 4 separates OPD-update settings from settings used only for process-reward estimation and segment filtering.

B.2 Training Data

Training uses student-generated responses and verifier-based process-reward rollouts to estimate intermediate solve probabilities and segment progress. Eligibility checks and conservative fallbacks determine when R2-OPD modifies the OPD token mask.

  • B.2 Training Data: All models train for one epoch on the deduplicated DAPO-Math-17K corpus with the same prompt set, and responses are generated on-policy by the current student.
  • B.2 Training Data: Each training example contains a mathematical prompt xi and verifier-compatible ground-truth answer gi.
  • B.3 On-Policy Response Generation: For each prompt xi, the student samples an on-policy response yi with a maximum length of 7,168 tokens.
  • B.3 On-Policy Response Generation: The teacher is queried on the same token contexts to compute support-restricted reverse-KL loss using the 16 highest-probability student tokens.
  • B.4 Model-Specific Reasoning Templates: DeepSeek and Qwen3 configurations use their respective family-specific reasoning templates, with Qwen3 set to thinking mode through enable_thinking=True.
  • B.5 Process-Reward Rollout Construction: Process rewards evaluate whether the student can reach the ground-truth answer from selected intermediate reasoning states using answer-eliciting rollouts.
  • B.5 Process-Reward Rollout Construction: The process reward for segment m is the finite difference PRm = ˆSm −ˆSm−1, where solve probabilities come from successful rollout fractions and terminal verification.
  • B.5 Process-Reward Rollout Construction: Responses lacking enough information for a reliable ranking conflict retain their original OPD supervision through conservative fallback behavior.

B.8 Training Procedure

R2-OPD generates student responses, estimates span-level progress, compares progress and distillation rankings, and selectively masks conflicting supervision before the OPD update.

  • Training Procedure: After filtering, the student receives a normalized OPD update using the remaining teacher supervision.The complete procedure operates at batch level and preserves teacher guidance when masking is not applicable.
  • Training Procedure: The student samples on-policy responses and initializes token-level masks to retain all supervision by default.Responses lacking a verified answer or enough segments bypass masking.
  • Training Procedure: Eligible responses are segmented into reasoning units, evaluated with answer-eliciting rollouts, and assigned progress rewards from successive estimated solution scores.Progress rewards are computed as differences between adjacent cumulative scores.
  • Training Procedure: Adjacent segments with same-sign progress rewards are merged into sign-consistent units before conflict detection.Zero-progress segments are absorbed during merging.
  • Training Procedure: The method ranks units by progress and segment-average KL loss, then identifies disagreements and masks selected units according to the masking budget.Responses with too few eligible units or no conflicts retain their supervision.

C.1 Case 1: Hyperbola-Constrained Rhombus

The hyperbola-constrained rhombus example contrasts incorrect boundary handling in the base model and OPD with R2-OPD’s verified solution of 480.

  • C.1 Case 1: Hyperbola-Constrained Rhombus: The problem asks for the greatest real number less than BD^2 for rhombi on x^2/20 − y^2/24 = 1 with diagonals intersecting at the origin.The reference answer is 480.
  • C.1 Case 1: Hyperbola-Constrained Rhombus: The base model and OPD both return 80 after admitting the boundary value t = 0 without establishing a valid rhombus.Their responses are labeled incorrect.
  • C.1 Case 1: Hyperbola-Constrained Rhombus: R2-OPD returns the verified answer 480 for the hyperbola-constrained rhombus problem.Its derivation is described as preserving the geometric constraints through a limiting argument.

C.2 Case 2: Product over Roots of Unity

The roots-of-unity example shows R2-OPD reaching the verified remainder 321, while the base model and OPD produce incorrect or incomplete outcomes.

  • C.2 Case 2: Product over Roots of Unity: The task asks for the remainder modulo 1000 of a product involving the 13th roots of unity and the factor 2 − 2ω^k + ω^{2k}.The reference answer is 321.
  • C.2 Case 2: Product over Roots of Unity: The base model incorrectly reduces the quadratic factor and returns 191.The erroneous reduction changes the product being evaluated.
  • C.2 Case 2: Product over Roots of Unity: OPD explores multiple algebraic routes but terminates after truncation, leading the evaluator to extract answer 1.Its response is labeled incomplete and incorrect.
  • C.2 Case 2: Product over Roots of Unity: R2-OPD factors the quadratic over 1 ± i, evaluates the conjugate cyclotomic-polynomial terms, and obtains 8321 ≡ 321 modulo 1000.The resulting answer matches the reference answer.
  • C.2 Case 2: Product over Roots of Unity: The example provides qualitative rather than causal evidence, illustrating competing failures to preserve constraints or complete a consistent derivation.R2-OPD reaches the verified answer in this case.
Loading 2608.19408v1…