HyperAIHyperAI

Command Palette

Search for a command to run...

批評誘導型オンポリシー蒸留による信頼できる教師信号の回復

Kaixiang Zhao Tianrun Yu Shangzhe Li Yuxiao Yang Amanda Hughes Weitong Zhang Taylor W. Killian

概要

オンポリシー蒸留(OPD)は、生徒モデルが自ら生成したロールアウトに対して、教師からトークン単位の教師信号を与えて学習する手法であり、生徒が生成したどのような途中経過からでも教師が正答に到達できるという仮定に基づく。我々はこの性質を教師の回復可能性と呼び、誤ったロールアウトではそれが成立しないことを示す。教師は誤りを認識しているにもかかわらず、生徒の誤りに追従してしまう。我々は、誤り認識を利用して回復可能性を回復するCriticOPDを提案する。教師が誤ったロールアウトを批評し、生徒はその誤りの後から解を書き直すことができる。批評を特権情報として含む生成結果は教師によって評価され、生徒は訓練時には批評を含まない改訂版を用いる。2つのQwen3生徒モデルと6つの数学的推論ベンチマークにおいて、CriticOPDはOPDを改善し、特に競技数学の問題で最大の向上を示した。コードは github.com/kzhao5/CriticOPD で入手可能である。

One-sentence Summary

Researchers from Brigham Young University and the University of North Carolina at Chapel Hill propose CriticOPD, a critique-guided on-policy distillation method that restores teacher recoverability by having the teacher identify errors, letting the student rewrite its solution, and scoring the revision with the critique as privileged information while the student trains on the revision without it; CriticOPD improves on-policy distillation on two Qwen3 students and six mathematical reasoning benchmarks, with the largest gains on competition problems.

Key Contributions

  • The paper introduces teacher recoverability in on-policy distillation and shows empirically that recoverability drops sharply on incorrect student rollouts, leaving many training tokens supervised from prefixes where the teacher rarely recovers correctly.
  • The paper proposes CriticOPD, in which the teacher critiques an incorrect rollout, the student rewrites from the last error using that feedback, and the critique serves only as privileged scoring context rather than as part of the student training signal.
  • Experiments on two Qwen3 student models and six mathematical reasoning benchmarks show that CriticOPD improves on-policy distillation, with the largest gains on competition problems, and code is available.

Introduction

On-policy distillation (OPD) trains a student on its own rollouts with token-level teacher supervision, reducing the occupancy mismatch that hurts offline distillation when inference prefixes differ from training prefixes. However, teacher supervision becomes unreliable after the student makes an error because the teacher often follows the student’s mistake rather than correcting it, and prior remedies either truncate supervision or let the teacher take over generation, which reintroduces occupancy mismatch. The authors propose Critic-OPD, where the teacher critiques an incorrect rollout and the student rewrites from the last error, while the critique is used only as privileged context for scoring during training. This improves teacher recoverability on revised trajectories, keeps every supervised token student-generated, and requires no feedback at inference.

Method

The authors introduce CriticOPD to address the issue where a teacher model fails to continue correctly after a student makes an error, even though the teacher can still recognize those errors. The method constructs training sequences based on the states the student reaches during inference, minimizing occupancy mismatch while ensuring reliable teacher supervision. The process involves a critique and rewrite mechanism followed by a specific training objective.

The process begins with the student generating a rollout yyy for a given problem xxx using a training prompt. If the rollout is incorrect or reaches the length limit without an answer, it is sent for critique and rewriting. Correct rollouts are trained using standard Online Policy Distillation (OPD).

The teacher acts as the critic. It receives the problem, the reference answer, and the student's solution formatted in numbered paragraphs. The teacher identifies errors in the solution, listing the paragraph index, the erroneous sentence, and a short fix for each. These errors are combined with the reference answer to form a critique ccc. The inclusion of the reference answer allows the teacher to critique problems it cannot solve itself. If no complete error is found in the critique, the rollout is trained with standard OPD.

To determine the cut point, the student retains its prefix up to the paragraph containing the last error identified in the critique and discards the remainder. The discarded part starts at position τ\tauτ, so the retained prefix is hτ=(x,y<τ)h_{\tau} = (x, y_{<\tau})hτ​=(x,y<τ​). This prefix contains all errors from the critique, ensuring the rewrite starts from the erroneous states where teacher recoverability typically declines.

The student then appends the critique ccc to the retained prefix using a feedback prompt and generates a new suffix, resulting in the revised rollout y~\tilde{y}y~​. The tokens before τ\tauτ remain unchanged (y~<τ=y<τ\tilde{y}_{<\tau} = y_{<\tau}y~​<τ​=y<τ​), and the student samples the suffix with the critique in its context. Unlike teacher takeover, the student writes every token of the training sequence.

As shown in the figure below:

The critique identifies errors in the original solution, and the student rewrites from the last error to reach the correct answer. The authors also evaluate the recoverability of the teacher on these rewrites compared to the original rollouts. Recoverability is estimated as the fraction of correct answers among eight continuations from each prefix.

As shown in the figure below:

The figure illustrates that the teacher is significantly more reliable on the rewrite than on the original rollout. At 75% of the continuation, the recoverability of the teacher is 0.14 on the original rollout and 0.67 on the rewrite. Similarly, the success rate of the student without the critique rises from 0.03 to 0.16.

The training objective combines standard OPD for correct rollouts and feedback distillation for rollouts with a valid critique.

For correct rollouts and those without a valid critique, the authors minimize the reverse KL divergence at each student-generated prefix hth_tht​:

