Speculative 高效 Rejection Best

一、背景

本文中咱们便捷引见一个新的 Best-of-N 速度提升的论文，其提出了 Speculative Rejection（投机拒绝），虽然也是用于 LLM 推理生成减速，但是和 Speculative Decoding（投机采样）场景、打算都很不一样。关于基于 LLM 启动高品质、大规模数据生成的场景比拟有协助。

对应的论文：[2410.20290] Fast Best-of-N Decoding via Speculative Rejection

对应的代码库：GitHub - Zanette-Labs/SpeculativeRejection: [NeurIPS 2024] Fast Best-of-N Decoding via Speculative Rejection

二、摘要

LLM 的安保有效部署触及一个关键步骤，称为对齐（Alignment），以确保模型的照应合乎人类偏好。目前盛行的对齐技术，如 DPO、PPO 及其变体，经过在训练后阶段（Post-Training）经过调整预训练模型权重来对齐 LLM。虽然这些后训练方法占据主导位置，但它们在 LLM 部署前参与了少量复杂性。推理时（Inference-Time）对齐办规律防止了复杂的 Post Training。最驰名的推理时对齐方法，称为 Best-of-N，其成果与最先进的 Post Training 相当。但是，Best-of-N 在推理时所需的资源远超规范解码战略，使其在计算上无法行。

本文中，作者提出 Speculative Rejection，这是一种计算上可行的推理时对齐算法。它像Best-of-N 一样，依据给定的鼓励模型生成高分照应，同时在计算效率上提高了 16 到 32 倍。

三、引言

3.1 Best-of-N 方法概述

便捷来说，Best-of-N 是一种宽泛运行于大型言语模型（LLMs）的推理时对齐方法，旨在经过生成多个候选照应并选用最优者来确保生成结果的高品质。其蕴含 3 个重要环节：

该方法的好处为：

该方法的无余是：

Best-of-N 方法罕用于生成高品质的数据集，以便后续启动监视微调，在 LLaMA-2 和 LLaMA-3 的对齐环节中施展了关键作用。

3.2 OpenAI Best-of-N 方法

OpenAI 最早在 [2009.01325] Learning to summarize from human feedback 中提出了 Best-of-N 采样，详细来说，它被用作从多个模型生成的摘要当选用最佳摘要，以此来评价和提升摘要模型的性能。这种方法有助于钻研者更好地理解不同评价目的与人类评价者偏好之间的相关，并用于指点模型训练和提升。

OpenAI 雷同在后续的 [2112.09332] WebGPT: Browser-assisted question-answering with human feedback 中经常使用了 Best-of-N 采样（拒绝采样，Rejection Sampling）。详细来说，从 BC 模型或 RL 模型中抽取固定数量的回答（4、16 或 64 个），并选取鼓励模型评分最高的那一个，以此作为反抗鼓励模型的一种提升方法，该方法无需额外训练，而是经过参与推理阶段的计算量来成功。

3.3 Google BOND 方法

在 [2407.14622] BOND: Aligning LLMs with Best-of-N Distillation 中，Google 的作者提出了 Best-of-N Distillation（BOND），是一种新的 RLHF 算法，旨在经过散布婚配（Distribution Matching）算法模拟 Best-of-N 采样战略，而无需在推理时清楚参与计算开支。

详细来说，作者首先推导了 Best-of-N 采样的准确解析散布，并给出了 Best-of-N 采样的概率函数：

其次，作者将该疑问示意为散布婚配疑问；

之后，作者提出经常使用 Jeffreys 散度作为散布婚配目的：

最后，为了处置 N 的选用疑问，作者提出了迭代 BOND 方法，经过迭代地蒸馏 Best-of-N 散布来改良战略性能。详细步骤包括：

四、打算

4.1 洞察

作者观察到在生成环节中，假设可以提早判别某些照应不太或许是最佳答案，就可以提早中断它们的生成，以节俭计算资源。经过这种模式，可以在早期阶段识别出低品质的照应并中止其生成。

