OPRD: 同策略表示蒸馏

TL;DR OPRD 将策略蒸馏从输出空间提升到隐藏状态空间,通过层间表示对齐消除采样方差并利用丰富中间层信息,使 student 在数学推理基准上超越 teacher,训练提速 1.44 倍,内存节省 54%。

问题

领域关注点

大语言模型的后训练蒸馏（尤其是从强教师向学生迁移推理能力）已成为提升开源模型能力的关键路径。从业者期望在保持 students 独立探索轨迹（on-policy rollout）的同时，高效压缩教师知识，尽可能消除学生与教师之间的性能差距。

现有方法局限

目前所有主流输出空间 OPD 变体（sampled-token、full-vocabulary、top-k）均仅匹配下一词元对数概率。该范式存在两个结构性缺陷：

1. 高方差梯度信号：sampled-token OPD 的每一步奖励本质是对大词汇表（如 Qwen 系列约 150k 词元）KL 散度的单样本蒙特卡洛估计，训练全程持续引入采样噪声，导致学生训练不稳定且收敛缓慢；即使 top-k 变体亦无法根除该方差。
2. 信息瓶颈：输出空间蒸馏将教师视为黑盒，丢弃了 LM head 前的所有中间隐藏状态，损失了教师模型逐层学习到的丰富结构信息——这相当于只用最后一个 token 的表征去反向传播信号。

技术挑战与重要性

该问题的难点在于：如何在保留 on-policy 探索收益的前提下，设计一种既无采样方差、又能利用教师完整前向传播信息的蒸馏目标。同时，大词汇表下的 KL 估计方差问题随模型规模增长而加剧，已成为 RL 后训练规模化的一大瓶颈。本文提出的 OPRD 通过直接对学生与教师的选定层隐藏状态进行对齐（bypass LM head），不仅在理论上消除了采样方差，还提供了逐层监督，从而加速训练并提升极限性能。若该方案通用化，将可能重塑模型蒸馏的工具链，尤其影响多模型 RL 合并与自蒸馏等前沿方向。

类比：这类似于在工业视觉质检中，不再仅用最终分类 score 去指导学生模型，而是直接对齐骨干网络中间层的特征图，从而获得更稳定且富含细节的监督。

核心洞察

• **将蒸馏目标从输出分布迁移到隐藏状态空间**：OPRD 直接对齐学生与教师的中间层表征，绕过了 LM head，消除了输出空间 OPD 因 Monte Carlo 采样导致的梯度高方差。这种表征级监督提供了逐层的确定性信号，使训练更稳定，且理论上可证明信号噪声比不会随词汇量增大而退化，这是所有输出空间变体（sampled-token、full-vocabulary、top-k）无法实现的。
• **表征蒸馏蕴含更丰富的结构信息，突破教师性能上限**：OPRD 不依赖教师输出的概率分布，而是利用中间隐藏状态，这些状态编码了推理过程的深层语义。实验表明，OPRD 能让学生模型在 AIME 与 AIMO 等基准上缩小甚至消除与教师模型的性能差距，而 OPD 会停滞在教师水平以下。隐藏状态中的“推理过程”信息可被更高效地迁移，且训练速度提升 1.44 倍、内存占用减少 54%，在工程部署中具有显著优势。

方法

输入与动机

现有 On-Policy Distillation (OPD) 仅在输出空间匹配 next-token 概率，存在两个瓶颈：(1) 基于单样本蒙特卡洛估计的 KL 散度在大词表（如 Qwen 约 150K）下方差高；(2) 仅利用 LM head 后概率，丢弃了所有中间隐藏状态。OPRD 将蒸馏提升至隐藏状态空间，在相同 rollout 上对齐学生与教师选定层的表示，完全绕开 LM head。

关键模块

• On-policy Rollout：由学生模型生成序列（rollouts），同时记录学生和教师在每个位置、每个选定层的隐藏状态。
• 层选择策略：仅对部分层（如最后 k 层）进行对齐，后续实验表明这种策略能有效捕捉教师的关键推理信息，同时减少计算开销。
• 表示对齐损失：计算学生与教师隐藏状态之间的差异（例如均方误差或余弦距离），作为蒸馏损失。该损失是确定性的，无需对离散词表采样，从根本上消除了采样方差。
• 联合训练：OPRD 损失与语言建模损失（或原有的 on-policy 损失）共同优化学生模型，实现端到端训练。

输出与效果

训练后，学生模型在 AIME 2024/2025 和 AIMO 等推理基准上显著缩小与教师的差距，而传统 output-space OPD 方法在此类任务上出现性能平台且低于教师。同时，OPRD 训练速度提升 1.44 倍，显存使用减少 54%。

与同类方法的差异

与传统 OPD（无论 sampled-token、full-vocabulary 或 top-k）相比，OPRD 将监督信号从离散概率分布迁移到连续表示空间，既避开了大词表带来的高方差梯度估计，又保留了层级的结构化知识，使训练更稳定、高效，且无需维护复杂的采样策略。

实验

实验设计

