On-policy distillation

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 2605.10889 §1. The paper contrasts this with the sparse outcome reward used by RL methods such as GRPO 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 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 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 2605.07865 §2.1 2605.07865 §1. Their fix subtracts a per-token control variate that the forward pass already provides 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 2509.20322 §3.1 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 2605.10889 §2.4. The paper frames this as an offline evaluation tool, not a training signal 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 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 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 2605.10889 §2.3. This is the only variant the paper's experiments score 2605.10889 §3.

Card: Unmasking On-Policy Distillation.

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 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 2605.10889 §2.3.

Card: Unmasking On-Policy Distillation.

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 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 2605.07865 §3.1 eq. 10. Detaching matters: the unbiasedness argument assumes the baseline is a constant in the gradient 2605.07865 §2.2 app. A.1.

The baseline has no dependence on the sampled token, so the expected gradient is unchanged 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 2605.07865 app. A.3 eq. 11. It predicts the largest reduction at tokens where student and teacher disagree most 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 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 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 2605.07865 §4.3 tab. 5.

Card: KL for a KL: On-Policy Distillation with Control Variate Baseline.

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 2605.10889 §2.3 eq. 2.10 app. B.5. The paper cites MiniLLM (Gu et al., ICLR 2024) without an arXiv id 2605.10889 §2.3.

Card: Unmasking On-Policy Distillation.

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 2509.20322 §3.1 2509.20322 §3.2. The paper gives no per-token loss in the cited sections. Iteration counts and dataset sizes are not stated 2509.20322 §3.1 2509.20322 §3.2.

Card: VisualMimic.

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 2605.10889 §2.2 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 2605.10889 §2.5. The score lies between −1 and +1, where positive means the distillation gradient points toward success 2605.10889 §2.4.

Card: Unmasking On-Policy Distillation.

When to use which

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 2605.07865 §6. The DAgger row comes from a robot paper, and its numbers are simulation results on a humanoid, not language-model benchmarks 2509.20322 tab. 2. The paper gives real-robot results only qualitatively 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 2605.10889 §2.5 2605.10889 §3. Teacher rankings from the AIME 2025 case studies rest on four questions and are indicative only 2605.10889 §4.4 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 2605.10889 tab. 1 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 2605.10889 app. E.3. The paper calls this an oracle, not a method 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 2605.10889 §2.3, so check it on a case with a known answer.

# 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)

# 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)

# 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 2605.07865 tab. 2, the diagnostic paper reports alignment scores 2605.10889 §4.1, and the robot paper reports simulated task success on a humanoid 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" 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 2605.10889 app. E.6 2605.10889 §4.1.
• The reverse-KL paper's tables report no variance across seeds or runs 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: 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: 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: DAgger distillation of privileged teachers into depth-only robot policies; a robot-control use of the on-policy idea, not a token-level method.