技術要約: Double Preconditioning (DoPr)

1. 問題提起: テスト時フィードバック (TTF) と損失・性能の乖離

本論文は、モデルが1ステップの教師あり目的関数（例：次トークン予測、ステップごとの模倣学習）で訓練される一方で、デプロイ時には自身の予測に基づいた逐次的なロールアウトを通じて運用されるという、現代のディープラーニング・アプリケーションにおける決定的な断絶に対処している。著者らはこの現象をテスト時フィードバック (Test-Time Feedback: TTF) と呼んでいる。

オートレグレッシブな言語モデリング、フローベースの生成モデリング、ロボット・ポリシー学習などのTTF設定では、モデルの予測に含まれる微小な誤差が時間の経過とともに累積する。これにより、デプロイ中に遭遇する状態の分布が、訓練分布から乖離してしまう（TTFシフトとして知られる現象）。

特定された核心的な問題は、標準的なオプティマイザ（例：Adam, Muon, Shampoo）は、訓練分布下での検証損失 ($L_{val}$) を最小化するように設計されている点である。しかし、本論文は、$L_{val}$ を最小化することが、必ずしもダウンストリーム・タスクの報酬 ($R_{test}$) や、モデル自身のデプロイ分布上で定義される理想的な損失 ($L_{ideal}$) を最小化することにはならないと主張している。具体的には、標準的なオプティマイザは、訓練分布には最適であるが、TTFのフィードバック・ダイナミクス下では失敗する特徴量を学習してしまうという、不完全な特徴学習 (imperfect feature learning) を誘発する可能性がある。これにより、検証損失が良好に収束している場合でも、ダウンストリームの性能が低下する。

2. 手法: Double Preconditioning (DoPr)

データ、アーキテクチャ、または訓練目的関数を変更することなくTTFシフトを軽減するために、著者らは Double Preconditioning (DoPr) を提案している。DoPrは、2つの異なるプリコンディショニング戦略を組み合わせた「プラグイン」型の最適化フレームワークである。

1. 活性化プリコンディショニング (Activation Preconditioning: AP): 層ごとの活性化の経験的共分散 ($\hat{\Sigma}_z \approx E[zz^\top]$) を用いて勾配をプリコンディショニングする、レイヤー単位のプリコンディショナー。

   • メカニズム: 生の勾配 $G$ に対して、活性化共分散行列の逆行列を乗算する：$M = G \cdot \hat{\Sigma}_z^{-1}$。
   • 理論的動機: 活性化が非等方的（不良条件）である場合、標準的な勾配降下法はバイアスのある特徴学習に陥ることを著者らは示している。APはこのバイアスを補正し、状態空間のあらゆる方向に対して一様な特徴学習 (uniform feature learning) を促進する。これにより、モデルは、訓練損失では軽視されるものの、TTF下での誤差蓄積を防ぐために不可欠な方向においても、ロバストな特徴を学習できるようになる。
   • 不変性の原理: 本論文は、APが入力分布のアフィン変換に対する不変性を誘導することを確立している。入力分布がアフィン変換された場合、APの更新ルールは層の出力を不変に保ち、分布シフトに対する最適化軌道を安定させる。

2. 勾配プリコンディショニング (Gradient Preconditioning: GP): APによってプリコンディショニングされた勾配に対して適用される、標準的な勾配ベースのプリコンディショナー（例：Adam, Muon, Signum）。

   • 役割: APは特徴学習の幾何学を扱うが、単独では数値的に不安定になる可能性がある。GPコンポーネントは、訓練プロセスを安定させ、収束を加速させ、モメンタムを提供する。
   • 統合: DoPrは、まず勾配にAPを適用し、その結果をGPに渡す。これにより、DoPrは既存のGP（AdamやMuonなど）の安定性とハイパーパラメータのスケーリング規則を活用しながら、APのTTF軽減特性を継承することができる。

アルゴリズムの流れ:
重み $W$ と活性化 $z$ を持つ層について：

1. バッチ勾配 $G = \nabla_W \mathcal{L}$ を計算する。
2. 経験的活性化共分散 $\hat{\Sigma}_z$ を計算する。
3. APを適用する：$M = G \cdot (\hat{\Sigma}_z + \gamma \text{tr}(\hat{\Sigma}_z)I)^{-1}$。
4. GPを適用する：$D = \text{GP}(M, \dots)$。
5. 重みを更新する：$W \leftarrow W - \eta D$。

