OPD 的 RL 视角：Least Square Policy Distillation 实现样本高效的 LLM 推理

TL;DR OPD 可视为 KL 正则化 RL；本文提出 LSPD，引入乐观探索与 off-policy 重用，提升样本效率和策略多样性，数学推理 +1.59 Avg@16。

问题

问题背景

当前 LLM 推理能力提升的核心路线之一是知识蒸馏，即用强教师模型引导小模型。on-policy distillation (OPD) 因学生模型从自身策略采样进行对齐，在数学推理等任务中表现突出。

现有方法局限

现有 OPD 存在两个关键局限：

• 优化目标层面：OPD 通常最小化学生与教师的 reverse-KL 散度，该目标天然倾向模式搜索，会导致学生策略多样性下降。在需要探索多种解题路径的推理任务中，容易过度模仿教师的高概率答案。
• 数据效率层面：OPD 的 rollout 数据仅用于单次更新，之后即被丢弃。在 LLM 场景下，每次 rollout 计算成本极高，这种 on-policy 约束导致样本利用率低，训练需要大量新鲜采样，限制扩展性。

为什么难/重要

将 OPD 形式化为 KL-regularized policy optimization 后，如何引入乐观探索与离线数据重用，同时保持理论上的 regret 保证，是一个兼具理论和实践价值的挑战。业界对样本高效且保持多样性的 LLM 推理蒸馏方法高度关注，因为这直接影响训练成本和模型泛化能力。

行业类比

这一挑战类似于在线强化学习中的经验回放机制，要在不牺牲策略改进质量的前提下，用更少的在线交互换取更稳定的学习过程。

核心洞察

• 将 OPD 中的 reverse-KL 目标与 KL-regularized policy optimization 等价起来，揭示了蒸馏本质上是一种策略优化过程。这为引入 RL 中的乐观探索和 off-policy 数据复用提供了理论依据，而传统蒸馏方法只做监督式拟合，忽略策略空间探索与样本效率。
• LSPD 通过最小二乘估计实现乐观策略更新，理论承诺 O(log K) 后悔界，实际在 Pass@k 评估中随 k 增大优势扩大，说明其保留策略多样性优于 baseline；同时完全 off-policy 变体仅用前 25% rollout 批次即可达到 vanilla OPD 性能，极大降低样本成本。

方法

输入与问题建模

LSPD 将 on-policy distillation (OPD) 视为 KL-regularized policy optimization 问题。输入包括：teacher 策略 π_teacher、student 策略 π_student、prompt 数据集 D。OPD 最小化 student 与 teacher 轨迹分布的 reverse KL 散度，等价于最大化 student 在 teacher 对数概率上的期望奖励，同时惩罚与 teacher 的 KL 偏离。该视角使蒸馏可以借用 RL 的探索与样本复用技术。

关键模块

1. Least-Square Policy Distillation 将策略更新看作基于值函数的最小二乘回归：从 rollout 轨迹中估计每个状态-动作对的优势值，并用 Huber-type penalty 增强对噪声 expert 反馈的鲁棒性。
2. Optimistic exploration 借鉴 value-based RL 的 upper confidence bound (UCB) 思想，在估计的策略质量上加乐观偏差，鼓励学生探索与 teacher 不同但可能更优的生成路径，从而保持策略多样性。
3. Off-policy data reuse 使用 replay buffer 存储历史 rollout（来自不同迭代的 student 分布），通过重要性采样或最小二乘拟合从这些离策略数据中学习，大幅提升 rollout 效率。
4. Lagrangian relaxation 将原约束优化目标转化为无约束 Lagrangian 形式，平衡蒸馏保真度与探索强度，使训练稳定可控。

输出与训练流程

每次迭代：student 从当前策略采样生成 rollout → 存入 replay_buffer → 用最小二乘法更新策略（结合乐观 bonus 调整目标）→ 返回更新后的 student。输出为训练好的 student 策略，在数学推理任务上 Pass@k 指标随 k 增大优势更明显，证明多样性得到保留。

与同类方法差异：相比 vanilla OPD 仅用当前 on-policy 数据进行反向 KL 最小化，LSPD 通过 optimistic exploration 与 off-policy least-squares 更新，首次在蒸馏框架中同时实现理论遗憾界（O(log K)）与显著样本效率提升。

实验

实验设计

论文在六个数学推理基准上评估 LSPD，覆盖多种 teacher–student 设置，并对比现有蒸馏基线。主要指标包括 Avg@16 与 Pass@k（k 最高 64）。同时测试了 off-policy 变体及带重放缓冲区的 LSPD-RB，以评估 rollout 效率。

