On-Policy Distillation Explained: White-Box, Self-Distillation & Black-Box Methods
This article provides a comprehensive analysis of On-Policy Distillation (OPD), covering white-box distillation with token-level probability matching, self-distillation using privileged context, black-box distillation via reward modeling, and industrial applications from MiMo, DeepSeek, and Qwen, while discussing OPD's limitations and generalization capabilities.
01. White-Box OPD: From Training Principles to Industrial Practice
White-box OPD starts with the student generating its own response. The teacher computes next-token probabilities on the student's prefixes, and the student updates its parameters using the probability difference.
Deriving Sequence-Level Update from Full Reverse KL
Let the teacher be a fixed distribution π_T and the student π_S. Using reverse KL (weights by student probability, penalizes student-high/teacher-low tokens), the full-sequence objective for a fixed input x and student-generated response y with prefix y_{<t} at step t is:
L = E_{y∼π_S}[ Σ_t log π_S(y_t|y_{<t},x) - log π_T(y_t|y_{<t},x) ]Defining the teacher signal Δ_t = log π_T(y_t|y_{<t},x) - log π_S(y_t|y_{<t},x), the objective becomes E[ Σ_t -Δ_t log π_S(y_t|...) ]. The gradient contains cross-position effects because parameter changes alter both current token probability and future prefixes visited. The sequence-level policy gradient is:
∇L = E_{y∼π_S}[ Σ_t (Σ_{k≥t} Δ_k) ∇log π_S(y_t|y_{<t},x) ]This matches MiniLLM §2.2 and Appendix A.2 [5].
Why Token-Level Approximation?
Sequence-level uses future cumulative return to update earlier choices. The difficulty: we want the average future performance of a choice, but a single trajectory only shows one outcome. A two-step geometry example illustrates: first step "set up equation" or "draw auxiliary line" (both 0.5 teacher/student). After "draw auxiliary line", two continuations AC (teacher 0.9, student 0.3, signal +1.10) and BD (teacher 0.1, student 0.7, signal -1.95). Student samples AC 30%, BD 70%; average future return ≈ -1.03 vs 0 for "set up equation", so average gradient should reduce "draw auxiliary line". But a single sample of AC gives +1.10, pushing the opposite direction — this is variance. Token-level approximation treats the generated prefix as fixed and only optimizes current teacher-student distribution difference, dropping future signals. It removes sampling noise but introduces bias relative to the full objective.
Local Gradient Computation: Sampled-Token, Full-Vocab, Top-K
At a fixed prefix, the local reverse KL is Σ_{v∈V} π_S(v) [log π_S(v) - log π_T(v)]. Its gradient can be estimated by sampled-token (one sample from π_S, weight implicit in sampling frequency) or computed exactly via full-vocab (sum over all tokens weighted by π_S). Both have same expected gradient; full-vocab eliminates current-token sampling variance. Top-K/Top-p truncates the vocab to a set S (chosen by teacher, student, or union). Renormalizing probabilities inside S changes the local objective and ignores probability mass outside S (e.g., teacher 10% outside, student 90% outside → truncated KL zero). Unbiased estimation requires tail correction.
Training Implementation: PPO-Style vs Direct Distribution Optimization
PPO-style reuses RL infrastructure (e.g., verl). The PPO policy loss L = -E[ min(r_t A_t, clip(r_t,1-ε,1+ε) A_t) ] keeps the same surrogate; just replace advantage A_t with teacher signal (current Δ_t for token-level, cumulative Σ_{k≥t} Δ_k for sequence-level). PPO can use sampled-token or full-vocab. Direct distribution optimization fixes student-generated prefixes and directly minimizes a divergence (reverse KL, forward KL, JSD) between teacher and student distributions. On a fixed prefix with local reverse KL weight and stop-gradient on teacher, the expected PPO gradient equals the full-vocab reverse KL gradient initially, but later updates diverge due to old weights and clipping.
MiniLLM and GKD: Design Choices
MiniLLM [7] keeps long-range terms via single-step decomposition: current step uses full-vocab reverse KL gradient; future steps use sampled cumulative return Σ_{k>t} Δ_k. It adds: (1) teacher-mixed sampling π_mix = (1-α)π_S + απ_T to reduce degenerate generations; (2) importance weight approximation using single-step ratio to correct for mixed sampling; (3) length-normalized future return to avoid penalizing longer generations; (4) PPO-style clipping on policy gradient term. These make the update a biased estimate of the original full reverse KL gradient.
GKD [8] follows direct distribution optimization. Student generates a batch of responses; prefixes are fixed; at each position the full teacher and student distributions are compared (forward KL, reverse KL, or JSD). JSD uses an intermediate distribution M = βπ_T + (1-β)π_S and averages KL(π_T||M) and KL(π_S||M). The paper leaves two choices to the practitioner: (1) which responses to train on (controlled by a selector); (2) which divergence to use. Forward KL emphasizes covering teacher modes; reverse KL concentrates on a subset; JSD interpolates.
Industrial OPD: Multi-Teacher Merging & Small Model Distillation
Multi-teacher merging : MiMo-V2-Flash [4] uses SFT → domain experts → MOPD merge. Math questions get math teacher, code questions get code teacher. Teacher signals and outcome rewards are combined additively in the advantage: A = Δ_t + λ A_outcome. DeepSeek-V4 [3] also trains domain experts then merges via OPD, choosing full-vocab reverse KL to reduce token sampling noise, with engineering optimizations: on-demand teacher weight loading, caching last-layer hidden states, custom KL kernel. Small model distillation : Qwen3 [2] first does off-policy distillation with teacher-generated answers (thinking/non-thinking modes), then OPD from the same strong teacher. Table 21 shows from the same off-policy distilled Qwen3-8B checkpoint, RL adds 17,920 GPU hours for AIME'24 67.6 → 74.4, while OPD adds only 1,800 GPU hours (≈1/10) and achieves higher scores (74.4 vs 67.6). Teacher training cost is separate.
02. Self-Distillation: Let a More-Informed Copy of Yourself Teach
Self-distillation exploits the context information gap : the same model, when given extra context (reference answer, execution feedback, corrected prompt, reusable experience, or accumulated deployment experience), can provide better token-level guidance. The extra context is privileged information only seen by the teacher during training. One round: (1) student generates y from normal input x; (2) teacher sees x + z (privileged info) and recomputes probabilities along y 's prefixes; (3) student updates via distribution matching (gradient only through student branch).
OPSD: Reference Answer in Teacher Context
OPSD [16] uses existing (question, answer) pairs. Student sees only question; teacher sees question + reference solution. Example: solving 2x=8, student writes x=8; teacher with reference solution prefers x=4 at that prefix. Teacher does not regenerate; one forward pass on (question, reference, student_prefix) yields probabilities. Teacher initialized from student, then frozen.
SDPO: Successful Attempts & Environment Feedback as Token-Level Supervision
SDPO [18] constructs privileged context from the student's own attempts: (1) successful attempt on the same problem — e.g., one passing code among three generations guides the failing ones; (2) environment feedback (error messages, failed unit tests) when no success exists. Example: counting non-negatives, student writes n > 0, test [-1,0,2] expects 2, gets 1. Teacher sees failure feedback, lowers probability of > and raises >= at the if n prefix. Feedback directly re-evaluates the original response without generating a corrected version.
Cursor: Corrective Prompt for a Specific Step
In long-horizon agents, an intermediate tool call may be wrong even if final outcome succeeds. Cursor Composer 2.5 [19] constructs a short corrective prompt for that step (e.g., list of available tools) and adds it to teacher's local context. Teacher re-evaluates the original erroneous call, lowering probability of the wrong tool name. Distillation loss is added only for that step alongside the full-trajectory RL objective.
OPCD: Distilling Reusable Prompts & Experience into Parameters
OPCD [20] focuses on reusable context c (optimized system prompt, summarized experience). Student sees x; teacher sees c + x. Experience c can be mined by having the model attempt a batch of problems, then summarizing patterns (e.g., "find unit-digit cycle, then modulo exponent") without requiring ground-truth answers. Distillation uses reverse KL along student-generated responses.
OEL: Continuous Learning from Deployed Interactions
OEL [21] connects OPCD to a deployment loop: model acts in user environment → server collects interaction logs → extracts experience → OPCD updates model → redeploy. Teacher's privileged context comes from actual interactions. Training re-generates only the current step from historical context h (student sees h, teacher sees h + experience), avoiding re-execution of the environment. On-policy here means the response comes from current student policy; old-task regression must still be monitored separately.
RLSD: Outcome Determines Sign, Teacher Adjusts Weights
RLSD [22] also gives teacher the ground-truth answer, but changes teacher's role: outcome verifier produces response-level advantage A (like GRPO); teacher compares per-token probabilities with student, converts difference to positive weight w_t = clip(π_T/π_S, 1/3, 3), multiplies onto A. For A>0, teacher-preferred tokens get more reinforcement; for A<0, teacher-dispreferred tokens get more suppression. Weight clipping keeps weights positive. Final token advantage is blended with original A. Unlike MiMo's additive combination, RLSD's multiplicative scheme prevents flipping the sign of the advantage, aiming to limit privileged information leakage (e.g., student learning to say "according to reference answer...").
03. Black-Box Distillation: OPD Without Teacher Probabilities
When teacher only returns text via API, we cannot compute token-level probability differences. Two methods follow: student generates → construct response-level reward from teacher text → RL update student . They keep on-policy sampling but optimize a teacher-text-derived reward, not white-box distribution matching.
GAD: Train a Discriminator on Teacher/Student Responses
GAD [24] trains a discriminator D (initialized from student + scoring head) to score (question, response) pairs. Teacher provides reference texts. Alternating steps: (1) student samples responses; (2) fixed D scores them, scores used as reward for GRPO update of student; (3) update D with Bradley-Terry pairwise loss on (teacher response, student response) pairs: min -log σ(D(y_T) - D(y_S)). Student reward is D(y_S), not backprop through discriminator loss. Risk: reward hacking — discriminator score may rise while true quality stagnates (e.g., favoring verbose, teacher-like phrasing). Continuous discriminator updates mitigate but don't eliminate this; independent task metrics are essential.
Rubric OPD (ROPD): Teacher Writes Rubric, Verifier Scores Items
ROPD [25] makes scoring criteria explicit. Rubricator reads question, multiple teacher responses, and current student responses, outputs a weighted rubric (e.g., "final answer 5050" weight 4, "valid summation method" weight 3, "arithmetic consistency" weight 2, "explains pairing" weight 1). Same rubric used for all responses to that question. Verifier blindly judges each response against each rubric item (teacher responses mixed in as quality anchors). Student reward = weighted pass rate. Only student rewards used for GRPO update. Rubric transparency helps audit what is rewarded, but verifier errors and gameable items (e.g., "mentions checking" vs "actually checks correctly") remain risks. Trade-off: GAD trains/maintains a scoring model; ROPD pays per-batch rubric generation and scoring cost.
Both convert teacher text to response-level reward, so no tokenizer alignment or teacher logits needed. But they still compute policy gradients through student token probabilities; the reward is not equivalent to the sequence-level distillation objective from full reverse KL. Success depends on whether the constructed reward reflects desirable behaviors.
04. Essence and Boundaries of OPD
Common post-training algorithms have known issues: SFT suffers train-inference mismatch; offline DPO limited by fixed preference pairs; RLHF with reward models risks reward hacking; RLVR needs reliable verifiers and gives sparse signal (only final correctness).
Why Strong Teacher May Fail to Teach
Three distinct failure modes: (1) Teacher cannot handle student's erroneous prefixes — teacher solves from scratch but fails to correct a student's half-written wrong derivation (e.g., student does 2+3=5 then 5×4 for 2+3×4; teacher continuing from that prefix also yields 20). Must test teacher continuation on actual student prefixes. (2) Training signal encourages useless behavior — student repeats "Wait, Wait..." and teacher assigns higher probability to those tokens, giving positive local signal, but repetition doesn't help solve the task. Signal correctness ≠ task progress. (3) Teacher's method is correct but student lacks execution skill — teacher suggests "differentiate then set to zero" for maximization, but student cannot differentiate correctly. Making student prefer a method ≠ teaching every execution step. Diagnosis: inspect samples with per-token distillation signals (as in Fig 4-1 from TML blog [32]) to see where errors start, whether teacher gives corrective signals, whether useless repetitions are rewarded. Then adjust: better teacher / privileged context for (1); add outcome reward for (2); SFT cold-start or stronger base for (3). Final validation on held-out tasks, not just KL decrease.
Single-Question Training Can Generalize
Thinking Machines Lab [33] and Rethinking OPD II [34] observe that training on a single fixed question improves performance on other questions. TML: Qwen3-8B-Base → RL teacher → original base student learns from that teacher on one question for 20 steps (256 responses/step) → approaches teacher on AIME'24. Rethinking OPD II: student never solves the hard training question, yet other test scores rise; at 1,000 steps single-question training achieves ~72% of full-data training gain. Explanation: even with fixed question, student explores different intermediate steps (e.g., 12×13 directly, 12×(10+3), correct split but wrong arithmetic 120+26). Teacher provides per-prefix next-token distributions; local signals on erroneous steps (e.g., low teacher probability on 26) push student away from those errors. If these updates improve behaviors used in other problems (arithmetic, decomposition), cross-task gains emerge. Rethinking OPD II's state analysis [35]: one question's trajectories cover 71.5% of hidden-state clusters visited by full-data training; 16 semantically diverse questions cover 98.9%. Implication: for a student with existing base, transferring part of teacher's post-training gains may need very few representative questions. Practical approach: validate transfer with a small representative set, then augment based on failure types on unseen tasks.
Conclusion
OPD is attractive: student samples around an input, improves on other tasks. Teams can verify gains with few representative inputs, then augment around real failures. Black-box route (teacher text → reward) is viable when probabilities unavailable. Self-distillation combined with environment feedback (answers, failed tests, interaction logs) lets the model's own privileged views guide learning, potentially a path toward recursive self-improvement. But boundaries are clear: white-box OPD optimizes distribution matching, not task completion directly. Student practices its own real generations, but teacher must guide on prefixes it rarely produces — guidance may be unreliable. Training closer to student ≠ more reliable teacher guidance; matching teacher ≠ better task performance. Whether using teacher probabilities or teacher-text-derived rewards, final validation must be on real tasks.
Signed-in readers can open the original source through BestHub's protected redirect.
This article has been distilled and summarized from source material, then republished for learning and reference. If you believe it infringes your rights, please contactand we will review it promptly.
Machine Learning Algorithms & Natural Language Processing
Focused on frontier AI technologies, empowering AI researchers' progress.
How this landed with the community
Was this worth your time?
0 Comments
Thoughtful readers leave field notes, pushback, and hard-won operational detail here.
