PaperlayerView exact revision on arXiv

Source-linked AI summary

When to Follow the Teacher: Advantage-Based On-Policy Distillation

Mahdi Sabbaghi, George Pappas, Hamed Hassani, Adel Javanmard

arXiv:2610.31563v1cs.LGcs.AIstat.ML

TL;DR

Standard OPD matches teacher actions along student-generated trajectories without distinguishing which actions help the student succeed. AOPD weights those actions by student-relative advantage and estimates the objective from student rollouts and terminal rewards; it improves over typical distillation and RL baselines across external-teacher and self-distillation settings.

  • Problem

    Standard OPD supervises teacher actions along student-generated trajectories without distinguishing which actions improve the student’s chance of success.

    Source: §1 · p. 1

  • Method

    AOPD weights teacher actions by their advantage under the current student and estimates the objective using student on-policy trajectories and terminal rewards.

    Source: §1 · p. 2

  • Results

    Across external-teacher and self-distillation settings, AOPD consistently improves over typical distillation and RL baselines.

    Source: §6 · p. 11

  • Takeaways & Limitations

    Effective distillation should optimize for student success rather than teacher agreement.

    Source: §6 · p. 11

    The theoretical analysis considers a limited update from the frozen student to an updated policy under a KL-movement constraint.

    Source: §5 · p. 9

Abstract

from arXiv · show

Online Policy Distillation (OPD) often outperforms offline methods such as Supervised Fine-Tuning (SFT) by training on trajectories sampled directly from the student. However, standard OPD trains the student to follow the teacher at every state of the trajectory, assuming that teacher guidance is always beneficial. We instead distinguish between what the teacher prefers and what actually helps the current student succeed. We introduce Advantage-Based On-Policy Distillation (AOPD), which weights teacher actions by their counterfactual advantage under the student. Teacher actions that improve the student's expected outcome are reinforced, while those that do not are downweighted. Directly computing these advantages requires trying alternative teacher actions at each step and running new rollouts from each. Instead, we show that the AOPD objective can be estimated using only standard student rollouts and their terminal rewards. Our estimator uses teacher importance weights on the student-sampled actions to provide unbiased advantage-weighted supervision, without any additional rollouts or reward evaluations. This allows AOPD to focus teacher supervision on actions that help the current student while preserving the rollout efficiency of standard OPD. Experiments on math, coding, and Medqa show gains over distillation and RL baselines in both external-teacher and self-distillation settings.

1 Introduction

Standard on-policy distillation supervises the student at states it visits, but still treats teacher-preferred actions as useful regardless of their effect on student success. AOPD weights teacher actions by their advantage for the current student and estimates that objective from student rollouts; the paper reports gains over distillation and RL baselines and a reward-improvement guarantee under its stated analysis.

1 Introduction

  • Offline distillation can mismatch student-generated states, and standard OPD addresses that mismatch but still matches the teacher without distinguishing actions by their effect on success.The paper frames the central question as how to weight teacher supervision to improve the student’s eventual task success.Source: §1 · pp. 1–2
  • Prior selective-supervision approaches use auxiliary signals or counterfactual trajectories but do not directly define usefulness by its impact on the student’s expected task reward.The cited examples include teacher uncertainty, student plausibility, auxiliary reward models, and intervention rollouts.Source: §1 · p. 2
  • AOPD weights teacher actions by their advantage under the current student, reinforcing actions with positive advantage and suppressing those with negative advantage.The same teacher action can help one student and be unhelpful for another.Source: §1 · p. 2
  • An unbiased importance-sampling estimator uses student on-policy trajectories and terminal rewards, without additional rollouts or auxiliary models.The approach incorporates outcome-level rewards into token-level teacher supervision without a value function, counterfactual continuations, or an external process reward model.Source: §1 · p. 2
  • Experiments cover multiple model sizes and math and coding tasks in external-teacher and self-distillation settings, with improvements over OPD, OPSD, and GRPO.The paper also examines whether teacher actions are uniformly beneficial in an ablation study.Source: §1 · p. 2
  • AOPD has a reward-improvement guarantee under an exact KL-proximal analysis, while forward-KL OPD can reduce reward even with a stronger teacher.The paper also gives a two-state counterexample showing that simply adding an RL objective to OPD does not guarantee reward improvement.Source: §1 · p. 2

