Fast Best-of-N Decoding via Speculative Rejection
作者: Hanshi Sun, Momin Haider, Ruiqi Zhang, Huitao Yang, Jiahao Qiu, Ming Yin, Mengdi Wang, Peter Bartlett, Andrea Zanette
分类: cs.CL
发布日期: 2024-10-26 (更新: 2024-10-31)
备注: NeurIPS 2024
💡 一句话要点
提出Speculative Rejection算法，加速Best-of-N解码，提升LLM推理时对齐效率。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大型语言模型 推理时对齐 Best-of-N解码 Speculative Rejection 计算效率 奖励模型
📋 核心要点
- 现有LLM对齐方法，如DPO和PPO，依赖于复杂的后训练过程，增加了部署难度。
- Speculative Rejection算法在推理时进行对齐，避免了后训练，同时生成符合人类偏好的高质量响应。
- 实验表明，Speculative Rejection算法在保证对齐效果的同时，计算效率提升了16到32倍。
📝 摘要（中文）
大型语言模型（LLM）的安全有效部署涉及一个关键步骤，即对齐，以确保模型响应符合人类偏好。主流对齐技术，如DPO、PPO及其变体，通过在后训练阶段改变预训练模型权重来实现LLM对齐。虽然这些后训练方法很普遍，但它们在LLM部署之前增加了相当大的复杂性。推理时对齐方法避免了复杂的后训练步骤，而是将生成偏向于与人类偏好对齐的响应。最著名的推理时对齐方法称为Best-of-N，其效果与最先进的后训练程序一样。不幸的是，Best-of-N在推理时比标准解码策略需要更多的资源，这使得它在计算上不可行。在这项工作中，我们介绍Speculative Rejection，一种在计算上可行的推理时对齐算法。它像Best-of-N一样，根据给定的奖励模型生成高分响应，同时计算效率提高了16到32倍。
🔬 方法详解
问题定义：Best-of-N解码方法虽然在推理时对齐方面表现出色，但其计算成本过高，严重限制了其在实际应用中的可行性。该论文旨在解决Best-of-N解码计算效率低下的问题，使其能够在资源受限的环境中部署。
核心思路：Speculative Rejection的核心思想是利用一个奖励模型来评估候选响应的质量，并基于此有选择地拒绝或接受这些响应。通过这种方式，算法能够专注于生成更有可能获得高分的响应，从而减少了不必要的计算。
技术框架：Speculative Rejection算法主要包含以下几个阶段：1) 使用LLM生成N个候选响应；2) 使用奖励模型对每个候选响应进行评分；3) 根据评分，决定接受或拒绝该响应。如果响应被拒绝，则重新生成新的候选响应，直到找到一个可接受的响应。整个过程迭代进行，直到生成满足要求的输出。
关键创新：Speculative Rejection的关键创新在于其有选择地拒绝低质量响应的机制。与Best-of-N盲目地生成大量候选响应并从中选择最佳响应不同，Speculative Rejection能够更有效地利用计算资源，从而显著提高解码效率。此外，该方法不需要对LLM进行后训练，简化了部署流程。
关键设计：Speculative Rejection算法的关键设计包括：1) 奖励模型的选择：奖励模型的质量直接影响算法的性能。需要选择一个能够准确反映人类偏好的奖励模型。2) 拒绝阈值的设置：拒绝阈值决定了哪些响应会被拒绝。需要根据实际情况调整拒绝阈值，以平衡生成质量和计算效率。3) 候选响应数量N的选择：N越大，找到高质量响应的可能性越高，但计算成本也会增加。需要根据资源限制选择合适的N值。
🖼️ 关键图片
📊 实验亮点
实验结果表明，Speculative Rejection算法在保证与Best-of-N相当的对齐效果的同时，计算效率提升了16到32倍。这意味着在相同的计算资源下，Speculative Rejection能够生成更多的对齐响应，或者在更少的计算资源下达到相同的对齐效果。
🎯 应用场景
Speculative Rejection算法可广泛应用于需要对LLM输出进行对齐的场景，例如对话系统、文本摘要、代码生成等。该算法能够显著降低推理成本，使得在资源受限的设备上部署对齐的LLM成为可能，从而加速LLM在各个领域的应用。
📄 摘要（原文）
The safe and effective deployment of Large Language Models (LLMs) involves a critical step called alignment, which ensures that the model's responses are in accordance with human preferences. Prevalent alignment techniques, such as DPO, PPO and their variants, align LLMs by changing the pre-trained model weights during a phase called post-training. While predominant, these post-training methods add substantial complexity before LLMs can be deployed. Inference-time alignment methods avoid the complex post-training step and instead bias the generation towards responses that are aligned with human preferences. The best-known inference-time alignment method, called Best-of-N, is as effective as the state-of-the-art post-training procedures. Unfortunately, Best-of-N requires vastly more resources at inference time than standard decoding strategies, which makes it computationally not viable. In this work, we introduce Speculative Rejection, a computationally-viable inference-time alignment algorithm. It generates high-scoring responses according to a given reward model, like Best-of-N does, while being between 16 to 32 times more computationally efficient.