vOPD

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

Accepted to NeurIPS 2026.

Paper · Blog · Citation

Overview

vOPD stabilizes on-policy distillation by subtracting a detached baseline from each token's teacher–student log-probability reward. The baseline is the negative reverse KL, computed from the existing forward pass. It reduces gradient variance without an extra critic or rollout, while preserving OPD's expected gradient. Approximating only the baseline with the student's top-k tokens also preserves this property.

Results

Math accuracy (%) from Tables 2 and 5 of the paper. Each score averages MATH500 and Minerva avg@8, and AMC23 and AIME24/25 avg@32; it is not pass@n. Top-k uses k = 20.

Student	OPD	Full-vocabulary OPD	vOPD top-k	vOPD full-vocabulary
Qwen3-0.6B-Base	17.5	20.6	20.0	21.1
Qwen3-1.7B-Base	29.8	33.0	33.0	33.1
Qwen3-4B-Base	42.9	45.1	45.3	45.4
Olmo-3-7B-Base	29.9	31.8	33.1	32.9

• Science: Qwen3-1.7B vOPD top-k improves SciKnowEval chemistry from 29.3% to 33.2% and GPQA-Diamond from 24.7% to 28.6% over OPD (Table 3).
• Speed: At 4B, vOPD top-k takes 32.88 s/step versus 77.66 for full-vocabulary OPD: 57.7% less time in the paper's single-H200 measurement (Figure 3).
• Stability: Gradient norms are 10–100× lower at 1.7B; top-k ablations retain gains across k = 5, 20, 50, and 100 (Figures 2 and 4).

Installation

conda env create -f environment.yml
conda activate opd
pip install flash-attn==2.8.3 --no-build-isolation

or

bash install.sh

Training

Run from the repository root after installation. This example trains Qwen3-1.7B-Base with Qwen3-1.7B as teacher on the English DAPO subset, using one GPU:

CUDA_VISIBLE_DEVICES=0 accelerate launch \
  --config_file configs/accelerate.yaml \
  --num_processes 1 \
  --module src.opd_train \
  training_config=configs/training/vopd-topk.yaml \
  dataset_name=dapo14k \
  model_name_or_path=Qwen/Qwen3-1.7B-Base \
  teacher_model_name_or_path=Qwen/Qwen3-1.7B \
  kl_top_k=20 \
  report_to=none

The YAML presets inherit shared settings from default.yaml. Command-line key=value arguments override them.

Method	Training config	Settings
vOPD, full-vocabulary baseline	vopd.yaml	use_baseline=true, kl_top_k=0
vOPD, top-k baseline	vopd-topk.yaml	Preset uses kl_top_k=100; override with 20 as above for the paper's main setting
Vanilla OPD	sample-opd.yaml	use_baseline=false
Full-vocabulary OPD	full-opd.yaml	distill_mode=gkd_rkl
Top-k OPD	full-opd.yaml	Add opd_top_k=20

kl_top_k controls vOPD's detached baseline; opd_top_k restricts the OPD loss itself. Generation sampling uses the separate top_k setting.

Citation

@misc{oh2026klklonpolicydistillation,
  title={KL for a KL: On-Policy Distillation with Control Variate Baseline},
  author={Minjae Oh and Sangjun Song and Gyubin Choi and Yunho Choi and Yohan Jo},
  year={2026},
  eprint={2605.07865},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2605.07865},
}