在根據某個獎勵微調生成式語言模型時，使用 KL 正則化的強化學習（KL-RL）來對齊生成式語言模型是一種常用框架。而 KL-RL 通常需要訓練一個獎勵模型，然后使用一個強化學習求解器。其它方法還包括直接偏好優(yōu)化、獎勵模型蒸餾、best-of-N  蒸餾的不同變體版本。

在度量 KL-RL 框架的效果時，常用的指標是已對齊模型相較于參照模型在給定任務上的勝率。

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

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

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

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

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

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

使用獎勵變換實現強化學習

為了解決以下定義 5 中的語言模型對齊問題，該團隊提出了一個通用框架。

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

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

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

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

其中 是該推理時間轉換策略下已校準的獎勵。

對該定理的證明請訪問原論文。而基于該定理，可以自然地得到一個迭代 EM 式算法，其可以根據 (7) 式使用固定的 R 更新 π，并可以根據 (9) 式使用固定的 π 來更新 R，直到收斂。

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

為了更高效地設計獎勵變換，該團隊研究了不執(zhí)行推理時間流程的情況。在這種情況下，在這種情況下，T (π) = π 且 

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

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

解決 InfAlign

使用已校準獎勵實現 KL-RL

對已校準獎勵 C_{r,π_ref} 的性質的討論請訪問原論文。接下來看如何在 KL 正則化的強化學習中使用這個已校準獎勵。

根據其性質，校準之后，基礎策略的輸出的獎勵分布獨立于獎勵模型與基礎策略本身。這樣一來，便可以設計一個僅關注推理時間流程 T 的變換函數 Φ，并將其用于已校準獎勵函數。

更確切地說，令 Φ : [0, 1] → ? 為一個變換函數，該團隊提出了以下獎勵函數：

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

推理感知型獎勵變換。對于給定的推理時間流程 T ，目標是推導或設計一個合適的變換 Φ，使得該解能在推理時間勝率 W^T 和與基礎策略的 KL 散度之間實現良好甚至最佳的權衡。

標準勝率（無推理時間流程）。當不使用推理時間流程時（即 T 是恒等映射），W^T 會約簡為標準勝率。將 Φ 設置為恒等變換能得到最佳的勝率與 KL 權衡曲線，注意 

該團隊考慮了一系列僅依賴于輸出的已校準獎勵的推理時間流程，這被稱為已校準流程（calibrated procedures）。然后他們探討了如何為這一系列變換設計合適的 Φ。下面先定義已校準流程。

接下來的結果表明，對于已校準推理時間流程，求解 (13) 式的已對齊策略的勝率和 KL 散度獨立于基礎策略和獎勵函數。

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

為 BoN 和 WoN 尋找更好的變換

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

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

定理 3 描述了 BoN 和 WoN 的性質。

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

該團隊還研究了不同的變換類型，并分析性地計算了使用定理 3 時通過調整 β 而得到的對齊曲線，即不同 β 下  的圖。具體涉及的變換包括標準勝率的最優(yōu)變換、指數函數和基于優(yōu)化的變換。對這些變換的詳細描述請參閱原論文。結果則見圖 1。

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

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

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

CTRL：校準和變換式強化學習

該團隊還提出了 Calibrate-and-Transform Reinforcement Learning（CTRL），即校準和變換式強化學習。這是一種用于推理時間勝率優(yōu)化問題的求解器。

回想一下，新提出的解決方案可以分成三個階段：獎勵校準、獎勵變換、標準的 KL-RL 求解器。前面已經重點介紹了獎勵變換，下面將關注近似經驗校準。再將其與獎勵變換相結合，可得到最終的 CTRL 算法，見算法 1。

經驗校準是這樣的，首先對于強化學習訓練數據中的每個提示詞 x，從參照模型 π_ref 采樣 K 個樣本 z_1, z_2, ..., z_K。然后將所有響應的獎勵 {r (x, z_1), r (x, z_2), ...r (x, z_K)} 排序，然后將 RLHF 訓練期間提示詞與響應對 (x, y) 的經驗式已校準獎勵分數配置為：

理想情況下，當 K → ∞ 時，經驗的已校準獎勵將收斂到真正的已校準獎勵，并且可以通過 PPO 用于強化學習訓練目標。但是，其成本可能很高，因為精確計算這個已校準獎勵需要在 KL-RL 求解器中對每個提示詞和每個 roll-out 進行采樣并存儲 K 個獎勵分數。

于是，該團隊提出了一種更為實用的方法，即通過在對數域中使用逐步函數來擴展校準曲線，進而近似求取該曲線。具體做法是，選取 p 個錨點 q_1, q_2, ...，其中在每個分位數 q_i ∈ (0, 1) 處都取得零校準誤差。算法 2 給出了簡單情況（p = 1，中位數）的算法。更復雜的情況請訪問原論文算法 3。

實驗結果

這里略過具體的實驗配置和過程，僅關注實驗結果。

獎勵模型通常沒有正確校準

該團隊首先測量了在 Anthropic helpfulness preference 數據集上訓練的獎勵模型的校準錯誤情況，具體做法是計算訓練分割的數據中 10 個隨機提示詞的 100 個參照 - 策略響應的分數。

然后，該團隊對這些分數進行排序，計算每個響應對應的排名，并將這些值繪制為散點圖，如圖 2（左）所示。如果模型經過完美校準，則每個提示詞的點都將位于 y = x 線上。但是，請注意，對于大多數提示詞，散點圖與 y = x 線有顯著偏差，并且這種偏差的程度因提示詞而異。

然后，該團隊測量了獎勵分數與其對應排名之間的絕對誤差（AE），并在圖 2（右）中繪制各種校準近似的 AE 的累積分布函數（CDF）。如果模型經過良好校準，則 AE 始終為零，因此 CDF 在零 AE 時達到一。

該團隊發(fā)現獎勵分數（identity）沒有校準（平均 AE：0.22），并且使用固定獎勵多項式變換函數（如平方根、立方、平方）不會降低校準誤差（平均 AE > 0.15）。但是，使用基于分位數的獎勵校準（quantile）可顯著降低校準誤差（平均 AE：0.02）。

已校準獎勵可提升標準勝率

從圖 3 的結果可以看到，相比于 IPO 和 BoNBoN，已校準獎勵優(yōu)化可以實現更好的勝率 - KL 權衡。

該團隊將此增益歸因于使用來自基礎模型的 m 個樣本能更高效地計算訓練數據的勝率，而不是依賴 KL-RL 期間現有的成對比較數據。

CTRL 可提升 BoN

對于 Anthropic 對話數據集中的有用性目標，該團隊的目標是通過已校準獎勵的指數變換來優(yōu)化已對齊模型的 Best-of-N 性能。

該團隊測量了相對于基礎策略模型的 Best-of-N（N =4）的勝率，見圖 4。

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

CTRL 可提升 WoN（BoN 越獄）

對于 Anthropic 對話數據集中的無害性目標，該團隊的目標是提高已對齊策略模型的最差 Worst-of-N 性能，以提高面對對抗攻擊的安全性。

在這里，該團隊使用負指數變換 t < 0。從圖 4 可以看到，與未校準的模型相比，基于每個提示詞的中位數獎勵的校準實現了高 4-9% 的 Worst-of-N 勝率。已校準獎勵的負變換優(yōu)于其它，其中 t = ?10 表現最佳。