Multi-Turn On-Policy Distillation with Prefix Replay
基于前缀重放的多轮在线策略蒸馏
🧭 Agent 多轮交互蒸馏 · 前缀重放
Microsoft Research · University of Amsterdam
用教师前缀重放替代环境交互，以步骤衰减平衡学生相关性与教师可靠性，实现加速免工具调用的Agent蒸馏
📌 工作参考价值
- 离线环境训练Agent：若业务中存在多轮对话推荐、工具调用Agent（如商品推荐对话、搜索推荐辅助），可用教师模型的历史交互轨迹作为离线重放池，学生只在重放前缀上生成动作，无需实时调用外部API或模拟器，大幅降低训练基建成本。
- 前缀分布控制策略：将多轮蒸馏转化为“有效前缀分布设计”的思路，可迁移到序列生成或对话生成的密度蒸馏，通过在轨迹不同位置施以衰减权重，抑制后期累积误差对监督信号的影响，改善学生模型的长程表现。
- 复用RL训练轨迹：教师进行在线策略强化学习（如GRPO）时产生的轨迹可直接作为蒸馏前缀池，无需额外采集费用，该“免费”数据复用模式可用于业务中已有RL训练管线的场景，提升样本效率。
- 多环境统一蒸馏：若需要训练一个同时处理多个垂直领域（如搜索、推荐、工具调用）的通用Agent，可将各环境的教师轨迹分别收集后合并为一个离线池，从而避免同时部署多个在线环境，降低系统工程复杂度。
动机
在线策略蒸馏（OPD）用于多轮Agent任务时，每次更新都需学生在线与环境交互并查询教师，导致大量工具调用和计算开销。而且学生自己的交互历史可能进入教师不可靠的区域（前缀陷阱），使蒸馏目标劣化。本文提出ReOPD，通过重用教师离线轨迹作为重放前缀，取消环境交互，并将前缀选择建模为可靠性感知的分布设计，以简单的步骤衰减调度平衡学生相关性与教师可靠性。
方法关键点
- 离前缀重放：从教师轨迹池采样前缀（包括历史动作与观察），学生仅对当前步生成动作，教师提供该步的目标分布；前一步均为教师记录，环境永不参与。
- 双边偏移分解：推导出理想目标与重放目标间的差距可分解为：学生占有率不匹配项（前缀与学生实际访问历史之间的TV距离）和教师可靠性项（教师在该前缀下的监督误差）。
- 几何桥梁与步骤衰减：通过几何桥梁连接学生与教师分布，导出最优前缀分布的形式为 (d^t_{\theta_{\text{old}}}^{\gamma_t} \cdot d^t_T^{1-\gamma_t})，其实现权重随步骤深度递减。进而用简单指数衰减 (\omega(t)=\kappa^t)（(\kappa=0.6)）替代高方差的每前缀重要性权重。
- 采样式训练：按衰减权重从轨迹池中采样训练位置，不产生额外加权损失，直接优化目标期望。
关键实验
- 任务：数学推理（Python工具，6个竞赛基准）和搜索问答（Wikipedia检索，7个QA基准），学生/教师模型均为Qwen3系列。
- 对比线上OPD与SFT。在数学任务上，ReOPD一致优于OPD（如4B教师/4B学生：57.2% vs. 55.1%），尤其在师生能力差距大时提升显著；搜索任务上当教师可靠时，两者持平（40.5% vs. 40.6%）。
- 多环境联合训练（数学+搜索）中，ReOPD在双领域均保持OPD水平，且训练时不需任何工具调用。
- 效率：ReOPD每步训练时间比OPD快至少4倍，工具调用次数降为0。
核心观点
多轮蒸馏的核心不是将数据变为“在策略”，而是为每一步选择有效的前缀分布，在学生对历史的占有率和教师监督的可靠性之间取得折中。