奖励偏移投机采样:一种高效的测试时弱到强对齐方法
Reward-Shifted Speculative Sampling Is An Efficient Test-Time Weak-to-Strong Aligner
AI总结:
本文提出奖励偏移投机采样(SSS)算法,利用对齐的小型草稿模型与未对齐目标模型间的分布偏移,在测试时以更低推理成本恢复RLHF最优解,实现高效的弱到强对齐。
摘要:
将大语言模型(LLMs)与人类偏好对齐已成为其发展中的关键步骤。近期的研究越来越关注测试时对齐,即在推理过程中分配额外计算资源以增强LLM的安全性和推理能力。然而,这些测试时对齐技术通常会产生大量的推理成本,限制了它们的实际应用。我们受到投机采样加速方法的启发,该方法利用一个小型草稿模型高效地预测未来token,以解决测试时对齐的效率瓶颈。我们提出了奖励偏移投机采样(SSS)算法,其中草稿模型与人类偏好对齐,而目标模型保持不变。我们从理论上证明,通过修改接受准则和奖励token分布,可以利用对齐的草稿模型与未对齐的目标模型之间的分布偏移,在实际未获得RLHF最优解的情况下恢复该最优解。在测试时弱到强对齐实验中,我们的算法以显著降低的推理成本取得了更优的黄金奖励分数,从而验证了其有效性和高效性。
英文摘要:
Aligning large language models (LLMs) with human preferences has become a critical step in their development. Recent research has increasingly focused on test-time alignment, where additional compute is allocated during inference to enhance LLM safety and reasoning capabilities. However, these test-time alignment techniques often incur substantial inference costs, limiting their practical application. We are inspired by the speculative sampling acceleration, which leverages a small draft model to efficiently predict future tokens, to address the efficiency bottleneck of test-time alignment. We introduce the reward-shifted speculative sampling (SSS) algorithm, in which the draft model is aligned with human preferences, while the target model remains unchanged. We theoretically demonstrate that the distributional shift between the aligned draft model and the unaligned target model can be exploited to recover the RLHF optimal solution without actually obtaining it, by modifying the acceptance criterion and bonus token distribution. Our algorithm achieves superior gold reward scores at a significantly reduced inference cost in test-time weak-to-strong alignment experiments, thereby validating both its effectiveness and efficiency.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。