Beyond Imitation: Filtering On-Policy Distillation by Reasoning Progress

arXiv:2608.194082026-08-21

A fix for AI models that get penalized by their teacher even when they're reasoning correctly

On-Policy Distillation (OPD) trains a student model by scoring its self-generated answers against a teacher model, but this scoring can unfairly punish reasoning steps that diverge from the teacher even when they're making real progress toward the correct answer. The researchers built R2-OPD, which separately measures whether a reasoning step is actually getting closer to the solution, and filters out teacher-based penalties whenever they conflict with genuine progress. On DeepSeek-R1-Distill-Qwen-1.5B, this improved average accuracy by 2.51 points and the chance of getting at least one correct answer out of four tries by 4.46 points, compared to standard OPD.

What they did

The team found that OPD's token-by-token comparison to a teacher model doesn't always reflect whether the student is actually making reasoning progress toward a correct answer
To measure real progress independently, they ran Neval=8 simulated rollouts from each intermediate reasoning point to estimate how likely the student was to reach the correct answer from there
Adjacent reasoning segments showing progress in the same direction were merged (sign-consistent merging) to reduce noise, then compared against the teacher-similarity ranking to find segments where the two disagreed, and masked the teacher-based penalty in those conflicting segments
On DeepSeek-R1-Distill-Qwen-1.5B with JustRL as teacher, R2-OPD beat standard OPD by 2.51 points on avg@4 and 4.46 points on pass@4, with the biggest gains on AIME math competition benchmarks
The method also transferred to a different model pairing, Qwen3-1.7B student with e3-1.7B teacher, showing it generalizes across model families

Figure 1: Distribution of evaluation outcomes for DeepSeek-R1-Distill-Qwen-1.5B and Qwen3-1.7B on DAPO (500 samples, max 7,168 response tokens).

Table 1: OPD performance on DeepSeek-R1-Distill-Qwen-1.5B with JustRL as the teacher model.
Method	AIME 24	AIME 25	Olympiad	Avg.
avg@4	pass@4	avg@4	pass@4	avg@4	pass@4	avg@4	pass@4
Student	22.50	43.33	23.33	36.67	43.19	58.31	29.67	46.10
Teacher	41.67	56.67	30.83	43.44	53.28	68.43	41.92	56.18
OPD (1)	28.33	50.00	22.50	30.00	46.86	62.10	32.55	47.37
E-OPD (17)	18.33	36.67	12.50	23.33	49.91	64.56	26.91	41.52
TIP-OPD (36)	17.50	33.33	11.67	20.00	48.24	62.83	25.80	38.72
IW-OPD (34)	20.83	36.67	17.50	26.67	43.59	59.09	27.31	40.81
Uni-OPD (13)	20.00	43.33	19.17	26.67	53.16	69.97	30.78	46.66
R2-OPD (Ours)	32.50	56.67	25.83	36.67	46.86	62.19	35.06	51.83

Figure 2: representative mismatch between reasoning progress and the OPD divergence signal. A reasoning regressive segment (J9) receives a smaller KL divergence under Equation (9) than segments that derive or recover the correct solution, respectively(J2 and J10).

Table 2: Transfer performance of R2-OPD on Qwen3-1.7B with e3-1.7B as the teacher model.
Dataset	Base	OPD	R2-OPD
avg@4	pass@4	avg@4	pass@4	avg@4	pass@4
AIME 24	24.17	30.00	22.50	36.67	25.00	40.00
AIME 25	18.75	20.00	27.50	33.33	25.83	36.67
Olympiad	50.10	63.87	53.82	67.09	54.31	67.91
Avg.	31.01	37.96	34.61	45.70	35.04	48.19

Figure 3: Overview of the proposed Reasoning-Progress-Aware Reward Filtering OPD. Adjacent reasoning spans with consistent progress signals are first merged. For each merged segment, we sum its process rewards and average its token-level KL divergences. Rank conflicts between the resulting segment-level signals are then used to identify unreliable spans, which are masked during the OPD policy update.

Table 3: Effect of sign-consistent segment merging on avg@4 accuracy. Δ denotes the improvement of the full R2-OPD over its no-merging variant.
Dataset	No Merge	R2-OPD	Δ
AIME 24	17.5	32.50	15.0
AIME 25	11.67	25.83	14.16
Olympiad	47.8	46.86	-0.96

Figure 4: Sensitivity of R2-OPD to the masking ratio q. A moderate ratio of q=30 provides the best overall performance, whereas overly aggressive masking substantially degrades accuracy on AIME 2024 and AIME 2025.

Table 4: Optimization, process-reward estimation, and filtering settings used in all R2-OPD experiments.
Setting	Value
Optimization and sequence settings
Training epochs	1
Optimizer	AdamW
Learning rate	5×10−6
Global batch size	64
Maximum prompt length	1,024 tokens
Maximum response length	7,168 tokens
KL support size	Student top-16 (H=16)
Process-reward estimation
Rollouts per evaluated boundary	Neval=8
Sampling temperature	0.7
Top-k	50
Top-p	1.0
Maximum rollout length	300 tokens
Segment filtering
Minimum sentences between boundaries	Smin=3
Minimum merged segments	nmin=3
Segment masking ratio	q=30%

Figure 5: PR–KL rank agreement during training with and without sign-consistent segment merging. Light lines show per-step values, and dark lines show 15-step moving averages. The two vertical axes use the same numerical span.

Table 5: Eligibility checks and conservative fallback actions used by R2-OPD.
Condition	Action
The string-level pre-check does not find gi in yi, or segmentation yields fewer than two segments.	Skip process-reward rollouts and retain all token-level OPD supervision.
The response has fewer than max⁡(3,nmin) merged segments.	Do not apply segment masking.
No strict adjacent pair is available after process-reward ranking.	Do not apply segment masking.
All candidate segments have zero inconsistency score.	Do not apply segment masking.
At least one candidate has a positive inconsistency score.	Mask up to the response-specific budget; retain every other token.

Table 6: Complete lexicon of discourse markers used to propose reasoning-segment boundaries.
Category	Matched markers
Reconsideration	wait, hold on, let me reconsider, hmm
Correction	actually
Verification	let me check
Alternative reasoning	alternatively

Why it matters

This work exposes a structural flaw in standard teacher-student distillation: it can suppress creative but correct reasoning paths simply for not matching the teacher's style. That matters directly for anyone training reasoning-focused AI models more cheaply and effectively, such as math or coding assistants.

Terms in this paper

On-Policy Distillation (OPD) · a training method where a student model's own generated answers are graded against a teacher model's outputs
reverse KL divergence · a measure of how different the student's probability distribution is from the teacher's, used as a training signal
process reward · a score measuring how much a reasoning step increases the estimated chance of reaching the correct final answer
sign-consistent merging · combining adjacent reasoning segments that show progress in the same direction to reduce measurement noise
avg@4 / pass@4 · avg@4 is the average accuracy over 4 attempts; pass@4 checks whether at least one of 4 attempts was correct

Original abstract (English)

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.

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

Read on arXiv

Latest papers

All papers →

Latest from METAL LAB

Figures: Chen Yang et al., arXiv:2608.19408, arxiv-nonexclusive