Global Convergence and Rich Feature Learning in $L$-Layer Infinite-Width Neural Networks under $μ$P Parametrization

作者: Zixiang Chen, Greg Yang, Qingyue Zhao, Quanquan Gu

发布时间: 2025-03-13

来源: arxiv

研究方向: 深度学习理论，神经网络优化，特征学习

主要内容

该论文研究了在无限宽度L层神经网络中，使用张量程序（TP）框架下的训练动力学。主要关注在最大更新参数化（µP）和轻微的激活函数条件下，随机梯度下降（SGD）如何使网络学习与初始值显著不同的线性独立特征。这些丰富的特征空间捕捉了相关数据信息，并确保训练过程的任何收敛点都是全局最小值。

主要贡献

1. 证明了在µP参数化下，多层感知器（MLPs）可以学习捕获任务相关信息的线性独立特征。

2. 通过分析相邻层之间的特征交互和高斯随机变量的性质，提供了对深度表示学习的新见解。

3. 通过在真实世界数据集上的实验验证了理论发现，证明了µP的独特能力，在同时实现有意义的特征学习的同时，保持特征丰富性。

研究方法

1. 张量程序（TP）框架

2. 随机梯度下降（SGD）

3. 分析高斯随机变量和特征交互

4. 谱分析

5. 输入维度缩放

实验结果

实验结果表明，与标准参数化、神经网络切线核（NTK）参数化和可积参数化相比，µP参数化能够同时实现有意义的特征学习和特征丰富性。实验结果表明，随着网络宽度的增加，非零特征值持续存在，这支持了µP的这种独特能力。

未来工作

将理论框架扩展到Transformer架构，特别是注意力机制，以了解现代语言模型中的特征学习。分析结构不变量，以提供关于收敛率（不仅仅是全局收敛）的新视角，可能为深度学习中的优化策略提供信息。研究特征非退化性如何影响泛化界限，可能为理解深度神经网络泛化属性提供更深入的理论基础。探索µP与更复杂的训练范式（如微调和自监督学习）的交互作用，可以进一步增强对实际设置中深度网络训练动力学的理解。