面向在线策略蒸馏的信任域行为混合

TL;DR 在线蒸馏中，TRB 用信任域内的教师近似行为替代早期学生采样并逐步退火，避免劣质前缀损害训练，在数学推理蒸馏中获得最优表现。

问题

知识蒸馏中的前缀分布不匹配

大语言模型压缩广泛采用 知识蒸馏 (knowledge distillation)，其中 on-policy distillation (OPD) 通过让学生模型从自身策略采样前缀并匹配教师分布，试图解决离线蒸馏中训练-推理时的 前缀分布偏移 (prefix mismatch)。然而，在训练初期，学生策略尚未收敛，采样出的前缀常包含低质量或错误的推理步骤，导致教师监督落在这类劣质样本上，损害收敛效率和最终性能。

现有方法的局限

• 离线蒸馏：使用固定的教师强制序列或离线数据集，前缀分布与推理时实际学生行为不一致，导致性能上限受限。
• 纯 OPD：回归到学生自身前缀虽然缓解了 mismatch，但早期 rollout 质量差，教师在此基础上的指导可能产生误导，类似强化学习中的“盲目探索”，训练不稳定且样本效率低。
• 现有 warmup 策略：如 SFT 预热、温度退火或间歇性教师注入，要么破坏 on-policy 特性重新引入分布偏移，要么需精细调参，难以在保持学生探索自主性的同时平滑过渡到高质量前缀。

技术挑战与重要性

该问题的核心难点在于 平衡学生探索与教师引导：完全依赖教师前缀（如同行为克隆）会丧失 on-policy 优势；完全放任学生探索则易陷入劣质局部区域。尤其在数学推理等长序列决策任务中，早期单个 token 的偏差可能导致整个 rollout 坍塌，后续监督完全失效。业界正将大模型能力压缩至小模型以削减推理成本，而该问题直接影响蒸馏后模型在复杂推理场景下的可靠性，因此寻找一种能够在不破坏 on-policy 特性的前提下提供高质量初始化的方法，具有很高的工程价值和学术关注度。

场景类比

这个问题类似于在 RLHF 管线中，先用监督微调给出合理初始策略以避免随机探索时的奖励稀疏——为 OPD 设计一个“受控探索”阶段，用教师行为软化学生前缀，确保 warmup 后学生能平滑接管训练。

核心洞察

• **信任区域行为混合 (TRB)**：在 on-policy 蒸馏的初期，以 KL 散度信任区域约束，将学生 rollout 策略替换为最接近教师的行为策略，从而避免低质量前缀引发训练崩溃。这与直接注入教师轨迹或 SFT warmup 不同，TRB 在 warmup 阶段依然保持 per-prefix 反向 KL 损失的 OPD 框架，不引入额外 loss 项或强制完全教师数据；TRB 通过退火预算平滑过渡到纯学生 rollout，实现了从稳定引导到自主学习的一致优化。
• **逐前缀闭式解与 token 级控制**：TRB 为每条前缀独立求解一个闭式行为策略，通过 bisection 算法在 token 级精确满足 KL 约束，无需手动调整混合比例。这区别于 veto 机制或固定 ε 混合，TRB 在每个生成步骤动态确定最教师-like 的 token 分布，既保留行为多样性又严格信任学生当前能力边界，最终在数学推理基准上取得平均最强，证明了理论引导的精确干预优于启发式调度。

方法

Trust-Region Behavior Blending (TRB) 是一种针对 On-Policy Distillation (OPD) 的 warmup 方法，旨在缓解早期学生策略前缀质量低导致的教师监督偏差。

输入与背景

OPD 使用学生自身策略采样前缀，并由教师在这些前缀上施加分布监督。但训练初始阶段学生策略较差，前缀质量低，教师监督容易落在弱或不恰当的轨迹上。TRB 在不改变 per-prefix reverse-KL OPD 损失的前提下，引入一个行为策略优化模块来改进前缀采集质量。

关键模块

1. Per-prefix 行为策略计算：对每个生成步骤的前缀，TRB 不直接使用学生策略采样下一个 token，而是求解一个带约束的最优化问题：在 KL 信任区域 内（约束为行为策略与学生当前策略的 KL 散度不超过预算 ε），寻找与教师分布最接近的行为策略。该优化问题具有闭式解（通过二分法求解，见附录 E），等价于在预算受限下将学生策略向教师方向倾斜。

2. OPD 损失保持：获得行为策略后，从前缀按此策略采样，并采用与标准 OPD 相同的 reverse-KL 损失（即最小化教师分布与学生预测分布之间的 KL 散度）来训练学生模型，无需修改损失函数或额外辅助目标。

3. 退火 warmup：KL 预算 ε 在 warmup 阶段从较大初始值逐步退火至零。初期行为策略近似教师，提供高质量前缀；随着预算收缩，行为策略平滑过渡至学生自身策略，warmup 结束后即恢复为标准 OPD 的纯学生 rollout。

输出与效果

经过 TRB warmup 后，学生策略在不增加推理开销的情况下，享有更好的初始前缀分布，从而提升最终蒸馏性能。在数学推理场景下，TRB 在平均性能上超过 temperature warmup、SFT warmup、Veto 及交错 teacher injection 等多种同类 warmup 方法。

差异点：区别于单纯混合教师分布或调整采样温度，TRB 显式地在信任区域内优化行为策略，以数学上最优的方式逼近教师，从而更有效地桥接学生与教师之间的策略差距。

实验

实验设计

