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

TL;DR

This paper introduces Alternating Gradient Flows (AGF), a new framework that models how two-layer neural networks learn features during training, capturing dynamics like neuron activation and feature acquisition.

Contribution

AGF provides a novel theoretical framework that approximates feature learning dynamics in two-layer networks, extending existing analyses and offering new insights into training behavior.

Findings

AGF matches experimental observations of loss plateaus and drops.

In linear networks, AGF converges to gradient flow as initialization vanishes.

In quadratic networks, AGF reveals Fourier features are learned in decreasing order.

Abstract

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.