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

技术摘要：广宽网络中的锐利特征学习相变与贝叶斯最优神经缩放律

问题设定
本文研究了从含噪查询中学习具有分层特征的单隐层教师神经网络的理论信息极限，具体背景是将知识迁移至更小的学生模型。该研究在广宽（extensive-width）机制下运行，其中教师网络的宽度 $k$ 与输入维度 $d$ 呈线性缩放关系（$k = \Theta(d)$）。这一设定介于研究较多的多索引机制（$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}$），意味着特征的重要性呈几何级数衰减。作者分析了两类学习者：

1. 贝叶斯学生：知晓数据生成机制（激活函数 $\sigma$、噪声方差 $\Delta$、读入向量 $v$），但不知晓具体的教师权重 $W_0$。它通过后验均值进行预测，达到贝叶斯最优误差（$\varepsilon_{BO}$）。
2. SGD 学生：一个宽度为 $k_s \leq k$ 的单隐层网络，使用 ADAM 进行训练。

本文解决的核心问题是：（Q1）给定固定数据集，哪些教师特征在统计上是可恢复的？（Q2）可恢复特征子集如何约束 $\varepsilon_{BO}$ 随数据量 $n$ 变化的缩放律？

方法论
作者采用了一种启发式的留一法（空腔）解耦论证，这是统计物理中的标准技术，用于推导学习过程的渐近锐利刻画。

• 解耦：该方法通过将除单个特征 $w_i$ 外的所有特征贡献视为有效噪声项，从而隔离单个特征 $w_i$ 的贡献。
• 高斯假设：作者假设“空腔场”（来自其他特征的有效噪声）表现为独立的高斯变量。这将多特征推断问题简化为一系列独立的标量广义线性模型（GLM）。
• 不动点系统：通过分析真实特征方向与后验估计之间的重叠，作者推导出一组封闭的不动点方程组。这些方程将特征 $i$ 的重叠 $q_i$ 与信噪比（SNR）及全局泛化误差 $\varepsilon$ 联系起来。
• 验证：理论预测通过哈密顿蒙特卡洛（HMC）模拟在不同参数区间进行了数值验证，即使在中等维度（$d \sim 10^2$）下也显示出极好的一致性。

主要结果

1. 锐利的特征学习相变阶梯：
   分析表明，特征的可学习性受控于一个锐利的相变。特征 $i$ 是可恢复的（重叠 $q_i > 0$）当且仅当其有效信噪比 $snr_i = \frac{(n/d)v_i^2}{\Delta + \varepsilon_{BO}}$ 超过一个依赖于激活函数的临界阈值 $\lambda_\sigma$。

   • 随着数据量 $n$ 的增加，特征按顺序跨越该阈值，从最强特征（$v_i$ 最大）开始。
   • 这导致特征重叠出现不连续的跳跃，并伴随泛化误差的相应下降，形成“阶梯”状，而非平滑过渡。

2. 有效宽度（$k_c$）与误差分解：
   在给定数据预算下，可学习特征的数量被定义为有效宽度 $k_c = |\{i : q_i > 0\}|$。贝叶斯最优误差分解为两项：
   $$\varepsilon_{BO} = \Theta\left(\frac{k_c d}{n}\right) + \Theta\left(\sum_{j > k_c} v_j^2\right)$$
   第一项代表可学习特征上的估计误差（每个特征按 $d/n$ 缩放），而第二项是不可学习特征带来的不可约偏差。

3. 统一的缩放律：
   对于指数 $\beta > 1/2$ 的幂律读入，本文推导出了两种不同的缩放机制，它们在关系式 $\varepsilon_{BO} = \Theta(k_c d / n)$ 下统一：

   • 特征学习机制（$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_s$ 接近理论有效宽度 $k_c$ 的 ADAM 训练学生，能够实现这些最优缩放律（存在微小的算法间隙）。

   • 如果 $k_s < k_c$，学生欠拟合，受限于其容量。
   • 如果 $k_s \gg k_c$，学生过拟合不可学习方向，导致性能下降。
   • 最优学生宽度根据特征分层随数据量缩放，为超参数迁移提供了理论基础。

意义与主张
本文声称首次提供了具有分层特征且二阶埃尔米特系数消失（$\mu_2 = 0$）的网络在广宽机制（$k = \Theta(d)$）下的贝叶斯最优学习的解析刻画。

• 统一性：它将此前截然不同的两种经验缩放机制（特征学习与细化）统一到一个由有效宽度 $k_c$ 驱动的单一框架中。
• 机制：它为实践中观察到的“阶梯”式学习动态提供了机制性解释，将其归因于特征可恢复性的顺序相变。
• 方法论贡献：该工作将空腔法从密集读入扩展到任意读入结构（包括稀疏和幂律），在恢复密集读入已知结果的同时，推广到了误差衰减率显著不同的稀疏设定。
• 实践洞察：推导出的有效宽度 $k_c$ 为基于数据预算调整学生模型规模提供了第一性原理的准则，表明模型容量应根据可学习特征的数量而非教师特征的总数进行缩放。

局限性
作者承认，其推导依赖于未证明的假设，特别是空腔场的高斯性以及重叠的集中性。此外，分析仅限于二阶埃尔米特系数消失（$\mu_2 = 0$）的激活函数，这包括奇函数激活，但排除了 ReLU 等其他激活函数（尽管论文指出，ReLU 可在其他工作中通过互补机制 $\mu_2 \neq 0$ 进行分析）。留一法归约是启发式的，但已通过数值验证。