技術概要：ラグランジュ報酬増幅（LARA）による安全な推論時アライメント

問題提起

人間の好みに大規模言語モデル（LLM）を適合させるための標準的な枠組みは、強化学習からの人間によるフィードバック（RLHF）です。しかし、従来のアライメントはポリシーの微調整（例：RLHFやDirect Preference Optimization (DPO)）に依存しており、これは計算コストが高く、安全性要件やユーザー層が変化するたびに再学習を行う必要があります。

推論時アライメント（Best-of-N リランキングやトークンレベルの報酬ガイド付きデコーディングなど）は、重みの更新なしに補助的な報酬信号を用いて凍結されたモデルを制御する手法として登場していますが、これらは現在、「安全なアライメント」のための原理的なメカニズムを欠いています。既存の手法は通常、単一のスカラー値を最適化します。有用性（helpfulness）と安全性（harmlessness）の目的が衝突する場合、実務家はこれら二つのトレードオフを手動で調整しなければなりません。さらに、これらの手法はデコーディング中に明示的な安全制約（例：期待コストの予算）を本質的に強制しないため、不安全な出力、あるいは有用性の不必要な低下を招く可能性があります。

手法：ラグランジュ報酬増幅（LARA）

著者らは、Safe RLHFの制約付き最適化構造を推論段階へと転送する、デコーダーに依存しないフレームワークであるLARAを提案しています。LARAはモデルの重みを更新する代わりに、スカラー双対変数を較正することで、増幅された報酬信号を作成します。

1. 理論的定式化

本手法は、参照ポリシー（$\pi_{ref}$）に対するKLダイバージェンス項によって正則化された、期待される有用性（$r$）を最大化しつつ、期待されるコスト（$c \le \tau$）を制約とするSafe RLHFの目的関数から始まります：
$$\max_{\pi} \mathbb{E}[r(x,y)] - \beta D_{KL}(\pi \parallel \pi_{ref}) \quad \text{subject to} \quad \mathbb{E}[c(x,y)] \le \tau$$

ラグランジュ双対性を用いて、著者らは安全制約を双対化します。これにより、無限次元のポリシー最適化問題が、非負の双対変数 $\lambda$（ラグランジュ乗数）に関する一次元の凸最適化問題へと変換されます。

• 増幅された報酬: 双対変数は、増幅された報酬 $r_\lambda(x, y) = r(x, y) - \lambda c(x, y)$ を定義します。
• 双対目的関数: 双対関数 $g(\lambda)$ は凸関数であることが示されています。その微分は、報酬 $r_\lambda$ によって誘導されるギブス・ティルテッド・ポリシー（Gibbs-tilted policy）の下での、安全予算 $\tau$ と期待コストの差に対応します。
• 最適性: 最適な $\lambda^*$ は相補性条件を満たし、制約がアクティブな場合には期待コストが予算 $\tau$ に等しくなることを保証しつつ、有用性を最大化します。

2. 実証的較正（Empirical Calibration）

LARAは、最適解を解析的に解く必要はありません。代わりに、小さなデータセットを用いて $\lambda^*$ を実証的に較正します。

1. データ: 少数のプロンプトをサンプリングし、凍結された参照ポリシーから複数のレスポンスを生成します。
2. 推定: コンパクトな区間 $[0, \Lambda]$ 上で実証的双対目的関数を最小化することにより、双対変数を推定します。
3. アルゴリズム: 双対勾配の凸性と単調性により、二分探索を用いて最適な $\lambda$ が効率的に発見されます。
4. 堅牢性: 較正とデプロイメントの間の分布の不一致に対処するため、著者らは較正セットをブートストラップし、推定された $\lambda$ 分布の上位分位数を保守的な選択肢として採用することを提案しています。

3. デコーディングへの統合

較正後、スカラー $\hat{\lambda}$ を用いて増幅された報酬 $r_{\hat{\lambda}}$ を計算します。このスコアは、既存の推論時アライメント手法への「ドロップイン」信号として渡されます。

