技術要約：形状化された深層非線形MLPにおけるベイズ推論

問題設定

深層学習理論は、モデルパラメータ（幅 $N$、深さ $L$）とデータセットサイズ（$P$）が同時に巨大であるとき、ニューラルネットワークがどのように予測を行うかを特徴付けることを目的としている。根本的な課題は、$N, L, P \to \infty$ の極限が必ずしも可換ではない点にある。既存のフレームワークは、固定された深さにおける無限幅極限（ニューラルタンジェントカーネルやニューラルネットワークガウス過程をもたらす）や、特定の線形ネットワーク・レジームに依存することが多い。しかし、これらのアプローチは、非線形ネットワークにおける深さの役割や、$L \propto N$ かつ $P \propto N$ である比例スケーリング・レジームにおける特徴学習とカーネル手法の相互作用を捉えることに失敗する場合が多い。

本研究で扱う具体的な問題は、訓練サンプル数（$P$）、入力次元（$N_0$）、隠れ層の幅（$N$）、および隠れ層の数（$L$）がすべて発散する結合スケーリング極限の下での、深層非線形多層パーセプトロン（MLP）におけるベイズ推論の厳密な特徴付けである。著者らは以下の点を明らかにしようとしている：

1. このレジームにおいて、予測事後分布はどのように振る舞うか。
2. 深さが恩恵をもたらす（ベイズモデル証拠を増加させる）条件は何か。
3. 非線形ネットワークは、データ依存のカーネル手法によって近似され得るか、またどのように近似されるのか。

手法

1. 形状化された活性化関数

$L, N \to \infty$ において非退化な挙動を保証するために、著者らは形状化された活性化関数 $\phi_s$ を採用している。基底となる非線形関数 $\phi_{base}$ に対し、活性化関数は $\phi_s(x) = s \phi_{base}(x/s)$ としてスケールされる。このメカニズムにより、多くの層にわたる非線形性の累積的な影響を、制御可能かつ線形効果と比較可能な状態に保つことができる。形状因子 $s$ は、比率 $t = L/N$ が初期化分布を一意に制御するように選択される。

2. ニューラル共分散SDE (NSDE)

核心となる技術的ツールは、ニューラル共分散確率微分方程式 (NSDE) である。著者らは、隠れ特徴の共役カーネル（グラム行列）$\Phi^{(\ell)}$ の層間における進化を分析している。

• $N, L \to \infty$ かつ $L/N \to t$ の極限において、カーネルの離散的な進化は拡散過程 $\Phi_t$ に収束する。
• このプロセスは、基底となる非線形性に依存するドリフト項 $b(\Phi_t)$ と、ブラウン運動行列 $B_t$ を用いて次のように記述される：
  $$d\Phi_t = \alpha_{act} \cdot b(\Phi_t) dt + \Phi_t^{1/2} dB_t \Phi_t^{1/2}$$
• 大きな $P$ を扱うために、タイムチェンジ $\tau = P t$ が導入され、「深さ時間」 $\tau^* = LP/N$ が定義される。拡散項とドリフト項のバランスをとるために、動力学が再スケールされる。

3. SDEによるベイズ分配関数

著者らは、予測事後分布をエンコードするベイズ分配関数 $Z_\beta(x_0, \kappa)$ を、共役カーネルの事前分布に関する期待値として定式化している。

• ガウス積分恒等式を用いることで、分配関数は次のように表される：
  $$Z_\beta(x_0, \kappa) \propto \int_{\mathbb{R}^P} \exp\left(-\frac{1}{2\beta}\|p\|^2 + ip^\top Y\right) \mathbb{E}_{prior}\left[\exp\left(-\frac{1}{2}v^\top \Phi v\right)\right] dp$$
  ここで $v = [p^\top, \kappa]^\top$ であり、$\Phi$ は最終層直前の共役カーネルである。
• 定理1によれば、結合極限において、MLP事前分布に関する期待値は、NSDEの解に関する期待値に収束する。これにより、深層MLPにおけるベイズ推論の問題は、$\Phi_\tau$ に関するSDEを解き、その結果得られる期待値を計算することへと還元される。

4. 摂動解析