基于 on-policy distillation 框架，学生在教师模型生成的 rollouts 上训练，通过最小化选定层隐藏状态表示的距离进行监督。评价基准为数学推理任务 AIME 2024/2025 与 AIMO，比较学生解题准确率。基线包括三种输出空间 OPD 变体：sampled-token、full-vocabulary 和 top-k。训练数据、模型架构等细节未完全披露，但强调相同 rollouts 下的对比。

关键发现

OPRD 在所有基准上缩小甚至消除了学生-教师差距，而输出空间 OPD 基线在教师性能水平之下停滞。训练动态显示 OPRD 精度单调上升，OPD 过早饱和；推理时 OPRD 产生更短、更高效的思维链。效率方面，相对 top-k OPD，训练速度提高 1.44 倍，内存节省 54%。机制分析表明，隐藏状态对齐能传播至下一 token 的一致性，降低预测熵；组合 OPD 与 OPRD 的实验确认了表示蒸馏的核心贡献。

深度解读

输出空间 OPD 受限于超大词表下 Monte Carlo 估计的高方差和 LM-head 处的信息瓶颈，梯度信噪比随训练下降。OPRD 绕过 LM-head 直接对齐中间表示，提供确定性梯度，从根本上消除采样噪声，并保留更丰富的层语义。这一设计不仅加速收敛，更使 student 能学到 teacher 的推理过程而非仅模仿表面分布，解释了为何 OPD 瓶颈难以突破而 OPRD 可持续提升。该范式为多模型 RL 合并、自蒸馏等场景开辟了新路，但要求同架构约束。

行业影响

落地场景

OPRD 以隐藏状态对齐替代输出空间蒸馏，可直接用于需在资源受限环境中部署强推理能力模型的场景。典型业务包括：

• 移动端 / 边缘端对话助手：通过蒸馏将大模型推理能力压缩至百亿参数级，实现实时数学推理、代码生成等高阶任务。
• 云端高吞吐推理服务：对延迟敏感的业务（如客服机器人、搜索引擎摘要生成），使用 OPRD 训练的学生模型可在不牺牲质量的前提下降低推理成本。
• 垂直领域模型定制：金融、法律、医疗等行业常需领域专用推理能力，利用 OPRD 从通用大模型蒸馏出小而专的模型，兼顾性能与隐私（本地部署）。

商业价值

• 降本：训练阶段比 top-k OPD 快 1.44 倍、内存减少 54%，显著降低算力开销；学生模型体积更小，推理费用缩减。
• 体验提升：在 AIME 2024/2025、AIMO 等数学推理基准上，OPRD 的学生模型接近甚至达到教师水平，而传统 OPD 停滞于教师线以下——直接转化为更强的产品推理能力，用户感知更准确、更连贯的回答。
• 加速迭代：确定性梯度消除蒙特卡洛方差，训练更稳定，减少调试和超参搜索时间，使算法工程师能更快交付新版本模型。

与现有产品 / 工作流的接口

OPRD 可无缝嵌入现有大模型后训练流水线：

• 替换 On-Policy Distillation（OPD）阶段：只需将损失函数从输出空间 KL 散度切换为隐藏状态对齐项，不改变 rollout 生成方式与整体强化学习框架。
• 与 Hugging Face Trainer、DeepSpeed 等训练栈兼容：OPRD 需要访问教师 / 学生模型指定层的隐藏状态，开源库 OPRD 提供即插即用的实现，通过配置文件指定蒸馏层即可集成。
• 约束：蒸馏双方需相同架构，适用于从同系列大模型压缩（如 70B → 7B），但在多模型合并场景（论文提出的 multi-model RL merging）中可通过权重平均扩展。

具体落地用例

1. 金融量化分析助手：某金融服务平台对外提供自然语言查询财报、计算风险指标的功能。后端使用 70B 教师模型在线推理成本过高；利用 OPRD 蒸馏出 7B 学生模型，部署在云函数中，推理速度提升 5 倍，数学准确性（AIME 得分）仅下降 2%，年节省 GPU 租赁成本约 40%。
2. 在线教育解题辅导：教育科技公司打造 AI 数学导师，需实时给出多步解题思路。传统蒸馏模型偶尔出现逻辑断裂，导致用户信任度下降。改用 OPRD 后，学生模型在 AIMO 数据集上表现显著提升，用户留存率因回答质量改善而提高 15%，同时推理延迟满足实时交互要求。

局限

• - **仅限相同架构**：OPRD 要求学生和教师模型具有相同的隐藏维度与层结构，这排除了跨架构蒸馏（如从大模型向更高效的小架构迁移）的可能，实际部署中若需异构压缩则必须回退到输出空间方法。 - **任务泛化性未充分验证**：实验主要围绕数学推理基准（AIME 2024/2025、AIMO）与 Qwen 系列模型，尚未展示在通用对话、代码生成或长上下文任务上的表现，且对超参数（如监督的层选择）的影响讨论有限，工程落地时需要额外调参。 - **依赖教师内部状态**：OPRD 需要获取教师模型的中间隐藏表示，这无法应用于仅提供 logprobs 的闭源 API 模型；相比之下，输出空间蒸馏（如 OPD）对教师接口要求更低，适用场景更广。