How Neural Reward Models Learn Features for Policy Optimization: A Single-Index Analysis

「How Neural Reward Models Learn Features for Policy Optimization: A Single-Index Analysis」という論文の詳細な技術的サマリーを以下に示す。

問題定義

報酬モデリング（RM）は、大規模言語モデルの整合性確保（RLHF などによる）において重要な構成要素であるが、標準的な教師あり予測とは根本的に異なる。KL 正則化付き方策最適化において、学習された報酬関数 $r(x)$ は、展開される方策 $\pi_\beta(dx) \propto e^{r(x)/\beta}\pi_{ref}(dx)$ を定義するために指数化される。したがって、システムの downstream における価値は、単に参照分布 $\pi_{ref}$ 下での予測精度に依存するのではなく、報酬が高く方策が集中する領域における誤差に依存する。これにより、方策に起因する分布シフトが生じる。

オフライン整合性の既存の理論的解析は、しばしば固定された既知の線形特徴を仮定しており、そのような分布シフト下でニューラルネットワークがどのように関連する特徴を発見するか（特徴学習）、および非線形報酬ランドスケープを近似するかという課題を回避してきた。本論文は、訓練目的が指数関数的報酬重み付けを含む場合に、2 層ニューラルネットワークがどのように特徴を学習し、非線形報酬関数を適合させるかを解析することで、このギャップに取り組む。

手法

著者らは、問題を**ガウス単一インデックスモデル（SIM）**の枠組み内で定式化する：

• データ生成: 入力 $x \sim \mathcal{N}(0, I_d)$、ラベル $y = r^*(x) + \zeta$。ここで真の報酬は $r^*(x) = \sigma^*(\langle \theta^*, x \rangle)$ である。$\theta^*$ は隠れた目標方向、$\sigma^*$ は上限付き多項式リンク関数、$\zeta$ は有界なノイズである。
• モデルアーキテクチャ: $N$ 個のニューロンを持つ 2 層ニューラルネットワーク：$r_a(x) = \frac{1}{N} \sum_{j=1}^N a_j \sigma(\langle w_j, x \rangle + b_j)$。
• 2 段階訓練プロトコル:
  1. 特徴回復（ステージ 1）: 第 1 層の方向 $w_j$ は、温度 $\beta_1$ を用いた報酬重み付け目的関数によるオンライン球面 SGD によって訓練される。第 2 層の係数 $a$ は初期化時に固定される。目標は、$w_j$ を隠れた方向 $\theta^*$ と整合させることである。
  2. 読み出し適合（ステージ 2）: 特徴が回復された後、方向 $w_j$ とバイアス $b_j$ は凍結される。第 2 層の係数 $a$ は、新しいサンプルを用いた重み付きリッジ回帰によって適合される。重み付け方式は、展開温度 $\beta_2$（$\beta_1$ とは異なる場合がある）を使用する。2 つの重み付け方式が解析される：
     • ラベル重み付け: 重みは $e^{y/\beta_2}$。
     • 代理重み付け: 重みは $e^{r_{proxy}(x)/\beta_2}$。ここで $r_{proxy}$ は凍結された代理報酬（例：特徴回復段階からのもの）である。

主要な貢献と結果

1. 報酬に起因する分布シフト下での特徴学習
本論文は、指数関数的報酬重み付け下での第 1 層の特徴回復を特徴づける。

• 信号変換: 指数関数的重み付けは、第 1 層が利用可能なエルミート信号を変化させる非線形ラベル変換として作用する。
• 指数の減少: 解析により、次元に依存しない $O(1)$ の閾値以上の任意の特徴学習温度 $\beta_1$ に対して、実効的な学習のボトルネックが**情報指数（IE）から生成指数（GE）**に減少することが示される。
• 回復保証: 上限付き多項式の仮定の下、一定割合のニューロンが高確率で隠れた方向 $\theta^*$ を回復する。弱回復の複雑さは生成指数によって支配され、$\tilde{O}(\beta_1^{2(I^*-1)} d^{(p_{gen}-1)\vee 1})$ としてスケーリングする。ここで $I^*$ は GE に到達するために必要な冪である。
• $\beta_1$ の重要性: 決定的なことに、$\beta_1$ の必要な下限は次元 $d$ に依存しない。これは、次元に依存する対数多項式温度を必要とし、それによって指数減少の恩恵を無効化していた以前のクリップド指数解析を改善するものである。

