技術要約: REAR: 報酬分解によるテスト時における好みの再整合（REAR: Test-time Preference Realignment through Reward Decomposition）

問題提起

大規模言語モデル（LLM）を多様でパーソナライズされたユーザーの好みに適合させることは、極めて重要な課題である。RLHF（人間からのフィードバックによる強化学習）のようなポストトレーニング手法は効果的であるが、コストの高いデータキュレーションと追加の学習を必要とするため、動的または特定のユーザーニーズに対して柔軟性に欠ける。テスト時スケーリング（Test-Time Scaling: TTS）は、推論時の追加計算を活用することで、学習を必要としない代替案を提供する。しかし、既存のTTS手法は、正解性が容易に判定できる検証可能なドメイン（数学やコーディングなど）に限定されていることが多い。回答の質は全体的なものであり、単純な検証可能な答えに還元できないため、好みの適合にTTSを適用することは困難である。現在のテスト時アライメント手法は、固定された計算オーバーヘッドを課したり、バリューヘッドの訓練を必要としたりすることが多く、そのスケーラビリティと効率性に制限がある。

手法

本論文では、テスト時の好みの再整合を**再整合問題（realignment problem）として定式化する新しいフレームワークであるREAR（REAlignment Reward）**を紹介する。核心となる洞察は、ベースモデルの学習目的が、推論時において特定のユーザーの好みを最適に満たさない可能性があるということである。REARは、基礎となる報酬関数を2つのコンポーネントに分解し、それらの重要性を動的に再スケーリングすることでこの問題に対処する。

1. 報酬の分解

著者らは、最大エントロピー強化学習（MaxEnt RL）の枠組みの中で、事前学習済みLLMの報酬構造を定式化している。状態 $s$（質問）と好み $x_p$ に対して、報酬 $r(s \oplus x_p, a)$ は以下のように分解できると仮定している：

• 質問に関連する報酬 ($r_0$): 質問に正しく答えるためにベースモデルが最大化する暗黙的な報酬。
• 好みに関連する報酬 ($r_p$): 回答が好み $x_p$ をどの程度満たしているかを捉える報酬コンポーネント。

合成報酬は、$r(s \oplus x_p, a) = r_0(s, a) + \alpha r_p(s \oplus x_p, a)$ とモデル化される。ここで $\alpha$ は固有のバランス係数である。

2. 再整合報酬 (REAR)

特定のユーザーに合わせてモデルを再整合するために、REARは推論中に好みのコンポーネントを再スケーリングするための制御可能な係数 $\hat{\alpha}$（または $\lambda = \hat{\alpha}/\alpha$）を導入する：
$$r_{REAR}(s \oplus x_p, a) = r_0(s, a) + \hat{\alpha} r_p(s \oplus x_p, a)$$

極めて重要な点は、著者らがこの報酬の計算効率の高い定式化を導出したことである。MaxEnt RLにおける最適方策とバリュー関数の関係を利用することで、$r_{REAR}$ はトークンレベルの方策対数確率の線形結合として表現できることを示している：
$$S_{REAR}(\tau) = \sum_{t=0}^{T} \gamma^t \left( (1-\lambda) \log \pi(a_t|s_t) + \lambda \log \pi(a_t|s_t \oplus x_p) \right)$$
このスコア関数は、質問のみのプロンプトと、好みを含むフルプロンプトという2つのコンテキスト下でのベースモデルの対数確率のみに依存する。これには、別の報酬モデルの訓練や、スコアリングのための追加のトークン生成は必要ない。

3. テスト時スケーリングとの統合

REARスコアは、以下の2つのTTSアルゴリズムに統合される：

• Best-of-N (BoN) サンプリング: $N$ 個の回答を生成し、最も高い $S_{REAR}$ スコアを持つものを選択する。
• Diverse Verifier Tree Search (DVTS): 木探索アルゴリズムであり、REARスコアを使用して分岐の拡張と選択をガイドし、好みの適合を最適化しながら多様性を促進する。

