When Both Layers Learn: Training Dynamics of Representing Linear Models via ReLU Networks
2026-06-03 • Machine Learning
Machine Learning
AI summaryⓘ
The authors study how gradient descent trains both layers of a simple neural network with ReLU activation to match a linear function, using data from a Gaussian distribution. They analyze why gradient descent can avoid getting stuck on tricky points in the loss landscape and prove it converges efficiently when starting from small random weights. Their work breaks down the training into three phases: aligning the hidden layer to the target, growing the weights while keeping alignment, and fine-tuning for fast convergence. They also provide mathematical guarantees and confirm their findings through experiments.
gradient descentone-hidden-layer neural networkReLU activationloss landscapenon-strict saddle pointsGaussian distributionsample complexityalignment phaseconvergence ratetrajectory analysis
Authors
Berk Tinaz, Changzhi Xie, Mahdi Soltanolkotabi
Abstract
In this paper, we study the gradient descent dynamics for jointly training both layers of a one-hidden-layer ReLU network to fit a linear target function. Concretely, we consider a realizable setting where inputs are drawn i.i.d. from a Gaussian distribution and labels follow a planted linear model. This stylized framework captures salient features of end-to-end training in inverse problems and certain auto-encoder models. Despite its apparent simplicity, the dynamics remain poorly understood, in part because the loss landscape contains multiple non-strict saddle points, making it unclear why gradient descent from random initialization reliably escapes bad stationary regions. We provide a detailed characterization of the optimization landscape and prove that gradient descent from a moderately small random initialization-simultaneously training both layers-converges to a global minimizer at a linear rate with order-wise optimal sample complexity. Our analysis tracks the trajectory through three phases: an alignment phase in which hidden weights progressively align with the planted direction while the output weights maintain the correct sign pattern; a growth phase in which the norms of both layers increase while preserving alignment; and a local refinement phase in which the aligned neurons rapidly converge to the planted direction, yielding fast local convergence. To rigorously show that GD avoids non-strict saddles, we develop trajectory-level control arguments for the end-to-end dynamics. In addition, we establish novel uniform concentration results that hold along the entire trajectory, and are essential for obtaining order-wise optimal sample complexity. We corroborate our theory with extensive experiments across a range of configurations.