Learning from Evolving Errors: Adaptive Iterative Repair for On-Policy Distillation

2026-10-01

AI Summary

This study addresses the issues of shortcut learning induced by reference solutions and the lack of error-correction supervision in online distillation. To this end, we propose the AIR-OPD framework, which incorporates an adaptive iterative repair mechanism and outcome-aware stage weighting, enabling the student model to progressively rectify errors under its own policy sampling. Furthermore, it achieves fine-grained error alignment by integrating a guided generator with privileged-context teacher supervision, and conducts online self-distillation training on the DAPO-Math dataset. Experimental results demonstrate that our method surpasses the strongest baseline by 3.6 points on mathematical benchmarks such as AIME, while maintaining robust generalization performance on out-of-distribution tasks.

Abstract

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.

Problem

• on-policy self-distillation
• solution-conditioned shortcut
• error-to-repair supervision
• knowledge distillation

Innovation

• On-policy distillation
• Adaptive iterative repair
• Error-to-repair supervision
• Stage weighting
• Guidance generator