2 Related Work

Related work spans fixed offline supervision, on-policy reinforcement learning, on-policy distillation, and selective teacher supervision. AOPD differs by defining teacher-action usefulness through its advantage for the current student’s eventual task reward.

Offline supervision

  • Offline supervision uses fixed demonstrations or teacher trajectories, which can create covariate shift and compounding errors when the learned policy visits unseen states.The mismatch is especially relevant in autoregressive reasoning, where early deviations can alter later steps.Source: §2 · p. 3

On-policy reinforcement learning

  • On-policy RL learns from current-policy trajectories and directly optimizes task reward, but sparse outcome supervision can be difficult when useful trajectories are rarely sampled.The passage lists PPO, RLOO, and GRPO among related policy-gradient methods.Source: §2 · p. 3

On-policy distillation

  • OPD combines student-generated trajectories with dense teacher supervision, while OPSD uses a privileged version of the student as teacher; both generally supervise uniformly across the trajectory.OPSD may condition on ground-truth solutions or other privileged information.Source: §2 · p. 3

Selective teacher supervision

  • Selective-supervision methods use uncertainty, compatibility, validity, verifier signals, or intervention rollouts, whereas AOPD weights teacher actions by student-relative advantage estimated from existing rollouts and terminal rewards.These prior criteria do not directly define usefulness by its effect on eventual task reward.Source: §2 · p. 3

3 Method: Advantage-based On-Policy Distillation

AOPD retains teacher supervision at student-generated histories but weights each teacher action by its effect on the student’s expected reward. Its estimator uses student rollouts and terminal rewards, and matched importance-weight clipping preserves the stated objective and guarantee.

3 Method: Advantage-based On-Policy Distillation

  • Forward-KL OPD reinforces teacher-supported actions at student-generated histories without distinguishing whether they improve the student’s final task outcome.The method develops AOPD from the forward-KL objective; the passage says forward- and reverse-KL objectives roughly reach the same performance in practice.Source: §3 · p. 4
  • AOPD weights each teacher action by its advantage under the frozen rollout student, comparing expected reward after that action with expected reward under the student’s own policy.The objective retains teacher supervision at the same student-generated histories.Source: §3 · p. 4
  • Unlike forward-KL, AOPD gives positive weights to teacher actions that improve expected reward and negative weights to those that reduce it.The advantage and history distribution are defined under the frozen rollout policy; the paper states that its objective increases average reward with a guarantee, whereas OPD’s sign depends on teacher-action advantage.Source: §3 · p. 4
  • A student-sampled rollout and its terminal reward estimate the objective instead of requiring separate student continuations after teacher interventions.The estimator’s terms estimate the action-value and student-value contributions whose difference gives the desired advantage.Source: §3 · p. 4
  • The estimator uses teacher-to-student importance ratios so the teacher evaluates actions already sampled by the student.The stated unbiased-estimator lemma assumes that any action with positive teacher probability has positive student probability at each visited history.Source: §3 · pp. 3–4
  • Clipping caps large importance weights, and applying the same truncation to both terms preserves the advantage-weighted objective and its reward-improvement guarantee.The experiments use a clipping constant of c = 4.Source: §3 · p. 5

4 Experiments

The experiments compare AOPD with distillation and RL baselines in external-teacher and self-distillation settings across math, coding, biology, and medical question answering. AOPD performs strongly across these settings, while teacher interventions can both help and harm student outcomes.

4.1 Experimental Setup

  • The evaluation spans external-teacher and privileged self-distillation across math, coding, biology, and medical question answering, using Qwen2.5 and Qwen3 models of different sizes.The reported tasks include held-out Open-Thoughts, AIME24, AMC, and HMMT for math; HumanEval, HumanEval+, and MultiPL-E for coding; and held-out biology and medical QA benchmarks.Source: §4 · p. 5, §4.1 · pp. 5–6

