论文

ResOPD:稀疏 在策略蒸馏的尾部残差

ResOPD: Tail Residualization for Sparse On-Policy Distillation

摘要

在策略蒸馏 (OPD) 在自生成轨迹上训练学生模型模型,但跨长推理轨迹传输密集的教师模型分布会产生令人望而却步的通信和记忆瓶颈。因此,实际系统依赖于稀疏的教师模型接口,通常传输采样的token分数或小的 Top-$k$ 分布。然而,这种稀疏设置面临着一个根本性的困境:采样的token估计器是无偏的,但受到严重的梯度方差的影响,而直接优化 Top-$k$ 目标会引入系统偏差。为了改善这种权衡,我们提出了 ResOPD(在策略 Distillation with Tail Residualization),它在 在策略采样下提供无偏全词汇反向 KL 梯度估计,在评估设置中的相同稀疏有效负载下显着减少方差。 ResOPD 将未观察到的词汇聚合成可观察的粗尾事件,计算其精确聚合梯度,并仅对细粒度尾内残差进行采样,这不需要额外的教师模型查询或前向传递。大量实验表明,ResOPD 显着降低了梯度方差,稳定了在线训练动态,并提高了评估设置中的下游性能。这些结果将 ResOPD 确立为一种高效、即插即用的方差缩减原语,适用于稀疏 在策略蒸馏。

ResOPD:稀疏 在策略蒸馏的尾部残差的原论文方法或结果图
图 2:ResOPD 概述。 (a) 当采样的学生模型-token落在头部时,ResOPD仅使用精确的粗梯度;尾部撞击时,它会添加残余修正。 (b) 与 HETS 相比,ResOPD 通过将尾部样本集中在 $\bm{\mu}_{B}$ 周围来抑制事件间方差,从而在执行轨迹步骤中实现与教师模型分布的出色对齐。