任意の $\tau^*$ に対してSDEを解くことは困難であるため、著者らは有効な深さが小さい $\tau = LP/N$ に対する摂動展開を行う。彼らは、$\tau=0$ における対数分配関数の $\tau$ に関する一次微分を計算する。これにより、線形ネットワーク成分（拡散のみ）と非線形カーネル成分（ドリフトのみ）の寄与を分離することが可能になる。

主な貢献と結果

1. 厳密な統一フレームワーク

本論文は、比例極限 $LP/N \in \Theta(1)$ の下での非線形深層MLPにおけるベイズ推論を分析するための、初の数学的に厳密なフレームワークを提供している。これは、無限幅カーネル手法と深層線形ネットワークの間の溝を埋めるものであり、任意の温度（$\beta$）において滑らかな活性化関数およびReLU活性化関数の両方をカバーしている。

2. 既存の結果の回収と一般化

• 深層線形ネットワーク: ドリフト項が消失する特殊なケースとして、深層線形ネットワーク（分配関数がMeijer-G積分を含む）の正確な結果を回収している。
• 無限幅極限: $N \gg LP$ の極限において、ベイズ推論が、Neural ODEから導出される決定論的な特徴マップを持つデータ非依存のカーネル手法に帰着するという結果を回収している。
• 弱非線形ネットワーク: 物理学の文献（例：[HZ24]）における弱非線形ネットワークに関する摂動的な知見を一般化しているが、それらは未証明の簡略化に依存しているのに対し、本研究は完全な数学的厳密性を持っている。

3. ベイズ証拠の摂動展開

著者らは、深さを増すこと（$\tau = LP/N$ を増やすこと）がベイズモデル証拠（周辺尤度）を改善するかどうかの単純な基準を導出している。

• 結論 2.1: 対数証拠の一次変化は、線形ネットワーク項と非線形カーネル項の和である。
• 結論 3.1 & 3.2: 特定のデータ分布（正規化され、ほぼ直交している入力）に対して、著者らは深さが有益となる条件を明示的な不等式として導出している。
  • 滑らかな活性化関数の場合、データ統計量と活性化関数の微分が、信号対雑音比（SNR）と活性化関数の曲率を含む特定の不等式を満たすとき、深さは有益となる。
  • ReLU活性化関数の場合、ラベルの分散がReLUの非対称性（$c_+ - c_-$）に関連する閾値を超える場合に限り、深さは有益となる。
• 相転移: これらの結果は、非線形性がデータ構造に対して「強すぎない」場合にのみ深さが有益となるという、相転移の存在を示唆している。

4. データ依存カーネル手法への等価性

• 結論 2.2: $\tau$ の一次のオーダーにおいて、テストデータに対するベイズ予測事後分布は、データ依存の特徴マップを持つカーネル手法と等価である。
• 無限幅極限では特徴マップは固定されているが、ここでの有効な特徴マップには、訓練データのラベルとテスト入力の直交成分に比例する補正項が含まれる。これは、「特徴学習」のベイズ的な意味における理論的正当性を与えている。すなわち、ネットワークは、事前特徴マップとは異なる、訓練セットに適応した特徴マップを学習している。

意義と主張

著者らは、ヒューリスティックな仮定に依存しがちな物理学的な導出に代わる、数学的に厳密な代替案を提示している。NSDEと摂動展開を用いることで、彼らは以下を実現している：

1. 深さの役割の明確化: 深さが普遍的に有益であるわけではなく、その有用性は、データを生成するプロセス、活性化関数、および温度の相互作用に依存することを実証した。
2. カーネル学習と特徴学習の統一: 深層の非線形レジームにおいても、予測事後分布はカーネル手法として見なせるが、その特徴マップは訓練データによって動的に調整されることを示した。
3. 分析ツールの提供: これまで扱いにくいと考えられていた、非線形ネットワーク（幅と深さが有限だが大きい場合）におけるベイズ統計（証拠、事後平均、分散）の計算を可能にするフレームワークを構築した。

本論文は、任意の $\tau$ に対する一般的なケースを解決したり、現実世界のデータセットを用いた実験的検証を提供したりすることを目的としているのではない。むしろ、確率微分方程式と比例スケーリング極限の観点から、深層ネットワークにおけるベイズ推論のメカニズムを理解するための理論的基礎を確立することを目指している。