在根據某個獎勵微調生成式語言模型時，使用 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 表現最佳。