When Both Layers Learn: Training Dynamics of Representing Linear Models via ReLU Networks

2026-06-03

AI Summary

This work investigates how single-hidden-layer ReLU networks, trained end-to-end with Gaussian inputs, escape non-strict saddle points and converge to the global optimum when fitting a linear target function. By analyzing the optimization trajectory through three distinct phases—alignment, growth, and local refinement—the study provides the first rigorous characterization of the dynamical mechanism that avoids non-strict saddles in full gradient-based training. The key contributions include establishing a novel uniform concentration inequality valid along the entire optimization path, achieving near-optimal sample complexity, and proving that gradient descent, initialized with modestly small random weights, converges linearly to the global minimum. These theoretical findings are corroborated across multiple experimental settings.

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.

Problem

• gradient descent dynamics
• ReLU networks
• non-strict saddle points
• global convergence
• training both layers

Innovation

• gradient descent dynamics
• ReLU networks
• non-strict saddle points
• trajectory analysis
• uniform concentration