Learning to Revise Reasoning with Segment-wise On-Policy Distillation

Yuxiang Zhang; Ding Cao; Shuting Cui; Lei Wang; Weijieying Ren; Tianxiang Zhao

cs.AIOct 2, 2026

Learning to Revise Reasoning with Segment-wise On-Policy Distillation

Authors: Yuxiang Zhang, Ding Cao, Shuting Cui, Lei Wang, Weijieying Ren, Tianxiang Zhao

Organizations: The Hong Kong University of Science and Technology (Guangzhou) · University of Science and Technology of China · Alibaba Group · Stanford University

Abstract

On-policy distillation (OPD) improves large language model reasoning by training students on their own rollouts with dense token-wise supervision from the teacher. However, token-wise OPD does not explicitly provide a coherent alternative reasoning step showing how the student's step could be revised to improve subsequent reasoning. Furthermore, this paradigm can become less effective when the student produces a degenerate reasoning prefix, as subsequent teacher supervision remains conditioned on that prefix and may reinforce poor reasoning patterns. In this work, we focus on learning reasoning revision with segment-wise OPD to rework intermediate reasoning steps and better support subsequent reasoning. Through controlled reasoning interventions, we find that replacing student segments with teacher redrafts improves subsequent reasoning accuracy. Therefore, we address the problem of turning teacher redrafts into explicit supervision for learning to revise reasoning. We propose Segment-wise On-Policy Distillation (Seg-OPD), which selects student segments based on an uncertainty metric and obtains corresponding teacher redrafts. Seg-OPD trains the student to prefer teacher redrafts over their paired student segments while retaining dense token-wise OPD supervision. Extensive experiments on mathematical reasoning and competitive programming tasks show that Seg-OPD-trained students achieve higher revision success rates than baselines. Seg-OPD consistently outperforms the compared state-of-the-art baselines in reasoning accuracy with an average relative improvement of 5.22% across diverse models and tasks. Code is available at https://anonymous.4open.science/r/Seg-OPD.

Figures & tables

Figure 1: (a). An example of reasoning revision. The student’s largest-coin-first attempt overshoots the target, whereas the teacher redraft reaches it with two equal-value coins. (b). Segment-wise supervision. The paired student segment and teacher redraft illustrate explicit supervision for learning reasoning revision. (c). Overall Performance. Reasoning accuracy (%) of Seg-OPD and baselines for a Qwen3-14B student on AIME25.

Figure 2: Effectiveness of teacher redrafts and segment selection. (a). Reasoning accuracy of initial/OPD-trained students with/without teacher redrafts on AIME24. (b). After replacing segments selected with random/pivot anchors with teacher redrafts, the reasoning accuracy (bars, left axis) and revision success rate (lines, right axis) of Qwen3-4B on AIME26.

Figure 3: Framework of Seg-OPD . Seg-OPD enables students to learn reasoning revision by pairing student segments with teacher redrafts and using preference learning to favor these alternative reasoning steps over the corresponding student segments. Simultaneously, Seg-OPD retains dense token-wise supervision outside the selected segments.

Method	Codeforces	TACO	HumanEval+	AIME24	AIME25	AIME26	MATH500	MINERVA
Student: DeepSeek-1.5B
OPD	33.30	28.30	73.25	43.09	36.67	33.33	86.45	37.22
E-OPD	34.90	29.30	78.84	43.31	35.83	30.83	86.05	37.59
G-OPD	33.75	30.06	78.67	50.00	39.17	31.67	86.85	38.69
OmniOPD	33.90	29.70	79.39	42.50	32.50	28.33	86.05	37.04
Seg-OPD	36.30	31.30	81.20	51.33	40.19	35.25	89.37	39.36

Table 1: Main results on mathematical reasoning and competitive programming. We report the mean Pass@1 (%) across diverse student models and datasets. Bold and underline denote the best and second-best results per evaluation setting.

Figure 4: Reasoning accuracy (dashed lines, left axis) and behavioral metric statistics (solid lines, right axis). Gray and gold horizontal lines indicate the initial behavioral metric values for the student and teacher models, respectively.

Variant	Codeforces	TACO	HumanEval+	AIME24	AIME25	AIME26	MATH500	MINERVA
Seg-OPD	45.89	43.65	85.89	64.50	63.00	52.53	89.10	48.85
Seg-OPD \S	42.60	38.07	72.53	54.16	59.18	51.54	75.49	45.23
Seg-OPD \T	40.27	33.25	71.89	56.90	57.25	52.39	70.05	46.27
Seg-OPD \N	37.58	31.75	70.69	42.96	56.89	49.20	73.29	40.67

Table 2: Ablation studies of Seg-OPD components with Qwen3-4B as the student.

Figure 5: Hyperparameter sensitivity analysis of Seg-OPD .

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Student Model	Audit results	Math (%)	Code (%)
Qwen3-4B	Redraft preferred	86.9	89.7
	Segment preferred	2.6	6.4
	Tie	10.5	3.9
Qwen3-8B	Redraft preferred	87.3	88.4
	Segment preferred	3.2	5.9
	Tie	9.5	5.7

Appendix

Table 3: Revision unit quality assessment.

