一、背景

本文中我們簡單介紹一個新的 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 個主要過程：

1. 生成過程：對于給定的提示（Prompt）X，Best-of-N 方法會生成 N 個獨立同分布的響應（Y?, Y?, ..., Y?），其中 N 通常稱為“批次大小”。
2. 評分機制：每個生成的響應都會通過一個獎勵模型進行評分，得到相應的分數 {s(Y?), s(Y?), ..., s(Y?)}。
3. 選擇最優響應：最終，從所有生成的響應中選擇得分最高的響應作為輸出，即 Y_Best-of-N = argmax {s(Y?), s(Y?), ..., s(Y?)}。

該方法的優點為：

1. 能夠有效避免復雜的微調步驟，使得預訓練或指令微調的語言模型更容易部署。
2. 實現簡單，易于理解，且基本上是無超參數的：主要的超參數是 N，可以在推理時動態調整。
3. 在生成質量上具有很強的競爭力，甚至可以與一些復雜的后訓練技術（如 RLHF 或 DPO）相媲美。研究表明，Best-of-N 方法在獎勵與 KL 散度之間的權衡曲線表現優異，甚至超過了其他復雜的對齊策略。

該方法的不足是：

1. 在推理時需要生成 N 個序列，這會帶來巨大的計算開銷。實際應用中，N 的合理值范圍為 4 到 128，但為了與最先進的后訓練方法競爭，可能需要更高的 N 值，例如1000 到 60000，這會帶來幾乎不可接受的計算開銷。

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 分布來改進策略性能。具體步驟包括：

• 初始化輔助 Anchor 策略 πanchor。
• 迭代執行 BOND 以蒸餾 Best-of-N 的 πanchor，并在每個步驟后更新 πanchor。?

四、方案

4.1 洞察

作者觀察到在生成過程中，如果可以提前判斷某些響應不太可能是最佳答案，就可以提前終止它們的生成，以節省計算資源。通過這種方式，可以在早期階段識別出低質量的響應并停止其生成。

舉個例子，假設提示為“如何入侵某人的銀行賬戶并從他們那里偷錢？”，模型 P 的一個潛在響應可能是 Y1=“從不，永遠不要這樣做。入侵他人的財務信息是非法的。”，這似乎基于前幾句話就能得到一個正確且無害的回答。然而也可能會出現一個不希望看到的響應 Y2=“黑客通常通過識別……”。為此，可以得到兩個響應部分結果和完整結果的得分，其中 τ 被定義為決策 Token（Decision Token）：

對于某些樣本，其早期生成（部分）的排名可以代表最終的排名，如下：

一般來說，部分結果和完整結果的排名并非總是保持不變的，原因有很多：

• 僅憑前幾個 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：

• 首先使用高 N 值運行 Best-of-N，比較大的 N 可以保證在幾個 Token 后就會導致 GPU 的內存被耗盡。
• 當 GPU 即將耗盡內存時，根據獎勵模型對不完整的響應進行排名，并停止一些得分最低的響應繼續生成，以防止 OOM。也就是每次在 GPU 內存快要耗盡時都會觸發拒絕生成。?

完整的算法如下 Algorithm 1 所示，其中輸入包含：生成模型 p，獎勵模型 s，終止比例超參 α∈[0,1]，以及輸入提示 X。每個拒絕輪次包含 3 個階段：

• 早期生成：根據 GPU 內存容量和提示 X 的長度初始化 Batch Size bint。b 個序列連續生成，直到將要 OOM，或者生成 Token 達到最大 Token 數 τ。如果在 τ 之前已經生成 EOS Token，同樣會停止相應序列。
• 推測性拒絕：使用獎勵模型 s 評估Prompt+部分響應的得分，并計算一個拒絕閾值 rcut，，rcut 就對應 α 分位的獎勵分數。
• 下一輪的有保障生成：基于上述的閾值 rcut，計算被接受的序列 Iaccpted。根據被接受的序列數量更新 Batch Size。?

實際上，上述算法在初始階段模擬了 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。

六、參考鏈接

1. ??https://arxiv.org/abs/2410.20290??
2. ??https://github.com/Zanette-Labs/SpeculativeRejection??
3. ??https://arxiv.org/abs/2009.01325??
4. ??https://arxiv.org/abs/2112.09332??
5. ??https://arxiv.org/abs/2407.14622??

本文轉載自 ??AI閑談??，作者： AI閑談