在根據(jù)某個(gè)獎(jiǎng)勵(lì)微調(diào)生成式語言模型時(shí)，使用 KL 正則化的強(qiáng)化學(xué)習(xí)（KL-RL）來對(duì)齊生成式語言模型是一種常用框架。而 KL-RL 通常需要訓(xùn)練一個(gè)獎(jiǎng)勵(lì)模型，然后使用一個(gè)強(qiáng)化學(xué)習(xí)求解器。其它方法還包括直接偏好優(yōu)化、獎(jiǎng)勵(lì)模型蒸餾、best-of-N  蒸餾的不同變體版本。

在度量 KL-RL 框架的效果時(shí)，常用的指標(biāo)是已對(duì)齊模型相較于參照模型在給定任務(wù)上的勝率。

但是，在推理時(shí)間很少會(huì)使用已對(duì)齊的模型，而是會(huì)通過一個(gè)推理時(shí)間流程來完成任務(wù)，比如 best-of-N 采樣、best-of-N 越獄、思維鏈推理、自我一致性。這樣一來，推理時(shí)間解碼過程與訓(xùn)練 KL-RL 目標(biāo)之間就不匹配了。

于是，問題來了：給定一個(gè)已知的推理時(shí)間流程，我們可以對(duì)齊模型，從而優(yōu)化相對(duì)于參照模型的推理時(shí)間勝率嗎？其中，為了求取推理時(shí)間勝率，需要通過該推理時(shí)間流程獲取每個(gè)模型的響應(yīng)并統(tǒng)計(jì)每個(gè)模型的樣本獲勝次數(shù)。

Google DeepMind 和 Google Research 近日的一篇論文嘗試解答了這個(gè)問題。他們發(fā)現(xiàn)，盡管很難直接優(yōu)化推理時(shí)間勝率，但可通過一組優(yōu)化目標(biāo)來獲取其最優(yōu)解。該團(tuán)隊(duì)將這個(gè)框架稱為 inference-aware alignment (InfAlign)，即推理感知型對(duì)齊。

• 論文標(biāo)題：InfAlign: Inference-aware language model alignment
• 論文地址：https://arxiv.org/abs/2412.19792

他們還進(jìn)一步證明，對(duì)于被 δ 限定的語言模型（其中隨著 δ → 0，所有輸出都被 δ 限定了上限），可通過對(duì)獎(jiǎng)勵(lì)使用一個(gè)特定的變換來求解 KL-RL，從而得到這個(gè)最優(yōu)解。

如此一來，針對(duì)推理時(shí)間勝率進(jìn)行優(yōu)化的難題就可以這樣解決了：設(shè)計(jì)一個(gè)適合特定推理時(shí)間流程的獎(jiǎng)勵(lì)變換，然后使用 PPO 等已有的優(yōu)化算法來求解 KL-RL。

使用獎(jiǎng)勵(lì)變換實(shí)現(xiàn)強(qiáng)化學(xué)習(xí)

為了解決以下定義 5 中的語言模型對(duì)齊問題，該團(tuán)隊(duì)提出了一個(gè)通用框架。

定義 5：令 T 為一個(gè)給定的推理時(shí)間流程，且 β > 0。那么，最大化推理時(shí)間勝率的優(yōu)化問題就可以表述成

他們提出的新方法的基礎(chǔ)是：基于獎(jiǎng)勵(lì)模型 r、推理時(shí)間流程 T、基礎(chǔ)策略 π_ref 設(shè)計(jì)一個(gè)新的獎(jiǎng)勵(lì)函數(shù) R；這樣一來，解決帶有已變換獎(jiǎng)勵(lì) R 的 KL 正則化強(qiáng)化學(xué)習(xí)問題就基本上能得到一個(gè)最優(yōu)解。更確切地說，已對(duì)齊策略就是以下優(yōu)化問題的解：

