An RL View of OPD: Least Square Policy Distillation for Sample-Efficient LLM Reasoning
Shangzhe Li, Yuxiao Yang, Tianrun Yu, Kaixiang Zhao, Xiaoyun Wang, Taylor W. Killian, Weitong Zhang
cs.LG
2026-09-29
LSPD recasts OPD as KL-regularized RL and matches log-probs with least squares. +1.59 Avg@16 on six math benchmarks; LSPD-RB matches vanilla OPD with 25% of rollout batches.
On-policy distillation (OPD) has the teacher score tokens on prefixes the student actually produced. That sits closer to the student's distribution than forward-KL training on teacher rollouts. Typical implementations wrap PPO or GRPO around the reverse-KL objective. Every update needs a fresh batch of long traces, old traces get dropped, and exploration stays thin.
Rollouts dominate the bill. Math traces often run to thousands of tokens, so regenerating after every policy step is expensive. The objective is already reverse KL. The updater is still an on-policy policy gradient, which never brings in the two tools value-based RL is built around: optimistic reward estimates and replaying old data.
Rewrite reverse KL as KL-regularized policy optimization. Define a teacher-induced reward R = log(π^E / π^ref). Minimizing reverse KL from student to teacher is the same as maximizing R minus a KL penalty to π^ref. With η = 1 this identity is OPD, not an extra regularizer. If π^ref is the initial student, the shape matches the reward-KL pair in RLHF, except the reward is the teacher-to-reference log ratio.
PPO follows the current reward estimate and does not treat uncertainty. LSPD fits that teacher-induced reward with least squares on accumulated student traces, takes the most optimistic value inside a statistical confidence set, and applies the closed-form KL-regularized token update. A Lagrangian relaxation of that optimistic program leaves a practical loss with two pieces:
Because the match is on stored prefixes, the traces need not come from the current policy. LSPD updates on the latest batch, once or several times, then drops it. LSPD-RB keeps a replay buffer of past query-response pairs and trains fully off-policy.
An idealized version (Algorithm 2) gets Õ(log K) reverse-KL regret under online exploration. The fast rate needs a bounded log ratio between teacher and student, meaning enough coverage. The code trains Huber plus entropy, not the confidence-set update the theorem studies.
Training set: DAPO-Math-17K. Six benchmarks: MATH-500, Minerva, Olympiad-Bench, AMC23, AIME24, AIME25. Three Qwen3 pairs: 8B → 4B-Base, 4B → 1.7B-Base, 1.7B → 0.6B-Base, thinking mode off on every teacher. Baselines: KD (teacher traces plus forward KL), OPD, and EOPD (OPD plus forward KL on high-entropy teacher tokens). Training uses temperature 1.0, top-p 1.0, max 7168 tokens, 4 samples per prompt; evaluation uses temperature 0.7 and top-p 0.95. LSPD-RB is trained for 10 steps; the others need at least 40 steps to settle.
Averages over 18 model-benchmark cells:
| Method | Avg@16 | Pass@16 |
| OPD | 29.61 | 51.03 |
| EOPD | 30.69 | 52.45 |
| LSPD | 31.60 | 53.30 |
| LSPD-RB | 31.51 | 54.86 |
Versus all distillation baselines, Avg@16 / Pass@16 move +1.59 / +1.87. Versus the strongest baseline, EOPD, LSPD is only +0.91 / +0.85. LSPD is best or tied-best Avg@16 in 11 of 18 cells and Pass@16 in 9 of 18. Taking the better of LSPD and LSPD-RB, that variant is top-2 in all 36 metric-setting pairs.
On 8B → 4B MATH-500, LSPD Avg@16 is 81.61 against EOPD 80.50 and OPD 79.31. On the same pair, AIME24 Pass@16 is 50.00 for LSPD-RB, 36.67 for EOPD, 30.00 for OPD. On AIME25 Pass@16 in that pair, LSPD scores 36.67, below OPD and EOPD at 40.00. On 4B → 1.7B MATH-500, LSPD is 70.86 versus EOPD 69.75 and OPD 69.13. The 1.7B → 0.6B numbers on the same set are 52.04 / 51.98 / 51.40: the smallest student barely separates from EOPD on per-sample accuracy.
Pass@k goes up to k=64 for 8B distilling 4B-Base. On AMC23 / AIME24 / AIME25, LSPD reaches Pass@32 of 85.54 / 43.33 / 46.67 and Pass@64 of 89.16 / 46.67 / 50.00. Three-benchmark averages: 58.51 / 61.94 versus EOPD 57.00 / 60.43. During training, student entropy for LSPD and EOPD sits near 0.5, above vanilla OPD.
Each rollout batch is 64 prompts × 4 responses. Raising inner updates N from 1 to 4 cuts steps-to-saturation from more than 50 to about 30; 16 and 64 add little. LSPD-RB runs 256 optimizer steps on the buffer after every new batch and saturates in about 10 rollout batches, which is the abstract's first 25% of rollout batches relative to vanilla OPD.
Dropping the entropy term costs 0.17 Avg@16 for LSPD and 0.75 for LSPD-RB, but 2.26 and 1.64 Pass@16. Pass@16 falls in all 12 method-benchmark cells. Entropy buys coverage; per-sample accuracy barely moves.
If reverse-KL on-policy distillation is already in the stack and rollouts are the bottleneck, the updater can be a least-squares log-prob match that trains off-policy, because the objective is already KL-regularized RL. The Avg@16 gap to EOPD is under one point. Rollout count falling from 40+ steps to about 10, and a Pass@k curve that still pulls away as k grows, are the pieces more likely to change a training run.
Even without a replay buffer, N=4 inner updates on the current batch cut rollout batches from 50+ to about 30. This is an incremental fix on OPD/EOPD. The study is Qwen3, math, non-thinking teachers. Code is public.
There is no standalone limitations section.
The theorem is about the optimistic confidence-set procedure (Algorithm 2). The experiments train Huber plus entropy. Õ(log K) does not transfer automatically. Assumption 6.1 needs a bounded log ratio; if teacher-student coverage is poor, the fast reverse-KL rate does not apply, and the paper cites empirical work where coverage matters.
Evaluation is Qwen3 and math only, with thinking turned off. No code, no instruction following, no thinking-mode teachers. LSPD-RB's 10 steps each run 256 gradient updates, so the comparison with OPD is in rollout batches, not wall clock or FLOPs. AIME has 30 problems; 50.00 versus 36.67 Pass@16 can move with sampling noise. Appendix B mentions out-of-domain numbers that are not in the extracted body. There is no GRPO-with-verifiable-reward baseline, only distillation methods.