主な貢献

1. 再整合の定式化: 本論文は、テスト時の好みの適合を再整合問題として定式化し、分解された目的関数から導出される計算効率の高い報酬であるREARを提案した。
2. 学習フリーのスコアリング: REARはベースモデルの内部確率のみから導出されるため、追加の学習、報酬モデルの訓練、または投機的生成ステップを必要としない。
3. スケーラブルなアルゴリズム: 著者らは、REARによってガイドされ、様々なサンプル予算に応じて生成性能を高めることができる2つのスケーラブルなテスト時アルゴリズム（BoNおよびDVTS）を開発した。
4. 汎用性: 本手法は、タスク固有の好みのテキストを注入することにより、数学的推論や視覚的タスクへの好みの適合を超えて一般化できることが示されている。

実験結果

著者らはQwenおよびLlamaモデルを用いて、いくつかのベンチマークでREARを評価した：

• 好みの適合ベンチマーク: PrefEval（明示的/暗示的な好みおよび選択肢をカバー）、Multifaceted Bench、および Ping-Pong（ロールプレイング）において、REAR誘導型手法（BoNおよびDVTSの両方）は、greedy decoding、Amulet、Linear Alignment (LA)、およびGenerative Reward Models (GenRM) を用いたBest-of-Nを含むベースラインを一貫して上回った。例えば、REARを用いたDVTSは、指標全体でgreedy decodingに対して平均11.6%向上した。
• 長文脈への堅牢性: REARは、コンテキスト長が増加すると著しく低下するGenRMと比較して、長文脈シナリオ（最大16kトークン）において優れた堅牢性を示した。
• 数学的推論: 適切な好みのテキスト（例：「厳格な数学者として振る舞え」）を伴う数学的タスク（MATH500, AIME, AMC）に適用した場合、REAR誘導型の選択は多数決（majority voting）を上回り、一貫したスケーリングゲインを示した。
• 視覚的タスク: 視覚的ハルシネーション検出のための MMHal-Bench において、REARはgreedy decodingやGenRMと比較して、総合スコアを向上させ、ハルシネーション率を減少させた。
• クロスモデル転移: 本手法は、モデルごとのチューニングなしにLlama-3.1-8B-Instructに一般化し、報酬の分解がQwenファミリーに特有のものではないことを確認した。
• 効率性: REAR誘導型手法は、外部の報酬モデルや反復的な最適化（Amuletなど）を必要とするベースラインよりも大幅に高速であった。これは、別個のモデルのロードを回避し、スコアリングのための単一のフォワードパスに依存するためである。

重要性と主張

本論文は、REARが、学習や外部の報酬モデルのオーバーヘッドなしに、テスト時にLLMをユーザーの好みに適合させるためのスケーラブルで一般的な解決策を提供すると主張している。報酬を質問と好みのコンポーネクションに分解することで、REARは、質問への回答と好みへの適合の間のトレードオフを動的に制御することを可能にする。

著者らは、REARを用いることで、従来の検証器が失敗する主観的でオープンエンドなドメイン（好みの適合やロールプレイングなど）へとテスト時スケーリングを拡張できることを強調している。また、デフォルトの $\lambda = 20$ が多くのタスクで堅牢に機能する一方で、本手法は再整合の強さを調整できることも述べている。本研究は、テスト時アライメントをより身近で効率的なものにし、検証可能なドメインにおけるTTSの成功と、主観的な好みの適合の複雑さとの間の溝を埋める一歩となるものである。

認められた限界: 著者らは、REARが調整可能なハイパーパラメータ（$\lambda$）をユーザーに露出させており、好みの分布が著しく異なるタスクでは、小さな検証スイープが必要になる可能性があると指摘している。さらに、TTSはファインチューニングよりは軽量であるが、最適な計算予算を特定することは今後の課題であるとしている。