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

• Vanilla OPD's token-wise corrections do not show how to revise a coherent reasoning step, while harmful student prefixes can degrade teacher guidance and reinforce unproductive continuations.Source: §1 · pp. 1–2
• Reasoning revision reworks an intermediate step while preserving the prefix, giving the student an alternative reasoning transition intended to redirect subsequent reasoning.Source: §1 · p. 2
• SEG-OPD selects student segments using entropy-based uncertainty and trains students to prefer paired teacher redrafts through a segment-wise preference loss while retaining dense token-wise OPD feedback.Source: §1 · pp. 2–3
• Controlled trajectory interventions found that teacher redrafts improve downstream reasoning, and high-uncertainty pivots provide more useful revision targets than random anchors.Source: §1 · p. 2
• 2.26% higher reasoning accuracy than the strongest baseline with Qwen3-14B on AIME25.Source: Figure 1 · p. 2, §1 · p. 2
• 5.22% average relative improvement in reasoning accuracy over compared state-of-the-art OPD methods; SEG-OPD also achieves higher revision success rates and fewer reasoning loops after training.Source: §1 · p. 3

2.1 ON-POLICY DISTILLATION

• SEG-OPD emphasizes transfer of teacher reasoning behavior by training students to prefer teacher redrafts over their own segments.OPD trains students on their own rollouts with teacher supervision; how teacher reasoning patterns transfer remains underexplored.Source: §2.1 · p. 3

2.2 LEARNING FROM REASONING INTERVENTIONS

• Unlike offline step-pairing or interleaved teacher generations, SEG-OPD pairs student segments with teacher redrafts and retains token-wise OPD on remaining tokens.Prior interventions include step-level reward labels (Lightman et al., 2024; Wang et al., 2024), preference learning over reasoning steps (Lai et al., 2024; Zhang et al., 2024; Yang et al., 2025b), and teacher guidance during on-policy training (Xu et al., 2026b; Huang et al., 2026; Zhao et al., 2026b).Source: §2.2 · p. 3

3.1 NOTATIONS

• The notation defines an input prompt x, response y, prefix q_t = (x, y_<t), and an OPD setup with a trainable student π_θ and fixed teacher π*.Source: §3.1 · p. 4

3.2 ON-POLICY DISTILLATION

• Token-wise OPD gives individual-token feedback but no coherent alternative reasoning step, motivating teacher redrafts as explicit revision targets compared with student segments.SEG-OPD uses conditional joint probabilities to compare each student segment with its teacher redraft while retaining dense token-wise supervision outside selected segments.Source: §3.2 · p. 4, Figure 3 · p. 5

3.3 UNDERSTANDING REASONING REVISION IN OPD

• Replacing an incorrect student segment with a teacher redraft increases subsequent reasoning accuracy for both initial and OPD-trained students.The experiment manually located an incorrect segment, generated a teacher redraft from the same prefix, and compared student continuations from each version.Source: §3.3 · p. 4
• At matched rollout positions, pivot anchors yield higher reasoning accuracy and revision success rates (RSR) than random anchors.RSR is the proportion of revision attempts in a complete trajectory that correct an earlier error with a substantively different alternative; GPT-5.6 assessed whether attempts met this criterion.Source: §3.3 · pp. 4–5

4.1 REVISION UNIT CONSTRUCTION

• Reasoning pivots anchor candidate segments, and SEG-OPD selects them from incorrect rollouts using predictive uncertainty, favoring segments with high mean entropy and low entropy dispersion.It expands high-entropy pivot positions to sentence boundaries, then greedily selects segments in descending score order, up to K per rollout; pivot-based segments showed higher revision success and post-redrafting accuracy than random-anchor segments.Source: §4.1 · pp. 5–6
• For each selected segment, the teacher rewrites it under the same prefix, and the leading sentence of a task-correct full draft becomes the redraft in the revision unit.Each unit comprises the prefix, student segment, and teacher redraft.Source: §4.1 · p. 6

4.2 REASONING REVISION LEARNING

