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

Feng Chen (Microsoft AI / Stanford University) · Surya Ganguli (Stanford) · Dhruva Karkada (University of California, Berkeley) · James Simon (University of California Berkeley) · Michael Deweese (University of California-Berkeley) · Daniel Kunin (University of California, Berkeley) · Giovanni Luca Marchetti (KTH Royal Institute of Technology, Stockholm, Sweden) · Nina Miolane (University of California, Santa Barbara)

active neuronsalternating gradient flowsattention-only linear transformersdiagonal linear networksdormant neuronsfeature learningfourier featuresfully connected linear networksloss minimization principal componentssaddle-to-saddle analysissingular modesstaircase-like loss curvetraining dynamicstwo-layer networksutility function

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, corresponding to an initialization at the origin. 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.