arXiv:2506.06489cs.LGstat.ML2025-06NeurIPS被引 13

提出交替梯度流模型,解释两层网络如何逐步学习特征。

Alternating Gradient Flows: A Theory of Feature Learning in Two-layer Neural Networks

  • 将特征学习拆解为激活新神经元与优化已有神经元的交替过程。
  • 预测了损失下降的时机与幅度,与多种网络架构实验结果一致。
  • 适用于线性网络、注意力模型及模加任务,揭示傅里叶特征学习顺序。

神经网络学习到的特征及其学习机制仍是未解之谜。本文提出交替梯度流(AGF),一种描述小初始化下两层网络特征学习动态的算法框架。已有研究显示,该场景中梯度流呈现阶梯状损失曲线:在神经元缓慢对齐有用方向的平台期后,出现快速增长的突降。AGF将此行为近似为两步交替过程:对休眠神经元最大化效用函数,对活跃神经元最小化代价函数。初始时所有神经元休眠,每轮激活一个新神经元,触发特征获取与损失下降。AGF量化了这些突降的顺序、时机与幅度,在多个常用架构中与实验匹配。我们证明AGF统一并拓展了全连接线性网络与仅注意力线性变换器中的鞍点到鞍点分析,分别对应奇异模式与主成分学习。在对角线性网络中,我们证明当初始化趋近零时AGF收敛至梯度流。将AGF应用于训练模加任务的二次网络,首次完整刻画训练动态,揭示网络按系数幅度递减顺序学习傅里叶特征。总体而言,AGF为理解神经网络特征学习提供了重要进展。

原文摘要 · Abstract (English)

What features neural networks learn, and how, remains an open question. In this paper, we introduce Alternating Gradient Flows (AGF), an algorithmic framework that describes the dynamics of feature learning in two-layer networks trained from small initialization. Prior works have shown that gradient flow in this regime exhibits a staircase-like loss curve, alternating between plateaus where neurons slowly align to useful directions and sharp drops where neurons rapidly grow in norm. AGF approximates this behavior as an alternating two-step process: maximizing a utility function over dormant neurons and minimizing a cost function over active ones. AGF begins with all neurons dormant. At each iteration, a dormant neuron activates, triggering the acquisition of a feature and a drop in the loss. AGF quantifies the order, timing, and magnitude of these drops, matching experiments across several commonly studied architectures. We show that AGF unifies and extends existing saddle-to-saddle analyses in fully connected linear networks and attention-only linear transformers, where the learned features are singular modes and principal components, respectively. In diagonal linear networks, we prove AGF converges to gradient flow in the limit of vanishing initialization. Applying AGF to quadratic networks trained to perform modular addition, we give the first complete characterization of the training dynamics, revealing that networks learn Fourier features in decreasing order of coefficient magnitude. Altogether, AGF offers a promising step towards understanding feature learning in neural networks.

神经网络特征学习梯度流两层网络

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。