• SEG-OPD trains students to prefer each teacher redraft over its paired student segment using a length-normalized preference margin and a Bradley–Terry objective.The loss penalizes small or negative revision margins more heavily and decreases as the margin grows.Source: §4.2 · p. 6
• Alongside segment-wise preference learning, SEG-OPD retains dense token-wise OPD on valid rollout tokens outside segments used in revision units.The segment mask only excludes those positions from the token-wise loss; the final objective jointly optimizes token-wise distillation and segment-wise preference losses.Source: §4.2 · pp. 6–7

5.1 EXPERIMENT SETUP

• Experiments pair DeepSeek-1.5B and Qwen3-4B, 8B, and 14B students with same-family, same-scale RL-trained teachers and evaluate mathematical reasoning and competitive programming.Training uses DeepMath-103K and Eurus-RL-Code; evaluations include AIME24/25/26, MATH500, MINERVA, Codeforces, TACO, and HumanEval+.Source: §5.1 · pp. 7–8

5.3 ANALYSIS OF STUDENT REASONING BEHAVIOR

• On HumanEval+, Qwen3-8B’s RSR peaks at 70.0% with SEG-OPD versus 44.6% with OPD, while accuracy peaks at 82.4% versus 76.6%.The initial teacher’s RSR is 75.0%; the authors interpret the higher student RSR as evidence that teacher redrafts help students revise reasoning more successfully.Source: Figure 4 · p. 8, §5.3 · p. 9
• For Qwen3-4B on AIME25, SEG-OPD achieves a 15.6% loop rate and 63.0% accuracy, versus 24.0% and 61.1% with OPD.The authors suggest that revision learning helps students move past stalled reasoning toward a solution.Source: Figure 4 · p. 8, §5.3 · p. 9
• For Qwen3-8B on AIME25, SEG-OPD reaches 59.7% accuracy versus 55.6% with OPD, with mean response lengths of 6.1K and 5.4K tokens, respectively.Across both length comparisons, SEG-OPD achieves higher final accuracy with comparable mean response lengths.Source: Figure 4 · p. 8, §5.3 · p. 9

5.2 SEGMENT-WISE OPD PERFORMANCE

• SEG-OPD yields a 5.22% average relative accuracy gain over baselines across all models and datasets.The comparisons cover OPD, G-OPD, E-OPD, and OmniOPD.Source: §5.2 · p. 8, §5.1 · p. 8
• Against OPD and OmniOPD, SEG-OPD improves accuracy in most settings, with gains of up to 8.24% and 8.83%, respectively, for DeepSeek-1.5B on AIME24.SEG-OPD uses student uncertainty to select reasoning pivots and trains students to prefer teacher redrafts over paired student segments.Source: §5.2 · p. 8

5.4 ABLATION STUDIES AND HYPERPARAMETER SENSITIVITY ANALYSIS

• Ablations show that reasoning-pivot anchors, student-tailored teacher redrafts, and segment-wise preference learning each contribute to SEG-OPD’s advantage.The variants use random anchors (SEG-OPD \S), teacher continuations from the prefix alone (SEG-OPD \T), or SFT on redrafts instead of preference loss (SEG-OPD \N); each comparison favors the full method.Source: Table 2 · p. 9, §5.4 · p. 9
• At λ = 2, Qwen3-8B AIME26 accuracy ends at 33.8%, versus 47.2% at λ = 1; at K = 2, accuracy falls from a 45.8% peak to 32.7% by training’s end.The results favor balancing revision supervision with token-wise distillation, as larger revision budgets can weaken sustained performance and training stability.Source: Figure 5 · p. 9, §5.4 · p. 9

6 CONCLUSION

• SEG-OPD outperforms competitive baselines in average reasoning accuracy across models and tasks, with higher revision success rates and fewer reasoning loops.The conclusion attributes these gains to learning students to prefer teacher redrafts over their paired segments.Source: §6 · p. 10

AI USE STATEMENT

