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

May 29, 2026

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

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%.

Yuxuan Jiang, Francis Ferraro

University of Maryland, Baltimore County

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

Ask PDF

Oct 4, 2026

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

On-policy distillation (OPD) improves the reasoning capabilities of small language models through token-level teacher supervision on student-generated trajectories. Yet can teachers that excel at solving problems independently also guide student reasoning effectively? Prior work shows that when student prefixes follow reasoning paths that differ from the teacher's own or contain errors, teachers can be less accurate when continuing from these prefixes than when solving problems independently. To this end, we propose Prep-OPD, which uses reinforcement learning (RL) before distillation to train the teacher to adapt to the student's existing reasoning state and correct course when errors arise. Training optimizes teacher continuations from fixed student prefixes using final-answer correctness as the reward. The prepared teacher then trains the student through trajectory guidance and token-level supervision. We evaluate Prep-OPD on eight mathematical reasoning benchmarks, using Qwen3-4B-Instruct-2507 as the teacher and Qwen3-0.6B and Qwen3-1.7B as students. With the 4B teacher and 1.7B student, Prep-OPD improves average accuracy over standard OPD and the strongest baseline, Relay-OPD, by 8.28 and 2.30 percentage points, respectively. Controlled experiments further show that teacher RL conditioned on student-generated prefixes yields higher student accuracy than problem-start teacher RL with and without handoff on Qwen3-1.7B. Reusing the same prepared teacher also improves Qwen3-0.6B.

Xiaoyu Ma, Haoyue Liu, Zhichao Wang +2

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)

On-Policy Distillation Teacher-Student Learning RL for Language Model Reasoning

Ask PDF

Aug 14, 2026

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

On-policy distillation (OPD) offers a promising way to transfer reasoning capabilities from stronger teacher models, but applying it to long-context reasoning teachers and short-context students introduces practical challenges, including tokenizer mismatch, teacher-student distribution mismatch, response length explosion, and training instability. In this work, we study this setting by transferring proof-reasoning capabilities from the long-context reasoning model SU-01 to short-context student models. To handle tokenizer differences, we perform OPD in a shared text space and align only tokens that occupy identical text spans under the student and teacher tokenizers. To mitigate the problem of excessive generation length and frequent truncation, we introduce a student reference KL loss and mask the advantages of special termination tokens such as </think> and <|im_end|>. This strategy constrains the student from drifting excessively from its initial policy, thereby mitigating the teacher-student distribution mismatch problem and fostering steady length growth. Experiments on both same-family and different-family student models, including Qwen3, Qwen3.5, Intern-S2, GLM-4.7, Gemma-4, show consistent gains in mathematical reasoning, especially natural-language math proving. Notably, Intern-S2-Preview improves by 21.2 points on ProofBench, reaching 55.2 and surpassing Gemini-2.5-Pro. It also improves on science benchmarks such as HLE and HiPhO, suggesting that OPD transfers reasoning capabilities that generalize beyond the mathematical training domain.

Haonan He, Haodi Lei, Yun Luo +13

SU-01 Team, Shanghai Artificial Intelligence Laboratory

Long-Context Reasoning in Language Models On-Policy Distillation Cross-Tokenizer Knowledge Distillation

Ask PDF

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