2. 重み付け予測から傾いた方策の後悔へ
著者らは、標準的な予測誤差の境界を超えて、傾いた方策の価値ギャップを評価する：
$$\mathcal{R} = \mathbb{E}_{x \sim \pi_{\beta^*}}[r^*(x)] - \mathbb{E}_{x \sim \hat{\pi}_{\beta_2}}[r^*(x)]$$
ここで $\pi_{\beta^*}$ は目標方策、$\hat{\pi}_{\beta_2}$ は学習された報酬によって誘発される方策である。

• 分解: 後悔は 3 つの項に分解される：
  1. 温度不一致（$T_{temp}$）: 目標温度 $\beta^*$ と展開温度 $\beta_2$ の差に依存する。
  2. 射影切り捨て（$T_{cut}$）: 積分可能性を確保するために解析を射影部分空間に制限することから生じる。
  3. 学習項（$T_{learn}$）: 適合された報酬の重み付け予測誤差に依存する。
• 低温挙動: 本論文は、傾いた方策下での報酬の分散が $\beta_2^2 / (2p_{max})$ としてスケーリングすることを示す明示的な境界を導出する。ここで $p_{max}$ は報酬の極大点の局所的次数である。

3. ラベル重み付けおよび代理重み付けの後悔境界
本論文は、両方の適合方式に対する明示的な価値ギャップ境界を提供する：

• ラベル重み付け境界: 理論的ベンチマークとして機能する。この境界は温度不一致、切り捨て、学習コストを分離し、方策価値率を非線形報酬極大の局所的形状と関連付ける。
• 代理重み付け境界: プラグインパイプラインにとってより実用的である。この境界は、代理誘発測度と参照分布との密度比を表す追加因子 $C_{0,R}(\beta_2)$ を導入する。
  • 代理不一致メカニズム: 代理報酬が不完全な場合、$C_{0,R}(\beta_2)$ は大きくなり得る（$1/\beta_2$ に対して多項式的にスケーリング）。これは近似誤差と有限サンプル誤差の両方を増幅し、「代理不一致」と報酬の過剰最適化の背後にある統計的メカニズムを特定する。
• 許容温度: 解析により、展開温度 $\beta_2$ の許容集合が得られる。より小さな $\beta_2$ を選択することは目標方策との一致を改善するが、指数関数的重み付けを通じて学習コストを増幅する。最適な選択はこれらのトレードオフをバランスさせるものである。

意義と主張

本論文は、制御された理論的設定内における報酬モデリングと方策価値の相互作用の基礎的な特徴付けを提供すると主張する。

• 特徴学習: 温度 $\beta_1$ が小さすぎない場合、報酬に起因する指数関数的重み付けが、実効的な学習指数を IE から GE に減少させることで特徴学習を促進し得ることを実証する。次元に依存する温度を必要とした以前の仮定を修正し、次元に依存しない $O(1)$ の温度で回復が可能であることを確立する。
• 評価指標: 報酬モデリングにおける関連する統計的ターゲットは、参照分布の予測誤差ではなく、傾いた方策の価値であると論じる。導出された境界は、価値ギャップを報酬ランドスケープの局所幾何学と展開温度の選択と明示的に結びつける。
• 代理のリスク: 代理報酬を重み付けに使用すること（代理重み付け）が、代理が不完全な場合、誤差を増幅し得る特定の統計的メカニズムを特定し、実務で観察される報酬の過剰最適化現象に対する理論的説明を提供する。

著者らは、ガウス SIM は理論的抽象化であるが、特徴学習と方策に起因する分布シフトの中核メカニズムを分離し、より豊かな入力分布や多インデックス構造に関する将来の作業の基盤を提供していると主張している。