• The authors used generative AI to illustrate the hiking metaphor and help identify and assess reasoning steps, then reviewed the AI-assisted work.They state that no other tasks requiring or recommending disclosure used generative AI, and accept responsibility for the final content.Source: AI USE STATEMENT · p. 11

ETHICS STATEMENT

• The authors state that SEG-OPD may be misused to enhance or adapt LLMs for malicious purposes and encourage safeguards and ethical guidelines.They state that the work conforms to the ICLR Code of Ethics.Source: ETHICS STATEMENT · p. 11

REPRODUCIBILITY STATEMENT

• The paper and supplementary materials provide implementation details, including datasets, models, baseline adaptations, training parameters, and prompts; code is available in an anonymous repository.The methodology is described in Section 4 and the implementation in Appendix A.Source: REPRODUCIBILITY STATEMENT · p. 11

A.1 DETAILS OF RSR EVALUATION

• RSR counts successful revision attempts divided by all revision attempts, and excludes steps that do not attempt revision.GPT-5.6 Sol identifies reasoning steps in complete rollouts and judges whether attempts correct earlier errors with substantively different alternatives; evaluation boundaries need not match training segments.Source: §A.1 · p. 15

A.2 DETAILS OF REVISION UNIT CONSTRUCTION

• Teacher redrafts were preferred in 86.9%–89.7% of audited units, versus 2.6%–6.4% for student segments across mathematics and programming pairs.The audit judged retained segments and redrafts under the same prefix for correctness and progress; ties remained in the denominator.Source: Table 3 · p. 17, §A.2 · p. 17

A.3 DETAILS OF MODELS

• Student models range from 1.5B to 14B parameters, and each reported student uses the RL-trained version of its base model as teacher.The experiments include public pretrained and fine-tuned checkpoints as well as locally trained models.Source: §A.3 · p. 17

A.4 DETAILS OF DATASETS

• Experiments span mathematical reasoning and competitive programming, training on 57,046 filtered DeepMath-103K problems and 25,276 Eurus-RL-Code problems.The datasets cover a range of topics and difficulty levels; evaluation datasets vary in source, scale, and difficulty.Source: §A.4 · p. 18

A.5 DETAILS OF COMPARED METHODS

• All baselines were reproduced in the same verl framework with shared student–teacher pairs, training data, and evaluation protocol, then tuned on a held-out validation split.Each method’s best validation configuration was selected; G-OPD with Qwen3-4B reached 61.2% on AIME25 in this setting, versus 56.1% reported in its original paper (Yang et al., 2026).Source: §A.5 · p. 18
• The comparisons cover OPD’s sampled-token reverse-KL supervision, E-OPD’s teacher-entropy-gated forward-KL term, G-OPD’s reward scaling, and OmniOPD’s agreement-weighted chunks.Implementations use teacher log-probabilities at response tokens for OPD (Yang et al., 2026), top-k teacher tokens for E-OPD (Jin et al., 2026), a 1.25 reward extrapolation coefficient and frozen initial-student reference for G-OPD (Yang et al., 2026), and ROUGE-1 F1 agreement with frozen-student regularization for OmniOPD (Zhou et al., 2026).Source: §A.5 · p. 18

A.6 OTHER IMPLEMENTATION DETAILS

• SEG-OPD and its baselines use PyTorch 2.10.0, verl 0.8.0, vLLM 0.18.0, and Transformers 5.3.0, with experiments run on eight NVIDIA H100 80 GB GPUs.The framework versions are attributed to Sheng et al. (2024), Kwon et al. (2023), and Wolf et al. (2020).Source: §A.6 · p. 19
• The shared setup uses AdamW without weight decay, a 5 × 10^-6 cosine-decay learning rate without warmup, 48 training steps, and validation every four steps.Default settings include seed 42, segment budget K = 1, and preference margin scale β = 0.01; rollout batch size is 256 for Qwen3-14B and 512 for other students.Source: §A.6 · p. 19
• Mathematical-reasoning and competitive-programming runs use different segment-wise loss weights for each student size.For DeepSeek-1.5B, Qwen3-4B, Qwen3-8B, and Qwen3-14B, λ is {0.75, 0.35, 0.33, 0.5} for math and {0.67, 0.9, 0.67, 0.75} for programming; response limits are 22,000 and 30,000 tokens, respectively.Source: §A.6 · p. 19
• Reasoning accuracy is reported as final-checkpoint mean Pass@1, averaging correctness over four sampled responses per problem and then over problems.Responses are generated at temperature 1.0 and top-p 1.0; peak values, when specified, denote the highest observed metric over displayed checkpoints.Source: §A.6 · p. 19