其中 R 是變換后的獎(jiǎng)勵(lì)函數(shù)。大致看的話，可能不太容易理解為什么 (6) 式可幫助求解 (5) 中的問題。但該團(tuán)隊(duì)會(huì)證明，對(duì)于任意給定的推理時(shí)間流程 T ，都存在一個(gè)可解決 (5) 的變換后獎(jiǎng)勵(lì) R。

定理 1（InfAlign 解決方案的特征）假設(shè) T 使得對(duì)于所有 x、y_1、y_2 都存在 ?T (π)(y1 | x)/?π(y_2 | x)，則就得到了最優(yōu)轉(zhuǎn)換獎(jiǎng)勵(lì) R，并且 (5) 式中的最優(yōu)策略 π? 必須滿足以下耦合方程：?x, y

其中 是該推理時(shí)間轉(zhuǎn)換策略下已校準(zhǔn)的獎(jiǎng)勵(lì)。

對(duì)該定理的證明請(qǐng)?jiān)L問原論文。而基于該定理，可以自然地得到一個(gè)迭代 EM 式算法，其可以根據(jù) (7) 式使用固定的 R 更新 π，并可以根據(jù) (9) 式使用固定的 π 來更新 R，直到收斂。

然而，這種算法有兩個(gè)缺點(diǎn)：首先，對(duì)于一般的語言模型，難以評(píng)估方程 (9) 或者評(píng)估效率很低，因?yàn)檫@需要在巨大甚至無限的輸出空間上評(píng)估策略；其次，尚不清楚這種算法是否能得到最優(yōu)解。

為了更高效地設(shè)計(jì)獎(jiǎng)勵(lì)變換，該團(tuán)隊(duì)研究了不執(zhí)行推理時(shí)間流程的情況。在這種情況下，在這種情況下，T (π) = π 且 

(9) 式便可簡(jiǎn)化為 R (x, y) = C_{r,π_ref} (x, y)，即 π_ref 下的 CDF 或已校準(zhǔn)獎(jiǎng)勵(lì)。

因此，定理 1 可以看作是這些結(jié)果與一般推理時(shí)間流程的泛化。這一觀察促使該團(tuán)隊(duì)考慮基于此已校準(zhǔn)獎(jiǎng)勵(lì)的一系列獎(jiǎng)勵(lì)變換，如下一節(jié)所述。我們將看到，對(duì)于這類已校準(zhǔn)推理時(shí)間流程（定義 6），可以通過一個(gè)實(shí)驗(yàn)性語言模型有效地評(píng)估此系列中的不同變換，從而找到優(yōu)良甚至最佳的變換。

解決 InfAlign

使用已校準(zhǔn)獎(jiǎng)勵(lì)實(shí)現(xiàn) KL-RL

對(duì)已校準(zhǔn)獎(jiǎng)勵(lì) C_{r,π_ref} 的性質(zhì)的討論請(qǐng)?jiān)L問原論文。接下來看如何在 KL 正則化的強(qiáng)化學(xué)習(xí)中使用這個(gè)已校準(zhǔn)獎(jiǎng)勵(lì)。

根據(jù)其性質(zhì)，校準(zhǔn)之后，基礎(chǔ)策略的輸出的獎(jiǎng)勵(lì)分布獨(dú)立于獎(jiǎng)勵(lì)模型與基礎(chǔ)策略本身。這樣一來，便可以設(shè)計(jì)一個(gè)僅關(guān)注推理時(shí)間流程 T 的變換函數(shù) Φ，并將其用于已校準(zhǔn)獎(jiǎng)勵(lì)函數(shù)。

更確切地說，令 Φ : [0, 1] → ? 為一個(gè)變換函數(shù)，該團(tuán)隊(duì)提出了以下獎(jiǎng)勵(lì)函數(shù)：

而我們希望已對(duì)齊策略是 KL-RL 問題的解。

