Beyond Imitation: Filtering On-Policy Distillation by Reasoning Progress

Abstract: On-policy distillation (OPD) has emerged as an effective framework for post-training LLMs 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.

• The paper introduces R²-OPD, which estimates segment-level solve-probability changes with Monte Carlo rollouts and masks reasoning spans where progress conflicts with reverse-KL distillation signals.
• R²-OPD improves aggregate performance over standard OPD by 2.51 avg@4 and 4.46 pass@4 points, reaching 35.06 avg@4 and 51.83 pass@4 across AIME 2024, AIME 2025, and OlympiadBench.
• The method works best with moderate filtering: a 30% masking budget improves results, while 50% masking removes useful teacher guidance, and its rollout cost, English segmentation, and benchmark dependence limit broader deployment.

Motivation and problem statement

On-policy distillation (OPD) trains a student on its own sampled trajectories under dense token-level reverse-KL supervision from a teacher, avoiding the exposure bias of static, teacher-generated distillation datasets [2402.xxxx-style setup; see ICLR2024_5be69a58]. The framework rests on an implicit assumption: that token-level agreement with the teacher is a valid proxy for reasoning quality. The paper's central empirical observation is that this proxy fails locally. A student-generated reasoning span can increase the probability of reaching a correct answer while deviating from the teacher distribution, in which case OPD assigns it an unfavorable (high-KL) reward; conversely, teacher-like spans need not make meaningful progress. Applying such uncalibrated signals uniformly risks suppressing productive but non-teacher-like reasoning pathways. The authors further note that outcome-level filtering is insufficient: their diagnostic shows roughly one third of responses contain correct reasoning paths that are truncated due to overthinking and would be discarded by trajectory-level correctness checks.

The R²-OPD method

Reasoning-Progress-Aware Reward Filtering for On-Policy Distillation (R²-OPD) uses process rewards not as an optimization target but as an independent reliability test on the distillation signal. Responses are partitioned into reasoning segments via discourse-marker matching with a minimum-sentence constraint ($S_{\min}=3$). At each boundary state, the solve probability is estimated from $N_{\mathrm{eval}}=8$ answer-eliciting Monte Carlo rollouts of the student itself, and the segment's process reward is the marginal change $\hat S_m - \hat S_{m-1}$. Trajectories lacking a verifiable ground-truth answer string or sufficient segmentation are passed through unfiltered, so the method never penalizes trajectories without process signals.

Three design elements structure the filtering:

• Sign-consistent merging. Adjacent segments whose process rewards share sign are merged into maximal runs. Because process rewards are finite differences of consecutive solve probabilities, the aggregated reward telescopes to depend only on endpoint estimates (Lemma 1), eliminating internal boundary estimation errors entirely.
• Segment-level divergence averaging. Token-level support-restricted reverse-KL losses are averaged within each merged segment. Under a weak-dependence covariance model with decay rate $\rho < 1$, Proposition 1 shows this average is unbiased with variance bounded by $O(1/L_j)$, trading temporal resolution for stability.
• Rank-based conflict detection and masking. Within each response, merged segments are ranked by aggregated progress and by average KL loss. When a higher-progress segment incurs a larger distillation loss — a local ranking conflict — the positive product of the two gaps accumulates into an inconsistency score. The top segments by this score, up to a budget of $q=30\%$ of segments (excluding the initial prefix segment), have their tokens masked from both the numerator and normalization of the OPD loss.

A deliberate architectural choice deserves emphasis: R²-OPD never replaces or augments the OPD objective. The process reward operates only as a gating signal, and masking with renormalization leaves the policy-optimization procedure otherwise unchanged. This distinguishes it from prior work that calibrates against outcome rewards [hou2026uni], reweights tokens by entropy or position [jin2026entropy6, xu2026tip, xie2026position], or optimizes process rewards directly.

Theoretical guarantees are modest but appropriate to the mechanism: they concern noise cancellation in the two signals being compared, not convergence properties of the filtered objective. The ranking-based comparison sidesteps cross-signal calibration, following preference-learning practice [NIPS2017_d5e2c0ad, ouyang2022training].

Empirical results

