Source-linked AI summary
Learning from Repaired Reasoning: Root-Cause-Guided On-Policy Distillation
Chenglei Shen, Haoyang Yao, Weijie Yu, Song Jin, Xiao Zhang, Jun Xu
TL;DR
Reference-based on-policy self-distillation may explain a correct solution without correcting errors in the student’s own reasoning. RC-OPD repairs student trajectories and tests repairs through continuation, assigning guidance to erroneous segments and valid prefixes. It achieves the highest mean accuracy among methods evaluated with Avg@4 across three benchmarks and three model sizes.
Problem
Reference-based guidance may provide a correct intermediate conclusion without addressing the student’s own error, leaving answer-level success compatible with unresolved reasoning.
Source: §1 · pp. 1–2Method
RC-OPD diagnoses and repairs local errors, tests repairs through student continuation, and uses distinct guidance for erroneous stages and their valid prefixes.
Source: §1 · p. 2, §3 · p. 3, §5.1 · p. 5Results
RC-OPD achieves the highest mean accuracy among methods evaluated with Avg@4 across three benchmarks and three model sizes.
Source: §6.2 · p. 7, §7 · p. 9Takeaways & Limitations
Analyses show better corrective alignment, stronger preservation of valid prefixes, and improved recovery.
Source: §7 · p. 9
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Reference-based hindsight can supply correct conclusions without correcting the student’s own reasoning, while uniform guidance can also constrain sound steps. RC-OPD uses local repairs, tests them through student continuation, and assigns supervision according to each segment’s role.
1 INTRODUCTION
- Reference-based guidance can leave the student’s inferential gap unresolved by supplying a correct intermediate conclusion without showing how its derivation follows from the student’s reasoning.This reasoning mismatch can produce a correct answer that relies on a borrowed reference segment rather than an established inference.Source: §1 · pp. 1–2
- Applying reference-conditioned guidance throughout a trajectory risks a distillation trap in which constraints on already-sound steps compete with correction of substantive errors.The challenge is to repair locally from the student’s retained prefix and test whether the student can continue from that repair.Source: §1 · p. 2
- RC-OPD diagnoses the earliest substantive error, repairs it into an anchor built on the retained prefix, and assigns distinct guidance to erroneous stages and valid prefixes.If continuation remains incorrect, iterative validation diagnoses further errors; chains that fail within the repair budget fall back to reference-conditioned distillation.Source: §1 · p. 2
- RC-OPD frames reasoning mismatch and the distillation trap as connected challenges and introduces iterative counterfactual validation to select local supervision based on successful student continuation.Its guidance uses failure diagnoses and corrective goals for erroneous segments, and reasoning chains to repaired intermediate results for valid prefixes.Source: §1 · p. 2
2 RELATED WORK
Related work develops on-policy distillation from student-generated prefixes, then explores how teacher guidance can be adapted in source, content, and scope. This paper studies jointly supporting local correction and valid reasoning with trajectory-specific hindsight.
2 RELATED WORK
- On-policy distillation methods query teachers along student-generated trajectories, and OPSD adds a reference solution as privileged hindsight for dense supervision at the student’s own prefixes.The passage situates Self-Distillation Fine-Tuning as a related use of demonstrations as teacher context in continual learning.Source: §2 · p. 3
- DASH adapts supervision allocation using the temporal evolution of teacher–student divergences, while ROSD constructs corrective reflection and targets the suffix after a first erroneous span.The passage characterizes this progression as linking the source of guidance to its content and scope.Source: §2 · p. 3
- The paper studies trajectory-specific hindsight that jointly supports local error correction and continued supervision of valid reasoning.This is presented as the next question after prior work on supervision source, content, and scope.Source: §2 · p. 3
3 PRELIMINARIES
The preliminaries describe on-policy distillation as teacher feedback at student-generated prefixes, then distinguish OPSD’s reference-conditioned teacher and uniform response-wide supervision. They also define stage-based trajectory units for RC-OPD’s local analysis.
3 PRELIMINARIES
- On-policy distillation queries the teacher at prefixes generated by the student, using student and teacher next-token distributions at each position.The teacher-to-student forward KL matches the full next-token distributions, and the loss is averaged over positions.Source: §3 · p. 3
- OPSD gives the student-derived teacher a reference solution as privileged hindsight, while the student generates from the problem alone and receives reference-conditioned supervision uniformly across response positions.The OPSD objective replaces the standard teacher distribution with the reference-conditioned teacher target; the supplied prose identifies stop-gradient in that target.Source: §3 · p. 3
- For local analysis, a trajectory is partitioned into contiguous token-span stages, with each diagnosis round identifying an Error Stage and its repaired Anchor Stage.RC-OPD assigns distinct hindsight to the Error Stage and its valid prefix; segment weights and loss aggregation are defined later in Section 5.3.Source: §3 · p. 3
4 ANALYSIS
With reference hindsight, answer accuracy rises sharply, but filtering correct responses with unresolved borrowing substantially reduces the scores. An example illustrates how accepting a reference conclusion can leave its inferential basis unreconstructed.
4 ANALYSIS
- The analysis compares a frozen Qwen3-4B on 30-problem AIME24 and AIME25 sets, using an external LLM to identify unresolved borrowing.Borrowing is counted when a necessary intermediate conclusion or final answer is adopted without derivation or independent verification in the complete response.Source: §4 · p. 4
- 23.4 and 25.0 percentage points: excluding correct responses with unresolved reference borrowing lowers AIME24 and AIME25 scores to 75.8% and 73.3%, respectively.Before filtering, reference hindsight raises answer accuracy from 71.7% to 99.2% on AIME24 and from 67.5% to 98.3% on AIME25.Source: Figure 2 · p. 4, §4 · p. 4
- AIME25 example: the model questions how the count 56 is obtained, then accepts it from the reference without reconstructing the counting argument.The passage says this illustrates the limitation of reference-based hindsight when reasoning exceeds the student’s current ability to follow and reproduce.Source: §4 · p. 4, Figure 3 · p. 4
5 RC-OPD: ROOT-CAUSE-GUIDED ON-POLICY DISTILLATION
RC-OPD builds supervision from local repairs of student reasoning, tests repaired continuations, and assigns distinct guidance to error stages and valid prefixes.
5 RC-OPD: ROOT-CAUSE-GUIDED ON-POLICY DISTILLATION
- RC-OPD combines diagnosis, iterative continuation testing, and differentiated teacher guidance, using local supervision after recovery or reference-based fallback when the repair budget is exhausted.The teacher is frozen and initialized from the student's base model; a separate diagnostic model supplies repairs and explanations.Source: §5 · p. 4, Figure 4 · p. 5
5.1 DIAGNOSE
- Diagnose locates the earliest substantive error, constructs an Anchor Stage for validation, and prepares distinct guidance for the Error Stage and valid prefix.The diagnostic model examines the trajectory with the problem and reference solution while following the student's reasoning route.Source: §5.1 · p. 5
= Diagnose
RC-OPD identifies local errors, tests repairs through student continuation, then assigns segment-specific or fallback teacher guidance and aggregates the selected losses.
= Diagnose
- Diagnose locates the earliest substantive error, replaces it with an Anchor Stage, and prepares an error explanation and viable-route reasoning chain for distinct segment guidance.The Anchor Stage connects to the valid prefix and derives a valid intermediate result.Source: = Diagnose · p. 5
5.2 ITERATIVE COUNTERFACTUAL VALIDATION
- Each round replaces the Error Stage and regenerates continuation from the repaired prefix with fixed student parameters; success selects chain supervision, while exhaustion falls back to the original rollout.The student receives only the problem and repaired prefix, and further diagnosis searches only the new continuation while earlier repairs remain unchanged.Source: §5.2 · pp. 5–6
5.3 DIFFERENTIATED DISTILLATION
- Local supervision conditions the teacher on Failure Reason & Goal at errors and COT to Anchor at valid prefixes; fallback uses reference hindsight on the original rollout.For an accepted chain, supervision is collected from each round, evaluated on the trajectory before that round's repair.Source: §5.3 · p. 6
- Local loss weights Error Stage positions 1 and valid-prefix positions λ>0, aggregates weighted losses over accepted chains, and normalizes batch loss by total supervision weight.Examples with no selected positions contribute zero to both batch sums.Source: §5.3 · pp. 6–7
6 EXPERIMENTS
Experiments compare RC-OPD with baselines, test its components, and analyze valid-prefix preservation, error-stage alignment, repair validation, and prefix-loss weight. RC-OPD leads the evaluated Avg@4 methods across three model sizes, while the analyses examine how its supervision behaves.
6.1 EXPERIMENT SETTINGS
- Experiments use Qwen3-1.7B, 4B, and 8B, train on OpenThoughts-Math-30K, and evaluate on AIME 2024, AIME 2025, and HMMT February 2025.The reported benchmark score is Avg@4, the mean correctness of four sampled responses per problem; Mean averages the three benchmark scores equally.Source: §6.1 · p. 7
6.2 OVERALL PERFORMANCE
- 9.66%, 5.78%, and 5.70% mean-accuracy gains over OPSD at Qwen3-1.7B, Qwen3-4B, and Qwen3-8B; RC-OPD leads Avg@4 methods at each size.Table 1 reports the main results across three benchmarks; RC-OPD results are averaged over three training seeds.Source: §6.2 · p. 7, Table 1 · p. 8
6.3 VARIANT COMPARISON
- Removing iterative validation, valid-prefix supervision, or reference fallback lowers mean accuracy relative to the full RC-OPD method.The largest decline occurs without Iterative Counterfactual Validation.Source: §6.3 · p. 8, Table 2 · p. 8
6.4 ANALYSES
- 44.7% lower valid-prefix KL from Base than OPSD indicates less divergence on all compared trajectories.Figure 5's left panel reports forward KL from Base; lower values indicate better preservation of valid reasoning.Source: Figure 5 · p. 8, §6.4 · p. 9
- 4.1% lower Error Stage KL than OPSD indicates closer alignment with corrective guidance; on 100 selected questions, final-answer accuracy is 54.75% versus 46.50% for OPSD and 30.25% for Base.The 8.25-point RC-OPD–OPSD difference applies only to the selected subset and does not establish an overall performance advantage.Source: Figure 5 · p. 8, §6.4 · p. 9
- Frozen-Base recovery rises from 51% to 85% and 95% over one to three repair rounds, while trained models score 41.94%, 44.17%, and 42.78% Avg@4 for budgets of one, two, and three rounds.Recovery gains diminish across rounds; two validation rounds yield the best training performance among the tested budgets.Source: Figure 6 · p. 9, §6.4 · p. 9
- Prefix weights 0, 0.1, 0.2, and 0.3 yield mean Avg@4 scores of 42.50%, 44.17%, 42.50%, and 40.56%, respectively.The 0.1 run scores highest among these selected runs, but differing training settings make the comparison descriptive rather than isolating the weight effect.Source: §6.4 · p. 9
7 CONCLUSION
The conclusion presents RC-OPD as combining diagnosis, iterative validation, and differentiated distillation to address reasoning mismatch and the distillation trap. It reports the highest mean accuracy among Avg@4 methods across three benchmarks and three model sizes.
7 CONCLUSION
- RC-OPD combines diagnosis, Iterative Counterfactual Validation, and Differentiated Distillation, and has the highest mean accuracy among Avg@4 methods across three benchmarks and three model sizes.The conclusion also reports better corrective alignment, stronger valid-prefix preservation, and improved recovery.Source: §7 · p. 9
AI USE STATEMENT
The authors report using generative AI tools during manuscript preparation and experimental work.
AI USE STATEMENT
- Generative AI assisted with drafting, revision, figure preparation, result interpretation, error diagnosis, local repair, and response evaluation.The authors state that they take full responsibility for the paper's accuracy, originality, and integrity.Source: AI USE STATEMENT · p. 10
A ADDITIONAL RESULTS AND EXAMPLES
These additional analyses illustrate unresolved reference borrowing, describe variant and error-stage alignment analyses, and report closer alignment with corrective guidance under RC-OPD.
B.1 REFERENCE-CONDITIONING ANALYSIS
- The analysis illustrates correct reference-conditioned answers and assesses unresolved borrowing by checking complete responses for adopted conclusions lacking derivation or independent verification.Responses were evaluated on AIME24 and AIME25; answer accuracy and filtered accuracy count correct answers, with the latter excluding unresolved borrowing.Source: §A.1 · p. 12, §B.1 · p. 13
A.1 EXAMPLES OF UNRESOLVED REFERENCE BORROWING
- Two correct-answer examples adopt reference counts without resolving their derivations: one accepts 56, and another accepts 451 while confused.The responses came from problems the Base response did not answer correctly; the examples are assessed from complete responses.Source: §A.1 · p. 12
B.2 VARIANT SETTINGS
- The variants remove Anchor Prefix supervision, student-continuation validation, or reference-fallback supervision while retaining specified components.Without validation, diagnosis and route checks still precede local supervision, but a student counterfactual continuation does not select the repair.Source: §B.2 · p. 13
B.3 ERROR-STAGE ALIGNMENT AND PREFIX PRESERVATION
- Of 88 archived answer-incorrect Base trajectories, 58 with identified first errors and local repairs entered Error Stage analysis, while 49 with valid prefixes entered prefix analysis.At Error Stage positions, forward KL was computed from the corrective teacher, excluding stage markers; scores were averaged over selected positions within each trajectory and equally across trajectories.Source: §B.3 · p. 13
B.4 CONTINUATION FROM IDENTICAL VALID PREFIXES
On identical valid prefixes, RC-OPD exceeds OPSD and Base in continuation accuracy; supplementary sections detail validation, loss implementation, and diagnosis prompts.
B.4 CONTINUATION FROM IDENTICAL VALID PREFIXES
- RC-OPD exceeds OPSD by 8.25 percentage points and Base by 24.50 points on identical valid-prefix continuations.The comparison reports final-answer accuracy for Qwen3-4B, with four continuations per problem.Source: Table 3 · p. 14, §B.4 · p. 14
B.5 ITERATIVE COUNTERFACTUAL VALIDATION SETTINGS
- Repair chains allow up to three rounds, diagnosing the first substantive error and testing a route-preserving local repair through student continuation.Separately, Qwen3-1.7B models trained with one, two, or three validation rounds are evaluated using mean Avg@4 across AIME24, AIME25, and HMMT25.Source: §B.5 · p. 14
B.6 DISTILLATION IMPLEMENTATION
- The implementation caps each vocabulary contribution to teacher-to-student forward KL before summation and excludes control tokens and stage-heading spans from supervision.The clipped surrogate can be negative and differs from the unmodified KL used for distributional measurements; an update is skipped when total supervision weight is zero.Source: §B.6 · p. 14
B.7 DIAGNOSIS PROMPTS
- Diagnosis preserves the viable route and valid prefix, repairs one logical step, and returns an anchor step, failure reason and goal, and derivation to the anchor.The prompt format can return unresolved or route_break instead of a repair when evidence is insufficient or the route cannot be preserved.Source: §B.7 · pp. 14–15