テクニカルサマリー：予測符号化ネットワークにおける無限の幅と深さの限界について

問題提起

予測符号化（Predictive Coding, PC）は、標準的なバックプロパゲーション（BP）に代わる生物学的に妥当な手法であり、ニューロンが重みを更新する前に局所的な予測誤差（エネルギー関数）を最小化します。近年の再パラメータ化によって、深層PCネットワーク（PCN）の学習の安定性は向上していますが、これらの手法を大規模な幅や深さにスケールアップさせるための理論的根拠は依然として不明です。具体的には、PCが無限の幅および無限の深さの極限において、BPのような安定性と特徴学習能力をスケールアップできるのか、あるいは標準的なPCNはモデルサイズが増大するにつれて不安定性（例：出力の爆発）に陥るのかが分かっていません。さらに、PCの勾配がこれらの極限においてBPの勾配に収束するための条件には、厳密な特性評価が必要です。

手法

著者らは、無限の幅（$N \to \infty$）および無限の深さ（$L \to \infty$）の同時極限におけるPCNの理論的および経験的分析を行っています。

1. 理論的枠組み:

   • 一般的なパラメータ化: 本研究では、線形多層パーセプトロン（MLP）および残差ネットワーク（ResNet）に対する一般的な「abcd」パラメータ化スキームを採用し、事前活性化（$a_\ell$）、重みの初期化分散（$b_\ell$）、学習率（$c$）、およびネットワーク出力（$d$）のスケーリング指数を定義しています。
   • 平衡エネルギー: 線形ネットワークに対して、著者らは「平衡エネルギー」（$F^*$：活動量に関してPCエネルギーを最小化したもの）を利用しています。彼らは、線形ネットワークにおいて、$F^*$ は標準的な平均二乗誤差（MSE）損失の再スケールされたバージョン、すなわち $F^*(\theta) = L(\theta) / s(\theta)$ （ここで $s(\theta)$ は重みに依存する再スケール因子）であることを示しています。
   • デシデラータ（要件）: 分析では、以下の3つの主要な基準（BPの文献から適応）を満たすパラメータ化を追求しています：
     1. 安定性: 層の事前活性化およびネットワークの予測が、$N \to \infty$ の際、消失も爆発もせず、有界であること。
     2. 特徴学習: 学習中（非レイジー学習）に特徴が非自明に進化すること。
     3. 深さの安定性（ResNetの場合）: 隠れ層の活性化が $L \to \infty$ において安定しており、各残差ブロックが特徴の進化に対して $\Theta(1/L)$ の寄与をすること。

2. 導出:

   • 著者らは、上記のデシデラータを満たすために必要な具体的なスケーリング指数を導出しています。
   • 再スケール因子 $s(\theta)$ および平衡エネルギーの勾配の漸近的挙動を分析しています。
   • これらの安定したパラメータ化の下で、幅が成長するにつれてPCの勾配がBPの勾配に収束することを証明しています。その誤差項は、深さと幅の比率（$L/N$）に依存します。

3. 経験的検証:

   • 実験は、線形および非線形アーキテクチャ（MLP、畳み込みニューラルネットワーク（CNN）、Transformerを含む）を用いて行われました。
   • モデルは、トイ・タスクからCIFAR-10、Fashion-MNIST、ImageNet-1kに至るデータセットを用いて訓練されました。
   • 本研究では、異なる幅と深さ（$L/N$）における、PC勾配（活動平衡時に計算）とBP勾配の間のコサイン類似度を測定しました。

主な貢献

1. 安定したパラメータ化の等価性:
   本論文は、PCにおいて安定性と非自明な特徴学習の両方を保証するパラメータ化の集合が、BPのそれと同一であることを証明しています。具体的には、線形MLPおよび残差ネットワークにおいて、「Mean-Field」または「Maximal Update」($\mu$P) パラメータ化が、安定性と特徴学習のデシデラータを満たす唯一の解です。これは、標準的なPCN（PyTorchのデフォルトのような標準的なパラメータ化を使用するもの）は本質的に不安定であり、幅が大きくなると出力が爆発することを意味しています。

