Kernel Renormalization in Bayesian Deep Neural Networks: the Equivalent Wishart Ansatz in the Proportional Regime

タイトル: ベイジアン深層ニューラルネットワークにおけるカーネル再正規化：比例領域における等価ウィシャール Ansatz

問題提起
深層ニューラルネットワーク（DNN）は現代の機械学習の中核をなすものであるが、その理論的理解は依然として困難であり、特に多層で非線形なアーキテクチャにおいて顕著である。「怠惰な」無限幅極限（ニューラルネットワークガウス過程、NNGP）は広幅のネットワークに対する厳密な記述を提供するが、DNN がカーネル法よりも優れた性能を発揮する原因となる有限幅の効果を捉えることはできない。逆に、特徴が著しく適応する「豊か」な学習領域では、しばしば $P \times P$ のカーネル行列のような高次元の順序パラメータが必要となり、計算的に扱いが困難となる。

本論文が扱う具体的なギャップは、比例領域におけるベイジアン DNN に対する低次元の有効理論の欠如である。ここで比例領域とは、訓練サンプル数 $P$ とネットワーク幅 $N$ が無限大に比例してスケールする領域（$P, N \to \infty$ かつ $\alpha = P/N$ を固定）を指す。この領域では有限幅の揺らぎは無視できず、既存の理論は深層ネットワークに対しては失敗する摂動的補正に依存するか、あるいは次元削減を提供しない結合した高次元方程式を解くことを要求する。

手法：等価ウィシャール Ansatz（EWA）
著者らは、**等価ウィシャール Ansatz（EWA）*に基づく近似ながら効果的な理論を提案する。中核的なアイデアは、深層非線形ネットワークにおける階層的経験カーネルの確率的揺らぎを、深層線形*ネットワークで見られる厳密な振る舞いに類似して、ウィシャール分布に従うものとしてモデル化することである。

1. カーネル階層: ネットワーク出力の事前分布は、各層 $\ell$ における経験カーネル $K^{(\ell)}_E$ の階層によって記述される。比例極限において、著者らは層 $\ell-1$ のカーネルが与えられたときの層 $\ell$ のカーネルの条件付き分布がウィシャール分布に従うと仮定する：
   $$K^{(\ell)}_E | K^{(\ell-1)}_E \sim \mathcal{W}_P\left( \frac{\Theta^{L-\ell+1}(K^{(\ell-1)}_E)}{N_\ell}, N_\ell \right)$$
   ここで $\Theta$ は NNGP カーネル関数、$N_\ell$ は幅である。

2. 次元削減: ウィシャール行列の性質（特に二次形式が $\chi^2$ 分布に従うという縮約性）を利用することで、著者らはパラメータ空間にわたる複雑な積分を、$L$ 個のスカラー順序パラメータ $q_\ell$（CNN の場合は行列 $q_\ell$）にわたる低次元積分へと還元する。

   • 全結合（FC）ネットワークの場合、有効作用はスカラー積 $Q = \prod q_\ell$ に依存する。
   • 畳み込みニューラルネットワーク（CNN）の場合、積層型等価ウィシャール Ansatzが導入され、パッチ間の相関を捉える低次元の行列順序パラメータ $Q$ が導かれる。

3. 有効作用: 分配関数は大偏差原理を通じて表現され、これらの順序パラメータとデータセットのみに依存する有効作用 $S(q)$ が得られる。再正規化されたカーネルは $K^{(R)}_Q = Q \Theta^L(C)$ で与えられ、ここで $C$ は入力グラム行列である。

4. 非ゼロ平均の扱い: 著者らは、比例領域において、ReLU のような非ゼロ平均を持つ活性化関数に起因する中心ウィシャール分布と非中心ウィシャール分布の区別は漸近的に無視できることを示す。したがって、平均ゼロの EWA は有効な記述として依然として妥当である。

主要な貢献

• 非線形 DNN に対する低次元理論: 本論文は、比例領域におけるベイジアン多層パーセプトロン（MLP）および CNN の最初の低次元近似記述を提供し、$P \times P$ の自己無撞着方程式に頼ることなく、非摂動的な有限幅効果を捉える。
• 等価ウィシャール Ansatz: 階層的経験カーネルの支配的な確率的揺らぎを捉えるメカニズムとして EWA を導入し、問題を $L$ 個（または小さな行列）の順序パラメータの決定へと還元する。
• CNN への拡張: 枠組みを畳み込みアーキテクチャに拡張し、空間的に相関したデータにおける有限幅効果を説明する階層的局所カーネル再正規化メカニズムを同定する。
• 比例領域における検証: 理論は、実世界データセット（MNIST、CIFAR-10）および合成データに対する広範な数値実験（MCMC サンプリング）を通じて検証され、深さ $L \sim 10$ までおよび様々な活性化関数（Erf、ReLU）を網羅している。

結果

• 定量的一致: EWA による訓練損失と汎化誤差の予測は、深さ $L \sim O(10)$ までのネットワークおよび様々なデータセットサイズにおいて、数値サンプリング実験と「全体的に非常に良好な一致」を示す。この理論は、広幅極限の NNGP が失敗する有限幅ネットワークの振る舞いを成功裡に追跡する。
• 系統的な偏差: 深さ $L$ と負荷 $\alpha$ の両方が大きい場合に、2 種類の偏差が観測される。
  1. 漸進的ドリフト: 高 $\alpha$ におけるガウスデータ実験では、EWA の誤差は深さとともに滑らかに増加する。これは、レート関数の幅における累積誤差または無視された方向的不均一性によるものと考えられる。
  2. 準安定性: CIFAR-10 上の ReLU ネットワークにおいて $L > 5$ かつ $\alpha \approx 2.5$ の場合、MCMC シミュレーションにおいて訓練誤差とテスト誤差の急激な急落が観測される。この現象は事後分布における準安定遷移として解釈されるが、滑らかな EWA 予測では捉えられておらず、極端な領域における Ansatz の破綻を示唆している。
• 平均場 vs 標準パラメータ化: この理論は、標準パラメータ化（SP）と平均場（$\mu$P）の両方において汎化性能を成功裡に予測する。豊か学習領域（$\mu$P）において、汎化の改善は主に予測器バイアスの強い適応ではなく、予測器分散の抑制によって駆動されていることが判明し、EWA の再正規化カーネルの描像と一致する。

意義と主張
著者らは、EWA が「怠惰な」無限幅極限と複雑で高次元の特徴学習領域の間のギャップを埋める扱いやすい有効理論を提供すると主張する。有限幅学習の記述をいくつかのスカラー（または低ランク行列）順序パラメータに還元することで、本論文は、データサイズと幅の相対的なスケーリング（$\alpha = P/N$）が、従来の「怠惰 vs 豊か」という二項対立よりも、有限幅学習にとってより基本的な組織原理であると論じる。

この研究は、中程度の深さと負荷において、非線形 DNN における特徴の複雑な適応は、カーネルの単純な大域的（CNN の場合は局所的）再正規化によって効果的に捉えられることを示唆している。これは特徴学習の理解のための基準を提供し、無限幅予測と有限幅の現象論の間の大きな隔たりは、この方向性均質な再正規化によって大部分が捉えられていることを示している。本論文は、EWA は近似であるが、多様なアーキテクチャとデータセットにわたる経験的学習曲線の予測におけるその成功は、比例領域におけるベイジアン深層学習の分析のための強力なツールであることを確立していると結論づけている。