---
kind: concept
slug: baseline-subtraction
title: Baseline subtraction
summary: >-
  Subtracting a baseline, usually the expected reward at the same context, from a policy-gradient reward to cut variance without changing the expected gradient. In on-policy distillation the closed-form baseline is the negative reverse KL; a top-k version is cheaper. Author-reported gains over plain distillation.
tags: [reinforcement-learning, variance-reduction, distillation]
papers: ["2605.07865"]
related: [policy-gradient, on-policy-distillation, knowledge-distillation, variance-reduced-gradient]
status: draft
updated_at: 2026-10-24
---

## 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 [arxiv:2605.07865 §2.2] [arxiv:2605.07865 app. A.1].
The value function, the expected reward at that context, is the canonical choice
[arxiv:2605.07865 §2.2] [arxiv: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 [arxiv: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 [arxiv:2605.07865 §2.1]. Its baseline is the negative
reverse KL between student and teacher at each position [arxiv:2605.07865 §3.1].
The paper says this needs no critic and no extra rollouts [arxiv: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 [arxiv:2605.07865 §2.1] [arxiv:2605.07865 eq. 3].
The authors state that this single-sample estimator has high variance, and that
training becomes unstable [arxiv:2605.07865 §1] [arxiv:2605.07865 §2.1]. Card:
[KL for a KL: On-Policy Distillation with Control Variate Baseline](/papers/2605.07865).

### 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 [arxiv:2605.07865 §3.1] [arxiv:2605.07865 eq. 9].
It uses the distributions already computed in the forward pass
[arxiv:2605.07865 §3.1]. The advantage becomes the reward plus this KL, with the
KL detached [arxiv:2605.07865 eq. 10]. The backward pass still touches only the
sampled token [arxiv:2605.07865 §3.1].

On Qwen3-1.7B the paper reports a math average of 33.1, against 29.8 for plain
OPD [arxiv: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
[arxiv:2605.07865 §1] [arxiv:2605.07865 fig. 3]. Card: [KL for a KL](/papers/2605.07865).

### 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 [arxiv:2605.07865 §3.2] [arxiv: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 [arxiv:2605.07865 §3.2]
[arxiv:2605.07865 eq. 15]. The main setting uses k = 20 [arxiv:2605.07865 §4.1].
On the 1.7B math average, k from 5 to 100 gives 32.3 to 33.4 [arxiv: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 [arxiv:2605.07865 tab. 2]
[arxiv:2605.07865 §4.2]. Card: [KL for a KL](/papers/2605.07865).

## When to use which

| Situation | Use | Why |
|---|---|---|
| You can afford a full-vocabulary pass and want the closed-form baseline | Full-vocabulary value baseline | Reports 33.1 against 29.8 for plain OPD on the 1.7B math average [arxiv:2605.07865 tab. 2]; the baseline stays unbiased by construction [arxiv:2605.07865 §3.1] |
| The vocabulary is large or wall-clock time is the limit | Top-k value baseline | On the 1.7B and 4B math averages it is within 0.1 points of full-vocabulary (33.0 vs 33.1; 45.3 vs 45.4) [arxiv:2605.07865 tab. 2]. Olmo-3-7B favours top-k by 0.2 points, which is inside the noise the paper's numbers allow [arxiv:2605.07865 tab. 2] |
| You can only get the sampled token's log-probability, not the teacher's distribution | No baseline, or a different method | Eq. 14 uses both models' distributions over the support, so the baseline is not available from sampled-token log-probabilities alone [arxiv:2605.07865 eq. 14] |
| You want to reduce noise by clipping gradients instead | Not a substitute | The authors state that simple gradient clipping cannot reproduce the effect [arxiv:2605.07865 §4.3] |
| The teacher is much larger than the student | Test it first | The paper tests one size-mismatched pair, a Qwen3-1.7B teacher with a 0.6B student, reporting 17.5 for OPD, 20.0 for top-k and 21.1 for full-vocabulary [arxiv:2605.07865 tab. 5] [arxiv:2605.07865 §4.2] |

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 [arxiv:2605.07865 §3.1] [arxiv: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.

```pseudocode
# 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()
```

```pseudocode
# 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 [arxiv:2605.07865 §2.2]. The official code detaches
  the gap, and the card notes that keeping the gradient through it changes the
  update [arxiv: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 [arxiv: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 [arxiv: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 [arxiv: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 [arxiv: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 [arxiv: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](/papers/2605.07865): 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.