• シーケンスレベル（例：Best-of-N）: 較正された $\lambda$ は、期待コスト制約付き問題の厳密な解を提供します。
• トークンレベル（例：報酬ガイド付き探索）: 較正された $\lambda$ は、制約付き目的関数から導出された原理的なヒューリスティックを提供しますが、トークンレベルの貪欲なデコーディングの性質上、集団レベルの制約充足を厳密に保証するものではありません。

主な貢献

1. 定式化: 本論文は、安全な推論時アライメントをSafe RLHFの双対として定式化し、問題を一次元の凸較正タスクへと簡略化しました。
2. 幾何学と効率性: 著者らは双対目的関数の幾何学的性質を特徴付け、その凸性と勾配の単調性を証明しました。これにより、効率的な二分探索による較正手順が可能になります。
3. 理論的保証: 推定値 $\hat{\lambda}_{N,K}$ に関する有限サンプル保証が提供されており、エラーがプロンプト数（$N$）とプロンプトあたりのサンプル数（$K$）によって制御されながら、母集団の最小化因子 $\lambda^*$ に集中することを示しています。
4. デコーダーに依存しないフレームワーク: LARAは、手動のスカラチューニングを原理的な安全制約付き構築に置き換える、既存の推論時アライメント手法と互換性のある較正レイヤーとして設計されています。

実験結果

著者らは、Llama-3.2-3B および Qwen-2.5-3B モデルを BeaverTails データセットを用いて評価しました。

• 制約充足: 実験により、較正された $\lambda^*$ が安全制約の境界に位置することが確認されました。$\lambda < \lambda^*$ の場合は安全予算を違反し、$\lambda > \lambda^*$ の場合は制約を満たすものの、有用性を不必要に低下させました。
• ベースラインとの比較:
  • 推論時手法: 推論時手法の中で、LARAを組み合わせた Best-of-N (BoN) が、最も優れた有用性と無害性のトレードオフを達成しました。これは、重みの更新なしに、有用性において SafeDPO を上回り、無害性においても競争力のある性能を示しました。
  • トークンレベル: トークンレベルの手法（ARGS, PARG, CD-Q）は、一般に BoN よりも有用性において劣る結果となりましたが、CD-Q は $\lambda$ の増加に伴い強い無害性の向上を示しました。
  • 微調整との比較: LARAを用いた BoN は、微調整ベースのベースライン（Safe-RLHF, SafeDPO）の性能に迫りましたが、微調整された手法の方が全体的なトレードオフにおいては依然として優れた性能を示しました。しかし、LARAはデプロイメントの柔軟性とコストにおいて大きな利点を提供します。
• 評価: 結果は、モデルベースの評価（BeaverTails 報酬/コストモデルを使用）と GPT-4o-mini ベースの評価の両方で一貫していましたが、GPT は報酬モデルよりも $\lambda$ の変動に対する感度が低い傾向にありました。

意義と主張

本論文は、LARAが安全な推論時アライメントのための手動のペナルティチューニングに代わる原理的な選択肢を提供すると主張しています。安全制約を双対化することで、LARAはモデルの重みを変更することなく、安全予算をデコーディングの目的関数へと変換します。

• シーケンスレベルのサンプラーに対して: LARAは厳密な期待コストの解釈を提供し、制約付き目的関数における最適なトレードオフを理論的に保証します。
• トークンレベルのデコーダーに対して: 集団レベルの制約を厳密に強制することはできないものの、制約付き目的関数の幾何学を尊重した、原理的なヒューリスティックを提供します。

著者らは、LARAが安全制約を訓練時のアライメントから実用的なデコーディング時の制御へと効果的に転送することに成功し、有用性と無害性のトレードオフを改善し、推論時手法を微調整ベースのアプローチの性能に近づけることができると結論付けています。本研究は、安全制約を推論時に効果的に管理できることを強調しており、繰り返しの再学習に伴う計算オーバーヘッドを回避できることを示しています。