AsyncOPD:在策略蒸馏可以有多陈旧?

arXiv cs.LG 论文

摘要

本文提出 AsyncOPD,一种完全异步的在策略蒸馏流程,用于大语言模型,系统研究了陈旧策略数据的影响,并提出了估计器设计,使训练吞吐量提升 1.6-3.8 倍,同时保持相当的准确率。

arXiv:2606.24143v1 Announce Type: new 摘要:在策略蒸馏(OPD)通过教师反馈指导学生在自身生成上进行训练,在大语言模型(LLM)后训练中越来越重要。然而,与强化学习(RL)类似,OPD 面临在策略系统瓶颈,因为对于推理任务,生成数据可能占据训练时间的主要部分。异步训练流程通过将数据生成与学习器更新解耦,可以缓解这一瓶颈,但会引入陈旧策略数据。虽然先前工作研究了异步 RL 中的陈旧数据问题,但其在 OPD 中的影响仍未充分探索。我们首次对异步 OPD 中的陈旧性进行系统研究,聚焦于一个实际场景:教师反馈通过局部 KL 损失实现,且跨完整词表的教师 logits 因存储或传输代价过高而不可行,因此需要有限的教师分数缓存。我们首先表明,KL 方向改变了陈旧数据问题:教师加权的正向 KL 对陈旧生成更鲁棒,而学生加权的逆向 KL 则易受影响。其次,针对这种易受影响的逆向 KL 情况,我们研究了用于稳定异步 RL 的方法是否能减轻 OPD 的陈旧性。在我们的实验中,这些方法并不比一个更简单的 OPD 特定替代方案更优:在学习器时刻用当前学生重新计算逆向 KL 信号。第三,我们分析了有限的教师分数缓存如何为稀疏和采样的逆向 KL OPD 估计器创建偏差-方差权衡。这激励了多样本蒙特卡洛(MC),它在保持 MC 可修正性的同时减少单样本方差。最后,我们提出并开源 AsyncOPD,一个基于这些估计器选择构建的完全异步 OPD 训练流程。实验表明,AsyncOPD 与严格的同步训练相比,训练吞吐量提升了 $1.6\times$ 到 $3.8\times$,同时达到相当的准确率。
查看原文
查看缓存全文

缓存时间: 2026/06/24 07:50

# AsyncOPD: 在线策略蒸馏可以有多陈旧?
来源:https://arxiv.org/html/2606.24143
Wonjun Kang1Kevin Galim∗1Seunghyuk Oh1Minjun Kang2Sanghyun Park2Donghoon Kim1Minjae Lee1Minseo Kim1Rishabh Tiwari3Yuchen Zeng4Hyung Il Koo1,2Kangwook Lee5,6

1FuriosaAI2Ajou University3UC Berkeley4Microsoft Research5KRAFTON6Ludo Robotics

代码:https://github.com/furiosa-ai/async-opd

###### 摘要

在线策略蒸馏(OPD)在教师反馈的指导下,使用学生自身的轨迹数据训练学生模型,并日益成为大语言模型(LLM)后训练的重要组成部分。然而,与强化学习(RL)类似,OPD面临在线策略的系统瓶颈,因为对于推理任务,轨迹生成可能占据训练时间的主导地位。异步训练流水线可以通过将轨迹生成与学习者更新解耦来缓解这一瓶颈,但这会引入陈旧策略数据。虽然先前的工作研究了异步RL中的陈旧数据,但其在OPD中的影响仍未得到充分探索。我们首次系统性地研究了异步OPD中的陈旧性问题,重点关注一个实际场景:教师反馈通过局部KL损失实现,而全词汇教师logits存储或传输成本过高,因此需要有限的教师分数缓存。我们首先表明,KL方向改变了陈旧数据问题:教师加权的前向KL对陈旧轨迹更鲁棒,而学生加权的反向KL则脆弱。其次,针对这种脆弱的反向KL情况,我们研究了旨在稳定异步RL的方法是否能缓解OPD的陈旧性。在我们的实验中,这些方法并不比一个更简单的OPD特定替代方案更好:在学习器时间使用当前学生重新计算反向KL信号。第三,我们分析了有限的教师分数缓存如何为稀疏和采样的反向KL OPD估计器创建偏差-方差权衡。这推动了多样本蒙特卡洛(MC)方法,该方法在减少单样本方差的同时保留了MC的可纠正性。最后,我们提出并开源了AsyncOPD,这是一个基于这些估计器选择构建的完全异步OPD训练流水线。实验表明,与严格同步训练相比,AsyncOPD将训练吞吐量提高了1.6倍到3.8倍,同时达到了可比的精度。

