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.

翻译：神经网络学习何种特征及其学习机制，至今仍是一个开放性问题。本文提出交替梯度流（AGF）这一算法框架，用于描述从小初始化训练的两层网络中特征学习的动态过程。先前研究表明，该机制下的梯度流呈现出阶梯状损失曲线：在平台期，神经元缓慢对齐至有用方向；在陡降期，神经元范数快速增长。AGF将这种行为近似为交替的两步过程：对休眠神经元最大化效用函数，对活跃神经元最小化代价函数。AGF初始时所有神经元均处于休眠状态。在每次迭代中，一个休眠神经元被激活，触发特征获取与损失下降。AGF量化了这些下降的次序、时机与幅度，其预测与多种常用架构的实验结果相符。我们证明AGF统一并扩展了全连接线性网络和纯注意力线性Transformer中现有的鞍点间分析框架——在这两类网络中，习得的特征分别为奇异模态和主成分。对于对角线性网络，我们证明在初始化趋近于零的极限情况下AGF收敛于梯度流。将AGF应用于执行模加运算的二次网络训练时，我们首次完整刻画了训练动态过程，揭示出网络按系数幅值降序学习傅里叶特征的规律。总体而言，AGF为理解神经网络的特征学习机制迈出了重要一步。