B.1 ADDITIONAL ANALYSIS ON BOUNDARY CONSTRUCTION RULES

• Period boundaries perform best overall and remain at or near the top late in training.The authors attribute this to preserving sentence-level context and a complete local inference with its explanation, clarifying the contrast used for preference learning.Source: Figure 6 · p. 20, §B.1 · p. 20
• The four rules extend selected reasoning pivots to the first period, comma, newline, or a fixed length of 2,048 tokens.The construction aims to preserve segment completeness while making segments and teacher redrafts informative and representative of the underlying reasoning process.Source: §B.1 · p. 20
• Comma boundaries generally trail periods, while newline boundaries approach period accuracy on Codeforces but remain weaker on HumanEval+.Comma cuts may separate an operation from its justification; line breaks may reflect layout rather than reasoning boundaries, despite prior use of double line breaks for segmentation (Zhang et al., 2025).Source: §B.1 · p. 20
• Fixed-length boundaries approach period-rule accuracy on MINERVA but underperform it on AIME26.Fixed token cuts may truncate a step or include later reasoning when student and teacher express the same step at different lengths.Source: §B.1 · p. 21

B.2 ADDITIONAL ANALYSIS ON REASONING BEHAVIOR

• Word ratio serves as a lexical proxy for potential reasoning transitions, while RSR assesses successful revisions.The predefined expressions include “wait,” “however,” “therefore,” “check,” “verify,” “correct,” “error,” “mistake,” and “incorrect”; related expressions have been used as indicators of reflection, backtracking, and self-verification (Huang et al., 2026).Source: §B.2 · p. 21
• SEG-OPD achieves higher late-training RSR than OPD in all four additional mathematical and competitive-programming settings.Successful revisions make up a larger proportion of revision attempts across these settings.Source: §B.2 · p. 21
• SEG-OPD generally reaches higher RSR earlier than OmniOPD, though late-stage gaps narrow on Qwen3-4B/AIME24 and Qwen3-8B/Codeforces.A clear late-training gap remains on Qwen3-8B/AIME26; OmniOPD weights student chunks by similarity to teacher continuations, while SEG-OPD learns to prefer teacher redrafts over paired student segments.Source: §B.2 · p. 21
• SEG-OPD has higher late-training word ratios than both baselines for DeepSeek-1.5B and Qwen3-14B, approaching the initial teacher’s level on DeepSeek-1.5B.Word ratio measures the share of English words in natural-language reasoning drawn from a predefined set of transition-related expressions; code is excluded.Source: §B.2 · p. 21

B.3 ADDITIONAL ANALYSIS ON TRAINING DYNAMICS OF SEG-OPD

• Token-loss gradient norms generally decrease while segment-loss norms fluctuate, and preference-loss norms are smaller than redraft-SFT norms at most steps in both tasks.The analysis reports pre-clipping gradient norms; the larger SFT norms occur with per-redraft token averaging.Source: Figure 8 · p. 22, §B.3 · p. 23
• The SFT variant’s policy entropy initially rises then drops sharply, while SEG-OPD maintains higher entropy in later training steps across both tasks.The analysis defines policy entropy as average next-token uncertainty; lower entropy indicates more concentrated predictions on sampled trajectories.Source: §B.3 · pp. 21, 23, Figure 8 · p. 22
• SEG-OPD scores above its SFT variant on all eight benchmarks, with accuracy gaps ranging from 3.33% on AIME26 to 21.54% on AIME24.The SFT variant uses direct teacher-redraft imitation while retaining token-wise OPD.Source: §B.3 · pp. 21, 23
• Token loss decreases in both tasks, while preference and SFT segment losses generally decline after early fluctuations; preference learning trains students to favor redrafts over paired segments.This contrast-based signal complements token-wise distillation, whereas the SFT variant directly imitates teacher redrafts.Source: §B.3 · pp. 21, 23

