Compute Aligned Training: Optimizing for Test Time Inference

以下は、Ousherovitch および Tewari による論文「Compute Aligned Training: Optimizing for Test Time Inference」の詳細な技術的サマリーです。

1. 問題定義：トレーニングと推論の不一致

本論文は、現代の大規模言語モデル（LLM）開発における根本的な断絶に対処しています：トレーニング目的は単一サンプルのパフォーマンス最適化向けに設計されているのに対し、デプロイメントはますますテスト時計算戦略（探索）に依存しているという点です。

• 現状： 教師あり微調整（SFT）や強化学習（RL）などの標準的なポストトレーニング手法は、単一の正しい出力の尤度（例えば、クロスエントロピーの最小化、または単一サンプルの期待報酬の最大化）を最適化します。
• 矛盾： 現代の推論戦略（例：Pass@N、多数決、Best-of-N）は、複数の候補解を生成し、それらを集約して最良の答えを見つけます。
  • 例（Pass@N）： モデルの正解確率が 50% から 90% に向上しても、Pass@100 メトリック（成功がほぼ確実な場合）では限界効用が逓減します。しかし、標準的なトレーニングは 100% の尤度を追求し続け、「易しい」サンプルに対して容量を浪費します。
  • 例（多数決）： モデルは競合他社に勝つために過半数（例えば 35% の確率）のみを必要とします。標準的なトレーニングは $p \to 1.0$ を目指し、堅牢な合意形成ではなく「圧勝」を促すため、モデルがアンサンブルとして動作を強要された際にパフォーマンスが低下する可能性があります。
• 結果： この不一致はポリシー崩壊（多様性の低下）を引き起こし、モデルがテスト時計算予算の増加を効果的に活用することを妨げます。

2. 手法：計算整合トレーニング（CAT）

著者らは、トレーニング目的をテスト時に使用される特定の推論戦略と整合させるフレームワークである**計算整合トレーニング（CAT）**を提案します。

核心的概念

著者らは、推論戦略を事後のヒューリスティックではなく、ベースポリシー $\pi_\theta$ を効果的なテスト時分布 $\tilde{\pi}_\theta$ へ変換する微分可能な演算子（$T$）として概念化します：
$$\tilde{\pi}_\theta(y|x) = T(\pi_\theta(\cdot|x), \phi)$$
ここで、$\phi$ はサンプル予算 $N$ などのハイパーパラメータを表します。

統合損失関数

CAT は、ベースポリシー（$\pi_\theta$）下での正解の負対数尤度を最小化するのではなく、効果的なシステムレベルポリシー（$\tilde{\pi}_\theta$）下での損失を最小化します：
$$L_{CAT}(\theta) = -\mathbb{E}_{(x, y^*)} [\log \tilde{\pi}_\theta(y^*|x)]$$

勾配スケーリング機構

全出力空間をサンプリングすることなく計算を可能にするため、著者らは対角勾配近似を導出しました。CAT 目的関数の勾配は、重み係数 $w(p, \phi)$ によって標準的な勾配をスケーリングすることで近似できることを示しています：
$$\nabla_\theta L_{CAT} \approx w(p, \phi) \cdot \nabla_\theta L_{CE}$$
ここで、$p$ は正解のベース確率であり、$w$ はテスト時戦略に対する $p$ 増加の限界効用を表します。

導出された主要なスケーリング因子：

1. Pass@N（SFT および RL）：
   • $w(p, N) = \frac{N p (1-p)^{N-1}}{1 - (1-p)^N}$
   • 挙動： 成功が既に保証されている易しいサンプル（高い $p$）の勾配を抑制し、限界利益がカバレッジを改善する「難しい」サンプルに学習を集中させます。
2. 多数決（SFT および RL）：
   • $w(p, N, k) \approx \frac{k \binom{N}{k} p^k (1-p)^{N-k}}{\sum_{i=k}^N \binom{N}{i} p^i (1-p)^{N-i}}$
   • 挙動： 勾配の質量を合意閾値（$p \approx k/N$）付近に集中させます。既に勝利している、あるいは絶望的に負けているサンプルを無視し、「転換点」にのみ焦点を当てます。
3. Best-of-N（RL）：
   • $w \propto N (P_{<y})^{N-1}$
   • 挙動： 「勝者総取り」のダイナミクスを生み出します。勾配はバッチ内の最高報酬サンプルのみを強化し、モデルが安全な平均値に収束するのではなく、分布の分散が高く報酬の高い尾部を探求することを促します。

