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

arXiv: https://arxiv.org/abs/2606.04476 · HTML: https://mlsanity.com/paper/2606.04476/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 46 of its 61 references that are in the map.
- First posted on arXiv 3 Jun 2026.

## Builds on (most cited first, 30 of 46)

- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network](https://mlsanity.com/paper/1609.04802/) (2016)
- [On Lazy Training in Differentiable Programming](https://mlsanity.com/paper/1812.07956/) (2018)
- [Gradient Descent Provably Optimizes Over-parameterized Neural Networks](https://mlsanity.com/paper/1810.02054/) (2018)
- [Gradient Descent Finds Global Minima of Deep Neural Networks](https://mlsanity.com/paper/1811.03804/) (2018)
- [Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks](https://mlsanity.com/paper/1901.08584/) (2019)
- [Deep Convolutional Neural Network for Inverse Problems in Imaging](https://mlsanity.com/paper/1611.03679/) (2016)
- [How to Escape Saddle Points Efficiently](https://mlsanity.com/paper/1703.00887/) (2017)
- [Escaping From Saddle Points --- Online Stochastic Gradient for Tensor Decomposition](https://mlsanity.com/paper/1503.02101/) (2015)
- [Non-convex learning via Stochastic Gradient Langevin Dynamics: a nonasymptotic analysis](https://mlsanity.com/paper/1702.03849/) (2017)
- [End-to-End Variational Networks for Accelerated MRI Reconstruction](https://mlsanity.com/paper/2004.06688/) (2020)
- [Gradient descent aligns the layers of deep linear networks](https://mlsanity.com/paper/1810.02032/) (2018)
- [Towards moderate overparameterization: global convergence guarantees for training shallow neural networks](https://mlsanity.com/paper/1902.04674/) (2019)
- [Neural Networks can Learn Representations with Gradient Descent](https://mlsanity.com/paper/2206.15144/) (2022)
- [When Do Neural Networks Outperform Kernel Methods?](https://mlsanity.com/paper/2006.13409/) (2020)
- [High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation](https://mlsanity.com/paper/2205.01445/) (2022)
- [Toward Deeper Understanding of Neural Networks: The Power of Initialization and a Dual View on Expressivity](https://mlsanity.com/paper/1602.05897/) (2016)
- [Theoretical insights into the optimization landscape of over-parameterized shallow neural networks](https://mlsanity.com/paper/1707.04926/) (2017)
- [Phase Retrieval via Wirtinger Flow: Theory and Algorithms](https://mlsanity.com/paper/1407.1065/) (2014)
- [Matrix Completion has No Spurious Local Minimum](https://mlsanity.com/paper/1605.07272/) (2016)
- [A Function Space View of Bounded Norm Infinite Width ReLU Nets: The Multivariate Case](https://mlsanity.com/paper/1910.01635/) (2019)
- [Low-rank Solutions of Linear Matrix Equations via Procrustes Flow](https://mlsanity.com/paper/1507.03566/) (2015)
- [Recovery Guarantees for One-hidden-layer Neural Networks](https://mlsanity.com/paper/1706.03175/) (2017)
- [A Geometric Analysis of Phase Retrieval](https://mlsanity.com/paper/1602.06664/) (2016)
- [Overparameterized Nonlinear Learning: Gradient Descent Takes the Shortest Path?](https://mlsanity.com/paper/1812.10004/) (2018)
- [Image Inpainting via Generative Multi-column Convolutional Neural Networks](https://mlsanity.com/paper/1810.08771/) (2018)
- [Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruction](https://mlsanity.com/paper/2106.15013/) (2021)
- [Phase Retrieval Under a Generative Prior](https://mlsanity.com/paper/1807.04261/) (2018)
- [A Hitting Time Analysis of Stochastic Gradient Langevin Dynamics](https://mlsanity.com/paper/1702.05575/) (2017)
- [Implicit Regularization in Nonconvex Statistical Estimation: Gradient Descent Converges Linearly for Phase Retrieval, Matrix Completion, and Blind Deconvolution](https://mlsanity.com/paper/1711.10467/) (2017)

## Built on by (most cited first, 1 of 1)

- [Removing spurious minima for planar features by skip connections](https://mlsanity.com/paper/2610.01728/) (2026)

## Abstract (arXiv)

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.