Sharp feature-learning transitions and Bayes-optimal neural scaling laws in extensive-width networks

技術的サマリー：広幅ネットワークにおける鋭い特徴学習遷移とベイズ最適なニューラルスケーリング則

問題設定
本論文は、ノイズの混入したクエリから階層的特徴を持つ一隠れ層教師ニューラルネットワークを学習する情報理論的限界を調査するものであり、特に小規模な学生モデルへの知識転移の文脈において研究されている。本研究は、教師の幅 $k$ が入力次元 $d$ に比例してスケーリングする（$k = \Theta(d)$）「広幅（extensive-width）領域」で動作する。この設定は、よく研究されているマルチインデックス領域（$k = \Theta(1)$）とニューラルタンジェントカーネル（NTK）領域（$k \gg d$）の間に位置し、実用的なネットワークの「大規模だが有限な」幅を捉えている。

教師ネットワークは、入力 $x_\mu \sim \mathcal{N}(0, I_d)$ から、加法性ガウスノイズを伴うラベル $y_\mu = v^\top \sigma(W_0 x_\mu)$ を生成する。読み出しベクトル $v$ は階層的構造（例えば、べき乗則減衰 $v_i \propto i^{-\beta}$）に従い、特徴の重要性が幾何学的に減衰することを意味する。著者は 2 種類の学習者を分析する：

1. ベイズ的学生：データ生成メカニズム（活性化関数 $\sigma$、ノイズ分散 $\Delta$、読み出し $v$）は知っているが、特定の教師重み $W_0$ は知らない。事後平均を通じて予測を行い、ベイズ最適誤差（$\varepsilon_{BO}$）を達成する。
2. SGD 学生：ADAM で訓練された幅 $k_s \leq k$ の一隠れ層ネットワーク。

主な問いは以下の通りである：(Q1) 固定されたデータセットが与えられたとき、どの教師特徴が統計的に復元可能か？また (Q2) 復元可能な特徴のサブセットは、データサイズ $n$ に対する $\varearepsilon_{BO}$ のスケーリング則をどのように制約するか？

手法
著者は、統計物理学で標準的な手法であるヒューリスティックな留め置（cavity）脱結合論法を用いて、学習過程の漸近的に鋭い特徴付けを導出する。

• 脱結合：この手法は、他の全ての特徴の寄与を有効ノイズ項として扱うことで、単一の特徴 $w_i$ の寄与を分離する。
• ガウス仮定：著者は、「空洞場（他の特徴からの有効ノイズ）」が独立したガウス変数として振る舞うと仮定する。これにより、多特徴推論問題は、独立したスカラー一般化線形モデル（GLM）の系列に還元される。
• 固定点系：真の特徴方向と事後推定との重なりを分析することで、著者は固定点方程式の閉じた系を導出する。これらの方程式は、特徴 $i$ の重なり $q_i$ を、信号対雑音比（SNR）および大域的汎化誤差 $\varepsilon$ と関連付ける。
• 検証：理論的予測は、様々なパラメータ領域におけるハミルトニアンモンテカルロ（HMC）シミュレーションを用いて数値的に検証され、中程度の次元（$d \sim 10^2$）においても優れた一致を示している。

主要な結果

1. 鋭い特徴学習遷移の階段：
   分析により、特徴の学習可能性は鋭い位相遷移によって支配されることが明らかになった。特徴 $i$ は、その有効 SNR $snr_i = \frac{(n/d)v_i^2}{\Delta + \varepsilon_{BO}}$ が活性化関数依存の臨界閾値 $\lambda_\sigma$ を超える場合にのみ、復元可能（重なり $q_i > 0$）である。

   • データサイズ $n$ が増加すると、特徴はこの閾値を順に越え、最も強い（最大の $v_i$ を持つ）ものから始まる。
   • これにより、特徴の重なりと対応する汎化誤差の低下において、滑らかな遷移ではなく、不連続なジャンプを伴う「階段」が生じる。