理論的保証

論文は、競争戦略（競合他社の確率増加が成功を損なう場合）において、対角近似が真の勾配の保守的な過小評価であることを証明しています。これにより、オプティマイザが間違った方向にステップすることなく、ポリシー崩壊に対する自然な安定化剤として機能します。

3. 主要な貢献

1. 統合演算子フレームワーク： テスト時戦略を微分可能な演算子として形式化し、任意の推論戦略（SFT および RL）とのトレーニング整合を導出する一般的なメカニズムを確立しました。
2. Pass@N を超えた拡張： 以前の研究（例：Direct Coverage Optimization）を多数決およびBest-of-Nに一般化し、SFT と RL の両方のパラダイムを網羅しました。
3. モダリティ非依存性： LLM だけでなく**タンパク質言語モデル（PLM）**へのフレームワークの適用性を示し、複雑なマルチモーダル報酬ランドスケープにおけるパフォーマンス向上を実証しました。
4. 分散分析： 勾配分散の増大に関する理論的および実証的分析を提供し、トレーニングを安定させるための正規化手法（例：対数加重推定量）を提案しました。

4. 実験結果

著者らは、MATH ベンチマーク（Mistral-7B 使用）およびタンパク質生成（ProtGPT2 使用）において CAT を評価しました。

A. 言語モデル（MATH ベンチマーク）

• Pass@N（SFT）： $N=64$ で CAT を用いてトレーニングされたモデルは、単一サンプル精度（Pass@1 が約 2.3% 低下）を犠牲にしましたが、標準的な SFT と比較して Pass@64 カバレッジで7.8% の絶対的改善を達成しました。
• 多数決（SFT）： CAT モデルは、すべての推論予算においてベースラインを一貫して上回りました。$N=16$ にてバイアスと分散を最適にバランスさせる「スイートスポット」が観察されました。
• 強化学習（GRPO）： 標準的な RL はポリシー崩壊に陥り、急速に飽和しました（Pass@32 $\approx$ 35.8%）。CAT トレーニングモデルは探索に必要な分散を維持し、Pass@32 で**40.0%**を達成しました。
• ウォームアップ効果： SFT ウォームアップ段階で CAT を使用すると、その後の RL パフォーマンスが大幅に向上し、モデルの推論経路を豊かにすることで単一ショット（Pass@1）精度さえも向上させました。

B. タンパク質生成（タンパク質言語モデル）

• 無条件生成（死の谷）： マルチモーダル報酬ランドスケープ（局所最適解 $h \approx 0.35$、大域最適解 $h \approx 0.75$、それらを「谷」の凝集が隔てている）において、標準的な RL は局所最適解に留まりました。Best-of-N CAT は高い分布分散を維持することで「谷」を越え、大域最大値を成功裏に見つけ出しました。
• 条件付き生成： CAT モデルは、入力と逆の疎水性を持つ配列を生成することを学習し、標準的な RL と比較して関係タスクとのスケーリングおよび整合性が優れていることを示しました。

5. 意義と含意

• パラダイムシフト： 本論文は、トレーニング目的は下流の推論手順を念頭に置いて設計されるべきであると主張しています。テスト時計算がモデル能力のスケーリングの主要なレバーとなるにつれ、トレーニングは単一サンプルではなく、集約された結果を最適化する必要があります。
• 効率性： CAT はトレーニング計算コストを増加させることなくこれらの利益を達成します。単に、サンプルのテスト時戦略に対する限界効用に基づいて既存の勾配を再重み付けするだけです。
• 安定性対探索： このフレームワークは、単一サンプルの精度と分布の多様性の間のトレードオフを原理的に提供し、モデルがモード探索行動に収束することなく、Best-of-N や多数決などの探索手順を効果的に活用できるようにします。
• 汎用性： LLM とタンパク質 LM の両方での有効性を証明することで、本論文は「計算整合トレーニング」がテスト時探索を利用する任意の生成モデルに対する普遍的な原理であることを示唆しています。

結論として、計算整合トレーニングは学習と探索の間のギャップを埋め、モデルが単一の答えを生成することに「賢い」だけでなく、計算予算を与えられた際に最良の答えを見つけることに「賢い」ように最適化されていることを保証します。