本文在两个数学推理蒸馏任务上验证 TRB 的有效性。基线方法包括 vanilla OPD、Veto、交错教师注入 (SKD)、温度预热、固定 ε 混合 和 SFT 预热。TRB 的核心是：在预热阶段，将学生 rollout 策略替换为在学生 KL 信任域内最接近教师的行为策略，保持原有的 per-prefix reverse-KL 损失不变；KL 预算通过退火逐步降至零，使训练平稳过渡到纯学生 rollout。

关键发现

• 在所有对比方法中，TRB 取得了最强的平均性能，验证了信任域行为混合能有效缓解 OPD 早期 rollout 质量低的问题。
• 预热期间的 KL 预算退火机制保证了训练最终回归标准的 on-policy 蒸馏，没有引入额外的推理开销或架构修改。
• 相比直接使用教师前缀（如 SKD）或仅靠温度调整，信任域约束确保了行为策略既不偏离学生分布过远，又提供更高质量的监督信号。

与基线对比的深度解读

TRB 区别于以往的 OPD 稳定化手段：它在每前缀级别动态求解闭式行为策略，而不是全局插值或启发式替换。这使得 teacher 监督能精确作用于学生信任域内的“最近教师点”，既避免了离线蒸馏的 prefix mismatch，又防止了早期弱 student rollout 破坏学习信号。与 SFT 预热相比，TRB 从第一轮迭代就保持 reverse-KL 的目标一致性，因此泛化能力更强。该工作为 on-policy 蒸馏的冷启动问题提供了理论优雅且易于实现的方案。

行业影响

落地场景

On-policy distillation (OPD) 通过让学生模型在自己的 rollout 前缀上学习教师分布，解决了离线蒸馏中的分布偏移问题。但早期学生策略差，生成的 prefix 质量低，导致教师监督浪费在弱前缀上。Trust-Region Behavior Blending (TRB) 在 warmup 阶段用最接近教师的策略（受限于学生 KL 信任区域）替换学生 rollout，从而提升早期学习质量，随后逐步退火回纯粹的学生 rollout。这特别适合需要将大模型能力压缩到小模型的场景，如：

• 边缘设备上的 智能客服 与对话系统，要求低延迟、低资源占用；
• 移动端的 代码补全 或 数学解题助手，需要小模型实时推理；
• 企业内部文档的 自动摘要 与 信息提取，通过小模型降低推理成本。

商业价值

• 降低推理成本：小模型部署成本远低于大模型，TRB 改善蒸馏质量，让小模型更接近教师性能，从而在不牺牲太多效果的情况下大幅削减算力开销。
• 提升用户体验：低延迟响应是交互式产品的核心诉求，高质量小模型可在端侧快速运行，避免网络延迟。
• 加速模型迭代：TRB 作为 warmup 方法，与现有 OPD 损失兼容，不需要重新设计训练流程，可快速在已有蒸馏管线中验证收益，缩短从大模型到生产级小模型的上线周期。

相比于直接离线蒸馏或标准 OPD，TRB 在数学推理等任务上取得最强平均性能，说明它能更稳定地传递教师能力。

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

TRB 的实现非常轻量：

• 训练早期替换 rollout 策略为求解一个闭式 KL 约束下最接近教师的行为策略。
• 引入退火系数 ε，从初始值逐步降至 0，回到标准 OPD。
• 损失函数无需改动，仍使用 per-prefix reverse KL 损失。

集成到现有蒸馏框架只需在数据加载时或采样阶段插入混合策略逻辑，无需修改模型结构或优化器。例如，在 Hugging Face Transformers 或自定义训练循环中，可以在前 N 步采样时调用一个 trb_sample 函数，该函数利用教师 logits 和学生 logits 计算混合分布（通过闭式解），然后采样 token。退火系数由外部调度器控制。现有 OPD 用户几乎可以零代码修改将其作为插件加入，迅速验证收益。

具体落地 Use Case

1. 电商智能客服：将大型对话模型蒸馏到小型模型，部署在客服终端设备上，用于实时回答用户询问。TRB 的 warmup 策略可以避免早期因学生回答错乱导致的训练不稳定，提升最终小模型的回答准确率与连贯性，从而减少人工干预。
2. 在线教育数学辅导：将擅长推理的教师模型蒸馏到学生模型，嵌入到移动应用或网页中，为学生提供解题步骤和讲解。低端设备上运行大模型成本高，TRB 帮助小模型更精准地模仿教师推理链，在保持高质量输出的同时大幅降低服务端推理成本。

局限

• **评估范围有限**：实验仅在两个数学推理数据集（GSM8K 和 MATH）上进行，方法在通用 NLP 任务、代码生成或多轮对话等场景下的效果未经验证。数学推理的 token 分布和错误模式与开放式生成有较大差异，信任区域的行为混合策略能否有效泛化仍存疑。
• **信任区域预算的敏感性**：TRB 引入了一个逐前缀的 KL 信任区域超参数（预算 B），并通过退火逐步收紧约束。论文未系统分析该预算对最终性能的鲁棒性，预算过小可能限制行为策略接近教师，过大则退火到纯学生策略的过程可能不稳定。不同学生-教师能力差距下，最优预算可能显著不同，实际部署需大量调参。
• **对教师策略的强依赖**：TRB 在 warmup 阶段需要教师策略能够为任意学生前缀生成高质量的下一个 token 分布，这对教师的容量和覆盖能力有较高要求。若教师本身在部分领域表现不佳，混合后的行为策略可能放大错误，反而削弱学生训练。此外，warmup 阶段的计算开销（需维护两个策略并求解优化问题）高于普通 OPD，资源受限时可能不适用。