2. 有効幅（$k_c$）と誤差分解：
   与えられたデータ予算において学習可能な特徴の数を、有効幅 $k_c = |\{i : q_i > 0\}|$ と定義する。ベイズ最適誤差は 2 つの項に分解される：
   $$\varepsilon_{BO} = \Theta\left(\frac{k_c d}{n}\right) + \Theta\left(\sum_{j > k_c} v_j^2\right)$$
   最初の項は学習可能な特徴に関する推定誤差（特徴あたり $d/n$ としてスケーリング）を表し、2 番目の項は学習不可能な特徴からの不可避なバイアスである。

3. 統合されたスケーリング則：
   指数 $\beta > 1/2$ のべき乗則読み出しに対して、論文は $\varepsilon_{BO} = \Theta(k_c d / n)$ という関係の下で統合される 2 つの異なるスケーリング領域を導出する：

   • 特徴学習領域（$n \ll k^{2\beta}d$）：ここでは、$k_c \sim (n/d)^{1/(2\beta)}$ となる。誤差は $\varepsilon_{BO} \sim n^{1/(2\beta) - 1}$ としてスケーリングする。この領域では、モデルは学習可能な特徴の数によって制限される。
   • 洗練領域（$n \gg k^{2\beta}d$）：全ての特徴が学習可能となる（$k_c = k$）。誤差は $\varepsilon_{BO} \sim n^{-1}$ としてスケーリングし、古典的なパラメトリック率を回復する。
     交差点 $n^\star = \Theta(k^{2\beta}d)$ は、最も弱い特徴が学習可能になる遷移点を示す。

4. SGD 訓練学生への示唆：
   実験結果は、理論的有効幅 $k_c$ に近い幅 $k_s$ を持つ ADAM 訓練学生が、これら最適スケーリング則を達成することを示している（小さなアルゴリズム的ギャップまで）。

   • もし $k_s < k_c$ なら、学生は容量によって制限され、未学習（underfitting）となる。
   • もし $k_s \gg k_c$ なら、学生は学習不可能な方向に過学習し、性能が劣化する。
   • 最適な学生幅は、特徴の階層に従ってデータサイズとともにスケーリングし、ハイパーパラメータ転送の理論的基盤を提供する。

意義と主張
本論文は、階層的特徴と消滅する第 2 ヘルミート係数（$\mu_2 = 0$）を持つネットワークについて、広幅領域（$k = \Theta(d)$）におけるベイズ最適学習の最初の解析的特徴付けを提供すると主張している。

• 統合：以前は別個であった 2 つの実証的スケーリング領域（特徴学習と洗練）を、有効幅 $k_c$ によって駆動される単一の枠組みに統合する。
• メカニズム：実務で観察される「階段」状の学習ダイナミクスに対するメカニズム的説明を提供し、それらを特徴の復元可能性における逐次的位相遷移に帰着させる。
• 方法的貢献：この仕事は、空洞法を密な読み出しから任意の読み出し構造（疎な構造やべき乗則を含む）へと拡張し、密な読み出しに関する既知の結果を回復しつつ、誤差減衰率が著しく異なる疎な設定へと一般化する。
• 実用的示唆：導出された有効幅 $k_c$ は、データ予算に基づいて学生モデルのサイズを決定するための第一原理的処方箋を提供し、モデル容量は教師特徴の総数ではなく、学習可能な特徴の数に応じてスケーリングすべきであることを示唆する。

限界
著者は、導出が空洞場のガウス性と重なりの一様集中という証明されていない仮定に依存していることを認めている。さらに、分析は第 2 ヘルミート係数が消滅する（$\mu_2 = 0$）活性化関数に限定されており、これには奇数関数の活性化が含まれるが、ReLU のような他の関数は含まれない（ただし、論文は ReLU は他の研究における補完領域 $\mu_2 \neq 0$ を通じて解析可能であると指摘している）。留め置還元はヒューリスティックであるが、数値的に検証されている。