论文
教师是方向,而不是目的地:在策略蒸馏中推断强化学习引起的表示残差
The Teacher Is a Direction, Not a Destination: Extrapolating RL-Induced Representation Residuals in On-Policy Distillation
摘要
on-policy蒸馏(OPD)训练学生在学生自己的轨迹上匹配老师的下一个词元分布,并产生了巨大的经验收益。广义变体允许学生通过推断输出空间中的隐式奖励来超越老师。然而,语言模型头各向异性地减弱了这种变化:教师隐藏状态中编码的大部分变化仅以其权重的一小部分传递至logits,并且输出空间外推所依赖的采样词元对数概率比注入了外推放大的噪声,从而使训练不稳定。我们观察到强化学习(RL)会相对于其基本检查点改变模型的内部表示,并且可以在每一层测量这种转变的方向。受这一观察的启发,我们提出了 RIDE(RL 引发的方向外推),它直接在表示空间中推断 RL 引发的变化:在每一层和标记位置,RIDE 计算教师与其预 RL 检查点之间的残差,并将学生的隐藏状态沿着该残差回归到超出教师的目标。以采样轨迹为条件,这种回归相当于最大化由以教师为中心的二次罚分下的残差定义的线性方向奖励,这明确了目标如何沿着 RL 诱导的方向移动学生,同时限制其与教师的偏差。在跨越不同规模、架构和预训练谱系的四对基础/强化学习教师对中,RIDE 在每对上都接近或超过了经过强化学习训练的教师,并且是唯一一种均值达到此目的的方法,并且它始终优于输出空间外推法,只要教师接近其基础,输出空间外推法就会降低学生的性能。项目页面:https://github.com/xixixixixxxx/RIDE.