推理感知型獎(jiǎng)勵(lì)變換。對(duì)于給定的推理時(shí)間流程 T ，目標(biāo)是推導(dǎo)或設(shè)計(jì)一個(gè)合適的變換 Φ，使得該解能在推理時(shí)間勝率 W^T 和與基礎(chǔ)策略的 KL 散度之間實(shí)現(xiàn)良好甚至最佳的權(quán)衡。

標(biāo)準(zhǔn)勝率（無推理時(shí)間流程）。當(dāng)不使用推理時(shí)間流程時(shí)（即 T 是恒等映射），W^T 會(huì)約簡(jiǎn)為標(biāo)準(zhǔn)勝率。將 Φ 設(shè)置為恒等變換能得到最佳的勝率與 KL 權(quán)衡曲線，注意 

該團(tuán)隊(duì)考慮了一系列僅依賴于輸出的已校準(zhǔn)獎(jiǎng)勵(lì)的推理時(shí)間流程，這被稱為已校準(zhǔn)流程（calibrated procedures）。然后他們探討了如何為這一系列變換設(shè)計(jì)合適的 Φ。下面先定義已校準(zhǔn)流程。

接下來的結(jié)果表明，對(duì)于已校準(zhǔn)推理時(shí)間流程，求解 (13) 式的已對(duì)齊策略的勝率和 KL 散度獨(dú)立于基礎(chǔ)策略和獎(jiǎng)勵(lì)函數(shù)。

基于上述定理，便可以通過關(guān)注易于計(jì)算和模擬的簡(jiǎn)單連續(xù)語言模型來對(duì)變換 Φ 進(jìn)行評(píng)估。下面，該團(tuán)隊(duì)使用了兩個(gè)常用的推理時(shí)間流程 best-of-N 和 worst-of-N 為示例，展示了該定理可以如何有效地評(píng)估不同 Φ 函數(shù)的推理時(shí)間勝率與 KL 散度權(quán)衡曲線，這可用于在實(shí)際場(chǎng)景中找到合適的變換 Φ。

為 BoN 和 WoN 尋找更好的變換

這一節(jié)將主要關(guān)注以下兩種推理時(shí)間流程：

• best-of-N 推理時(shí)間流程 (BoN)
• worst-of-N 推理時(shí)間流程 (WoN)

定理 3 描述了 BoN 和 WoN 的性質(zhì)。

通過調(diào)整式 (13) 中的 β，可以得到一條對(duì)齊曲線，該曲線繪制了不同對(duì)齊策略的推理時(shí)間勝率和 KL 散度偏差。這樣，便能比較不同變換函數(shù) Φ 的性能。

該團(tuán)隊(duì)還研究了不同的變換類型，并分析性地計(jì)算了使用定理 3 時(shí)通過調(diào)整 β 而得到的對(duì)齊曲線，即不同 β 下  的圖。具體涉及的變換包括標(biāo)準(zhǔn)勝率的最優(yōu)變換、指數(shù)函數(shù)和基于優(yōu)化的變換。對(duì)這些變換的詳細(xì)描述請(qǐng)參閱原論文。結(jié)果則見圖 1。

以上結(jié)果證明了在執(zhí)行對(duì)齊時(shí)考慮推理時(shí)間過程的重要性。

該團(tuán)隊(duì)發(fā)現(xiàn)具有不同 t 的指數(shù)變換適用于不同的推理時(shí)間流程，這將是該團(tuán)隊(duì)在實(shí)驗(yàn)中的重點(diǎn)。接下來，該團(tuán)隊(duì)將研究在理想化的連續(xù)語言模型上找到的好變換是否可以泛化用于現(xiàn)實(shí)世界場(chǎng)景。

但在繼續(xù)進(jìn)行實(shí)驗(yàn)之前，還必需一種實(shí)用的算法來解決推理時(shí)間 KL-RL 優(yōu)化問題。

CTRL：校準(zhǔn)和變換式強(qiáng)化學(xué)習(xí)

該團(tuán)隊(duì)還提出了 Calibrate-and-Transform Reinforcement Learning（CTRL），即校準(zhǔn)和變換式強(qiáng)化學(xué)習(xí)。這是一種用于推理時(shí)間勝率優(yōu)化問題的求解器。

