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
把OPD的reverse-KL重写成KL正则策略优化，提出最小二乘蒸馏LSPD。六项数学推理Avg@16平均高1.59；回放版用四分之一rollout批次追平普通OPD。
On-policy distillation（OPD）让老师在学生自己生成的前缀上给 token 级监督，比拿老师轨迹做 forward-KL 更贴学生分布。现有实现多半套 PPO、GRPO：每轮都要重新采样，旧轨迹几乎用不了，探索也弱。
蒸馏最贵的是 rollout。数学推理轨迹动辄数千 token，学生每更新一次就要再滚一批。OPD 的目标已经是 reverse-KL，算法却仍按策略梯度在跑，没有把 value-based RL 里「乐观估计奖励、反复吃旧数据」那套带进来。
把 reverse-KL 原样改写成 KL 正则策略优化。老师相对参考策略的 log 比当作奖励：R = log(π^E / π^ref)。最小化学生相对老师的 reverse-KL，等价于最大化这个 R，再减去对 π^ref 的 KL。η=1 时这就是 OPD 本身，没有额外正则项。参考策略取成初始学生，结构和 RLHF 的 reward-KL 一样，奖励由老师似然比直接给出。
PPO 跟着当前奖励估计走，不显式处理不确定性。LSPD 改用累积学生轨迹，最小二乘拟合这个老师诱导奖励，在统计置信集里取最乐观的点，再按 KL 正则闭式解更新 token 分布。对乐观约束做拉格朗日放松后，训练目标剩两块：
匹配发生在存储前缀上，轨迹不必来自当前策略。LSPD 对当前 batch 更新若干次后清空；LSPD-RB 把历史 query-response 放进 replay buffer，完全 off-policy 抽 batch。
理想化算法在在线探索下给出 Õ(log K) 的 reverse-KL regret，K 是更新轮数。这条快速率要求老师与学生的 log 比有界，覆盖不够就不成立。实际代码用的是 Huber 加熵，理论和实现之间隔了一层松弛。
训练数据是 DAPO-Math-17K。六个基准：MATH-500、Minerva、Olympiad-Bench、AMC23、AIME24、AIME25。三组 Qwen3 师生对：8B→4B-Base、4B→1.7B-Base、1.7B→0.6B-Base，老师关闭 thinking。对照 KD（老师轨迹加 forward-KL）、OPD、EOPD（在高熵老师 token 上再加 forward-KL）。训练采样 temperature 1.0、top-p 1.0、最大 7168 token、每题 4 条；评测 temperature 0.7、top-p 0.95。LSPD-RB 只跑 10 个训练 step，其余方法至少 40 step 才收敛稳定。
18 组设置平均：
| 方法 | 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 |
相对全部蒸馏基线，Avg@16 / Pass@16 平均 +1.59 / +1.87。相对最强基线 EOPD，LSPD 只高 0.91 / 0.85。LSPD 在 11/18 组拿到 Avg@16 最好或并列最好，Pass@16 是 9/18；两个变体合在一起，36 个指标-设置组合里更好的那个始终进前二。
8B→4B 的 MATH-500 上，LSPD Avg@16 为 81.61，EOPD 80.50，OPD 79.31。同设置 AIME24 的 Pass@16，LSPD-RB 到 50.00，EOPD 36.67，OPD 30.00。8B→4B 的 AIME25 Pass@16，LSPD 是 36.67，低于 OPD 与 EOPD 的 40.00。4B→1.7B 的 MATH-500，LSPD 为 70.86，对 EOPD 69.75、OPD 69.13。1.7B→0.6B 同一项是 52.04 / 51.98 / 51.40，最小学生的单条准确率和 EOPD 几乎贴在一起。
Pass@k 做到 k=64（8B 蒸 4B-Base）。AMC23 / AIME24 / AIME25 上 LSPD 的 Pass@32 为 85.54 / 43.33 / 46.67，Pass@64 为 89.16 / 46.67 / 50.00；三项平均 Pass@32/64 为 58.51/61.94，EOPD 为 57.00/60.43。训练中 LSPD 与 EOPD 的策略熵都维持在约 0.5，高于普通 OPD。
每个 rollout batch 是 64 道题、每题 4 条回复。当前 batch 上把更新次数 N 从 1 加到 4，饱和大约从超过 50 step 降到约 30 step；N 再加到 16、64，收益变小。LSPD-RB 每收集一批就在 buffer 上做 256 次更新，约 10 个 rollout batch 饱和，对应摘要写的「前 25% rollout 批次追平普通 OPD」。
熵消融：去掉熵项，Avg@16 只掉 0.17（LSPD）和 0.75（LSPD-RB），Pass@16 掉 2.26 和 1.64。12 组 Pass@16 全部下降。熵主要买覆盖，单条准确率几乎不动。
如果已经在做 reverse-KL 的 on-policy 蒸馏，而且 rollout 是瓶颈，这篇给了一个合法的 off-policy 替换：目标本来就是 KL 正则 RL，不必死守 PPO 式更新。相对 EOPD 的 Avg@16 提升不到 1 分。采样次数从 40+ 降到约 10，Pass@k 随 k 增大仍能拉开，这两项对训练配方更有用。
即便不建 replay buffer，只在当前 batch 上做 N=4 次更新，rollout 批次也能从 50+ 降到约 30。这是 OPD/EOPD 上的渐进改进，范围锁在 Qwen3、数学、非 thinking 老师。代码已公开。
论文没有单独的局限节。
理论管的是带置信集的乐观更新（Algorithm 2），实验用的是 Huber 加熵。Õ(log K) 不会自动落到实用损失上。Assumption 6.1 还要求 log 比有界，师生覆盖差时这条速率本身不成立，论文也引用了覆盖影响 reverse-KL 蒸馏的经验结果。
评测只有 Qwen3 家族和数学题，老师关闭 thinking，没有代码、指令跟随或 thinking 模式老师。LSPD-RB 每个 rollout batch 对应 256 次梯度步，和 OPD 对齐的是采样次数，墙钟和 FLOP 没有对照。AIME 只有 30 题，Pass@16 上 50.00 对 36.67 的差距，采样噪声不会小。附录提到了域外评测，抽取正文里没有数字。对照里没有带可验证奖励的 GRPO，只和蒸馏方法比。