一、背景

本文中我們簡單介紹一個新的 Best-of-N 速度優(yōu)化的論文，其提出了 Speculative Rejection（投機拒絕），雖然也是用于 LLM 推理生成加速，但是和 Speculative Decoding（投機采樣）場景、方案都很不一樣。對于基于 LLM 進行高質(zhì)量、大規(guī)模數(shù)據(jù)生成的場景比較有幫助。

對應(yīng)的論文：[2410.20290] Fast Best-of-N Decoding via Speculative Rejection

對應(yīng)的代碼庫：GitHub - Zanette-Labs/SpeculativeRejection: [NeurIPS 2024] Fast Best-of-N Decoding via Speculative Rejection

二、摘要

LLM 的安全有效部署涉及一個關(guān)鍵步驟，稱為對齊（Alignment），以確保模型的響應(yīng)符合人類偏好。目前流行的對齊技術(shù)，如 DPO、PPO 及其變體，通過在訓(xùn)練后階段（Post-Training）通過調(diào)整預(yù)訓(xùn)練模型權(quán)重來對齊 LLM。盡管這些后訓(xùn)練方法占據(jù)主導(dǎo)地位，但它們在 LLM 部署前增加了大量復(fù)雜性。推理時（Inference-Time）對齊方法則避免了復(fù)雜的 Post Training。最著名的推理時對齊方法，稱為 Best-of-N，其效果與最先進的 Post Training 相當(dāng)。然而，Best-of-N 在推理時所需的資源遠(yuǎn)超標(biāo)準(zhǔn)解碼策略，使其在計算上不可行。

本文中，作者提出 Speculative Rejection，這是一種計算上可行的推理時對齊算法。它像Best-of-N 一樣，根據(jù)給定的獎勵模型生成高分響應(yīng)，同時在計算效率上提高了 16 到 32 倍。

三、引言

3.1 Best-of-N 方法概述

簡單來說，Best-of-N 是一種廣泛應(yīng)用于大型語言模型（LLMs）的推理時對齊方法，旨在通過生成多個候選響應(yīng)并選擇最優(yōu)者來確保生成結(jié)果的高質(zhì)量。其包含 3 個主要過程：

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

該方法的優(yōu)點為：

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

該方法的不足是：

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

Best-of-N 方法常用于生成高質(zhì)量的數(shù)據(jù)集，以便后續(xù)進行監(jiān)督微調(diào)，在 LLaMA-2 和 LLaMA-3 的對齊過程中發(fā)揮了關(guān)鍵作用。

3.2 OpenAI Best-of-N 方法

OpenAI 最早在 [2009.01325] Learning to summarize from human feedback 中提出了 Best-of-N 采樣，具體來說，它被用作從多個模型生成的摘要中選擇最佳摘要，以此來評估和優(yōu)化摘要模型的性能。這種方法有助于研究者更好地理解不同評估指標(biāo)與人類評估者偏好之間的關(guān)系，并用于指導(dǎo)模型訓(xùn)練和優(yōu)化。

OpenAI 同樣在后續(xù)的 [2112.09332] WebGPT: Browser-assisted question-answering with human feedback 中使用了 Best-of-N 采樣（拒絕采樣，Rejection Sampling）。具體來說，從 BC 模型或 RL 模型中抽取固定數(shù)量的回答（4、16 或 64 個），并選取獎勵模型評分最高的那一個，以此作為對抗獎勵模型的一種優(yōu)化方法，該方法無需額外訓(xùn)練，而是通過增加推理階段的計算量來實現(xiàn)。

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 采樣策略，而無需在推理時顯著增加計算開銷。

具體來說，作者首先推導(dǎo)了 Best-of-N 采樣的精確解析分布，并給出了 Best-of-N 采樣的概率函數(shù)：

其次，作者將該問題表示為分布匹配問題；

之后，作者提出使用 Jeffreys 散度作為分布匹配目標(biāo)：

最后，為了解決 N 的選擇問題，作者提出了迭代 BOND 方法，通過迭代地蒸餾 Best-of-N 分布來改進策略性能。具體步驟包括：

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

四、方案

4.1 洞察

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

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

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

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

