Feature Learning Dynamics in Infinite-Depth Neural Networks

深度神经网络在实践中取得了显著成功，然而对于特征在训练过程中如何演化的机制性理解仍不完整，尤其是在大深度极限下。对于采用 depth-$\mu$P 缩放的 ResNets，先前的工作将层索引 $\ell$ 视为连续时间 $t_\ell = \ell/L$，从而得到训练动力学的 SDE 描述。一个关键的未解决问题是，反向传播通过其转置 $W_\ell^\top$ 重用每个前向权重矩阵 $W_\ell$，从而在前向特征和后向梯度之间产生相关性，其行为和在特征学习中的作用仍不清楚。我们研究了在 depth-$\mu$P 下单层 ResNets 中的这种权重重用前向-后向耦合。利用条件高斯表示，我们在取任何网络极限之前，明确地将权重重用引起的耦合项与解耦的高斯涨落分离开来。在初始化时，我们证明该耦合是一种有限宽度效应，并以 $O(n^{-1})$ 的速率消失，且在整个深度上一致。然而，在训练过程中，SGD 诱导出一个非平凡的前向-后向相关项，该项在无限宽度极限下仍然存在。关键的深度效应是，在 depth-$\mu$P 缩放下，这个存留项在深度上是高阶的，并且随着 $L\to\infty$，其在各层上的累积贡献变得可以忽略不计。这种深度诱导的抑制催生了神经特征动力学（NFD），这是一个具有解耦后向权重的前向-后向 SDE 系统，它保留了训练过程中生成的特征-梯度协方差结构。在非退化假设下，我们证明了有限网络训练动力学收敛到其 NFD 极限，且具有 $O(L^{-1})$ 的深度离散化误差，而重用-权重耦合项具有更快的 $O(L^{-2})$ 衰减。这些结果为 depth-$\mu$P 下单层 ResNets 的特征学习动力学提供了一个严格的无限深度极限。