Towards a theory of inference-time alignment with unknown rewards

技術的要約：未知の報酬を用いた推論時アライメントの理論に向けて

問題定式化
本研究は、生成モデルの応答を、モデルパラメータを更新することなく（RLHFとは異なり）、テスト時に報酬信号に基づいて選択または誘導するパラダイムである**推論時アライメント（inference-time alignment）の統計的学習の基礎に取り組むものである。著者らはこれを、未知の報酬を仮定した弱から強への学習（weak-to-strong learning）**問題として定式化している。

設定は以下の通りである：

• 参照ポリシー ($\pi_{ref}$): 「弱学習器」（例：事前学習済みモデルやSFTモデル）であり、候補となる応答を生成する。これは、任意の入力 $x$ に対して「良い」応答（報酬 $r^*(x, y) = 1$）を生成する確率が一定である $\gamma$ であると仮定される。
• 報酬クラス ($\mathcal{R}$): 二値の報酬関数 $r: \mathcal{X} \times \mathcal{Y} \to \{0, 1\}$ の集合。真の報酬 $r^*$ は未知であるが、$\mathcal{R}$ 内に存在する。
• データ生成: 学習データは $m$ 個の独立同一分布（i.i.d.）サンプルで構成される。各サンプルには、入力 $x_i$ と、$\pi_{ref}(\cdot|x_i)$ から抽出された $k$ 個の応答の「スレート（一連の候補）」$\{y_i^{(1)}, \dots, y_i^{(k)}\}$ が含まれる。学習者は、良い応答のインデックスの集合 $A_{r^*}(z_i) = \{j : r^*(x_i, y_i^{(j)}) = 1\}$ を観測する。
• 目標: テスト入力 $x$ と $\pi_{ref}$ からの $k$ 個の候補のスレートが与えられたとき、高い確率（$1-\epsilon$）で $r^*(x, y^{(j)}) = 1$ となるインデックス $j$ を選択する学習アルゴリズムを設計すること。

著者らは、既存の報酬モデルへのアクセスを仮定している既存の理論文献とは異なり、報酬関数自体が有限の選好データから推論されなければならない、ゼロからの学習として扱うことで、自身の研究を区別している。

手法と核となる概念
本論文では、報酬クラスの学習可能性を特徴付けるための新しい組合せ次元を導入している：

1. アライメント次元 ($d_{ALN}(\mathcal{R})$):
   入力の集合 $\{x_1, \dots, x_d\}$ が $\mathcal{R}$ によって**アライメント・シャッター（alignment shattered）**されているとは、任意のバイナリベクトル $b \in \{0, 1\}^d$ に対して、$r_b(x_i, y_i^{(b_i)}) = 1$ かつ $r_b(x_i, y_i^{(1-b_i)}) = 0$ を満たす報酬関数 $r_b \in \mathcal{R}$ が存在するような、対となる応答 $(y_i^{(0)}, y_i^{(1)})$ が存在することを指す。アライメント次元は、このようなシャッターされた集合の最大サイズである。

2. 非支配的ナタラジャン次元 (Non-Dominating Natarajan Dimension):
   学習可能性を証明するために、著者らは部分概念クラスに対して用いられる**ワン・インクルージョン・グラフ（OIG）**アルゴリズムを、誘導された集合値受理クラスへと関連付けている。彼らは、この「非支配的ナタラジャン次元」を定義し、これがアライメント次元によって上界に抑えられることを証明している。

3. 学習アルゴリズム（非支配的ペアに対するトーナメント）:
   提案される学習手順は、推移的（transductive）であり、テスト時に動作する：

   • テスト用のスレート $Z$ に対して、アルゴリズムは、互いに包含関係にない（$A \not\subseteq B$ かつ $B \not\subseteq A$）空でないインデックスの部分集合のペア $(A, B)$ すべてを検討する。
   • 各ペアについて、真の受理集合が $A$ または $B$ であったインスタンスを含む学習サブセットを構築する。
   • このサブセット上でOIGアルゴリズムを実行し、真の受理集合が $A$ か $B$ かを予測する。
   • このプロセスは「トーナメント」を形成し、可能な受理集合を連鎖（chain）へと絞り込んでいく。最終的な予測は、生き残った集合の共通部分に含まれる任意の応答インデックスである。

主要な結果
本論文は、アライメント学習可能性の完全な特徴付けを確立している：

• 定理1: 報酬クラス $\mathcal{R}$ がアライメント学習可能であるための必要十分条件は、そのアライメント次元 $d_{ALN}(\mathcal{R})$ が有限であることである。
  • 上界（定理2）: $d_{ALN}(\mathcal{R}) < \infty$ である場合、提案されたトーナメント・アルゴリズムは、次元、$\epsilon$、$\delta$、および $\gamma$ の有限関数であるサンプルサイズ $m$ と推論時スレートサイズ $k$ において、$(\epsilon, \delta)$-学習可能性を達成する。
  • 下界（定理3）: $d_{ALN}(\mathcal{R}) = \infty$ である場合、いかなる学習アルゴリズムも有限のサンプルでアライメント学習可能性を達成することはできない。証明では、確率法を用いて困難な事例を構築し、無限の次元性が、学習者が真の報酬関数を区別できない分布の存在を意味することを示している。

意義と主張
著者らは、本研究を一般的なアライメント問題に対する厳密な統計的学習理論への先駆的な一歩として位置づけている。主な貢献に関する主張は以下の通りである：

• 新しいフレームワーク: 報酬が未知であり、データから完全に学習されるという設定において、推論時アライメントを統計的学習問題として定式化した最初の研究である。これは、優れた報酬推定へのアクセスを仮定する先行研究とは異なる。
• 理論的特徴付け: アライメント次元の導入は、二値分類におけるVC次元に類似した、学習可能性のための必要十分条件を提供する。これは、複数の候補から一つの良い応答を選択するという特定の制約に適合するように調整されている。
• 既存のパラダイムとの区別: 著者らは、自身の設定がマルチクラス学習（一つのラベルが正しい）やリスト学習（ラベルのリストが出力される）とは異なることを明確にしている。彼らの設定では、複数の応答が正しい可能性があるが、学習者はそのうちの一つを正常に特定しなければならない。
• 範囲に関する謙虚な姿勢: 著者らは限界についても認めており、現在の定式化が二値報酬と実現可能な設定を想定していることを指摘している。彼らは、自身の学習手順が非標準的であること（テスト時の計算とタプルの順序付けの仮定に依存していること）を明示しており、単一応答の学習データ、実数値報酬、または選好ベースの学習（ブラッドリー・テリー・モデル）への拡張は、今後の研究課題であるとしている。

本論文は、アライメント次元を基本的な複雑性尺度として確立することが、推論時アライメントがいつ、どのように可能であるかを理解するための理論的基盤を提供し、より体系的な原則の開発を導く可能性があると結論付けている。