关键发现

• LSPD 在 Avg@16 上平均较基线提升 +1.59 点。
• Pass@k 评估显示，随着 k 增大，LSPD 的性能优势更明显，表明其更好地保留了策略多样性。
• 完全 off-policy 的 LSPD 仅使用前 25% 的 rollout batches，即可达到 vanilla OPD 的同等性能，大幅节省采样成本。

与基线对比的深度解读

传统 OPD 受限于 on-policy 更新，rollout 数据利用效率低。LSPD 从 value-based RL 引入乐观探索与 off-policy 数据复用，将 OPD 重新解释为 KL 正则化策略优化。其理论分析给出 O(log K) 的 regret 界，支持探索的高效性。实际中，LSPD 不仅指标更高，还能在减小 rollout 规模的情况下保持性能，对工程实践中的算力与数据预算优化具有直接参考价值。

行业影响

落地场景

LSPD 可直接用于 推理型大模型的蒸馏压缩，尤其适合需要高多样性与长链推理的场景：

• 教育解题助手：将大型数学推理模型蒸馏为中小型学生模型，部署到在线答疑、作业批改产品中。
• 企业级智能客服：处理需要多步逻辑计算的问题（如账单拆分、套餐最优选择），LSPD 训练出的模型在 Pass@k 指标上随 k 增大优势更明显，适合提供多个候选答案供人工或规则筛选。
• 内容平台的代码/数学辅助工具：为开发者或创作者提供轻量级代码解释、公式推导服务，降低推理延迟。

商业价值

核心价值在于大幅降低推理模型蒸馏的样本成本与计算开销：

• 降本：LSPD 的完全 off-policy 变体仅使用前 25% 的 rollout 批次即可达到 vanilla OPD 相当的性能，数据采集成本减少约 75%。
• 体验提升：通过乐观探索保留策略多样性，模型生成的候选答案质量更高，可提高用户对推理结果的满意度，减少二次提问。
• 加速迭代：off-policy 数据重用允许同一批轨迹被反复学习，缩短模型更新周期，利于快速响应业务需求变化。

与现有工作流的接口

LSPD 可无缝集成到当前 LLM 训练栈中：

1. 替换 on-policy distillation (OPD) 中的 rollout 采样步骤，将 reverse-KL 目标改为 Least-Square 形式，无需改动教师模型或推理框架。
2. 作为 RLHF/DPO 流程的前置蒸馏步骤，先用 LSPD 从大教师模型获得一个较强的 student policy，再进行偏好对齐，减少对齐阶段的探索成本。
3. 与 vLLM / SGLang 等推理引擎结合，利用 off-policy replay buffer 实现异步数据采集与训练，提高 GPU 利用率。

论文开源代码：UNCSciML/LSPD，可直接参考实现细节。

具体落地 use case

• 金融分析报告生成：蒸馏具备复杂数值推理能力的 LLM，用于自动生成财务数据解读、风险提示等，LSPD 的多样性保持能力可输出多种分析角度供分析师选择。
• 电商优惠策略推荐：在购物车结算场景中，需要计算跨店铺满减、优惠券叠加的最优组合，LSPD 训练的小模型能以低延迟提供多个可行方案，提升转化率。

局限

• 理论保证存在理想化假设：论文的 regret bound 依赖 **optimistic least-squares** 和 **online exploration** 的理想化设定，而实际实现采用 **Huber 损失**、**Lagrangian 松弛** 和 **KL 正则化** 进行近似，两者之间存在明显 gap。因此，理论分析无法直接解释 LSPD 在实际 LLM 蒸馏中的行为，限制了其指导意义。此外，理论分析仅覆盖在线探索设置，对 off-policy 重用的理论保障不足。
• 实验评估范围较窄：主要基准集中在数学推理（如 GSM8K、MATH 等六个数据集），缺少对代码生成、指令遵循等通用能力的测试。模型规模未明确，但推测以 7B-13B 为主，未验证在大规模模型上的效果。同时，与 SOTA 蒸馏方法（如 **ReST**、**VinePPO**）的系统对比不足，难以全面评估方法的竞争优势。
• off-policy 变体的性能增益有限：尽管 **LSPD-RB** 仅用 25% 的 rollout batches 达到与 vanilla OPD 相当的性能，但并未显著超越，说明样本效率提升的代价可能是性能上限未突破。此外，方法引入多个超参数（如探索系数、buffer 容量、正则化强度），论文未提供敏感性分析，实际部署调参成本较高。