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

In plain words

A language model can learn from a stronger teacher model by writing its own answers and having the teacher grade each word it chose. That signal is useful, but a single sampled word gives a noisy estimate, so training can swing wildly from one batch to the next.

This paper keeps the cheap single-word signal and adds a correction term, the same kind of "typical value" subtraction that reinforcement learning uses to calm its updates. The correction is the disagreement between student and teacher at that position, and it comes for free from numbers the training step already computes. When the student and teacher strongly disagree, the correction is large and cancels most of the harsh penalty on that word, so the updates stop being dominated by a few extreme cases.

Because the correction is treated as a fixed number, the expected direction of learning is unchanged. The authors report better accuracy than plain distillation on maths and science benchmarks, with lower training time than the full-vocabulary alternative, and they also show a cheaper version that only checks the most likely words.

TL;DR

On-policy distillation (OPD) trains a student on its own samples with a single-token estimator, which is cheap but noisy. vOPD subtracts a detached baseline, the negative reverse KL between student and teacher at each position, which the forward pass already provides. The gradient stays unbiased and its variance drops; a top-k version of the baseline costs less. On Qwen3 and Olmo-3 students, the paper reports average gains over vanilla OPD and parity with full-vocabulary OPD (author-reported).

Problem

On-policy distillation minimises the reverse KL between student and teacher using samples from the student §2.1 eq. 1. In practice it is optimised as policy gradient with a single sampled token per position and a per-token reward that is the teacher's log-probability minus the student's §2.1 eq. 2 eq. 3. The authors state that this single-sample estimator has high variance and that training becomes unstable §2.1 §1.

The two common fixes are costly or biased. Full-vocabulary OPD sums the KL over the whole vocabulary, which needs a backward pass per vocabulary entry §2.1 eq. 5. Top-k OPD restricts the KL to the student's top-k tokens, which changes the objective and biases the gradient §2.1 eq. 7.

The paper's assumptions: a stronger teacher is available, the teacher's per-token log-probabilities are accessible, and the student samples its own responses. The method is evaluated only on models up to 7B parameters §6.

Method

Starting point. The OPD gradient is a sum over tokens of the reward times the gradient of the sampled token's log-probability §2.1 eq. 3. A baseline that does not depend on the sampled token can be subtracted without changing the expected gradient, and the best practical choice is the value function §2.2 eq. 8 app. A.1 app. A.2.

Step 1: the value function of OPD has a closed form. Averaging the per-token reward over the student's own distribution gives the negative reverse KL at that context §3.1 eq. 9. It needs only the student and teacher distributions already computed in the forward pass; no critic and no extra rollouts §3.1.

Step 2: the advantage. The vOPD advantage is the reward plus the reverse KL at the same context, with the KL treated as a constant during backpropagation §3.1 eq. 10. The backward pass still touches only the sampled token §3.1.

Step 3: top-k baseline. The KL in the baseline can be computed on the student's k most likely tokens, with both student and teacher renormalised over that set §3.2 eq. 14. The baseline still depends only on the two models and the support, not on the sampled token, so the gradient remains unbiased for any k §3.2. The reward itself is not restricted §3.2 eq. 15.

Why it lowers variance. The paper derives an approximate reduction in per-token gradient variance proportional to the squared reverse KL times the expected squared norm of the score app. A.3 eq. 11. The derivation uses a weak-correlation approximation that the paper motivates but does not prove app. A.3. It implies the largest reduction at high-disagreement tokens §3.1 §4.3.

Effect on the advantage. Because the KL is non-negative, the vOPD advantage is never below the raw reward; positive rewards are almost unchanged, and large negative rewards at high-KL contexts are pulled toward zero §4.3 fig. 1.

Algorithm, as in the official code. Implementation details are in the Paper vs code section; the pseudocode below follows the paper's equations.

# vOPD loss for one batch of student rollouts  [eq. 10] [eq. 14]
# Input:  prompts; responses sampled from the student (B, T)
# Output: scalar loss

def vopd_loss(prompts, responses, student, teacher, top_k):
    # Forward both models on prompt + response  [§3.1]
    student_logits = student(prompts, responses)     # (B, T, V)
    with no_grad():
        teacher_logits = teacher(prompts, responses)  # (B, T, V)

    # Log-probability of the sampled token under each model  [§2.1] [eq. 3]
    student_logp = gather(log_softmax(student_logits), responses)  # (B, T)
    teacher_logp = gather(log_softmax(teacher_logits), responses)  # (B, T)

    # Per-token reward: teacher minus student log-probability  [eq. 3]
    reward = teacher_logp - student_logp            # (B, T)

    # Baseline: reverse KL(student || teacher) at each context  [eq. 14]
    kl = reverse_kl(student_logits, teacher_logits, top_k)  # (B, T)

    # Advantage, detached so the baseline is a constant  [eq. 10]
    advantage = (reward + kl).detach()              # (B, T)

    # Policy-gradient surrogate on the sampled token  [eq. 10]
    # Maximising advantage * log-prob, so the loss is its negative
    return -(advantage * student_logp).sum()

