用奖励调整的推测采样,让大模型推理更安全且更快。
Reward-Shifted Speculative Sampling Is An Efficient Test-Time Weak-to-Strong Aligner
- 用小模型预判未来词,大幅降低推理开销。
- 在测试阶段实现弱到强对齐,奖励得分更高且耗时少。
- 适合需要高效安全推理的应用场景。
将大语言模型(LLMs)与人类偏好对齐已成为其发展中的关键步骤。近年来研究越来越多关注测试阶段对齐,即在推理时分配额外计算资源以提升模型安全性与推理能力。然而,这类方法常伴随显著的推理成本,限制了实际应用。受推测采样加速机制启发,该工作提出奖励调整的推测采样(SSS)算法:使小的草稿模型与人类偏好对齐,而目标模型保持不变。理论证明,通过调整接受准则和奖励分配,可利用对齐草稿模型与未对齐目标模型之间的分布差异,无需直接获得强化学习与人类反馈(RLHF)最优解即可恢复其效果。实验表明,该方法在测试阶段弱到强对齐任务中显著降低推理成本的同时取得更优黄金奖励分数,验证了其有效性和高效性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。