B.4 CASE STUDY OF REASONING CORRECTION

• In the AIME26 case, the trained student recognizes that 16 has distinct factors 2 and 8, excludes only squares of primes, and answers 70 instead of the initial student’s 44.Both responses initially exclude all perfect squares, but the trained student revises that insufficient rule.Source: §B.4 · p. 24
• In the TACO #797C case, the trained student rejects its next-character stack rule on “acdb,” instead comparing the stack top with the minimum remaining character and producing “abdc.”Its final program passes all validation tests, while the initial student’s output is “acbd.”Source: §B.4 · p. 25

C TRAINING ALGORITHM

• Each iteration selects up to K student segments using uncertainty scores, trains for Nstep steps, then resamples the on-policy rollout and segments.Selection is based on Eq. (5), and the rollout is resampled after each training iteration.Source: C TRAINING ALGORITHM · p. 26
• Training constructs accepted revision units from teacher redrafts, computes revision margins and both losses, then updates the student parameters.The losses are segment-wise preference and token-wise OPD losses, optimized jointly under Eq. (9).Source: C TRAINING ALGORITHM · p. 26
• The teacher and reference policies remain fixed; if no revision unit is accepted, Ls=0 and all valid rollout tokens are retained for token-wise OPD.Source: C TRAINING ALGORITHM · p. 26

D COMPLEXITY ANALYSIS

• 1.73–1.93× OPD training time versus 3.39–3.89× for OmniOPD: Seg-OPD was faster across four student–teacher pairs and remained below twice OPD’s time.The timings include full teacher-draft generation and preference learning, and are normalized to OPD for each pair.Source: Table 6 · p. 26, D COMPLEXITY ANALYSIS · p. 27
• OPD, OmniOPD, and SEG-OPD have per-iteration bounds of O(L(1 + L)), O(L(1 + L) + Kn(L+c)(1+L+c)), and O(L(1 + L) + K(L+D)(1+L+D)), respectively.The analysis fixes model sizes, batch size, and optimization passes, and assumes dense attention with KV-cached decoding; boundary extraction adds only a linear scan.Source: D COMPLEXITY ANALYSIS · p. 26
• With fixed sampling budgets and D = O(L), all three methods have quadratic sequence-length upper bounds, but differ in teacher-generation and optimization work.OmniOPD uses K = 10, n = 10, and c = 50; SEG-OPD uses K = 1 and redrafts selected failed rollouts.Source: D COMPLEXITY ANALYSIS · pp. 26–27
• SEG-OPD’s practical cost includes generating full teacher drafts before extracting the retained redraft; OmniOPD generates up to 5,000 teacher tokens across 100 continuations per rollout.SEG-OPD’s added generation work depends on how many rollouts are redrafted and their full draft lengths; extraction is negligible relative to inference.Source: D COMPLEXITY ANALYSIS · p. 27

E DISCUSSIONS AND FUTURE WORK

• SEG-OPD learns reasoning revision through explicit teacher redrafting, and the authors encourage responsible use and safeguards against malicious adaptation.The paper reports empirical analysis across diverse model architectures and downstream tasks.Source: E DISCUSSIONS AND FUTURE WORK · p. 27
• Future work could extend SEG-OPD to multiple teachers, including black-box models, using preference supervision for redrafts and a white-box teacher for token-wise OPD logits.The segment-wise objective relies on preference supervision rather than teacher logits.Source: E DISCUSSIONS AND FUTURE WORK · p. 27