# Reverse KL over the full vocabulary or the student's top-k  [eq. 14]
# Input:  student_logits, teacher_logits (B, T, V); top_k or None
# Output: kl (B, T)

def reverse_kl(student_logits, teacher_logits, top_k):
    # Restrict both models to the student's top-k set if requested  [§3.2]
    if top_k is not None:
        support = topk_indices(student_logits, top_k)  # (B, T, k)
        student_logits = gather(student_logits, support)
        teacher_logits = gather(teacher_logits, support)

    # Renormalise both distributions over the support  [eq. 14]
    student_logp = log_softmax(student_logits)      # over the support
    teacher_logp = log_softmax(teacher_logits)      # over the support

    # Sum over the support of p_student * (log p_student - log p_teacher)
    student_p = exp(student_logp)
    return sum(student_p * (student_logp - teacher_logp), axis=-1)

Key equations (paper's notation, reproduced in our own symbols):

a_t = r_t + \mathbb{D}_{\mathrm{KL}}(\pi_\theta \,\|\, \pi_T) eq. 10

\hat b_t = -\mathbb{D}_{\mathrm{KL}}(\bar\pi_\theta \,\|\, \bar\pi_T) eq. 14

Hyperparameters

Results

Author-reported, on the paper's own setup. Math numbers are the paper's average of four Avg@n columns (MATH500, Minerva, AMC23, AIME24/25). Science numbers come from the ablation table. Each row compares against vanilla OPD.

Full-vocabulary OPD, the costly reference, scores 33.0 (1.7B), 45.1 (4B), 31.8 (Olmo-3-7B) and 20.6 (0.6B) on the math average tab. 2 tab. 5. Top-k OPD, which biases the gradient, scores 30.2 (1.7B), 42.8 (4B) and 29.5 (Olmo-3-7B) tab. 2; the paper reports only marginal gains for it §4.2.

Cost, from the introduction: vOPD reduces wall-clock time by up to 57.7% against full-vocabulary OPD §1 fig. 3. Gradient norms are 1–2 orders of magnitude lower than vanilla OPD at 1.7B §4.3 fig. 4. The top-k baseline is nearly insensitive to k on the math average §4.3: the 1.7B ablation ranges from 32.3 (k = 100) to 33.4 (k = 5) tab. 5.

Paper vs code

Code: https://github.com/holi-lab/vOPD, linked in the footnote of the paper's abstract. It is a single-repo training and evaluation package built on TRL, at commit d4146d5 (dated 2026-10-08, later than the arXiv v1 of 2026-05-08). Its README says the paper is accepted to NeurIPS 2026. Its license is not stated: there is no LICENSE file, so code_license is recorded as not stated and no code is quoted.

The core logic is in src/opd_trainer.py. The advantage is the teacher-minus- student log-probability of the sampled token plus the reverse KL over the full vocabulary or the top-k set, detached before the loss opd_trainer.py. The top-k branch gathers the top-k indices from the student's logits, renormalises student and teacher over that set, and computes the KL there opd_trainer.py. The advantage is added to inputs["advantages"] before the TRL loss is called opd_trainer.py opd_trainer.py.

Discrepancies found:

• Top-k size in the preset. The paper uses k = 20 for the main vOPD top-k setting §4.1. The preset vopd-topk.yaml sets kl_top_k: 100 vopd-topk.yaml; the README says to override it to 20 for the paper's main setting.
• Evaluation top-p. Table 4 gives evaluation top-p 0.9 tab. 4. The math evaluation route uses 1.0 default.yaml. The MCQ route uses 0.95 with top-k 20 default.yaml.
• Teacher thinking mode. The paper says Qwen teachers run with extended thinking disabled app. B.1. The code disables it only through the chat-template path opd_train.py. The default config sets use_chat_template: null default.yaml, so default training uses the raw math prompt. Whether the paper's runs used a chat template is not stated.
• Training prompt for science. The paper's scientific prompt asks for a JSON answer app. B.2. The training path always uses the math template unless a chat template is requested opd_train.py util.py. The JSON instruction appears only in the MCQ evaluation code evaluate_mcq.py. Whether the paper used it for training is not stated.
• Learning rate for Olmo-3. The paper gives 2e-5 tab. 4. The repo has no Olmo config; the default is 1e-5 default.yaml.
• Science epochs. The paper gives 10 epochs for science tab. 4. The default is 1 default.yaml; no science config is in the repo.
• Batch size. The paper gives 64 tab. 4. The repo uses 1 sequence per device with 64 accumulation steps default.yaml. The accelerate file sets num_processes: 4 accelerate.yaml. Inference: with 4 processes the effective batch would be 256 sequences. The README example uses one process.
• Gradient clipping. The paper states no clipping. The accelerate file sets gradient_clip: 1 under DeepSpeed accelerate.yaml. Whether DeepSpeed honours that key was not verified. This matters for §4.3, which says that simple gradient clipping cannot reproduce the effect §4.3.
• Schedule, weight decay, dropout, seed. The paper states none of these. The repo uses cosine decay, zero warmup and zero weight decay default.yaml, LoRA dropout 0.05 default.yaml, and seed 42 default.yaml.
• Extra options not in the paper. A reward-clipping option (reward_clip_lambda, off by default) and an OPD top-k option (opd_top_k, off by default) exist in the code opd_trainer.py default.yaml. The paper does not use them.
• Loss internals not in repo. The policy-gradient loss (_compute_loss), the advantage helper, the single-step loss and kd_temperature come from TRL, which is not in the repo opd_trainer.py opd_trainer.py. Clipping, importance ratio and token normalisation were not verified.
• Repo version. The repo commit is later than the arXiv v1 the card cites. Experiment scripts in scripts/ do not set the paper's model pairs; the README example is the only end-to-end command.

Pitfalls

• The baseline must be detached. The paper's unbiasedness argument assumes the KL term is a constant in the gradient §2.2 app. A.1. The code computes the gap with autograd and then detaches it opd_trainer.py. Keeping the gradient through the KL changes the update.
• Top-k baseline ignores teacher mass outside the student's top-k. Both models are restricted to the student's set, so the teacher's own likely tokens outside that set have no effect on the baseline §3.2 opd_trainer.py.
• Temperature enters the baseline. The code divides both logit tensors by kd_temperature before any log-probability is computed, so the KL baseline depends on that value opd_trainer.py. The paper defines no temperature, and the value is set by TRL, not stated in the repo.
• Full-vocabulary memory. The full-vocabulary branch materialises full log-probability tensors of shape (batch, length, vocabulary) for both models opd_trainer.py. The paper's claim of no extra backward pass is correct, but the forward memory cost is not stated in the paper and is large for a 150k-token vocabulary §2.1.
• The variance formula is approximate. Eq. 11 uses "≈" and a weak-correlation assumption app. A.3. Treat the variance claim as a heuristic that the paper motivates with a scatter plot fig. 5, not as a measured bound.
• Sign convention. Because the KL is non-negative, the vOPD advantage is shifted upward relative to the raw reward §4.3. An implementation that flips the sign of the KL term, or of the reward, reverses this effect.
• Noise in the reported numbers. The tables give no variance across seeds or runs. The AIME24/25 pass@32 values are multiples of 1/60 (inferred from the values, not stated), and AMC23 values are multiples of 1/40 (inferred). Differences of about one point, such as 33.0 vs 33.1, should not be read as a ranking. The "Avg." column is the mean of the four Avg@n columns; checking the OPD row for 1.7B gives 29.8.
• Truncation. Training responses are capped at 2048 tokens tab. 4. What happens to truncated responses, and whether they enter the KL term, is not stated.
• Teacher setup. The paper's teacher is the post-trained checkpoint, distilled into the base model §4.1. Changing the teacher changes the baseline magnitude as much as the reward; the paper tests only one size-mismatched pair §4.2.
• Paper and code scale. The paper's experiments stop at 7B §6. The README gives an end-to-end command only for the 1.7B run.

Lineage

• Builds on: vanilla on-policy distillation, as in GKD [arXiv:2306.13649] and MiniLLM [arXiv:2306.08543]; the single-sample policy-gradient estimator in §2.1.
• Builds on: the control-variate baseline from policy-gradient RL §2.2, and the value function as the canonical baseline §2.2 app. A.2.
• Compared against: full-vocabulary OPD (the GKD-style full KL) §2.1; top-k OPD, which restricts the KL and biases the gradient §2.1 [arXiv:2604.13016] (cited in the paper as li2026rethinkingopd).
• Superseded by: not known from content already in this wiki.