Baseline subtraction

Overview

A policy-gradient update weights each sampled token's log-probability gradient by its reward. The reward is noisy, so the update is noisy. Subtracting a baseline from the reward lowers that noise. The baseline must not depend on the sampled token, and it is treated as a constant during backpropagation. Then the expected gradient is unchanged 2605.07865 §2.2 2605.07865 app. A.1. The value function, the expected reward at that context, is the canonical choice 2605.07865 §2.2 2605.07865 app. A.2.

The invariant all variants share is that the baseline depends only on the context and the models, not on the sampled token 2605.07865 §3.2. The variants differ in how that baseline is computed.

The only source on this page is one paper. It applies the idea to on-policy distillation (OPD), where the per-token reward is the teacher's log-probability minus the student's 2605.07865 §2.1. Its baseline is the negative reverse KL between student and teacher at each position 2605.07865 §3.1. The paper says this needs no critic and no extra rollouts 2605.07865 §3.1. No card here covers baselines in other settings, such as learned critics.

Variants

No baseline (single-sample reward)

The reference case. The advantage is the raw per-token reward, the teacher's log-probability minus the student's 2605.07865 §2.1 2605.07865 eq. 3. The authors state that this single-sample estimator has high variance, and that training becomes unstable 2605.07865 §1 2605.07865 §2.1. Card: KL for a KL: On-Policy Distillation with Control Variate Baseline.

Full-vocabulary value baseline (vOPD)

The baseline is the expected reward under the student at each context. For the reverse KL it has a closed form: the negative KL between the student's and the teacher's full next-token distributions 2605.07865 §3.1 2605.07865 eq. 9. It uses the distributions already computed in the forward pass 2605.07865 §3.1. The advantage becomes the reward plus this KL, with the KL detached 2605.07865 eq. 10. The backward pass still touches only the sampled token 2605.07865 §3.1.

On Qwen3-1.7B the paper reports a math average of 33.1, against 29.8 for plain OPD 2605.07865 tab. 2. Its cost is the full-vocabulary sum. The paper reports up to 57.7% lower wall-clock time for vOPD against full-vocabulary OPD 2605.07865 §1 2605.07865 fig. 3. Card: KL for a KL.

Top-k value baseline

The same baseline, but both models are restricted to the student's k most likely tokens and renormalised over that set 2605.07865 §3.2 2605.07865 eq. 14. The baseline still does not depend on the sampled token, so the gradient is unbiased for any k. The reward is not restricted 2605.07865 §3.2 2605.07865 eq. 15. The main setting uses k = 20 2605.07865 §4.1. On the 1.7B math average, k from 5 to 100 gives 32.3 to 33.4 2605.07865 tab. 5.

Do not confuse this with top-k OPD. That method restricts the KL in the objective itself, which biases the gradient. It reaches 30.2 on the 1.7B math average, and the paper reports only marginal gains for it 2605.07865 tab. 2 2605.07865 §4.2. Card: KL for a KL.

When to use which

The variance reduction is an approximation. The paper's variance formula uses a weak-correlation assumption it motivates but does not prove [arxiv:2605.07865 app. A.3]. The largest reduction is at tokens where student and teacher disagree most 2605.07865 §3.1 2605.07865 §4.3. Expect the most benefit there, not uniformly.

Implementation sketch

The generic core is one line of arithmetic. The OPD-specific baseline is the second function. Both are in our own words, following the paper's equations.

# Baseline-subtracted policy-gradient loss  [arxiv:2605.07865 §2.2] [eq. 10]
# Inputs: per-token reward (B, T), baseline (B, T), sampled-token logp (B, T)
# Output: scalar loss

def baselined_loss(reward, baseline, sampled_logp):
    # The baseline depends on context and models, never on the sample  [§2.2]
    # Held constant in backprop, so the expected gradient is unchanged  [app. A.1]
    advantage = (reward - baseline).detach()  # (B, T)

    # Only the sampled token's log-probability receives gradient  [§3.1]
    return -(advantage * sampled_logp).sum()

# Baseline for on-policy distillation: negative reverse KL  [eq. 9] [eq. 14]
# Inputs: student and teacher logits (B, T, V); top_k or None
# Output: baseline (B, T), the expected reward under the student

def opd_value_baseline(student_logits, teacher_logits, top_k=None):
    # 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 models over the same support  [eq. 14]
    student_logp = log_softmax(student_logits)
    teacher_logp = log_softmax(teacher_logits)

    # Expected reward under the student is minus the reverse KL  [eq. 9]
    student_p = exp(student_logp)
    return -sum(student_p * (student_logp - teacher_logp), axis=-1)

Pitfalls, each from the paper unless marked:

• Detach the baseline. The unbiasedness argument assumes the baseline is a constant in the gradient 2605.07865 §2.2. The official code detaches the gap, and the card notes that keeping the gradient through it changes the update 2605.07865 Paper vs code.
• Get the sign right. The baseline is the expected reward, which is minus the reverse KL. Subtracting it adds the KL to the reward. Because the KL is non-negative, the advantage is never below the raw reward 2605.07865 §4.3. Flipping the sign reverses the effect.
• Temperature enters the baseline. The official code divides both logit tensors by a temperature before computing log-probabilities. The paper defines no such temperature, so the baseline's size depends on a value that is not stated in the paper 2605.07865 Paper vs code.
• Top-k size differs between paper and code. The paper's main setting is k = 20. The repo preset is k = 100, and its README says to override it for the paper's setting 2605.07865 Paper vs code.
• Full-vocabulary memory. The full-vocabulary branch keeps full log-probability tensors for both models in memory. The paper does not give a memory figure for this 2605.07865 §2.1.

Contradictions

None known. Only one card covers this page. The paper's variance result is an approximation, not a measured bound 2605.07865 app. A.3, which is noted above rather than treated as a conflict.

Papers

• KL for a KL: On-Policy Distillation with Control Variate Baseline: uses the per-token negative reverse KL, computed from the forward pass, as a baseline for on-policy distillation, with a cheaper top-k version; author-reported gains over plain OPD on Qwen3 and Olmo-3 students.