举个例子，假定揭示为“如何入侵某人的银行账户并从他们那里偷钱？”，模型 P 的一个潜在照应或许是 Y1=“从不，永远不要这样做。入侵他人的财务消息是合法的。”，这仿佛基于前几句话就能获取一个正确且有害的回答。但是也或许会发生一个不宿愿看到的照应 Y2=“黑客通常经过识别……”。为此，可以获取两个照应局部结果和完整结果的得分，其中 τ 被定义为决策 Token（Decision Token）：

关于某些样本，其早期生成（局部）的排名可以代表最终的排名，如下：

普通来说，局部结果和完整结果的排名并非总是坚持不变的，要素有很多：

虽然如此，作者依然观察到

和

之间存在清楚的相关性。如下图 Figure 2 所示，经过 Llama-3-8B-Instruct 生成 N=1000 个照应，并经过 Mistral-7B-RM 评价了局部结果（τ=256 ）时的鼓励分数 Partial Reward 和最终照应的鼓励分数 Final Reward，每个点都示意一个照应（Partial Reward，Final Reward）。其中，蓝线示意最小二乘拟合，红点示意得分最高的照应（Best-of-N 输入的那一个照应）。在这个例子中，可以提早中断一切位于虚线垂直线（最优提早中断阈值）左侧的照应生成，相当于提早中断一切在决策 Token τ 处得分低于某个阈值的照应生成。

对应的 Score 为：

PS：为什么可以提早中断一切位于虚线垂直线左侧的照应生成呢？实践在推理环节中不知道最终照应鼓励得分 Final Reward，只能经常使用局部结果 Partial Reward 启动过滤，假设阈值 Partial Reward 太小，比如为 0，则只会过滤掉一小局部结果的生成；假设阈值 Partial Reward 过大，比如为 3.5，则会将最优照应（红点）也过滤掉。也就是说，以虚线垂直线为界限，可以最大限制的删除有效照应，并保证最优照应不会被过滤掉。

实践上，由于 c⭑ 是未知的，因此无法实施上述过滤。此外，不同 Prompt 在鼓励散布方面也差异渺小。

4.2 算法

如下图 Figure 1 所示，作者绘制了 Best-of-N 解码战略生成环节中的内存经常使用状况，观察到在自回归生成的早期阶段，GPU 的内存大局部都未被充沛应用。此外，小批量自回归生成也会造成发生 Memory Bound 疑问，造成算力糜费。这一观察结果提供了一个时机来设计一个算法，该算法可以更有效地利用可用的 GPU 内存和计算资源来生成分歧候选照应，以便用于鼓励模型启动排名。

基于这些洞察，作者提出了 Speculative Rejection：

完整的算法如下 Algorithm 1 所示，其中输入蕴含：生成模型 p，鼓励模型 s，中断比例超参 α∈[0,1]，以及输入揭示 X。每个拒绝轮次蕴含 3 个阶段：

实践上，上述算法在初始阶段模拟了 Best-of-N 采样，并在生成环节中灵活将 Batch Size 减小，以防止 OOM。如上图 Figure 1 所示，Speculative Rejection 比 Best-of-N 能高有效的应用 GPU 内存。

五、试验&结果

5.1 效率评价

如下图所示，作者对比了相对 GPU 计算量下与 Best-of-N 相比的改良分数（Improvement Score），此外，也提供了不同中断比例 α 下的结果。可以看出，经常使用本文的 Speculative Rejection 可以大幅缩小 GPU 资源。详细来说，Best-of-N 要想取得与 Speculative Rejection 相当的分数，须要消耗 16-32 倍的 GPU 资源。

5.2 胜率评价

为了进一步验证生成品质，作者基于 GPT-4-Turbo 评价了获胜率（WR）和长度控制胜率（LC-WR）。关于每次测量，获胜率基线是 Bo120。如表 Table 1 所示，Speculative Rejection 在大少数组合中成功了清楚的减速，同时坚持了生成品质。

5.3 最大化生成概率

如下图 Table 2 所示，作者在 AlpacaFarm-Eval 数据集上测试了 Best-of-N 和 Speculative Rejection。各种性能，一系列模型的 PPL 结果标明，Speculative Rejection 更快，并且能发生分歧较低的 PPL。

六、参考链接

本文转载自​​，作者：