回想一下，新提出的解決方案可以分成三個(gè)階段：獎(jiǎng)勵(lì)校準(zhǔn)、獎(jiǎng)勵(lì)變換、標(biāo)準(zhǔn)的 KL-RL 求解器。前面已經(jīng)重點(diǎn)介紹了獎(jiǎng)勵(lì)變換，下面將關(guān)注近似經(jīng)驗(yàn)校準(zhǔn)。再將其與獎(jiǎng)勵(lì)變換相結(jié)合，可得到最終的 CTRL 算法，見算法 1。

經(jīng)驗(yàn)校準(zhǔn)是這樣的，首先對(duì)于強(qiáng)化學(xué)習(xí)訓(xùn)練數(shù)據(jù)中的每個(gè)提示詞 x，從參照模型 π_ref 采樣 K 個(gè)樣本 z_1, z_2, ..., z_K。然后將所有響應(yīng)的獎(jiǎng)勵(lì) {r (x, z_1), r (x, z_2), ...r (x, z_K)} 排序，然后將 RLHF 訓(xùn)練期間提示詞與響應(yīng)對(duì) (x, y) 的經(jīng)驗(yàn)式已校準(zhǔn)獎(jiǎng)勵(lì)分?jǐn)?shù)配置為：

理想情況下，當(dāng) K → ∞ 時(shí)，經(jīng)驗(yàn)的已校準(zhǔn)獎(jiǎng)勵(lì)將收斂到真正的已校準(zhǔn)獎(jiǎng)勵(lì)，并且可以通過 PPO 用于強(qiáng)化學(xué)習(xí)訓(xùn)練目標(biāo)。但是，其成本可能很高，因?yàn)榫_計(jì)算這個(gè)已校準(zhǔn)獎(jiǎng)勵(lì)需要在 KL-RL 求解器中對(duì)每個(gè)提示詞和每個(gè) roll-out 進(jìn)行采樣并存儲(chǔ) K 個(gè)獎(jiǎng)勵(lì)分?jǐn)?shù)。

于是，該團(tuán)隊(duì)提出了一種更為實(shí)用的方法，即通過在對(duì)數(shù)域中使用逐步函數(shù)來擴(kuò)展校準(zhǔn)曲線，進(jìn)而近似求取該曲線。具體做法是，選取 p 個(gè)錨點(diǎn) q_1, q_2, ...，其中在每個(gè)分位數(shù) q_i ∈ (0, 1) 處都取得零校準(zhǔn)誤差。算法 2 給出了簡(jiǎn)單情況（p = 1，中位數(shù)）的算法。更復(fù)雜的情況請(qǐng)?jiān)L問原論文算法 3。

實(shí)驗(yàn)結(jié)果

這里略過具體的實(shí)驗(yàn)配置和過程，僅關(guān)注實(shí)驗(yàn)結(jié)果。

獎(jiǎng)勵(lì)模型通常沒有正確校準(zhǔn)

該團(tuán)隊(duì)首先測(cè)量了在 Anthropic helpfulness preference 數(shù)據(jù)集上訓(xùn)練的獎(jiǎng)勵(lì)模型的校準(zhǔn)錯(cuò)誤情況，具體做法是計(jì)算訓(xùn)練分割的數(shù)據(jù)中 10 個(gè)隨機(jī)提示詞的 100 個(gè)參照 - 策略響應(yīng)的分?jǐn)?shù)。

然后，該團(tuán)隊(duì)對(duì)這些分?jǐn)?shù)進(jìn)行排序，計(jì)算每個(gè)響應(yīng)對(duì)應(yīng)的排名，并將這些值繪制為散點(diǎn)圖，如圖 2（左）所示。如果模型經(jīng)過完美校準(zhǔn)，則每個(gè)提示詞的點(diǎn)都將位于 y = x 線上。但是，請(qǐng)注意，對(duì)于大多數(shù)提示詞，散點(diǎn)圖與 y = x 線有顯著偏差，并且這種偏差的程度因提示詞而異。

