論文要約：期待報酬予測とモデルルーティングへの応用

1. 背景と問題定義

大規模言語モデル（LLM）の分野において、**報酬モデル（Reward Models）**は、特定のプロンプトに対する LLM の応答を評価・スコアリングする標準的なツールとして広く利用されています。従来の報酬モデルは、固定されたプロンプトに対して複数のモデル（または同一モデルからの複数サンプル）が生成した応答を比較し、どちらが優れているかをランク付けすることを目的としています。

しかし、本研究が扱う核心的な問題は、**「応答レベルの報酬モデルを、プロンプトに対して特定のモデルがどの程度適しているかを事前に評価する『モデルレベルの報酬』へと昇華（lift）できるか」**という点です。

具体的には、応答 $y$ に対して報酬 $r(x, y)$ を与える関数を、プロンプト $x$ に対してモデル $\pi$ が生成する応答の期待報酬 $E_{Y \sim \pi(Y|x)}[r(x, Y)]$ を予測する関数へと変換できるかが問われています。

• 課題: 生成モデルの挙動がプロンプト依存で予測困難であること、および多くの報酬モデルが相対評価（Bradley-Terry モデルなど）に基づいており、絶対的なスコア値自体が定義されていない（任意の定数加算が可能）ため、期待値の予測が理論的に可能か不明瞭であること。
• 目的: プロンプトを入力とし、特定のモデルがそのプロンプトに対して生成する応答の「期待される報酬」を、実際の応答生成を行わずに（事前的に）高精度に予測すること。

2. 手法（Methodology）

2.1 期待報酬予測（Expected Reward Prediction: ERP）

本研究では、プロンプトの埋め込み表現（Embedding）から期待報酬を直接予測するシンプルなアプローチを提案します。

1. データ構築:

   • 各プロンプト $x_i$ に対して、対象モデル $\pi$ から $K$ 個（実験では $K=32$）の応答をサンプリングします。
   • 各応答に報酬モデル $r$ を適用し、得られたスコアの平均 $\hat{ER}_\pi(x_i)$ を「正解ラベル（ターゲット）」として定義します。
   • これにより、$(x_i, \hat{ER}_\pi(x_i))$ の対からなる学習データセットを構築します。

2. モデル構造:

   • 入力: プロンプト $x$ の固定長ベクトル表現（埋め込み）。ここでは、事前学習済みの汎用埋め込みモデル gte-large-en-v1.5 を使用します。
   • 予測モデル: 埋め込みベクトルを入力とし、期待報酬を出力するリッジ正則化付き線形モデル（Linear Probe）を訓練します。
   • 学習目標: 線形モデルのパラメータ $\theta$ を、以下の二乗誤差と正則化項を最小化するように最適化します。
     $$\arg\min_{\theta} \mathbb{E}_{X}[(\theta^\top v(X) - \hat{ER}_\pi(X))^2] + \beta\|\theta\|^2_2$$

3. モデルルーティングへの応用:

   • 複数のモデルプール $\Pi = \{\pi_i\}$ がある場合、各モデルに対して個別に期待報酬予測器を訓練します。
   • 入力プロンプト $x$ に対して、各モデルの予測期待報酬から計算コスト $c(i)$ を差し引いた値（コスト調整済み報酬）を最大化するモデルを選択します。
   • 最適化問題: $\max_\rho \mathbb{E}[r(x, Y) - \lambda c(\rho(x))]$

2.2 実験設定

• データセット: open-perfectblend（一般チャット、指示追従、数学推論、コード推論の 4 カテゴリから構成）。
• 対象モデル: Llama3.1-IT (8B, 70B), Gemma2-IT (9B, 27B), Gemma1-IT (7B)。
• 報酬モデル: OpenAssistant-RM, GRM-2B-RM, InternLM-RM の 3 種類。

3. 主要な結果（Results）

3.1 期待報酬の予測可能性

