Critique-Guided On-Policy Distillation

Why People Wanted It

On-policy distillation asks a teacher to supervise continuations from states the student actually reaches. But after the student makes an error, the teacher may no longer be reliable at continuing from that state. Simply handing the trajectory to the teacher to repair changes who wrote the training tokens; ignoring the bad region leaves the student without useful guidance. So the real question is: can we keep the student's trajectory, make the problematic part recoverable, and give it supervision that is reliable without requiring feedback at inference?

The Core Move

CriticOPD has the teacher identify errors and actionable fixes, lets the student rewrite from the last identified error with that critique in context, and has the teacher score the rewritten suffix with the critique visible while the student is scored without it.

The key distinction is that the critique changes the teacher's target context, not the student's inference-time context. No new closed-form objective is needed here: the source gives the sampled-token advantage and states that the overall objective combines standard OPD with feedback distillation.

Quiz

In CriticOPD's feedback distillation, which contexts are used to score a rewritten suffix token?

The student sees the critique; the teacher does not.

Both teacher and student see the critique.

The teacher sees the critique; the student does not.

Neither sees the critique.

$ check

Why the Last Error Sets the Cut Point

Here is the part that gets technical: the retained prefix is chosen to preserve the student's history up to the paragraph containing the last error, while discarding everything after it. The critique lists each error's paragraph index, erroneous sentence, and short fix, and also includes the reference answer. If there is no complete error in the critique, the rollout instead receives standard OPD.

The retained prefix is denoted $h_\tau=(x,y_{<\tau})$. In words, it contains the problem and the original student tokens before the cut point. The rewrite begins at the last identified error so that it starts from the erroneous states the critique addresses; the prefix itself is not rewritten.

Why the Student Must Write the Rewrite

Now the version for readers who want more of the machinery: the student appends the critique to the retained prefix and samples a new suffix. The tokens before the cut remain unchanged, and the student—not the teacher—writes every token in the revised training sequence. The concrete example in the paper has two identified errors: the rewrite starts at the later one and uses the critique to reach the correct answer.

This is a deliberate boundary. The teacher supplies an error diagnosis and later a score, but does not author the replacement suffix. Why score that student-authored suffix with feedback still attached?

Why Feedback Changes the Teacher Target, Not the Student Input

Here is the technical distinction that matters: the teacher evaluates each rewritten suffix token with the critique in context, while the student's log-probability is computed from the same rewritten prefix without the critique. The source gives the sampled-token advantage as $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)$. In words, it compares the teacher's probability with critique visible against the student's probability without critique.

Only rewritten suffix tokens receive this feedback supervision; the retained prefix still contains the errors that motivated the cut. Because the suffix was sampled with critique context but the student's probability is computed without it, the update is a surrogate, not an unbiased gradient of the corresponding objective. Correct rollouts and rollouts without a valid critique use standard OPD; CriticOPD's overall loss combines that term with feedback distillation, with each rollout entering one term.

The distinction lets the student learn a continuation from privileged guidance without needing that guidance at inference. What tradeoff comes with asking the student to learn from a suffix sampled under context it will not have later?

Minimal Causal Code

This loop makes the context asymmetry and suffix-only update explicit. teacher_score and student_logprob stand in for model calls; they are not optimized for speed.

for rollout in batch:
    if rollout.is_correct or not rollout.has_valid_critique:
        for token in rollout.tokens:
            loss += opd_loss(rollout.prefix_before(token), token)
        continue

    prefix = rollout.tokens_before(rollout.last_error_cut)
    critique = rollout.critique
    suffix = student.sample_suffix(prefix, extra_context=critique)

    for token in suffix:
        teacher_lp = teacher.logprob(prefix, token, extra_context=critique)
        student_lp = student.logprob(prefix, token)
        advantage = teacher_lp - student_lp
        loss += feedback_distillation_loss(student_lp, advantage)
        prefix = prefix + [token]

What You Gain, and What You Lose

The method keeps the student's own rollout states and has the student author the rewritten sequence, while critique context gives the teacher a more useful basis for scoring the suffix. It also leaves the student without critique at inference, so the training signal and inference context are intentionally different. That makes the feedback update a surrogate rather than an unbiased gradient, and the method depends on critiques that identify complete errors and fixes. The design is not simply standard OPD made faster: it changes the rewrite, the teacher's context, and which suffix tokens receive feedback supervision.

Quiz

Why does feedback supervision apply only to the rewritten suffix?

The critique is removed before the rewrite begins.

The retained prefix contains errors at states where teacher recoverability is low.

The teacher writes the suffix instead of the student.

The student did not generate the retained prefix.

$ check