4.2 External-Teacher Distillation

  • On external-teacher math, AOPD achieves the highest pass@1 on Open-Thoughts, AIME24, and AMC, while coding benchmark winners vary by metric.GRPO has the highest HumanEval pass@10, and OPD (FKL) performs best on MultiPL-E; AOPD improves during training on TACO while OPD makes little evaluation-reward progress.Source: §4.2 · p. 6, Table 1 · p. 7, Figure 2 · p. 7

4.3 Self-Distillation

  • In self-distillation math, AOPD outperforms the baselines on all three benchmarks for Qwen3-1.7B and achieves 52.1% on AIME24 with Qwen3-8B.For Qwen3-1.7B, AOPD reaches 35.0% on AIME24, 12.7 points above GRPO and 19.4 above OPSD; for Qwen3-8B, its AIME24 score is 10.4 points above OPSD.Source: §4.3 · p. 7, Table 2 · p. 8
  • With Qwen2.5-7B-Instruct, AOPD reaches 74.8% on MedQA and 69.8% on MedMCQA, exceeding OPSD by 2.7 and 3.2 percentage points, respectively.AOPD also achieves the highest biology accuracy on SciKnowEval Biology and MMLU-Pro with Qwen2.5-7B, and best HumanEval+ scores with both coding models; its MultiPL-E advantage over GRPO is specific to Qwen3-4B-Instruct.Source: §4.3 · pp. 7–8, Table 3 · p. 8
  • AOPD outperforms GRPO and OPSD across biology, medical question answering, and coding, and leads several external-teacher math and coding benchmark comparisons.In external-teacher math, it improves AIME24 by 2.4 percentage points and AMC by 3.3 points over the baselines; Table 3 reports its cross-domain self-distillation results.Source: §4.2 · p. 6, §4.3 · pp. 7–8, Table 3 · p. 8

4.4 Ablation: Is the teacher always useful?

  • At the highest-disagreement position, external-teacher interventions help in 9.6% of cases and harm in 9.8%; privileged self-teacher interventions help more often than external-teacher ones at random positions.At random positions, the helpful rates are 8.6% for the privileged self-teacher and 6.9% for the external teacher.Source: §4.4 · p. 8, Table 4 · p. 9

5 Theory: Reward Improvement with Advantage-based Distillation

The theory compares single, KL-proximal updates from a frozen student, showing that AOPD has a nonnegative reward effect under stated conditions while forward-KL OPD’s effect can have either sign. It also examines clipped supervision, self-distillation, and adding an RL objective to OPD.

Reward Improvement from a Single Proximal Update

  • For sufficiently small η, AOPD’s exact proximal update has a nonnegative reward effect, while the sign of the forward-KL OPD update remains indeterminate.The guarantee is for a single update in the proximal analysis.Source: §5 · p. 10

Proximal Update Conditions

  • AOPD’s small-step update preserves a valid probability distribution by limiting probability removed from any negative-advantage action to its current mass.The analysis holds the teacher, rollout distribution, and student advantage fixed at their current-student values.Source: §5 · pp. 9–10

Surpassing the Teacher

  • When the student equals the teacher, a sufficiently small AOPD update can still improve the policy and achieve reward at least as high as the teacher’s.The corollary frames the teacher as action-level supervision rather than a fixed target.Source: §5 · p. 10

Reward Improvement with Clipping

  • Clipped teacher weights can use the same reward-improvement argument when the truncated weights are nonnegative and proximal coefficients remain positive.The stated clipped update assumes c ≥ 1 and 0 < η < 1/c.Source: §5 · p. 10

RL plus FKL-OPD

  • Adding an explicit RL reward term to forward-KL OPD does not guarantee improvement: a two-step task can have a higher-reward teacher yet a reward-decreasing update.The counterexample applies for every fixed λ > 0 and sufficiently small η.Source: §5 · p. 10

