技術要約：合意形成に基づく最適化によって訓練される2層ニューラルネットワークの平均場モデル

問題定義

ニューラルネットワークの訓練は、本質的に経験リスク関数を最小化することを目的とした最適化問題である。確率的勾配降下法（SGD）やその適応型変種（Adamなど）といった勾配ベースの手法が標準となっているが、これらは目的関数の非凸性により、局所解に陥りやすいという課題がある。さらに、マルチタスク学習では、相反する勾配のバランスを取ることや、異なるタスクに対して個別のモデルやアンサンブルを訓練する際の膨大なメモリオーバーヘッドの管理が求められる。

本論文では、2層ニューラルネットワークの訓練において、相互作用する粒子系に基づく勾配フリーのグローバル最適化手法である**合意形成に基づく最適化（Consensus-Based Optimization: CBO）**の使用可能性を調査する。著者らの目的は以下の通りである：

1. CBOの性能をAdamと比較すること。
2. CBOとAdamの両方の強みを活用するためのハイブリッドアプローチを開発すること。
3. メモリオーバーヘッドを削減するために、CBOをマルチタスク学習に適応させること。
4. 2つの領域（無限ネットワーク幅 $M \to \infty$ および無限粒子数 $N \to \infty$）における訓練ダイナミクスの厳密な平均場モデルを構築すること。

手法

1. 最適化フレームワーク

本研究では、2層ニューラルネットワーク $\hat{g}(x; \theta)$ の訓練を、経験リスク $\hat{R}(\theta)$ を最小化することとして定式化している。

• Adam: 勾配の1次および2次のモーメント推定に基づき適応的な学習率を用いる、ベースラインとなる勾配ベースのオプティマイザとして使用される。
• Consensus-Based Optimization (CBO): $N$ 個の粒子 $\theta_n$ のアンサンブルが確率微分方程式に従って進化する粒子ベースの手法である。ダイナミクスは、合意点 $V^k$（低リスクを好む粒子による加重平均）へのドリフト項と、探索のための拡散項で構成される。
• ハイブリッド手法: AdamステップとCBOステップを組み合わせた新しい更新則：
  $$\theta_n^{k+1} = \theta_n^k - \gamma \Delta t (\text{Adam update}) + (1-\gamma) (\text{CBO update})$$
  ここで $\gamma \in [0,1]$ はそのバランスを制御する。
• Multi-Task CBO: 単一の粒子アンサンブルを複数のタスク間で再利用する定式化。個別のアンサンブルを維持する代わりに、粒子を特定のタスクに割り当てることで、メモリオーバーヘッドを削減する。合意点は、共有された粒子セットを用いてタスクごとに計算される。

2. 平均場定式化

著者らは、ニューラルネットワークの無限幅極限を扱うために、CBOを**最適輸送（Optimal Transport: OT）**の枠組みの中で再定式化している。

• 無限幅 ($M \to \infty$): ネットワークを $\mathbb{R}^{M(d+2)}$ 内の点として表す代わりに、ネットワークをパラメータ空間上の確率測度 $\mu$ として表現する。CBOのダイナミクスは、ワッサースタイン空間 $P_2(\mathbb{R}^{d+2})$ 上に持ち上げられる。合意点は、ワッサースタイン距離の二乗の加重和の最小化因子として定義される**バリセンター（重心）**となる。
• 無限粒子 ($N \to \infty$): 粒子アンサンブルのダイナミクスは、ワッサースタイン空間上のワッサースタイン空間（Wasserstein-over-Wasserstein）上の測度の空間へと持ち上げられる。著者らは離散時間平均場モデルを導出し、粒子の分散が単調に減少することを証明している。

3. 数値実験

著者らは、以下の3つのタスクを用いて手法を検証している：

1. 正弦波近似（Sine Approximation）: CBOとAdamを比較する1次元回帰問題。
2. MNIST分類: CBO、Adam、およびハイブリッド手法を比較するマルチクラス分類タスク。
3. マルチタスク学習: Multi-Task CBO戦略を用いて、100個のシフトした正弦波を同時に近似する。
4. 正方形近似（平均場検証）: ネットワーク幅 ($M$) と粒子数 ($N$) を変化させ、平均場極限への収束を確認するための実験。

主な結果

パフォーマンス比較

• 正弦波回帰: CBOはAdamよりも低い最終的な経験リスクを達成し、より高い安定性を示したが、全粒子に対するフォワードパスが必要であるため、反復あたりの計算コストは高くなった。
• MNIST分類: AdamはCBOよりも速く収束し、より低いリスクを達成した。しかし、ハイブリッド手法は両者を上回り、Adam単独よりも速く収束した。このハイブリッドアプローチにより、（純粋なAdamでは不安定になるような）大きな学習率の使用が可能になった。これはCBO成分がダイナミクスを安定化させたためである。
• Multi-Task CBO: この手法は、単一の粒子アンサンブルを用いて100の異なるタスクにおけるリスクを最小化することに成功した。中央値および最小の経験リスクの両方が単調に減少しており、粒子の再利用が、個別のアンサンブルを維持するメモリコストなしにマルチオブジェクティブ最適化を効果的に処理できることを裏付けている。

平均場への収束

• 分散の減衰: 理論的分析（命題3）によれば、無限粒子極限において、粒子アンサンブルの分散は $(1-\Delta t)^{2k}$ として幾何級数的に減衰し、合意が保証される。
• 実証的検証: 正方形近似に関する数値実験では、ネットワーク幅 $M$ と粒子数 $N$ の両方が増加するにつれて経験リスクが単調に減少し、平均場領域付近で飽和することが示された。

意義および主張

本論文は、以下の貢献と意義を主張している：

1. ハイブリッド最適化: CBOとAdamを組み合わせることで、特にMNISTのような非凸なランドスケープにおいて、どちらの手法よりも高速かつ安定した堅牢なオプティマイザが得られることを示している。
2. メモリ効率の高いマルチタスク学習: Multi-Task CBOの定式化は、粒子を再利用することで、個別のモデルを訓練する場合と比較してメモリオーバーヘッドを大幅に削減し、関連する複数のタスクに対してモデルを訓練するための実用的な方法を提供する。
3. 理論的統一: 本研究は、以下の点を通じて粒子ベースの最適化とニューラルネットワーク理論を橋渡ししている：
   • CBOを最適輸送の枠組みで再定式化し、無限幅のニューラルネットワークを連続的な測度として扱うことを可能にした。
   • これらの無限幅ネットワーク上のCBOダイナミクスの離散時間平均場モデルを導出した。
   • 平均場極限において、粒子集団の分散が単調に減少することを証明した。
4. 限界と今後の方向性: 著者らは、分散の減衰は合意を保証するものの、それがOT設定におけるグローバルな最小化への収束を厳密に保証するわけではない（この関連性は未解決の問題である）と控えめに述べている。また、現在のフレームワークは2層ネットワークに限定されており、OT定式化には現在、厳密な拡散項が欠けている（実際にはノイズの追加に依存している）。今後の課題として、より深いアーキテクチャへの拡張と、厳密な収束証明の確立が挙げられている。

要約すると、本論文は、CBOがニューラルネットワーク訓練において、勾配降下法の実行可能かつグローバルに収束する代替手段となり得ることを提示している。それは、安定性とマルチタスク効率において特有の利点を持ち、そのダイナミクスはワッサースタイン空間における平均場極限を通じて厳密に分析可能である。