2. バックプロパゲーションへの収束:
   著者らは、これらの安定したパラメータ化の下で、ネットワークの幅が深さよりも十分に大きくなる（$L/N \to 0$）とき、PCによって計算される勾配がBPの勾配に収束することを実証しています。

   • 固定された深さにおける線形MLPでは、PC勾配とBP勾配の相対誤差は $O(N^{-1})$ でスケールします。
   • 同時極限における線形残差ネットワークでは、相対誤差は $O(L/N)$ でスケールします。
   • この収束は目的関数と勾配の両方に成立しており、これは、広い極限において、PCが事実上BPを実装していることを意味します。

3. 「高速なサドル・トゥ・サドル（Saddle-to-Saddle）」領域の不安定性:
   本論文は、PCが（深く狭いネットワークにおいて）BPよりも速い「サドル・トゥ・サドル」ダイナミクスを示すという、以前観察された領域が、標準パラメータ化（SP）に依存していることを特定しています。著者らは、この領域が幅と深さの両方に対して不安定であり、安定性のデシデラータに違反していることを示しています。したがって、このスピードアップは、大規模モデルにおける数値的不安定性と引き換えになっています。

4. 経験的な汎用性:
   理論的分析は線形ネットワークに対して導出されたものですが、実験の結果、異なるオプティマイザ（Adam）および損失関数を用いて訓練された非線形アーキテクチャ（CNN、Transformer）においても、高い勾配アライメントが維持されることが示されました。ただし、これはネットワークが十分に広く、かつ活動平衡に達している場合に限られます。

結果

• 理論的: 導出により、出力のスケーリング指数 $d=1/2$ が、BPと同様にPCにおける非レイジーな特徴学習にとって極めて重要であることが確認されました。再スケール因子 $s(\theta)$ は、$N \to \infty$（MLPの場合）または $L/N \to 0$（ResNetの場合）において1に近づき、PCの目的関数がMSE損失に収束することを引き起こします。
• 実験的:
  • 勾配のアライメント: 幅が増加するにつれて、特にResNetやTransformerのように $L/N$ が小さい場合において、PC勾配とBP勾配のコサイン類似度は1.0に近づきます。
  • 非線形モデル: 32層の残差MLP、10層のCNN、および12ブロックのTransformerにおいて、高いアライメントが観察されました。
  • 標準パラメータ化の失敗: 標準的なパラメータ化（SP）を用いたモデルは、幅が増加するにつれて出力の爆発と不十分な勾配アライメントを示し、理論的な不安定性を裏付けました。

意義と主張

本論文は、PCと共にスケール可能なパラメータ化に関する最初の原理的な制約を提供したと主張しています。その主な意義は以下の2点にあります：

1. スケーラビリティの制約: PCを大規模なモデルにスケールさせるためには、BPと同じパラメータ化規則（具体的には $\mu$P）を採用しなければならないことを確立しました。これは、近年の深層PC研究で使用されている「ヒューリスティックな」再パラメータ化が、実行可能であるためにはこれらの理論的限界に適合する必要があることを示唆しています。
2. BPの生物学的妥当性: これらの結果は、脳（おそらく、深さよりもはるかに広いネットワークである）が、局所的な学習規則のみを使用して、バックプロパゲーションのグローバルなクレジット割り当てをどのように実装し得るかというメカニズムを示唆しています。もし脳が $L/N \to 0$ の領域で作動しているのであれば、PCは理論的にBPへと収束でき、非局所的な情報の伝播なしに効率的な学習を可能にします。

著者らは、理論的証明が線形ネットワークと正確な活動収束に依存していることを認め、限界についても謙虚な姿勢を保っています。また、非線形ネットワークにおける活動平衡に達するための計算コストが実用的なボトルネックであることも認めていますが、将来的にアナログハードウェアの実装がこれに対処できる可能性を示唆しています。彼らは、PCがすべての領域においてBPより優れていると主張しているのではなく、適切なパラメータ化の下で、広い極限において両者が数学的に等価になることを主張しています。