6 Conclusion

AOPD weights teacher actions by their advantage under the current student, reinforcing useful guidance without extra counterfactual rollouts or auxiliary value models. Across external-teacher and self-distillation settings, it improves over typical distillation and RL baselines.

Conclusion

  • Across external-teacher and self-distillation settings, AOPD consistently improves over typical distillation and RL baselines.It weights teacher actions by student-relative advantage and requires no additional counterfactual rollouts or auxiliary value models.Source: §6 · p. 11

A Additional Experiments

The additional experiments examine group size and importance-weight clipping. Larger groups yield little benefit for AOPD, while clipping improves its stability and performance.

Group size effect on OPD and AOPD

  • Increasing group size yields only negligible reward improvement for AOPD, and that gain does not transfer to better benchmark numbers.The main experiments use group size G = 4 for AOPD and OPD.Source: A Additional Experiments · p. 16, Figure 4 · p. 16

Clipping avoids performance fluctuations

B Proofs: Reward Improvement

The proofs relate each update’s local student advantage to its reward effect. They show that AOPD’s squared-advantage term is nonnegative and that this local property carries through to final reward under the stated support conditions.

Forward-KL reward effect

  • Forward-KL OPD’s reward effect depends on the teacher’s average student advantage, which can have either sign.The proof drops terms independent of the policy before optimizing at each history.Source: B Proofs: Reward Improvement · pp. 17–18

AOPD local reward effect

  • For sufficiently small η, AOPD’s expected student advantage is nonnegative because its numerator is an expectation of squared advantages and its denominator is positive.The proximal objective separates across histories.Source: B Proofs: Reward Improvement · pp. 18–19

From local advantage to final reward

  • The local result gives a nonnegative final-reward change, with strict improvement when a reached history has a teacher-supported action with nonzero student advantage.The proof applies the local result along complete responses and takes expectations over responses.Source: B Proofs: Reward Improvement · p. 19

Reward improvement with clipping

  • The clipped-objective proof substitutes truncated teacher weights into the AOPD proximal update and applies the same performance-difference argument.It requires the resulting proximal probabilities to remain positive.Source: B Proofs: Reward Improvement · pp. 19–20

B.1 Surpassing the teacher

Setting the student equal to the teacher in Theorem 5.1 gives the strict-improvement condition, while AOPD can continue improving beyond the teacher by reweighting its actions according to task advantage.

B.1 Surpassing the teacher

  • Setting S = T in Theorem 5.1 yields the strict-improvement condition.The section then gives the corresponding FKL update.Source: §B.1 · p. 20
  • When the student matches the teacher, FKL makes no further update.The FKL update is introduced after setting S = T.Source: §B.1 · p. 20
  • AOPD can continue improving beyond the teacher by reweighting teacher actions according to their task advantage.Source: §B.1 · p. 20

B.2 Adding RL does not guarantee improvement

The section analyzes how an FKL-plus-RL update can reduce reward and gives a two-step counterexample, despite the RL contribution being nonnegative.

B.2 Adding RL does not guarantee improvement

  • The analysis expands the proximal update around the student policy, whose closeness to that policy is maintained by the KL penalty.Its first-order direction combines a reward-improving term with a term moving the student toward the teacher.Source: §B.2 · pp. 20–21
  • The RL contribution is nonnegative, but the FKL contribution can make the update reduce reward.For a fixed λ > 0, the constructed task admits sufficiently small ε and η for which the reward change is negative.Source: §B.2 · pp. 20–22
  • In the two-step counterexample, the student chooses the difficult route with probability ε and succeeds there with probability ε, while the teacher chooses it with probability 1 − ε.The teacher’s success probability is 1 − ε on either route.Source: §B.2 · p. 21
  • The teacher has higher expected reward, yet the student’s expected teacher advantage at the difficult-route history is negative for 0 < ε < 1/3.The safe-route history has zero expected teacher advantage.Source: §B.2 · pp. 21–22

Related papers

Loading 2610.31563v1…