• 高い予測精度: 線形モデルを用いた期待報酬予測は、テストデータにおいて驚くほど高い精度（$R^2$ 値）を達成しました。
  • 全体的な $R^2$ は 0.46〜0.59 の範囲にあり、プロンプトカテゴリ内およびカテゴリ間でも予測が可能であることを示しています（Table 1, Figure 2）。
  • 数学や指示追従などの特定のタスクでは予測が難しいケースもありましたが、全体的な傾向として「プロンプトの埋め込みからモデルの期待性能を線形に近似できる」ことが実証されました。
• 単純性の有効性: 複雑なニューラルネットワークではなく、オフ・ザ・シェルフの埋め込みと線形回帰だけで十分な性能が得られることが示されました。

3.2 モデルルーティングの性能

• ベースラインとの比較: 提案手法（ERP）は、以下のベースラインを上回る性能を示しました。
  • 固定モデル: 常に特定のモデルを使用する手法。
  • カテゴリ別最適モデル: プロンプトカテゴリごとに平均性能が最も良いモデルを固定で選択する手法（Oracle）。
  • Zooter 法: 複数モデル間の勝率分布を直接予測する複雑な手法。
• パレート最適性: 報酬と計算コストのトレードオフ（Regret-Cost Tradeoff）において、ERP ベースのルーティングは、カテゴリ情報を持たない状態でも、カテゴリごとの最適モデル（Oracle）と同等かそれ以上の性能を発揮しました（Figure 1b, Figure 6）。
• ペアワイズ比較との同等性: 2 モデル間の勝敗を予測するタスクにおいて、ERP ベースの予測器は、ペアワイズ比較データで訓練されたロジスティック回帰モデルと同等の AUROC を達成しました（Figure 3）。これは、期待報酬の予測がモデル間の相対的な優劣を捉えられていることを意味します。

4. 主要な貢献（Key Contributions）

1. 期待報酬予測の実証: 応答レベルの報酬モデルから、プロンプトに対するモデルの期待報酬を事前予測することが可能であることを実証しました。これは、報酬モデルが実用的な実装（事前学習済み LLM の微調整）に基づいているため、プロンプトごとの絶対的なスコアスケールが学習可能であることを示唆しています。
2. シンプルかつスケーラブルなルーティング手法: 複雑なペアワイズ比較学習や、複数モデル間の勝率分布の直接予測（Zooter 法など）を必要とせず、各モデルごとに独立して線形予測器を訓練するだけで、高性能なモデルルーティングを実現できることを示しました。
3. 拡張性の証明: 新しいモデルがプールに追加された場合、既存のモデルの予測器を再訓練する必要がなく、新しいモデルに対してのみ単一の予測器を訓練すれば済みます。データ要件がモデル数の「2 乗」ではなく「1 次」で済むため、スケーラビリティに優れています。
4. 既存手法の解釈: 複雑なルーティング手法（例：Zooter）が機能する背景には、本質的に「期待報酬の予測可能性」があることを示し、そのメカニズムを解明しました。

5. 意義と結論（Significance）

本研究は、LLM の推論時におけるリソース最適化（モデルルーティング）に対して、理論的かつ実用的な新たな視点を提供しました。

• コスト削減と品質向上: 計算コストと応答品質のバランスを最適化し、大規模なモデルプールから各プロンプトに対して最適なモデルを動的に選択するシステムを構築可能にしました。
• 応用範囲の拡大: モデルルーティング以外にも、システムプロンプトの動的変更や、モデルのホットスワッピングなど、推論時の様々な操作に応用できる可能性があります。
• 報酬モデルの理解深化: 報酬モデルが単なる相対評価だけでなく、プロンプトとモデルの組み合わせに対する絶対的な「期待値」を内包している可能性を示唆し、今後の報酬モデルの設計やトレーニング戦略への示唆を与えています。

結論として、**「プロンプトの埋め込みから、特定の LLM が生成する応答の期待報酬を高精度に予測できる」**という発見は、LLM エコシステムの効率化と、より賢い推論システムの実現に向けた重要な一歩です。