LOPD(y)=∑t=1∣y∣KL(πθ(⋅∣ht)∥πT(⋅∣ht))\mathcal{L}_{\mathrm{OPD}}(y) = \sum_{t=1}^{|y|} \mathrm{KL}\left(\pi_{\theta}(\cdot \mid h_t) \| \pi_T(\cdot \mid h_t)\right)LOPD​(y)=t=1∑∣y∣​KL(πθ​(⋅∣ht​)∥πT​(⋅∣ht​))

This uses the sampled-token advantage.

For rollouts with a valid critique, the training focuses on the revised sequence y~\tilde{y}y~​. Let h~t=(x,y~<t)\tilde{h}_t = (x, \tilde{y}_{<t})h~t​=(x,y~​<t​) denote its prefix without the critique. The teacher scores the rewrite with the critique ccc in its context, while the student does not see it. The feedback distillation loss is defined as:

Lfb(y~;c)=∑t=τ∣y~∣KL(πθ(⋅∣h~t)∥πT(⋅∣h~t,c))\mathcal{L}_{\mathrm{fb}}(\tilde{y}; c) = \sum_{t=\tau}^{|\tilde{y}|} \mathrm{KL}\left(\pi_{\theta}(\cdot \mid \tilde{h}_t) \| \pi_T(\cdot \mid \tilde{h}_t, c)\right)Lfb​(y~​;c)=t=τ∑∣y~​∣​KL(πθ​(⋅∣h~t​)∥πT​(⋅∣h~t​,c))

Only the rewritten suffix receives supervision because the retained prefix contains errors at states where the teacher has low recoverability. Each suffix token receives the sampled-token advantage:

Atfb=log⁡πT(y~t∣h~t,c)−log⁡πθ(y~t∣h~t)A_t^{\mathrm{fb}} = \log \pi_T(\tilde{y}_t \mid \tilde{h}_t, c) - \log \pi_{\theta}(\tilde{y}_t \mid \tilde{h}_t)Atfb​=logπT​(y~​t​∣h~t​,c)−logπθ​(y~​t​∣h~t​)

Since the suffix is sampled with the critique but the student's log-probabilities are computed without it, this update serves as a surrogate rather than an unbiased gradient.

The overall objective minimizes Lours=LOPD+Lfb\mathcal{L}_{\mathrm{ours}} = \mathcal{L}_{\mathrm{OPD}} + \mathcal{L}_{\mathrm{fb}}Lours​=LOPD​+Lfb​, where each loss sums over its respective rollouts in the batch. Compared with standard OPD, the teacher target changes from πT(⋅∣ht)\pi_T(\cdot \mid h_t)πT​(⋅∣ht​) to πT(⋅∣h~t,c)\pi_T(\cdot \mid \tilde{h}_t, c)πT​(⋅∣h~t​,c), under which the teacher is much more recoverable. The student learns without the critique and requires no feedback at inference.

Experiment

The study evaluates critique-guided on-policy distillation using Qwen3-1.7B and Qwen3-0.6B students distilled from a Qwen3-4B teacher on DAPO-Math-17k, with assessment on AIME, AMC, MATH500, OlympiadBench, and Minerva Math, plus ablations and reliability analyses. Teacher continuation accuracy drops after student errors because the teacher tends to follow incorrect prefixes, but its ability to flag valid errors remains stable. CriticOPD improves over OPD and related baselines, with the critique being the most important component, and it helps the teacher assign token-level credit to corrections; gains concentrate on post-error states where the student learns to recover from its own mistakes, while outside mathematics performance is largely preserved.

CriticOPD improves average mathematical reasoning accuracy over the OPD baseline at both tested Qwen3 scales, with the largest gains on competition-style benchmarks including AIME and AMC. It outperforms FastOPD and RelayOPD and is the best student method on most mathematical benchmarks at both scales. Ablations show that answer-only or generic feedback is less effective, and removing the critique causes the largest performance drop. CriticOPD achieves the best student accuracy on most mathematical benchmarks at both tested scales and improves over OPD at both scales. The largest gains appear on AIME 2024, AIME 2025, and AMC 2023, indicating stronger improvements on competition-style reasoning. Feedback ablations show the critique is essential; removing it causes a large accuracy drop, while answer-only guidance falls below the OPD baseline.

The ablation isolates the main components of the feedback and training setup. The critique is the most important feedback element, while removing the reference answer or using generic feedback causes smaller average declines. The feedback also matters more for scoring the suffix than for which model writes the suffix, indicating the critique improves both the rewritten suffix and the reliability of teacher scoring. Removing the critique produces the largest average drop among the feedback changes. Removing the reference answer or replacing the feedback with generic text leads to smaller average declines than removing the critique. Using a teacher-written suffix instead of a student-written suffix has a smaller effect than scoring the suffix without feedback. Scoring the suffix without the feedback causes a larger average drop, indicating that the critique supports reliable teacher scoring.

CriticOPD improves mathematical reasoning accuracy over the OPD baseline at both tested Qwen3 scales, with the largest gains on competition style benchmarks such as AIME and AMC, and it outperforms FastOPD and RelayOPD as the best student method on most mathematical benchmarks. Ablations show that the critique is the most important feedback component, since removing it causes the largest performance drop, while answer only or generic feedback and teacher written suffixes have smaller effects. Feedback also matters more for scoring the suffix than for which model writes it, indicating that the critique improves both the rewritten suffix quality and the reliability of teacher scoring.


AIでAIを構築

アイデアからローンチまで — 無料のAIコーディング支援、すぐに使える環境、最高のGPU価格でAI開発を加速。

AI コーディング補助
すぐに使える GPU
最適な料金体系

HyperAI Newsletters

最新情報を購読する
北京時間 毎週月曜日の午前9時 に、その週の最新情報をメールでお届けします
メール配信サービスは MailChimp によって提供されています