• 僅憑前幾個 Token 可能無法評估整個響應(yīng)的得分，因為生成過程可能以意想不到的方式繼續(xù)。
• 獎勵模型通常被訓(xùn)練用于評估完整響應(yīng)。

盡管如此，作者依然觀察到

和

之間存在顯著的相關(guān)性。如下圖 Figure 2 所示，通過 Llama-3-8B-Instruct 生成 N=1000 個響應(yīng)，并通過 Mistral-7B-RM 評估了部分結(jié)果（τ=256 ）時的獎勵分?jǐn)?shù) Partial Reward 和最終響應(yīng)的獎勵分?jǐn)?shù) Final Reward，每個點都表示一個響應(yīng)（Partial Reward，F(xiàn)inal Reward）。其中，藍(lán)線表示最小二乘擬合，紅點表示得分最高的響應(yīng)（Best-of-N 輸出的那一個響應(yīng)）。在這個例子中，可以提前終止所有位于虛線垂直線（最優(yōu)提前終止閾值）左側(cè)的響應(yīng)生成，相當(dāng)于提前終止所有在決策 Token τ 處得分低于某個閾值的響應(yīng)生成。

對應(yīng)的 Score 為：

PS：為什么可以提前終止所有位于虛線垂直線左側(cè)的響應(yīng)生成呢？實際在推理過程中不知道最終響應(yīng)獎勵得分 Final Reward，只能使用部分結(jié)果 Partial Reward 進行過濾，如果閾值 Partial Reward 太小，比如為 0，則只會過濾掉一小部分結(jié)果的生成；如果閾值 Partial Reward 過大，比如為 3.5，則會將最優(yōu)響應(yīng)（紅點）也過濾掉。也就是說，以虛線垂直線為界限，可以最大限度的刪除無效響應(yīng)，并保證最優(yōu)響應(yīng)不會被過濾掉。

實際上，由于 c? 是未知的，因此無法實施上述過濾。此外，不同 Prompt 在獎勵分布方面也差異巨大。

4.2 算法

如下圖 Figure 1 所示，作者繪制了 Best-of-N 解碼策略生成過程中的內(nèi)存使用情況，觀察到在自回歸生成的早期階段，GPU 的內(nèi)存大部分都未被充分利用。此外，小批量自回歸生成也會導(dǎo)致出現(xiàn) Memory Bound 問題，導(dǎo)致算力浪費。這一觀察結(jié)果提供了一個機會來設(shè)計一個算法，該算法可以更有效地利用可用的 GPU 內(nèi)存和計算資源來生成一致候選響應(yīng)，以便用于獎勵模型進行排名。

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

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

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

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

實際上，上述算法在初始階段模擬了 Best-of-N 采樣，并在生成過程中動態(tài)將 Batch Size 減小，以防止 OOM。如上圖 Figure 1 所示，Speculative Rejection 比 Best-of-N 能高有效的利用 GPU 內(nèi)存。

五、實驗&結(jié)果

5.1 效率評估

如下圖所示，作者對比了相對 GPU 計算量下與 Best-of-N 相比的改進分?jǐn)?shù)（Improvement Score），此外，也提供了不同終止比例 α 下的結(jié)果?？梢钥闯?，使用本文的 Speculative Rejection 可以大幅減少 GPU 資源。具體來說，Best-of-N 要想獲得與 Speculative Rejection 相當(dāng)?shù)姆謹(jǐn)?shù)，需要消耗 16-32 倍的 GPU 資源。

5.2 勝率評估

為了進一步驗證生成質(zhì)量，作者基于 GPT-4-Turbo 評估了獲勝率（WR）和長度控制勝率（LC-WR）。對于每次測量，獲勝率基線是 Bo120。如表 Table 1 所示，Speculative Rejection 在大多數(shù)組合中實現(xiàn)了顯著的加速，同時保持了生成質(zhì)量。

5.3 最大化生成概率

如下圖 Table 2 所示，作者在 AlpacaFarm-Eval 數(shù)據(jù)集上測試了 Best-of-N 和 Speculative Rejection。各種配置，一系列模型的 PPL 結(jié)果表明，Speculative Rejection 更快，并且能產(chǎn)生一致較低的 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??

本文轉(zhuǎn)載自 ??AI閑談??，作者： AI閑談