3. 主な貢献

• 理論的洞察: 本論文は、TTFシフトと不完全な特徴学習を厳密に関連付けている。標準的なオプティマイザは、検証損失は低いものの、特徴のサブスペース整合性が乏しい解に収束する可能性があり、それがTTFを悪化させることを証明している（命題3.2および4.1による）。
• DoPrフレームワーク: 方向性（一様な特徴学習を保証するためのAPが担当）と適応性/安定性（GPが担当）を分離する、新しい最適化パラダイムを導入した。
• ドロップイン互換性: DoPrは既存のオプティマイザのドロップイン置換として設計されている。モデルのアーキテクチャやデータのキュレーションを変更する必要はない。
• ハイパーパラメータの転移: DoPrはGPコンポーネントを保持しているため、AdamやMuonといった基本となるGPから導出された既存のハイパーパラメータ・スケーリング規則（学習率やウェイト減衰のための $\mu$P スケーリングなど）がDoPrに直接転移することを示している。これは大規模訓練を容易にする。

4. 実験結果

著者らは、連続制御、ロボティクス、および言語モデリングのタスクにおいてDoPrを評価している。主な知見は以下の通りである：

• 連続制御 (Humanoid-v5): DoPrは、ベースラインのオプティマイザ（AdamW, Muon, Signum, AdaMuon）と比較して、終端報酬を一貫して向上させた。特筆すべきは、このダウンストリーム性能の向上は、訓練または検証損失の改善と必ずしも相関していない点である。いくつかのケースでは、DoPrは同等またはわずかに劣る損失収束を示しながら、より高い報酬を達成した。
• ロボット・ポリシー学習 (Robomimic): 「Tool-Hang」や「Transport」のような困難なタスクにおいて、DoPrのバリアントはベースラインを上回る成功率を示した。ここでも、タスクの成功における利得は、訓練損失の一様な改善を伴うものではなかった。
• 言語モデル (数学的推論):
  • 3B SFT: DoPr-AdamWは、トークンレベルの訓練損失が同等または高くなっているにもかかわらず、様々な学習率において、AdamWと比較してGSM8Kにおけるピーク精度を達成した。
  • 8B SFT: より大規模なファインチューニングにおいて、DoPr-AdamWは、ベースラインのAdamWと比較して、検証損失とダウンストリーム性能（GSM8K, MATH-500）の間の相関を強めている。ベースラインのAdamWモデルは、高い学習率において検証損失は改善しているものの、タスク性能が低下する傾向が見られたが、DoPrは堅牢な性能を維持した。
  • タスク間の競合: DoPrはタスク間の競合を軽減する。GSM8K-CoTの最適学習率は、MATH-500においても最適に近い値となり、これはベースラインのAdamWでは観察されなかった特性である。

5. 重要性と主張

本論文は、最適化は、データキュレーションやアーキテクチャ設計とは異なる、TTF設定における誤差蓄積を軽減するための新しい設計軸であると主張している。

• 損失と性能の分離: 主要な意義は、TTF設定におけるダウンストリーム性能の向上には、検証損失の改善は必須ではなく（むしろ伴わないこともある）、むしろそれが必要であるということを示した点にある。これは、「訓練損失を最小化することがモデルの品質の唯一の代理指標である」という従来の常識に挑戦するものである。
• レバーとしての特徴学習: 本研究は、TTF下でのロバスト性のための決定的な要因は、損失収束の速度ではなく、「一様な特徴学習」（APによって保証される）であると位置づけている。
• 実用的な有用性: DoPrは、モデルが自身の予測をロールアウトしなければならない現実世界のデプロイメント・シナリオにおいて、モデルの堅牢性を向上させる、低オーバーヘッドの実用的な介入策（「プラグイン」）を提供する。

著者らは、最適化における「方向性」と「適応性」の分離が、特に訓練目的がデプロイの現実と乖離するタスクにおいて、依然として豊かな研究領域であることを示唆し、まだ「オプティマイザのフリーランチ」が存在するのではないかと問いかけている。