When to Follow the Teacher: Advantage-Based On-Policy Distillation
Standard on-policy distillation forces a student to copy a teacher at every token, even when the teacher's preferred actions hurt the student's eventual success. This paper proposes Advantage-Based On-Policy Distillation (AOPD), which weights each teacher token by its counterfactual advantage under the current student policy. The key insight is that imitation should optimize for student task reward, not teacher agreement, and the authors derive an unbiased estimator that needs only standard student rollouts and terminal rewards.
AOPD is a well-motivated and theoretically grounded refinement of on-policy distillation. The shift from uniform teacher imitation to advantage-weighted supervision is conceptually clean, and the paper provides both a formal proximal-update analysis and empirical gains on math and reasoning tasks. That said, the practical estimator relies heavily on importance-sampling ratios that require strong support assumptions and clipping heuristics, and the empirical benefits are uneven across domains—sometimes trailing simpler distillation or RL baselines in coding and biology.
The ablation in Section 4.4 directly validates the central premise: single-token teacher interventions at high-disagreement positions help in only 9.6% of cases and harm in 9.8%, confirming that teacher-preferred actions are not uniformly beneficial. The proximal theory (Theorem 5.1) crisply separates AOPD from forward-KL OPD by showing that AOPD's exact local update guarantees non-negative reward improvement $\mathcal{R}(S_{\mathrm{AOPD}}^{+}) \geq \mathcal{R}(S)$, whereas standard OPD's reward effect has indeterminate sign. The unbiased estimator (Lemma 3.1, Equation 3.6) is elegant because it requires no extra rollouts, value networks, or process reward models—only terminal rewards and per-token importance weights.
The theoretical guarantees assume exact proximal updates and a small-step condition ($0 < \eta < 1/c$ with $c=4$) that may not hold under neural-network training with Adam and LoRA adapters. The estimator's unbiasedness relies on a full-support assumption—"$T(a \mid h) > 0$ implies $S(a \mid h) > 0$ at every visited step"—which the authors address only by disabling top-p/top-k for AOPD, a protocol difference that weakens comparison fairness. Empirically, AOPD is not uniformly dominant: in external-teacher coding it trails OPD (FKL) on MultiPL-E (32.5 vs. 33.3), and in self-distillation with Qwen3-4B-Instruct on SciKnowEval it drops to 32.0 vs. 44.0 for OPSD. The single-token ablation uses only 512 responses with greedy continuation, giving suggestive but limited statistical power.
The evidence supports the core claim most strongly in math: AOPD improves AIME24 by 2.4 points and AMC by 3.3 points over baselines with Qwen3-1.7B, and in self-distillation it lifts Qwen3-8B AIME24 from 41.7% (OPSD) to 52.1%. However, the comparisons are not fully controlled: AOPD and OPD use group size $G=4$ while GRPO and KDRL use $G=8$, which matters because variance reduction scales with group size. Moreover, the paper reports no confidence intervals or standard deviations, making it impossible to assess whether modest gains are statistically significant. The mixed coding results—where GRPO wins on HumanEval pass@10 and OPD wins on MultiPL-E—suggest AOPD's advantage-weighting is domain-dependent.
The authors release code and report concrete hyperparameters: LoRA $(r=64, \alpha=128)$, learning rate $2 \times 10^{-5}$, clipping constant $c=4$, and temperature $T=0.8$ for AOPD. Still, reproduction is complicated by method-specific sampling: AOPD disables top-p and top-k to preserve support, whereas baselines use $T=1.1$, top-p $=0.95$, and top-k $=20$—a confound for anyone benchmarking wall-clock or sample efficiency. An independent reproduction would also need to extract per-token probabilities from both student and teacher for every step of every rollout, which is feasible but computationally nontrivial. The paper does not specify optimizer settings beyond the learning rate or the exact compute budget in GPU-hours, leaving some gaps for replication.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.