Learning from Repaired Reasoning: Root-Cause-Guided On-Policy Distillation
Authors: Chenglei Shen, Haoyang Yao, Weijie Yu, Song Jin, Xiao Zhang, Jun Xu
Organizations: Gaoling School of Artificial Intelligence, Renmin University of China · School of Software and Microelectronics, Peking University · School of Information Technology and Management, University of International Business and Economics
On-policy self-distillation (OPSD) uses reference solutions as privileged hindsight to supervise student-generated reasoning trajectories. However, reference-based guidance may explain a correct solution without addressing why the student's own reasoning fails. This reasoning mismatch between the guidance provided and the correction needed can encourage the student to borrow correct conclusions while leaving its reasoning errors unresolved. Moreover, applying the same hindsight throughout the trajectory risks a distillation trap, where unnecessary constraints on valid reasoning compete with correction of substantive errors. To address these issues, we propose Root-Cause-Guided On-Policy Distillation (RC-OPD), which uses repairs of the student's own reasoning to provide guidance that addresses its specific errors while building on valid progress. For each failed attempt, RC-OPD locates the earliest substantive error, develops a local correction, and uses the corrected intermediate result as an anchor for the valid prefix. An iterative diagnosis--repair--continuation process tests the repairs through student continuation, identifying further errors within a fixed repair budget. For repair chains that reach a correct answer, root--cause--guided distillation uses failure diagnoses and corrective goals to supervise the erroneous segments, while anchor-guided distillation supports the corresponding valid prefixes with reasoning chains leading to the repaired intermediate results. We evaluate RC-OPD across multiple datasets and model scales. Extensive experiments and analyses show that it mitigates reasoning mismatch and the distillation trap, yielding substantial performance gains.
Figures & tables
Figure 1: Reasoning mismatch (top): a Borrowed Reference Segment bypasses the Unreachable transition, creating apparent reachability. Distillation trap (bottom): Reference-conditioned KL on every step also constrains already-sound steps.
Figure 2: Answer accuracy (Avg@4) before and after reference conditioning. Filtering excludes answer-correct responses with unresolved reference borrowing while retaining the original evaluation denominator. Arrows illustrate the proposed reasoning-mismatch mechanism.
Figure 3: An AIME25 response excerpt. Red highlights mark reference borrowing and acceptance; the ellipsis denotes omitted text. Highlighting is editorial.
Figure 4: Overview of RC-OPD. (a) Diagnose produces three outputs: (1) the repaired Anchor Stage , (2) Failure Reason & Goal , and (3) COT to Anchor . (b) Iterative Counterfactual Validation retains repaired stages and regenerates the continuation with the student (blue dashed boxes). Checking the final answer selects valid prefixes and Error Stages for Local supervision upon recovery, or the original trajectory τ0 for Fallback supervision upon budget exhaustion. (c) Differentiated Distillation conditions the teacher on COT to Anchor for the valid prefix and Failure Reason & Goal for the Error Stage under Local supervision, and on the reference under Fallback supervision.
Qwen3-1.7B
Qwen3-4B
Qwen3-8B
Method
AIME 2024
AIME 2025
HMMT 2025
Mean
AIME 2024
AIME 2025
HMMT 2025
Mean
AIME 2024
AIME 2025
HMMT 2025
Mean
Base
47.50
36.67
20.83
35.00
71.67
67.50
41.67
60.28
78.33
65.00
41.67
61.67
SFT †
48.40
36.30
22.70
35.80
70.20
62.30
43.40
58.63
72.30
64.20
42.90
59.80
GRPO †
51.10
38.30
23.70
37.70
75.60
68.10
44.40
62.70
76.40
68.90
46.70
64.00
OPSD
54.17
39.17
27.50
40.28
75.83
65.00
46.67
62.50
78.33
64.17
47.50
63.33
EOPD
56.67
35.83
22.50
38.33
75.00
70.83
39.17
61.67
75.83
66.67
46.67
63.06
Table 1: Main results across Qwen3 model sizes (%). Mean is the unweighted average across AIME 2024, AIME 2025, and HMMT February 2025. † denotes SFT and GRPO results reported by Zhao et al. ( Zhao et al., 2026a ) . SFT and GRPO use a different evaluation protocol and are included for context only. All other entries report Avg@4; bold marks the highest Avg@4 mean at each size. RC-OPD results are averaged over three independent training seeds.
Variants
Avg@4 (%)
#
Method
AIME2024
AIME2025
HMMT2025
Mean
①
w/o Anchor supervision
57.50
42.50
27.50
42.50
②
w/o Validation
54.17
39.17
28.33
40.56
③
w/o Fallback supervision
57.50
42.50
24.17
41.39
④
Full RC-OPD
54.17
47.50
30.83
44.17
Table 2: Comparison of RC-OPD and its training variants. “w/o” denotes the omitted component. Mean averages the three benchmarks. Variant settings are described in Appendix B.2 .
Figure 5: Valid-prefix preservation (left) and error-stage alignment (right). Blue squares denote OPSD and red circles denote RC-OPD ; yellow stars mark the respective reference distributions at zero divergence. Both panels report forward KL from the marked reference to the evaluated model; lower is closer. The right panel uses an explicit axis break.
Figure 6: Effect of validation rounds on frozen-Base recovery (red, left axis) and final Avg@4 of models trained with the corresponding validation budget (blue, right axis).
Figure 7: Sensitive Analysis on valid prefix weight λ .
Figure 8: Correct answers with unresolved reference dependence. Excerpts from the supplied references and actual model responses. Blue identifies reference conclusions; red identifies their adoption in the response. Mathematical and Markdown formatting is normalized for display. The missing-link descriptions summarize the full-response evidence review.
Model
Correct / Total
Accuracy (%)
Base
121 / 400
30.25
OPSD
186 / 400
46.50
RC-OPD
219 / 400
54.75
Appendix
Table 3: Final-answer accuracy from identical valid prefixes (Qwen3-4B, four continuations per problem).
Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.
On-policy self-distillation (OPSD) supplies dense token-level feedback on trajectories sampled from the student's own policy, a richer training signal than the outcome-level rewards of reinforcement learning. This feedback comes from a teacher conditioned on a full reference solution unavailable to the student. The reference solution specifies the target but not how to move from the student's current error toward it, creating a solution-conditioned shortcut risk. We introduce AIR-OPD, an adaptive iterative repair framework for on-policy distillation that provides error-to-repair supervision. Given a failed response, a guidance generator synthesizes repair guidance for the current error. The student samples an on-policy retry with this guidance. If the retry remains incorrect, the generator produces new repair guidance for the newly observed error. At each round, a fixed teacher receives the guidance as privileged context and supervises the student on an error-aligned region of its latest failed response. Outcome-aware stage weighting favors early repair stages and credits stages whose immediate retry passes verification. We train AIR-OPD on the DAPO-Math-17K dataset and evaluate on AIME24, AIME25, and HMMT25, alongside out-of-distribution tests on MMLU-Pro and GPQA. We examine two guidance sources, self-guidance from the current student policy and external guidance from a larger model. For both Qwen3-4B and Qwen3-8B, AIR-OPD attains the best mathematical-reasoning averages, improving over the strongest baseline by up to 3.6 points, while preserving base-model performance on the out-of-distribution benchmarks.
Rui Li, Liyang He, Zheng Zhang +3
University of Science and Technology of China · Nanyang Technological University
On-policy self-distillation (OPSD) improves reasoning by training a model on its own trajectories with dense token-level supervision from a privileged copy conditioned on a reference solution. Existing methods typically use a single teacher distribution conditioned on the complete reference. We identify a limitation of this formulation: a reference is only one realization of the reasoning process. Although it provides evidence about the answer, it also encodes an ordering, phrasing, and solution path. Full-reference conditioning can therefore produce correct targets that depend excessively on information unavailable to the student at inference. We show that replacing one valid reference with another, while holding the problem and student state fixed, can substantially alter the teacher distribution, a phenomenon we call reference-induced target instability. To address this instability, we introduce Monte Carlo Privileged Distillation (MCPD), which models privileged context as a distribution over partial reference views and distills the marginal teacher distribution. MCPD evaluates the teacher under multiple views and averages the resulting next-token distributions. We show that this probability-space average is the exact target induced by minimizing the expected teacher-to-student KL divergence over privileged views. Across mathematical and scientific reasoning benchmarks, we show that MCPD produces more reference-consistent supervision and consistently improves over full-reference OPSD.
Khouloud Saadi, Di Wang
Department of Computer Science, King Abdullah University of Science and Technology, KAUST, Thuwal, Saudi Arabia
On-policy distillation is a powerful way to transfer reasoning ability from a strong teacher to a smaller student: the student samples trajectories from its own policy, and the teacher provides dense token-level supervision on the states the student actually visits. However, this supervision is not always reliable: a teacher can assign high likelihood to plausible but incorrect solutions, or low likelihood to correct student solutions that follow different reasoning paths. Unconditionally distilling the teacher can therefore reinforce bad modes or erase useful student behavior. To address these limitations, we introduce RG-OPD: Reward-Gated On-Policy Distillation that uses verifier feedback to decide when teacher logits should be trusted. RG-OPD bridges sparse verifier rewards and dense teacher logits, preserving token-level supervision while filtering misleading teacher signals. Across reasoning and coding benchmarks, RG-OPD produces stronger distilled students, outperforming both vanilla reverse-KL distillation and the recent TSD-KD baseline. At 1K generation length, RG-OPD improves over reverse-KL by 2.9 points and over TSD-KD by 4.9 points; in the long-generation setting, it improves over the untuned student by 8.2 points. Our code is available at https://github.com/UoC-tail/RG-OPD.
Mohammad Sadegh Akhondzadeh, Vijay Lingam, Atula Tejaswi +3