基于前缀重放的多轮在线策略蒸馏 Multi-Turn On-Policy Distillation with Prefix Replay
复用教师轨迹做离线多轮 agent 蒸馏,零工具调用、4倍加速且不掉点。
前置知识
On-Policy Distillation (OPD,在线策略蒸馏)
知识蒸馏的在线形式:学生模型在自己采样得到的前缀上生成动作,再由教师给出逐 token 的目标分布进行监督。与传统的离线蒸馏(在教师生成的轨迹上做 SFT)不同,OPD 让学生看到自己推理时真正会遇到的状态,从而缓解曝光偏差导致的误差累积,监督信号也比 RL 的标量奖励密集得多。
本文的核心研究对象,ReOPD 就是要在保留 OPD 优势的同时去掉昂贵的实时环境交互。
Prefix Replay(前缀重放)
多轮交互蒸馏中,不再让环境实时生成观察,而是直接复用预先收集好的教师轨迹里的前缀(含历史动作与环境观察)。学生在某个截断前缀上生成当前动作,由教师监督。整个学生训练不调用真实环境,把昂贵的在线交互转化为可复用的离线资源,并可让一个学生从多个异构环境合并出的统一池中联合学习。
这是 ReOPD 区别于传统 OPD 的关键机制,理解它才能理解为什么 ReOPD 能大幅降低训练成本。
GRPO(Group Relative Policy Optimization)
DeepSeek 提出的 PPO 变体,用于 LLM 的强化学习后训练。它用同一 prompt 下多条采样的相对优势作为 baseline,省去独立 critic 网络,配合可验证奖励在数学推理等任务上效果显著。本文教师模型用 GRPO 训练,其训练产生的 rollout 轨迹被直接复用为 ReOPD 的前缀池。
理解教师轨迹从哪来——它们是 GRPO 训练的免费副产品,使 ReOPD 的前缀池几乎零额外采集成本。
Covariate Shift 与误差累积(Compounding Error)
模仿学习的经典问题:学生只在教师访问过的状态上训练,一旦推理时犯一个教师不会犯的早期错误,就会进入训练时从未见过的状态分布,后续错误沿轨迹不断累积放大。DAgger 等方法通过在学生自身诱导的状态分布上训练来缓解。本文把这个结构推广到多轮 LLM agent 场景,并指出“完全学生在线”也会引入反向的可靠性偏移。
这是 ReOPD 理论分析中“前缀陷阱”时间层的根源,也是 OPD 相对 SFT 的核心动机。
Agentic LLM(智能体大模型)
把 LLM 当作与环境多轮交互的智能体:每一轮模型输出动作(如调用 Python、发起搜索检索),环境返回观察,模型再据此继续推理。典型任务包括带工具的数学推理和基于检索的问答。这种多轮结构让训练成本远高于单轮生成,因为每一步都可能触发昂贵的工具调用或检索。
本文研究的正是这种多轮 agent 场景的蒸馏,环境交互的高成本是 ReOPD 要解决的核心痛点。
研究动机
现有方法在多轮 agent 蒸馏中陷入两难。传统的离线蒸馏(SFT/序列级蒸馏)数据效率高,但学生在教师访问过的前缀上训练,推理时一旦犯早期错误就会漂移到从未见过的状态,误差沿轨迹累积(covariate shift / exposure bias)。完全在线的 OPD 能让学生在自己诱导的状态分布上学习,但代价极高:每次参数更新都要让学生在真实环境里重新滚动生成多轮轨迹,并在每个访问到的历史处查询教师。对于数学推理(每次最多 16 次 Python 调用)和搜索问答(需部署约 80GB 显存的检索系统)这类 agentic 任务,每个 step 都要并发约 32 个工具进程。当要蒸馏的环境数量增加(如同时做数学+搜索),OPD 的部署复杂度和资源消耗随环境数快速增长,几乎不可扩展,这是本文要打破的瓶颈。
本文的目标是本文的具体目标是:在去掉学生训练期间所有实时环境交互的前提下,仍保留 OPD 的两大核心好处——学生相关性(student relevance)和教师密集的逐 token 监督。可量化的指标包括:学生训练阶段做到零工具调用(zero tool calls),每个 rollout 至少比 OPD 快 4 倍;同时在数学推理和搜索两类 agentic 任务上,让蒸馏出的学生精度达到或超过完全在线的 OPD;并让一个学生能从多个异构环境的不同教师那里联合蒸馏,而无需同时在线部署所有环境。最终把昂贵的 agent-environment 交互转化为可复用的离线资源,实现跨工具、跨任务、跨环境的可扩展蒸馏。
与已有工作不同的是,作者抓住了一个被既有 OPD 工作忽视的关键点:在多轮场景下,“让前缀完全学生在线”并不自动最优。这是因为存在一个“前缀陷阱”的双重分布偏移——把前缀推向学生会让数据更贴近学生真实分布(降低学生占用偏移),但同时可能把教师查询到它本身很少访问的后期历史上,使教师的目标分布不再可信(教师可靠性偏移)。这两股力量随轨迹深度此消彼长,所以最优前缀分布应在学生相关性与教师可靠性之间做位置相关的平衡,而非一味追求 on-policy。这是 ReOPD 把多轮 OPD 重新表述为“可靠性感知的前缀分布设计”的独特切入点,也是它能在师生差距大时反而超过完全在线 OPD 的理论根源。
核心方法
直觉上,ReOPD 像是把“环境”从训练循环中摘出来,换成一本教师预先写好的“剧本”(教师轨迹池)。学生不再亲自跑到环境里摸爬滚打,而是被放在剧本的某一页(某个截断前缀,含历史动作和真实观察),只在这一页上自己写下一个动作,由教师批改。这样既保留了“学生在自己的动作上被监督”的在线性,又彻底免掉了环境调用。技术上 ReOPD 由两个部件组成:(1) 前缀来自预收集的教师轨迹,学生在被监督的 step $t$ 生成自己的动作,教师在该前缀上的条件分布 $\pi_T(\cdot \mid x, h_t)$ 作为逐 token 蒸馏目标;(2) 一个按位置衰减的采样调度 $\omega(t; \kappa) = \kappa^t$,$\kappa \in (0,1]$,越靠后的高偏移位置被采样越少。唯一自由旋钮是衰减陡度 $\kappa$(实验默认 $\kappa = 0.6$),可由师生每步 KL 差距 $\bar{c}$ 直接读出:$\kappa \approx \exp(-\gamma_t \bar{c})$。
核心创新是把多轮 OPD 重新表述为“可靠性感知的前缀分布设计”,而非简单的“让它 on-policy”。作者证明了双侧分解界(Proposition 1):实际目标与理想交互目标之差被两个项控制——$2B \cdot \mathrm{TV}(d_t^{\theta_{old}}, \rho_t)$(学生占用偏移)和 $\mathbb{E}_{\rho_t}[\epsilon^{\theta_T}_t]$(教师可靠性误差)。把前缀完全推向学生($\rho_t = d_t^{\theta_{old}}$)能把第一项归零却放大第二项;完全用教师前缀($\rho_t = d^T_t$)则相反。最优解是两者的几何桥接 $\rho^\star_t \propto (d^{\theta_{old}}_t)^{\gamma_t}(d^T_t)^{1-\gamma_t}$。这推出一个反直觉但可验证的结论:当师生能力差距大(如 8B→4B、30B→4B 数学)时,“更不 on-policy”的教师锚定前缀反而更好。这与传统“越 on-policy 越好”的直觉根本不同,是全文最重要的理论洞见。
方法步骤详情
完整流程对应 Algorithm 1:(1) 初始化 $\theta_{old} \leftarrow \theta$,准备教师轨迹池 $D_T$(复用教师 GRPO rollout,零成本)。(2) 每个迭代从池中所有位置里按 $p_t \propto \kappa^t$ 采样被监督位置,早段低偏移位置被采更多。(3) 对每个位置 $(x, h_t)$:前缀 $h_t$ 全部来自教师轨迹(动作与观察都是教师记录,不调用环境);学生在该前缀上自回归生成动作 $A_t$,沿该动作逐 token 累计蒸馏损失 $\sum_j D_{KL}(\pi_\theta(\cdot|h_t,a^{<j}_t) \,\|\, \pi_T(\cdot|h_t,a^{<j}_t))$,其中上下文 $a^{<j}_t$ 是学生自采,故该步真 on-policy。(4) 在累计损失上更新 $\theta$,令 $\theta_{old} \leftarrow \theta$ 进入下一轮,对同一教师前缀重新评估学生。采样与加权两种实现无偏等价,实验统一用采样。
技术新颖性
与已有技术的本质区别有三点。其一,相比 MiniLLM、GKD、DistiLLM 等单轮 OPD,ReOPD 首次把 on-policy 蒸馏扩展到多轮 agent 的固定教师轨迹池场景,核心设计对象从“教师目标分布”扩展到“在哪个前缀分布上查询教师”。其二,相比 DAgger、scheduled sampling 这类靠学生实时滚动环境来缓解 covariate shift 的方法,ReOPD 不重新滚动环境,而是用教师前缀加可靠性感知调度在离线池上近似学生占用,把“on-policy”压缩到单个被监督 step。其三,相比 rejection sampling、RAFT 这类只接受/拒绝整条轨迹的方法,ReOPD 把每条轨迹当成可查询教师条件分布的“位点”,有效前缀分布成为核心变量。理论上它给出首个把 OPD 目标差分解为“学生占用项 + 教师可靠性项”的分解界,并预测出可被实验直接验证的“两段制”行为。
实验结果
核心发现分四块。(1) 数学推理(Table 4):ReOPD 在所有教师配置下都超过 OPD——4B→4B 平均 55.1→57.2;8B→4B 平均 51.0→53.7(AIME24 +8.4 点);30B→4B 51.1→52.5;30B→8B 56.5→56.8。师生差距越大提升越大,与教师可靠性视角吻合。(2) 搜索(Table 5):师生同族、教师对学生诱导历史仍可靠,ReOPD 基本持平 OPD(4B 教师 40.5 vs 40.6;8B 教师 39.0 vs 39.1)。(3) 多环境(Table 6):一个 4B 学生同时从两域教师蒸馏,数学 55.3、搜索 41.0,均持平 OPD 且无需同时在线部署两套环境。(4) 效率(Figure 1/6/8):学生训练零工具调用,每个 rollout 快至少 4.2 倍(数学 175s→40s,搜索 169s→39s);OPD 需 32 并发进程与约 80GB 显存,ReOPD 全省。Figure 5 验证偏重早段最优,Figure 7 证明教师自身生成的前缀最好,Table 7 表明 RL 池与定常池几乎等价。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| AIME24 数学推理(8B 教师→4B 学生) | 准确率 (%) | 36.7 (ReOPD) | 28.3 (OPD) | +8.4 个点 |
| 六项数学竞赛平均(4B 教师→4B 学生) | 平均准确率 (%) | 57.2 (ReOPD) | 55.1 (OPD) | +2.1 个点 |
| 七项 QA 平均(4B 教师→4B 学生) | 平均准确率 (%) | 40.5 (ReOPD) | 40.6 (OPD) | -0.1(基本持平,可靠教师区间) |
| 单步 rollout 耗时(数学环境) | 秒/rollout | 约 40s (ReOPD) | 约 175s (OPD) | 4.2× 加速,且零工具调用 |
局限与改进
作者承认的局限:(1) 方法依赖一个预收集、带观察记录的教师轨迹池,池的覆盖度与质量直接决定学生能学到什么;(2) 可靠性代理只是“一阶近似”为按位置衰减的权重,比较粗糙——它只刻画前缀有多深,没有直接度量某条前缀离“学生相关历史与教师可靠支撑的重叠区”有多远,learned/data-dependent 的可靠性估计可能更优;(3) 双侧界用“教师支撑”作为可靠性的代理,更紧、可直接估计的可靠性度量及在线自适应调度是未来工作。我自己的观察:所有实验都用 Qwen3 系列,跨模型族的迁移未验证;$\kappa = 0.6$ 是全局固定值,论文虽声称无需逐任务调参,但 Figure 5b 显示性能对 $\kappa$ 较敏感,实际部署可能仍需小范围搜索;AIME 仅 30 题、平均@8/@4 的高方差让部分提升处于统计噪声边缘;搜索任务上 ReOPD 几乎无增益,说明在“教师本就可靠”的窄差距场景,方法的理论优势难以转化为实际收益。
独立分析的弱点
弱点一:可靠性调度过于粗糙。$\omega(t) = \kappa^t$ 只用位置索引,忽略同一步内不同前缀的真实偏移。改进:直接用学生-教师似然比 $\mathrm{br}_t = \prod_{s<t} \pi_{\theta_{old}}(a_s)/\pi_T(a_s)$(论文式 7 已给出但嫌方差高),或训练小型可靠性预测器在线打分。弱点二:固定 $\kappa$ 不自适应。不同任务、师生对的 $\bar{c}$ 不同,理想 $\kappa = \exp(-\gamma_t \bar{c})$ 应随差距动态调整;改进:在线估计 $\bar{c}$ 并自动设 $\kappa$。弱点三:教师池覆盖决定上限,教师 RL 没采到的题型学生永远学不到;改进:主动采集加覆盖度检测或针对性补采。弱点四:仅在数学与搜索两类环境验证,浏览器、长程 API、多模态等更高复杂度场景的前缀陷阱可能更严重,方法是否仍持平未知;改进:扩展到更多环境,尤其超长 horizon 任务。
未来方向
作者明确提出的方向:(1) 用可学习或数据相关的可靠性估计替代粗糙的 step-decay 调度;(2) 设计更紧、可直接估计的教师可靠性度量,以及能在线自适应的调度或权重。基于本文成果可延伸的方向:把几何桥接参数 $\gamma_t$ 做成逐 step 自适应(目前是把整个深度剖面塌缩成单 $\kappa$),在长 horizon 任务上应更精细;把 ReOPD 与 RL 后训练结合——在教师/学生都持续 RL 时让前缀池动态更新,探索“蒸馏与 RL 交替”的课程式训练;推广到多教师、多模态、异构动作空间(如图像+文本观察)的 agent 蒸馏;研究在线版 ReOPD——学生偶尔触发少量真实环境回放来校正池的偏移,形成混合精度蒸馏。另一个有价值的方向是把“教师可靠性偏移”与“奖励 hacking 检测”联系起来,因为两者本质都是“目标信号在 off-support 处不可信”。
复现评估
复现友好度较高。代码、模型、数据均开源(baohaoliao/ReOPD,baohaol/reopd[cs.LG]),配有项目主页。训练基于开源的 Qwen3 系列、Slime 框架和标准 GRPO/SFT/OPD 流程。数据集全部公开:DAPO 数学训练集 6.4K、NQ+HotpotQA 搜索 6.5K、六个数学评测(共 1547 题)、七个 QA 评测(共 51713 题)、2018 Wikipedia dump 加 E5 retriever。算力门槛适中:8×H100,ReOPD 训练 3 小时内完成(200 步)。最大复现难点在 agentic 环境部署——数学需并发 Python 沙箱,搜索需约 80GB 显存部署 embedding 与索引;但 ReOPD 学生训练本身不依赖这些(这正是其卖点),只需一次性为教师采集轨迹即可纯离线训练。$\kappa = 0.6$、学习率 1e-6、200 步、batch 256×1、温度 1.0 等超参在 Table 1/2 给全。整体属于“仔细按论文配置即可复现”的级别。
论文图表
左侧 OPD:随环境数量增加,需同时部署所有环境(数学+Python、搜索/QA、浏览器、代码/API 等),复杂度随环境数增长;右侧 ReOPD:各环境的教师轨迹可分别采集后合并进统一离线池,学生训练无需任何在线环境。
这张图说明了 ReOPD 的可扩展性卖点——多环境联合蒸馏时把 N 套在线环境压成 1 个离线池,是方法工程价值的关键论据。