Gradient descent dynamics explain feature learning in simple two neuron models
Gradient Descent on Two ReLU Neurons: Global Landscape and Bifurcation Dynamics
Machine Learning
Summary
Understanding how simple neural networks learn features is a key puzzle in AI. The authors study a very basic setup with two simple neurons and show that learning happens in phases: first, the neurons come together in an average direction, then they split and specialize to different tasks. They prove the math behind these steps and explain how the learning landscape looks, showing that only the best and saddle points appear with both neurons active. This helps explain why even simple neural networks break symmetry and learn specialized features over time.
What this means in practice
- •For machine learning engineers: Guide the design of training setups by understanding dynamics in simple two-neuron models to better handle feature specialization.
- •For automated ai system developers: Inform development of algorithms that rely on small neural modules by clarifying how training dynamics unfold in minimal settings.
A theory result. No direct application yet.
Authors
Binghua Li, Mengzhe Li, Denny Wu, Tianhao Wang
Abstract
Understanding how gradient descent learns features remains a central challenge in neural network theory. We study this question in a minimal multi-index setting that already exhibits rich multi-phase gradient dynamics: a well-specified two-neuron ReLU teacher-student model under isotropic Gaussian inputs, trained by population gradient flow from small initialization. We show that the dynamics are organized by two directions: the "easy" bisector direction, which carries the leading signal, and the "hard" splitting direction, which governs specialization. To characterize the loss structure and learning behavior, we analyze the population squared-loss landscape and show that every critical point with both student weights nonzero is either a global minimum or a saddle. We then track the gradient-flow trajectory, showing that student neurons first collapse toward a bisector saddle before escaping and specializing to teacher neurons. Our results provide both a landscape and dynamics account of the multi-phase symmetry-breaking behavior that arises in a simple multi-index model and under standard algorithmic and architectural choices.