## 1 引言

参见图注图1:异步OPD的估计器设计。(a) 密集KL是全词汇参考,但在异步OPD中,完整的教师logits缓存存储或传输成本高昂。(b) 稀疏top-k在陈旧性下暴露了支撑集不匹配:前向KL是教师支撑的,但反向KL是学生支撑的,可能需要缓存中未评分的动作。(c) 单样本蒙特卡洛通过重要性采样在期望上是可纠正的,但方差高;我们的估计器在学习器时间重新计算A_θ,并使用多样本MC来降低方差。

在线策略蒸馏(OPD)[20 (https://arxiv.org/html/2606.24143#bib.bib20),6 (https://arxiv.org/html/2606.24143#bib.bib3),1 (https://arxiv.org/html/2606.24143#bib.bib4)]和强化学习(RL)[28 (https://arxiv.org/html/2606.24143#bib.bib30),26 (https://arxiv.org/html/2606.24143#bib.bib28)]已成为提高大语言模型(LLM)推理能力[7 (https://arxiv.org/html/2606.24143#bib.bib29)]的核心后训练方法,包括数学[17 (https://arxiv.org/html/2606.24143#bib.bib27)]和编码[27 (https://arxiv.org/html/2606.24143#bib.bib32)]。OPD使用教师提供的密集词元级反馈,在学生自身的轨迹上训练学生[12 (https://arxiv.org/html/2606.24143#bib.bib2)],而RL则从轨迹的奖励反馈中学习。OPD为LLM后训练提供了一条高效且有效的途径,尤其适用于较小的学生模型[24 (https://arxiv.org/html/2606.24143#bib.bib11)]。最近的工作表明,OPD并不局限于将大教师蒸馏到小学生:它还支持在线策略自蒸馏[32 (https://arxiv.org/html/2606.24143#bib.bib31)]以及来自与学生规模相当的领域专用教师的多教师蒸馏[2 (https://arxiv.org/html/2606.24143#bib.bib26),21 (https://arxiv.org/html/2606.24143#bib.bib33)]。

OPD和RL继承了在线策略的系统瓶颈:每次学习者更新必须等待被训练模型生成新的轨迹[5 (https://arxiv.org/html/2606.24143#bib.bib34)]。对于推理任务,这些轨迹很长且昂贵,因此同步训练往往等待生成而非更新模型,导致学习者资源利用率不足。异步RL[14 (https://arxiv.org/html/2606.24143#bib.bib10),4 (https://arxiv.org/html/2606.24143#bib.bib1)]通过将轨迹生成与学习者更新解耦来缓解这一瓶颈:轨迹工作器持续生成数据,而学习者在较早的轨迹上进行更新,从而提高了训练效率和硬件利用率[23 (https://arxiv.org/html/2606.24143#bib.bib9),34 (https://arxiv.org/html/2606.24143#bib.bib22),18 (https://arxiv.org/html/2606.24143#bib.bib23)]。类似的流水线可以应用于OPD,通过并行运行学生轨迹、教师评分和学习者更新[19 (https://arxiv.org/html/2606.24143#bib.bib5)]。

然而,异步执行引入了陈旧策略数据,从这些数据中学习可能会降低模型质量[3 (https://arxiv.org/html/2606.24143#bib.bib7)]。这产生了权衡:更激进的异步提高了训练吞吐量,但也增加了轨迹与学习之间的策略延迟。因此,先前关于异步RL的工作研究了如何稳定地学习陈旧策略数据[4 (https://arxiv.org/html/2606.24143#bib.bib1),33 (https://arxiv.org/html/2606.24143#bib.bib8),10 (https://arxiv.org/html/2606.24143#bib.bib24)]。然而,这些思想和陈旧数据解决方案是否适用于OPD仍未得到充分探索,因为OPD的实际实现暴露了不同的反馈接口。教师反馈通常通过局部KL损失实现,这需要教师对学生访问的前缀上的动作进行评分。由于全词汇教师logits存储或传输成本高昂(尤其是在异步流水线中),教师分数通常仅缓存有限的动作集(图1 (https://arxiv.org/html/2606.24143#S1.F1))。一旦学习器接收到教师评分的缓存,它可以在缓存的动作上重新计算当前学生的对数概率,但无法恢复从未被评分的动作的教师分数。这引出了三个问题,构成了我们研究的结构:(i) 异步OPD在陈旧性下如何表现,(ii) 异步RL思想和陈旧数据解决方案是否适用于OPD,以及(iii) 有限的教师分数缓存如何塑造OPD估计器的设计。

首先,我们研究KL方向如何影响陈旧性。在带有缓存教师分数的异步OPD下,相同的陈旧轨迹缓存可能对不同KL目标产生不同影响。如图1 (https://arxiv.org/html/2606.24143#S1.F1)所示,前向KL是教师加权的,对陈旧轨迹更鲁棒,而反向KL是学生加权的,当当前学生动作落在评分缓存之外时会变得脆弱。因此,在后续的陈旧性分析中,我们专注于反向KL OPD。

其次,聚焦于反向KL情况,我们询问旨在稳定异步RL的方法是否也能缓解OPD的陈旧性。这一比较很自然,因为OPD中的反向KL允许使用一种RL风格的策略梯度替代项,其中教师-学生对数比率充当词元级优势。因此,我们评估了PPO风格裁剪[16 (https://arxiv.org/html/2606.24143#bib.bib35)]、解耦PPO[4 (https://arxiv.org/html/2606.24143#bib.bib1)]和M2PO[33 (https://arxiv.org/html/2606.24143#bib.bib8)]。在我们的实验中,它们并不比一个更简单的OPD特定替代项更好:在学习器时间使用当前学生重新计算反向KL词元级优势,而不进行裁剪。

第三,我们回到教师缓存约束,并研究由此产生的稀疏和采样反向KL OPD实现的偏差-方差权衡。陈旧学生top-k支撑集提供确定性覆盖,但存在支撑集不匹配,因为它们可能遗漏当前top-k目标所需的动作,并且在陈旧支撑集内重新加权无法恢复缺失的教师分数。单样本蒙特卡洛(MC)通过来自陈旧轨迹策略的重要性可纠正样本避免了这种固定支撑集不匹配,但遭受高方差。这推动了多样本MC方法,该方法在每一步解码时缓存并对多个陈旧策略样本进行教师评分,在减少单样本方差的同时保留了MC可纠正性。

最后,我们将这些发现实例化为AsyncOPD,一个完全异步的OPD流水线,重叠了学生轨迹、教师评分和学习者更新。在Qwen3-Base模型上,与严格同步训练相比,AsyncOPD将训练吞吐量提高了1.6倍到3.8倍,同时保持了可比的准确性。我们的贡献如下:

- •我们首次通过OPD特定的教师缓存约束视角,系统研究了异步OPD中的陈旧性问题。
- •我们表明KL方向改变了陈旧数据问题:前向KL对陈旧轨迹相对鲁棒,而反向KL由于是学生加权而脆弱(第4节 (https://arxiv.org/html/2606.24143#S4))。
- •我们确定最有效的反向KL策略梯度替代项是使用学习器时间重新计算的优势,而不进行裁剪,并且先进的异步RL替代项并不比此选择更好(第5节 (https://arxiv.org/html/2606.24143#S5))。
- •我们表明陈旧学生top-k支撑集存在支撑集不匹配,而单样本MC虽可纠正但方差高;这推动了多样本MC(第6节 (https://arxiv.org/html/2606.24143#S6))。
- •我们提出并开源了AsyncOPD,一个完全异步的OPD训练流水线,并展示了在保持OPD质量的同时提高了训练效率(第7节 (https://arxiv.org/html/2606.24143#S7))。

## 2 相关工作

#### 在线策略蒸馏

在线策略蒸馏(OPD)使用学生自身的轨迹训练学生,同时利用教师对访问的前缀提供密集的词元级反馈[20 (https://arxiv.org/html/2606.24143#bib.bib20),12 (https://arxiv.org/html/2606.24143#bib.bib2)]。GKD[1 (https://arxiv.org/html/2606.24143#bib.bib4)]引入了词元级KL公式,而MiniLLM[6 (https://arxiv.org/html/2606.24143#bib.bib3)]研究了序列级反向KL变体。Li等人[11 (https://arxiv.org/html/2606.24143#bib.bib12)]研究了词元级OPD的训练动态和不稳定配置的配方。TIP[22 (https://arxiv.org/html/2606.24143#bib.bib21)]通过学生熵和教师-学生散度刻画了每词元的重要性。G-OPD[25 (https://arxiv.org/html/2606.24143#bib.bib6)]将词元级OPD解释为密集KL约束的RL,并扩展了奖励缩放。这些工作阐明了OPD作为一种有效的后训练目标,但假设轨迹、教师评分和学习者更新保持同步。

#### 异步强化学习

在同步RL流水线中,训练通常等待批次中最长的轨迹完成,导致学习者资源闲置。异步RL通过将轨迹生成与学习者更新解耦来提高硬件利用率。Async RLHF[14 (https://arxiv.org/html/2606.24143#bib.bib10)]重叠生成和学习,使得在学习者训练较早样本的同时产生新样本。StreamRL[34 (https://arxiv.org/html/2606.24143#bib.bib22)]进一步将RLHF流水线分解为流式阶段。AReaL[4 (https://arxiv.org/html/2606.24143#bib.bib1)]将轨迹工作器与训练工作器完全解耦,以实现连续异步执行。Laminar[18 (https://arxiv.org/html/2606.24143#bib.bib23)]使用细粒度权重同步实现轨迹级异步。然而,异步RL必须从陈旧策略数据中学习。解耦PPO[4 (https://arxiv.org/html/2606.24143#bib.bib1)]通过将陈旧轨迹的行为策略与锚定PPO[16 (https://arxiv.org/html/2606.24143#bib.bib35)]更新的近端策略分离,稳定了异步RL训练。M2PO[33 (https://arxiv.org/html/2606.24143#bib.bib8)]使用二阶矩重要性权重约束稳定陈旧更新,A-3PO[10 (https://arxiv.org/html/2606.24143#bib.bib24)]通过陈旧性感知插值降低解耦PPO开销。

#### 异步在线策略蒸馏

VeRL[19 (https://arxiv.org/html/2606.24143#bib.bib5)]实现了逐步OPD调度器,通过将轨迹延迟固定为一个或两个学习者步骤,重叠学生轨迹、教师评分和学习者更新。这些调度器确立了异步OPD的实际可行性,但未说明OPD估计器在陈旧教师评分缓存下的行为。KDFlow[29 (https://arxiv.org/html/2606.24143#bib.bib25)]通过将教师推理与学习者训练解耦并传输教师隐藏状态来提高LLM蒸馏的系统效率,但针对同步OPD,并将异步执行留作未来工作。我们直接研究这种缺失的异步OPD机制,并根据由此产生的估计器选择构建AsyncOPD。

## 3 预备知识:在线策略蒸馏

#### OPD设置

在每个解码时间步,我们将已访问的前缀ss视为局部状态,并将下一个词元aa视为动作。设q(a∣s)q(a\mid s)表示教师策略,pθ(a∣s)p_\theta(a\mid s)表示当前学生策略。遵循先前关于词元级OPD的工作[12 (https://arxiv.org/html/2606.24143#bib.bib2),25 (https://arxiv.org/html/2606.24143#bib.bib6),11 (https://arxiv.org/html/2606.24143#bib.bib12)],我们将局部损失应用于生成的输出词元,并在固定的前缀状态ss下分析由此产生的目标。OPD可以用不同的散度定义;前向KL和反向KL是两个标准选择[1 (https://arxiv.org/html/2606.24143#bib.bib4)]。

#### 前向KL OPD

在固定前缀ss下,前向KL OPD是教师加权的:

DF(θ;s)=KL(q(⋅∣s)∥pθ(⋅∣s))=∑a∈Vq(a∣s)(logq(a∣s)−logpθ(a∣s)). (1)D_F(θ;s) = KL(q(⋅|s) ‖ p_θ(⋅|s)) = ∑_{a∈V} q(a|s)(log q(a|s) - log p_θ(a|s)). (1)

在固定前缀ss下,梯度为∇θDF(θ;s)=−∑a∈Vq(a∣s)∇θlogpθ(a∣s).∇_θ D_F(θ;s) = -∑_{a∈V} q(a|s) ∇_θ log p_θ(a|s).

#### 反向KL OPD

在相同前缀下,反向KL OPD是学生加权的:

DR(θ;s)=KL(pθ(⋅∣s)∥q(⋅∣s))=−∑a∈Vpθ(a∣s)(logq(a∣s)−logpθ(a∣s)). (2)D_R(θ;s) = KL(p_θ(⋅|s) ‖ q(⋅|s)) = -∑_{a∈V} p_θ(a|s)(log q(a|s) - log p_θ(a|s)). (2)

微分并利用Ea∼pθ(⋅∣s)[∇θlogpθ(a∣s)]=∇θ∑apθ(a∣s)=0E_{a~p_θ(⋅|s)}[∇_θ log p_θ(a|s)] = ∇_θ ∑_a p_θ(a|s) = 0,得到

∇θDR(θ;s)=∑a∈Vpθ(a∣s)(logpθ(a∣s)−logq(a∣s)+1)∇θlogpθ(a∣s).∇_θ D_R(θ;s) = ∑_{a∈V} p_θ(a|s)(log p_θ(a|s) - log q(a|s) + 1) ∇_θ log p_θ(a|s).

相似文章

ShortOPD:通过短到长在策略蒸馏恢复剪枝后的大语言模型

arXiv cs.LG

ShortOPD提出了一种短到长的在策略蒸馏方案,通过将训练集中在有效前缀上,恢复用于自由形式生成的剪枝后大语言模型,相较于未恢复模型实现了高达9倍的改进,并以四分之一的训练时间达到了长视界蒸馏的效果。

学会预见:揭示 On-Policy 蒸馏效率的解锁机制

arXiv cs.CL

本文研究了大型语言模型中 On-Policy 蒸馏(OPD)效率背后的参数级机制,将其归因于模块分配和更新方向上的早期“预见性”。本文提出了 EffOPD,一种即插即用方法,可在不损害最终性能的情况下将 OPD 训练速度提高 3 倍。

基于前缀重放的多轮在线策略蒸馏

Hugging Face Daily Papers

本文提出 ReOPD,一种用于 LLM 智能体的在线策略蒸馏方法,该方法重用预先收集的教师轨迹作为重放前缀,无需新的环境交互即可提高效率和准确性。