The primary configuration distills DeepSeek-R1-Distill-Qwen-1.5B from JustRL-1.5B on DAPO-Math-17K, evaluated on AIME 2024, AIME 2025, and OlympiadBench. R²-OPD achieves 35.06 avg@4 and 51.83 pass@4 aggregate, exceeding standard OPD by 2.51 and 4.46 points respectively, with the largest gains on the AIME benchmarks (e.g., AIME 2024 avg@4 of 32.50 versus OPD's 28.33). Against the strongest baseline, Uni-OPD, the margins are 4.28 avg@4 and 5.17 pass@4 points. Notably, Uni-OPD remains stronger on OlympiadBench (53.16 vs. 46.86 avg@4), so the benefit of progress-aware filtering is benchmark-dependent rather than uniform.

Method	AIME 24 avg@4	AIME 25 avg@4	Olympiad avg@4	Avg avg@4	Avg pass@4
Student	22.50	23.33	43.19	29.67	46.10
Teacher	41.67	30.83	53.28	41.92	56.18
OPD	28.33	22.50	46.86	32.55	47.37
E-OPD	18.33	12.50	49.91	26.91	41.52
TIP-OPD	17.50	11.67	48.24	25.80	38.72
IW-OPD	20.83	17.50	43.59	27.31	40.81
Uni-OPD	20.00	19.17	53.16	30.78	46.66
R²-OPD	32.50	25.83	46.86	35.06	51.83

On the heterogeneous Qwen3-1.7B / e3-1.7B pair, avg@4 gains are modest (+0.43 aggregate) but pass@4 improves consistently across all three benchmarks (45.70 → 48.19 aggregate), suggesting transfer across model families with effect size varying by metric. The paper attributes the weak AIME performance of entropy-based baselines partly to interaction with the response-length constraint: preserving uncertain positions may encourage lengthy or unfinished reasoning. It also observes that R²-OPD produces more concise answers, though the training-length confound relative to Uni-OPD's original settings means this claim should be read cautiously.

Two ablations isolate the components:

• Masking ratio. $q=30\%$ performs best; $q=50\%$ degrades AIME accuracy sharply (19.17 on AIME 24, 12.15 on AIME 25), confirming that aggressive masking removes useful teacher guidance. This is an explicit trade-off rather than a free win.
• Sign-consistent merging. Removing merging raises PR–KL rank disagreement dramatically: smoothed rank agreement improves from roughly 0.20–0.27 (below the 0.5 random baseline) to 0.55–0.73 when merging is applied. Downstream, merging contributes +15.00 avg@4 on AIME 2024 and +14.16 on AIME 2025, with a slight −0.96 on OlympiadBench. That the no-merge rank agreement falls below chance is itself evidence that raw fine-grained signals are too noisy for reliable conflict detection.

Qualitative case studies on AIME problems illustrate the failure mode concretely: in one example, standard OPD inherits the base model's boundary-value error (yielding 80 instead of 480), while the R²-OPD student preserves the geometric constraints through a limiting argument. The authors appropriately flag these as qualitative rather than causal evidence.

Limitations and open questions

Several limitations are conceded directly in the paper. Process-reward estimation adds nontrivial rollout cost ($M-1$ boundaries × 8 rollouts per eligible response), and the authors explicitly defer reducing its cost and variance to future work. Segmentation relies on a small hand-crafted lexicon of discourse markers ("wait", "actually", "alternatively"), which ties boundary detection to English self-reflection conventions and may not generalize to other languages or reasoning styles. The masking ratio $q$ is a sensitive hyperparameter, and the strong degradation at $N_{\mathrm{eval}}=8$0 indicates the method can discard valuable supervision if mis-tuned. Evaluation is confined to 1.5–1.7B mathematical reasoning models with a single epoch of training; generalization to larger models and non-mathematical domains remains untested, as does the interaction between progress-aware filtering and longer training horizons. Finally, because Uni-OPD outperforms R²-OPD on OlympiadBench, the conditions under which span-level progress conflicts dominate over outcome-calibrated supervision are not characterized.

Conclusion

R²-OPD reframes dense teacher feedback in OPD as a signal to be audited rather than uniformly applied. By constructing a teacher-independent reference for reasoning progress through Monte Carlo solve-probability estimates, merging sign-consistent spans to suppress estimation noise, and masking only segments where distillation penalties contradict estimated progress, the method improves aggregate reasoning performance over standard OPD — most substantially on AIME (avg@4 gains of +15.00 and +14.16 attributable to merging alone) while preserving the standard OPD optimization pipeline. The contribution is best understood as a selective-supervision mechanism whose effectiveness depends on reliable segment-level signals, moderate masking budgets, and domain-verifiable outcomes.