Abstract
Despite deep neural networks' powerful representation learning capabilities, theoretical understanding of how networks can simultaneously achieve meaningful feature learning and global convergence remains elusive. Existing approaches like the neural tangent kernel (NTK) are limited because features stay close to their initialization in this parametrization, leaving open questions about feature properties during substantial evolution. In this paper, we investigate the training dynamics of infinitely wide, $L$-layer neural networks using the tensor program (TP) framework. Specifically, we show that, when trained with stochastic gradient descent (SGD) under the Maximal Update parametrization ($\mu$P) and mild conditions on the activation function, SGD enables these networks to learn linearly independent features that substantially deviate from their initial values. This rich feature space captures relevant data information and ensures that any convergent point of the training process is a global minimum. Our analysis leverages both the interactions among features across layers and the properties of Gaussian random variables, providing new insights into deep representation learning. We further validate our theoretical findings through experiments on real-world datasets.
Abstract (translated)
尽管深度神经网络具有强大的表示学习能力，但关于如何同时实现有意义的特征学习和全局收敛的理论理解仍然不明确。现有的方法，如神经切片核（NTK），由于在这种参数化中特征保持接近初始化状态而受到限制，从而留下了有关特征在显著演变过程中属性的问题。在这篇论文中，我们使用张量程序（TP）框架研究了无限宽、$L$层神经网络的训练动力学。具体来说，我们展示了当使用随机梯度下降法（SGD）和激活函数的一些温和条件下，在最大更新参数化（$\mu$P）下训练时，SGD使这些网络能够学习与初始值显著不同的线性独立特征。这一丰富的特征空间捕捉到了相关的数据信息，并确保了训练过程中的任何收敛点都是全局最小值。我们的分析利用了层间特征之间的相互作用以及高斯随机变量的属性，为深度表示学习提供了新的见解。我们进一步通过在现实世界数据集上的实验验证了理论发现。
URL
https://arxiv.org/abs/2503.09565
