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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

解決 InfAlign

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

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

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

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

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

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

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

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

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

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

為 BoN 和 WoN 尋找更好的變換

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

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

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

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

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

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

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

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

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

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

回想一下，新提出的解決方案可以分成三個階段：獎勵校準、獎勵變換、標準的 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 時達到一。

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

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

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

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

CTRL 可提升 BoN

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

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

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

CTRL 可提升 WoN（BoN 越獄）

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

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