揭示深度神经网络特征学习的普适规律,发现权重与梯度关系随网络深度变化。
On the Neural Feature Ansatz for Deep Neural Networks
- 提出神经特征假设(NFA),描述深层网络第一层权重与输入梯度的关系
- 证明网络深度为L时,比例指数α=1/L,且在权重衰减下对非平衡初始化也成立
- 通过反例说明非线性结构不满足该假设,验证理论边界
理解特征学习是建立深度神经网络数学基础的重要开放问题。神经特征假设(NFA)认为:训练后,深层网络第一层权重的格拉姆矩阵与输入梯度外积的平均值(AGOP)的幂次α>0成正比。在梯度流动态与平衡初始化假设下,已证明两层线性网络中NFA在训练全程成立,且α=1/2(Radhakrishnan et al., 2024)。本文将此结果扩展至L≥2层网络,证明NFA在α=1/L时成立,揭示了NFA的深度依赖性。此外,我们证明在非平衡初始化下,若引入权重衰减,NFA仍可渐近成立。同时,我们提供了反例,表明某些带非线性激活的网络架构即使完美拟合训练数据,也不满足NFA。我们通过多种优化算法、权重衰减率和初始化方案的数值实验,全面验证了理论结果。
原文摘要 · Abstract (English)
Understanding feature learning is an important open question in establishing a mathematical foundation for deep neural networks. The Neural Feature Ansatz (NFA) states that after training, the Gram matrix of the first-layer weights of a deep neural network is proportional to some power $α>0$ of the average gradient outer product (AGOP) of this network with respect to its inputs. Assuming gradient flow dynamics with balanced weight initialization, the NFA was proven to hold throughout training for two-layer linear networks with exponent $α= 1/2$ (Radhakrishnan et al., 2024). We extend this result to networks with $L \geq 2$ layers, showing that the NFA holds with exponent $α= 1/L$, thus demonstrating a depth dependency of the NFA. Furthermore, we prove that for unbalanced initialization, the NFA holds asymptotically through training if weight decay is applied. We also provide counterexamples showing that the NFA does not hold for some network architectures with nonlinear activations, even when these networks fit arbitrarily well the training data. We thoroughly validate our theoretical results through numerical experiments across a variety of optimization algorithms, weight decay rates and initialization schemes.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。