Correlation flow governs learning at criticality

技術要約：臨界点における相関フローが学習を支配する

問題提起
深層ニューラルネットワークの初期化は、情報の伝播および勾配の伝播が可能かどうかを決定する上で極めて重要である。経験的な成功がネットワーク規模の拡大を推進してきた一方で、初期化の特性と学習ダイナミクスを結びつける統一的な理論は依然として未解明である。具体的には、無限に広く、かつ無限に深いネットワークの極限において、相関の伝播、エンドツーエンドのヤコビアンのスペクトル、およびニューラル・タンジェント・カーネル（NTK）の相互作用が、初期化の選択（ガウス分布か直交分布か）によってどのように条件付けられるのかという理解が不足している。先行研究では、「臨界性」（重みとバイアスの分散が特定の平面上にある点）によって、相関が指数関数的ではなく代数的に減衰することが確立されているが、この相関フローと結果としての学習ダイナミクス（NTKによって制御される）との正確な関係は、まだ十分に特徴付けられていない。さらに、この極限において有限サイズ補正を制御する上での直交初期化の役割については、厳密な理論的根拠が必要である。

手法
著者らは、幅 $N$ を無限大とした後に深さ $L$ を無限大とする「逐次極限（sequential limit）」（$N \to \infty$ を行った後に $L \to \infty$ とする）における全結合ニューラルネットワークを分析するために、平均場理論とランダム行列理論を組み合わせて用いている。

1. 理論的枠組み:

   • 平均場領域: 著者らは、深さが増すにつれて前活性化（pre-activation）の共分散と相関がどのように変化するかを示す決定論的な再帰関係を分析している。彼らは、順序領域（相関 $\to 1$）、カオス領域（相関 $\to$ 定数 $<1$）、および新たに特定された分散消失フェーズ（分散 $\to 0$ となるフェーズ）の3つの領域を特定している。
   • 臨界点: 彼らは、消失分散フェーズの境界と、「カオスの縁（edge of chaos）」のライン（$\sigma_w = 1/\phi_1$ かつ $\sigma_b = 0$、ここで $\phi_1$ は活性化関数の第1次テイラー係数）の交点に臨界点を位置づけている。
   • ヤコビアン解析: 自由確率論を用いて、層ごとの、およびエンドツーエンドのヤコビアンの特異値スペクトルを分析している。彼らは、ガウス初期化と直交初期化の両方において、ヤコビアン・スペクトルの平均と分散のスケール則を導出している。
   • NTKの導出: 彼らは、部分的なヤコビアンによって引き戻された共分散カーネルの和としてNTKを表現している。臨界点におけるこれらの項の漸近挙動を分析することで、NTKと固定点相関との関係を導出している。

2. 有限サイズ解析:

   • 著者らは、経験的カーネルに対する主要な一次の有限サイズ補正（$O(L/N)$）を導出している。彼らは、これらのゆらぎを「ウィック欠陥（Wick defects）」（前活性化の非ガウス性に由来するもの）と「重みの非直交性」の効果に分解している。
   • ガウス初期化と直交初期化の下で、これらの誤差がどのように蓄積するかを比較している。

3. 数値的検証:

   • 理論的な予測は、サイン関数（および一部のtanh）を用いた全結合ネットワークを用いて検証されている。実験では、多様体伝播、相関フロー、ヤコビアン・スペクトル、およびNTKのコンディショニングを、様々な深さと幅にわたって追跡している。

主要な貢献と結果

1. 相関伝播には臨界性が不可欠である:
   著者らは、非自明な相関の無限深への伝播は、臨界点においてのみ可能であることを証明している。他のすべての領域では、相関は1に崩壊する（順序領域）、入力とは無関係な定数に収束する（カオス領域）、あるいは分散が指数関数的に消失して数値的な崩壊を引き起こす。臨界点においては、相関は代数的に収束し、入力依存性を保持する。

2. エンドツーエンド・ヤコビアンの代数的消失:
   臨界性が「力学的等長性（dynamical isometry）」（ヤコビアンのスペクトルが1の付近に集中すること）を可能にするという期待に反して、著者らは、臨界点においてはエンドツーエンド・ヤコビアンが深さに対して代数的に消失する（$\Xi \propto 1/L$）ことを証明している。これは、直交初期化を用いても発生する。したがって、逐次極限における臨界点では、力学的等長性は存在しない。

3. NTKと出力相関の等価性:
   本論文の中心的理論的結果は、無限深の極限における臨界点において、NTKが正確に固定点相関行列に比例する（具体的には、$\Theta_L \to \frac{1}{A\sigma_w^2} \rho_\infty$）ことを証明したことである。これにより、情報伝播（相関）と学習ダイナミクス（NTKスペクトル）の間の直接的なつながりが確立された。この等価性は、ヤコビアンが消失している場合でも、NTKの定義で使用されるスケーリングされたオーバーラップ・ヤコビアンによって相殺されるため、成立する。

4. 直交初期化の役割:
   直交初期化は力学的等長性を回復させることはできないが、有限サイズ効果を制御する上で極めて重要な役割を果たす。

   • ガウス初期化の下では、相関とNTKにおける有限サイズのゆらぎは $O(\sqrt{L/N})$ でスケールし、深さとともに蓄積して、不良条件のスペクトルをもたらす。
   • 直交初期化の下では、重みの非直交性に由来する主要なゆらぎが抑制される。残るゆらぎ（ウィック欠陥）は $O(1/\sqrt{N})$ でスケールし、深さに依存しない。これにより、逐次極限の予測がより深いネットワークにおいても成立するようになる。

5. 活性化関数の制約:
   情報が自明な固定点に崩壊することなく伝播するためには、活性化関数が $\phi_2 = 0$（すなわち、最低次において奇関数であること）を満たさなければならないことを著者らは示している。もし $\phi_2 \neq 0$ であれば、初期化に関わらず相関は崩壊する。

意義と主張
本論文は、深層ネットワークにおける情報伝播、勾配制御、および学習ダイナミクスを結びつける統一的な枠組みを提供すると主張している。その主な意義は以下の通りである：

• 相関と学習の統一: NTKスペクトル（学習の速度と方向を決定するもの）が、臨界点における入力の相関構造によって漸近的に決定されることを示している。これは、ネットワークを通じて入力の類似性を保持することが、なぜ学習に不可欠であるかを説明している。
• 臨界初期化の正当化: 臨界点での初期化が、入力依存の相関を保持しつつ、明確な（ただし消失する）ヤコビアン構造を維持できる唯一の選択肢であることを示すことで、臨界初期化に対して厳密な理論的正当性を与えている。
• 直交初期化の明確化: 直交初期化が、深いネットワークにおいてNTKの予測可能性を破壊する制御不能な有限サイズゆらぎの増大を抑制するための必要条件であることを示すことで、「凍結された（線形）」領域と「特徴学習」領域の間の緊張関係を解決している。
• 実用的含意: これらの結果は、非常に深いネットワーク（例：SIRENs、暗黙的ニューラル表現）において、臨界初期化と直交重みの組み合わせが、安定した学習ダイナミクスと制御されたスペクトル・バイアスを確保するための原理的なアプローチであることを示唆しており、ネットワークが入力データセットの相関構造の主モードに沿って学習することを保証する。

著者らは、これらの結果が逐次極限および特定の活性化関数の仮定（有界、微分可能、$\phi_2=0$）に依存していることを注記しており、これらを共同極限（$L/N \to c$）や非有界な活性化関数（ReLUなど）へ拡張することは、今後の課題であるとしている。