然后，該團(tuán)隊(duì)測(cè)量了獎(jiǎng)勵(lì)分?jǐn)?shù)與其對(duì)應(yīng)排名之間的絕對(duì)誤差（AE），并在圖 2（右）中繪制各種校準(zhǔn)近似的 AE 的累積分布函數(shù)（CDF）。如果模型經(jīng)過良好校準(zhǔn)，則 AE 始終為零，因此 CDF 在零 AE 時(shí)達(dá)到一。

該團(tuán)隊(duì)發(fā)現(xiàn)獎(jiǎng)勵(lì)分?jǐn)?shù)（identity）沒有校準(zhǔn)（平均 AE：0.22），并且使用固定獎(jiǎng)勵(lì)多項(xiàng)式變換函數(shù)（如平方根、立方、平方）不會(huì)降低校準(zhǔn)誤差（平均 AE > 0.15）。但是，使用基于分位數(shù)的獎(jiǎng)勵(lì)校準(zhǔn)（quantile）可顯著降低校準(zhǔn)誤差（平均 AE：0.02）。

已校準(zhǔn)獎(jiǎng)勵(lì)可提升標(biāo)準(zhǔn)勝率

從圖 3 的結(jié)果可以看到，相比于 IPO 和 BoNBoN，已校準(zhǔn)獎(jiǎng)勵(lì)優(yōu)化可以實(shí)現(xiàn)更好的勝率 - KL 權(quán)衡。

該團(tuán)隊(duì)將此增益歸因于使用來自基礎(chǔ)模型的 m 個(gè)樣本能更高效地計(jì)算訓(xùn)練數(shù)據(jù)的勝率，而不是依賴 KL-RL 期間現(xiàn)有的成對(duì)比較數(shù)據(jù)。

CTRL 可提升 BoN

對(duì)于 Anthropic 對(duì)話數(shù)據(jù)集中的有用性目標(biāo)，該團(tuán)隊(duì)的目標(biāo)是通過已校準(zhǔn)獎(jiǎng)勵(lì)的指數(shù)變換來優(yōu)化已對(duì)齊模型的 Best-of-N 性能。

該團(tuán)隊(duì)測(cè)量了相對(duì)于基礎(chǔ)策略模型的 Best-of-N（N =4）的勝率，見圖 4。

可以看到，與未在有用性目標(biāo)上校準(zhǔn)的模型相比，基于每個(gè)提示詞的中位數(shù)獎(jiǎng)勵(lì)的校準(zhǔn)實(shí)現(xiàn)了高 8-12% 的 Best-of-N 勝率。已校準(zhǔn)獎(jiǎng)勵(lì)的指數(shù)變換優(yōu)于其他。該團(tuán)隊(duì)發(fā)現(xiàn)，指數(shù)因子 t = 10 的效果最佳。此外，該團(tuán)隊(duì)表明 N 值不同時(shí)，也存在這些增益。

CTRL 可提升 WoN（BoN 越獄）

對(duì)于 Anthropic 對(duì)話數(shù)據(jù)集中的無害性目標(biāo)，該團(tuán)隊(duì)的目標(biāo)是提高已對(duì)齊策略模型的最差 Worst-of-N 性能，以提高面對(duì)對(duì)抗攻擊的安全性。

在這里，該團(tuán)隊(duì)使用負(fù)指數(shù)變換 t < 0。從圖 4 可以看到，與未校準(zhǔn)的模型相比，基于每個(gè)提示詞的中位數(shù)獎(jiǎng)勵(lì)的校準(zhǔn)實(shí)現(xiàn)了高 4-9% 的 Worst-of-N 勝率。已校準(zhǔn)獎(jiǎng)勵(lì)的負(fù)變換優(yōu)于其它，其中 t = ?10 表現(xiàn)最佳。