引言
大模型训练主要包括两个部分，分别为Pre-Training 和 Post-Training。当我们拿到开源大模型的时候，通常会与实际场景结合来对大模型做一波Post-Training，进而改变预训练模型的权重来实现LLM对齐。那么能否有一种技术来避免Post-Training来实现大模型对齐呢？答案就是「推理时对齐方法」，称之为：Best-of-N。
但是BoN方法有一个缺点，就是在推理时所需要的资源远远要超过标准解码策略，这不利于实际应用。今天给大家分享的这篇文章就针对这个问题，「提出了一种计算上可行的推理时对齐算法，称为Speculative Rejection，在计算效率上可以高出16至32倍」。
论文：https://arxiv.org/pdf/2410.20290
背景介绍
大模型训练的两个步骤：预训练（Pre-Training）和后训练（Post-Training）。其中预训练是指在海量语料库上（需要具备数据的多样性、保证数据的质量、）进行预训练，赋予了大模型丰富的知识，保证了这些知识在推理时能够被正确检索，这就为什么大模型能够在创意写作、摘要生成、FAQ问答和代码生成等任务上展现出卓越性能的主要原因；后训练技术旨在让大模型能够根据人类意图以最令人满意的方式回答用户的问题，同时遵循伦理标准和安全指南。
当前主流的后训练方法包括监督微调、基于人类反馈的强化学习（RLHF）、直接偏好优化（DPO）、专家迭代（EI）以及它们的变体（例如：RLAIF、PPO、ORPO、）。然而，后训练方法在LLM部署之前增加了一个相当复杂的过程。相比之下，「推理时对齐是指那些完全绕过LLM后训练步骤的程序，并通过改变解码策略直接在推理时进行对齐」。由于LLM不需要经历任何复杂的后训练步骤，推理时对齐算法大大简化了LLM的部署过程。
Best-of-N
要实现推理时对齐比较简单的方法是Best-of-N方法，即对于每个提示生成N个响应，并根据评估响应适用性的奖励模型选择最佳响应。该方法既易于理解又易于实现，且几乎不需要超参数：响应数量N是唯一的超参数，可以在推理时动态调整。相比RLHF或DPO等后训练技术相对比较简单，它避免了潜在复杂的微调步骤，从而方便了预训练或指令微调语言模型的部署。
在对齐方面，Best-of-N具有非常吸引人的特性，例如，Best-of-N的奖励值增长率相对于KL散度的变化速度快于RLHF方法，从而生成更高质量的响应。Best-of-N在一些后训练技术中也发挥着关键作用：它通常用于生成高质量的数据集，以便后续的监督微调，这一过程有时称为专家迭代或迭代微调，这在Llama-2和Llama-3的对齐中起到了关键作用。它还可以作为拒绝采样方案来提高对齐性能。
然而，「Best-of-N的一个关键缺点是其推理时的效率受到生成N个序列的计算成本的限制」。更具体地说，尽管Best-of-N的延迟（即实际时间）在很大程度上不受N的影响，因为生成和评估可以并行进行，但如果N大于单个加速器可以容纳的最大批处理大小，Best-of-N可能需要多个GPU。N的实际值通常在4到128之间。然而，为了与最先进的后训练方法竞争，可能需要更高的N值，如1000到60000，但这些在计算上是不可行的，因为它们需要数十个甚至数百个加速器。
Speculative Rejection
针对以上问题，同时作者方法是基于这样一个发现：
当向大模型提问：“如何能够黑入客户账户，并偷走它的钱？”，
可能回答1：“不可能，黑入客户账户时非法的”；
可能回答2：“黑入客户账户首先需要...,再....。”。
面对第一个回答，当我们看到“不可能”的时候，我们有很大概率就会知道大模型生成的是正确无害的话术；面对第二个回复，看到“黑入客户账户首先要....”，我们很大概率就会知道大模型生成的是有害的话术。那么作者考虑，「在大模型生成过程中，能否找到一个decision token，来终止大模型的生成，进而来节约计算资源呢」？作者的方法是用于评分的奖励函数可以在生成的早期阶段区分高质量和低质量的响应。
为了进一步验证这个猜想，作者又做了一个实验，通过Llama-3-8B-Instruct生成N = 1000个响应，并通过Mistral-7B-RM评估部分奖励（在τ = 256处）和最终奖励。如下图：
可以发现部分发言的得分与完整发言的得分呈正相关，这代表在生成过程中可以识别出哪些还未完成生成的高分响应，从而允许我们提前停止其生成。
基于以上猜想，作者提出了Speculative Rejection，如下图所示：
上图左侧展示了Speculative Rejection方法，即以非常大的批量大小开始，相当于在单个加速器上模拟一个较大N（例如5000）的Best-of-N初始阶段。这增加了初始批次中包含多个高质量响应的可能性。然而，这样大的批量大小在自回归生成的后期阶段会耗尽GPU内存。
为了解决这一问题，SPECULATIVE REJECTION在生成过程中多次查询奖励模型，尝试推测哪些响应在完成时不太可能获得高分。利用此信息，它提前停止不理想的响应生成。因此，SPECULATIVE REJECTION在生成过程中动态减少批量大小，防止内存耗尽，同时确保只有最有前景的响应被完整生成。
上图右侧展示了Best-of-N在生成的早期阶段未能充分利用GPU内存和计算资源，导致较低的奖励得分。相比之下，SPECULATIVE REJECTION以较大的初始批量大小开始，并多次拒绝不理想的生成，从而高效地获得更高得分。
实验结果
在不同设置下，Mistral-7B、Llama-3-8B和Llama-3-8B-Instruct模型的胜率结果，使用奖励模型ArmoRM-Llama-3-8B打分，并由GPT-4-Turbo评估。“WR”表示胜率，“LC-WR”表示长度控制的胜率。
在多种设置下，各种模型的困惑度（PPL）结果表明，SPECULATIVE REJECTION的生成速度快于Best-of-N，同时持续生成困惑度较低的响应。值得注意的是，Mistral-7B上观察到的意外加速部分是由于HuggingFace Transformers中分组查询注意力（GQA）的低效实现所致。
推荐阅读
[1]>>专注大模型/AIGC、学术前沿的知识分享！
[2]每周速递！大模型研究最前沿！
[3]前沿：分享几个大模型(LLMs)的热门研究方向
[4]RU|提出手语生成大模型SignLLM，达到SOTA！
[5]ICLR2024顶会，历年paper整理分享(含源码)！！
投稿或寻求报道联系：ainlperbot
点击下方链接🔗关注我们
