--- kind: concept slug: on-policy-distillation title: On-policy distillation summary: >- A student writes its own rollouts and a teacher scores each token, a dense signal instead of one outcome reward. Variants differ in how they use the teacher's probabilities, and a reverse-KL baseline cuts the variance of the one-token signal. Trade-offs: teacher probabilities on every rollout, and a teacher choice that depends on student size and task. Based on three papers. tags: [distillation, on-policy, reasoning, variance-reduction, diagnostics] papers: ["2509.20322", "2605.07865", "2605.10889"] related: [knowledge-distillation, policy-gradient, baseline-subtraction, rlhf, preference-optimization, large-reasoning-models, sim-to-real-transfer] status: draft updated_at: 2026-10-31 --- ## Overview In on-policy distillation the student generates the responses it is trained on, and a teacher assigns credit to the student's choices along the way. The teacher's distribution over next tokens is compared with the student's at each node of the rollout tree, so the training signal is dense, one value per token, instead of one reward per answer [arxiv:2605.10889 §1]. The paper contrasts this with the sparse outcome reward used by RL methods such as GRPO [arxiv:2605.10889 §1]. The language-model variants share the same per-node structure. Each writes the distillation gradient as a function of the student's probability vector and a teacher-derived quantity at that node [arxiv:2605.10889 §2.3] [eq. 2.11]. They differ in whether they use the full teacher distribution, only the sampled token's probability, or a term that depends on the rest of the trajectory [arxiv:2605.10889 §2.3] [app. B.4] [app. B.5]. A second line of work keeps the cheap sampled-token estimator and reduces its noise. The one-token estimator is high-variance, and the authors of the reverse-KL baseline report that training becomes unstable with it [arxiv:2605.07865 §2.1] [arxiv:2605.07865 §1]. Their fix subtracts a per-token control variate that the forward pass already provides [arxiv:2605.07865 §3.1]. Outside language models, the same on-policy idea appears in robot control. There the student is trained on states its own policy reaches, with DAgger, and the teacher is a policy with access to privileged state [arxiv:2509.20322 §3.1] [arxiv:2509.20322 §3.2]. The distillation target is a control policy, not a per-token distribution, and the card describes no token-level KL term. The paper also proposes a diagnostic that scores a teacher configuration without training on it. It measures, per token, whether the distillation gradient points in the same direction as an estimated "ideal" gradient for success [arxiv:2605.10889 §2.4]. The paper frames this as an offline evaluation tool, not a training signal [arxiv:2605.10889 app. C]. Whether the teacher's guidance helps depends on the student and the task. The paper's central finding is that the guidance aligns better with success on incorrect rollouts than on correct ones [arxiv:2605.10889 §4.1]. ## Variants ### Full-distribution GKD gradient Each node uses the student's probabilities and the teacher's probabilities over the sufficiently visited tokens of that node. The objective is the KL divergence from student to teacher, written as KL(student || teacher), and its gradient at each node is the student-weighted deviation of the per-token log-ratio from its mean [arxiv:2605.10889 §2.3] [eq. 2.6] [eq. 2.8]. The paper analyses this form under the name GKD, which it attributes to Agarwal et al. (ICLR 2024) without an arXiv id [arxiv:2605.10889 §2.3]. This is the only variant the paper's experiments score [arxiv:2605.10889 §3]. Card: [Unmasking On-Policy Distillation](/papers/2605.10889). ### Single-sample estimator The estimator used in the "On-Policy Distillation" write-up by Thinking Lab uses only the token the student actually sampled at each node. Its weight is the teacher's log-probability of that token minus the student's, minus one. In expectation it gives the full-distribution gradient with the opposite sign [arxiv:2605.10889 §2.3] [eq. 2.9] [app. B.4]. The paper names this estimator but does not give an arXiv id for the write-up [arxiv:2605.10889 §2.3]. Card: [Unmasking On-Policy Distillation](/papers/2605.10889). ### Reverse-KL baseline (vOPD) Keeps the single-sample estimator and subtracts a baseline. The baseline is the negative reverse KL between student and teacher at the same context. It is computed from the student and teacher distributions the forward pass already produces, so no critic and no extra rollouts are needed [arxiv:2605.07865 §3.1]. The advantage for the sampled token is the usual teacher-minus-student log-probability plus that KL, with the KL detached [arxiv:2605.07865 §3.1] [eq. 10]. Detaching matters: the unbiasedness argument assumes the baseline is a constant in the gradient [arxiv:2605.07865 §2.2] [app. A.1]. The baseline has no dependence on the sampled token, so the expected gradient is unchanged [arxiv:2605.07865 §2.2] [eq. 8]. The authors derive an approximate per-token variance reduction that grows with the squared reverse KL. That derivation rests on a weak-correlation assumption the authors motivate but do not prove [arxiv:2605.07865 app. A.3] [eq. 11]. It predicts the largest reduction at tokens where student and teacher disagree most [arxiv:2605.07865 §4.3]. Because the KL is non-negative, the advantage is never below the raw reward; large negative rewards at high-disagreement tokens are pulled toward zero [arxiv:2605.07865 §4.3] [fig. 1]. A top-k version computes the KL only over the student's k most likely tokens, renormalising both models over that set. The baseline still ignores the sampled token, so the gradient stays unbiased for any k. The reward itself is not restricted [arxiv:2605.07865 §3.2] [eq. 14] [eq. 15]. The paper reports that the baseline is nearly insensitive to k on the math average, with k between 5 and 100 [arxiv:2605.07865 §4.3] [tab. 5]. Card: [KL for a KL: On-Policy Distillation with Control Variate Baseline](/papers/2605.07865). ### Trajectory-level REINFORCE form MiniLLM is written in a REINFORCE-style form. Its per-node factor depends on the full downstream trajectory, not only on the local distributions [arxiv:2605.10889 §2.3] [eq. 2.10] [app. B.5]. The paper cites MiniLLM (Gu et al., ICLR 2024) without an arXiv id [arxiv:2605.10889 §2.3]. Card: [Unmasking On-Policy Distillation](/papers/2605.10889). ### Teacher-labelled student rollouts (DAgger) Used for a humanoid robot, not for language models. A teacher with privileged state (object positions, contact forces) is trained first with PPO. A student that sees only depth images and proprioception is then distilled from it with DAgger, in two places: a motion tracker and a task-specific keypoint generator [arxiv:2509.20322 §3.1] [arxiv:2509.20322 §3.2]. The paper gives no per-token loss in the cited sections. Iteration counts and dataset sizes are not stated [arxiv:2509.20322 §3.1] [arxiv:2509.20322 §3.2]. Card: [VisualMimic](/papers/2509.20322). ### Gradient alignment diagnostic Not a training loss. At a node, the score is the cosine similarity between the ideal gradient for success and the distillation gradient, both restricted to the same set of tokens [arxiv:2605.10889 §2.2] [arxiv:2605.10889 §2.4] [eq. 2.12]. The ideal gradient is built from estimated success probabilities per token, computed by resampling rollouts from under-visited prefixes [arxiv:2605.10889 §2.5]. The score lies between −1 and +1, where positive means the distillation gradient points toward success [arxiv:2605.10889 §2.4]. Card: [Unmasking On-Policy Distillation](/papers/2605.10889). ## When to use which | Situation | Use | Why | |---|---|---| | Teacher probabilities are available for all tokens at each node | Full-distribution GKD gradient | The gradient uses the teacher's probabilities over the node's token set [arxiv:2605.10889 §2.3] [eq. 2.8] | | Only the teacher's log-probability of the sampled token is available | Single-sample estimator | The weight uses only the sampled token [arxiv:2605.10889 §2.3] [eq. 2.9] | | Sampled-token training is noisy, and a critic or extra rollouts are not wanted | Reverse-KL baseline | The baseline comes from the existing forward pass. On Qwen3-1.7B the paper reports a math average of 33.0 for vOPD top-k against 29.8 for vanilla OPD [arxiv:2605.07865 tab. 2] | | Full-vocabulary KL is too costly per step | Reverse-KL baseline with top-k | The paper reports wall-clock time cut by up to 57.7% against full-vocabulary OPD [arxiv:2605.07865 §1] [fig. 3] | | Full-vocabulary KL is affordable and the top-k bias is a concern | Full-vocabulary KL with the baseline | The paper's full-vocabulary variant scores close to top-k on the same models [arxiv:2605.07865 tab. 2]. Plain top-k OPD, which biases the gradient, gains little over vanilla OPD [arxiv:2605.07865 §4.2] | | Credit should depend on the rest of the response | Trajectory-level form | The per-node factor depends on the downstream trajectory [arxiv:2605.10889 eq. 2.10] | | Deciding between teachers or teacher contexts before a full training run | Alignment diagnostic | It scores configurations per token and per question [arxiv:2605.10889 §2.4] | | Choosing a teacher for a new student size or task | Run the diagnostic per task | Teacher rankings change with student size and task [arxiv:2605.10889 §4.2] [tab. 1] | | The teacher sees state the student cannot observe, and the student must act from sensors in a simulator | Teacher-labelled student rollouts (DAgger) | The student is trained on its own visited states with a teacher distilled into it [arxiv:2509.20322 §3.2]. In simulation the vision student falls well below the state teacher on Push Box distance, 37 ± 28 m against 152 ± 36 m, and matches it on some rows, such as Lift Box height [arxiv:2509.20322 tab. 2] | The reverse-KL baseline rows come from one paper's setup: Qwen3 students from 0.6B to 4B and Olmo-3-7B, on math and chemistry benchmarks [arxiv:2605.07865 §4.1]. The paper's experiments stop at 7B-scale models [arxiv:2605.07865 §6]. The DAgger row comes from a robot paper, and its numbers are simulation results on a humanoid, not language-model benchmarks [arxiv:2509.20322 tab. 2]. The paper gives real-robot results only qualitatively [arxiv:2509.20322 §4.2]. The diagnostic is costly. It needs roughly 45K–200K targeted rollouts per question, and the full suite took about 72 A100-days [arxiv:2605.10889 §2.5] [arxiv:2605.10889 §3]. Teacher rankings from the AIME 2025 case studies rest on four questions and are indicative only [arxiv:2605.10889 §4.4] [arxiv:2605.10889 app. A.1]. For the teacher, the paper's results do not give one best choice. On BoolQ with 0.6B, the top teacher in the table is the student given one correct solution (Self-1C, 0.047). With 1.7B, the table's highest value is Qwen3-8B (0.053), and the paper marks Self-Sum-1C (0.050) in bold alongside it. On MMLU the top teacher differs again: the 32B-generated correct solution for 0.6B, and Self-Sum-1C for 1.7B [arxiv:2605.10889 tab. 1] [arxiv:2605.10889 §4.2]. The paper argues for per-task diagnostics over a fixed recipe [arxiv:2605.10889 §4.4]. Selective distillation can improve the signal, but the gain uses the true alignment, which is unknown at training time. Keeping only tokens with positive alignment raised the mean signal for 0.6B BoolQ from 0.007 to 0.070 while keeping about half the tokens [arxiv:2605.10889 app. E.3]. The paper calls this an oracle, not a method [arxiv:2605.10889 app. E.3]. ## Implementation sketch The common core is the per-node student distribution, a teacher quantity over the same tokens, and the cosine score used by the diagnostic. The variants differ only in how the teacher quantity enters the gradient. The reverse-KL baseline adds one detached term to the single-sample advantage. The sign convention of the distillation gradient is not stated unambiguously in the paper [arxiv:2605.10889 §2.3], so check it on a case with a known answer. ```pseudocode # Full-distribution distillation direction at one node [arxiv:2605.10889 §2.3] # Input: student_probs, teacher_probs over the node's tokens Output: vector def gkd_direction(student_probs, teacher_probs): # Per-token log-ratio between student and teacher [eq. 2.7] log_ratio = log(student_probs) - log(teacher_probs) # Student-weighted mean of the log-ratio at this node [eq. 2.8] mean_log_ratio = sum(student_probs * log_ratio) # Negated loss gradient; the paper's sign convention is not fully stated return -student_probs * (log_ratio - mean_log_ratio) ``` ```pseudocode # Advantage with a reverse-KL baseline at one sampled token [eq. 10] [eq. 14] # Input: teacher and student log-probs of the sampled token, KL at the node # Output: scalar advantage used in a policy-gradient update def vopd_advantage(teacher_logp, student_logp, node_kl): # Per-token reward: teacher minus student log-probability [eq. 3] reward = teacher_logp - student_logp # Baseline is the negative reverse KL, so subtracting it adds the KL back; # the KL is treated as a constant in the gradient [§2.2] [app. A.1] return reward + stop_gradient(node_kl) ``` ```pseudocode # Alignment score at one node [arxiv:2605.10889 §2.4] [eq. 2.12] # Input: ideal and distill vectors over the same tokens Output: in [-1, 1] def alignment_score(ideal, distill): # Cosine similarity; undefined when either vector is zero [§2.4] dot = sum(ideal * distill) return dot / (norm(ideal) * norm(distill)) ``` ## Contradictions The papers do not give conflicting numbers: they test different things. The reverse-KL paper's results are downstream accuracy on its own Qwen3 and Olmo-3 setups [arxiv:2605.07865 tab. 2], the diagnostic paper reports alignment scores [arxiv:2605.10889 §4.1], and the robot paper reports simulated task success on a humanoid [arxiv:2509.20322 tab. 2]. None of these measures the others. Points to keep apart: - The reverse-KL paper reports that top-k OPD, which restricts the KL and biases the gradient, gains only marginally over vanilla OPD [arxiv:2605.07865 §4.2]. The diagnostic paper does not evaluate top-k restriction, so this claim has no counterpart there. - Within the diagnostic paper, the introduction says the pattern of higher alignment on incorrect rollouts "holds across all settings and metrics" [arxiv:2605.10889 §1]. The appendix table gives a p-value of 0.123 for the mean cosine on 1.7B MMLU, and §4.1 itself calls that gap not significant [arxiv:2605.10889 app. E.6] [arxiv:2605.10889 §4.1]. - The reverse-KL paper's tables report no variance across seeds or runs [arxiv:2605.07865 tab. 2], so differences of about one point between variants, such as the 1.7B top-k and full-vocabulary rows, are not a ranking. ## Papers - [Unmasking On-Policy Distillation](/papers/2605.10889): a training-free per-token diagnostic for on-policy distillation teachers, with results on Qwen3-0.6B and 1.7B. - [KL for a KL: On-Policy Distillation with Control Variate Baseline](/papers/2605.07865): a detached reverse-KL baseline for the single-sample on-policy distillation gradient, with a cheaper top-k form; author-reported gains on Qwen3 and Olmo-3. - [VisualMimic: Visual Humanoid Loco-Manipulation via Motion Tracking and Generation](/papers/2509.20322): DAgger distillation of privileged teachers into depth-only robot policies; a robot-control use of the on-policy idea, not a token-level method.