Mathematical Reasoning
Student Model	Teacher Model	Algorithm
DeepSeek-1.5B	JustRL-DeepSeek-1.5B ( He et al., 2025a )	GRPO
Qwen3-4B	Qwen3-4B-Non-Thinking-RL-Math ( Yang et al., 2026 )	GRPO
Qwen3-8B	Qwen3-8B-DAPO ( Cai et al., 2025 )	DAPO
Qwen3-14B	Qwen3-14B-Math-RL ( Cai et al., 2026 )	DAPO
Competitive Programming

Appendix

Table 4: Student–teacher pairs. The algorithm column refers to teacher RL training algorithms.

Evaluation Dataset	Size	Description
Training on DeepMath-103K ( He et al., 2025b )
AIME24		American Invitational Mathematics Examination from year 2024 to 2026.
AIME25	30
AIME26 ( Cai et al., 2026 )
MATH500 ( Lightman et al., 2024 )	500	MATH subset with varied difficulty
MINERVA ( Lewkowycz et al., 2022 )	272	Undergraduate math and science

Appendix

Table 5: The statistics of datasets in our work.

Figure 6: Reasoning accuracy of Seg-OPD with four different construction rules.

Figure 7: Additional analysis of reasoning behavior of Seg-OPD , OmniOPD, and OPD. The first five panels report word ratio, and the remaining four report revision success rate (RSR), ordered from left to right and top to bottom. Horizontal lines indicate the initial behavioral metric values for the student and teacher models, respectively.

Figure 8: Training dynamics of training a Qwen3-4B on DeepMath-103K (first row) and Eurus-RL-Code (second row) . Left column and middle columns compare the gradient norms and values of the token loss, segment loss, and segment loss (SFT) used in Seg-OPD \N variant, respectively. Right column compares the policy entropy of Seg-OPD and its SFT variant. We record the mean and standard deviation across five runs with different random seeds.

Student	Teacher	OPD	OmniOPD	Seg-OPD
DeepSeek-1.5B	JustRL-DeepSeek-1.5B	$1\times$	$3.67\times$	$1.89\times$
Qwen3-4B	Qwen3-4B-Non-Thinking-RL-Math	$1\times$	$3.89\times$	$1.93\times$
Qwen3-8B	Qwen3-8B-DAPO	$1\times$	$3.41\times$	$1.73\times$
Qwen3-14B	Qwen3-14B-Math-RL	$1\times$	$3.39\times$	$1.87\times$

Appendix

Table 6: Relative training time on DeepMath-103K, normalized to OPD within student–teacher pair.

LLM Reasoning

Large language model reasoning studies how models solve multi-step problems, verify intermediate steps, and use computation or tools reliably.

See all 931 papers

Bridging Reasoning Trajectories in On-Policy Distillation via Near-Future Guidance
May 29, 2026Yuxuan Jiang, Francis FerraroCoT Distillation Language Model Distillation
University of Maryland, Baltimore County
How Should Teachers Be Prepared? RL on Student-Induced States for On-Policy Distillation
Oct 4, 2026Xiaoyu Ma, Haoyue Liu, Zhichao Wang +2On-Policy Distillation Teacher-Student Learning
School of Science and Engineering, The Chinese University of Hong Kong, Shenzhen · Shenzhen Future Network of Intelligence Institute (FNiI-Shenzhen) · Guangdong Provincial Key Laboratory of Future Networks of Intelligence, CUHK(SZ)
SimpleOPD: Simple Tokenizer-Agnostic On-Policy Distillation for Long-Context Reasoning
Aug 14, 2026Haonan He, Haodi Lei, Yun Luo +13Long-Context Reasoning in Language Models On-Policy Distillation
SU-01 Team, Shanghai Artificial Intelligence Laboratory

May 29, 2026J/K move · Enter open · S save

Bridging Reasoning Trajectories in On-Policy Distillation via Near-Future Guidance

Yuxuan Jiang, Francis Ferraro

University of Maryland, Baltimore County

CoT Distillation Language Model Distillation On-Policy Distillation Teacher-Student Learning Knowledge Distillation

Open paper AskPDF arXiv

On-Policy Distillation (OPD) improves large language model reasoning by training a student model on trajectories sampled from its own policy under teacher supervision. Although OPD operates on trajectories, its learning signal remains token-level: it identifies deviations through high-loss tokens and repairs them through local reverse-KL correction. We show that this "trajectory-sampled but token-learned" mechanism cannot reliably bridge student trajectories toward teacher trajectories. About 30% of high-loss tokens fall into the low-divergence regime, indicating that many are surface-form mismatches rather than real reasoning forks. Moreover, even truly divergent tokens are difficult to repair with isolated token-level supervision, since reasoning failures often unfold as short-horizon distributional drift. We propose Trajectory-aware OPD (TOPD), which uses near-future trajectory information to identify real divergent states and distribute guidance across multiple future tokens. Experiments show that suppressing non-divergent high-loss tokens improves standard OPD from 47.8% to 48.2% average accuracy, while TOPD further improves performance to 52.2%, with gains on AIME24 from 60.0% to 63.3% and AIME25 from 46.7% to 53.3%.

Abstract

LLM Reasoning

Explore similar work

Bridging Reasoning Trajectories in On-Policy Distillation via Near-Future Guidance

How Should Teachers Be Prepared? RL on Student-Induced States for On-Policy Distillation

SimpleOPD: Simple Tokenizer-Agnostic On-Policy Distillation for Long-Context Reasoning