mlsanityData to 7 Oct 2026

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

arXiv 2606.04476 ↗First posted 3 Jun 2026cs.LG math.OC math.STcited by 1

Not on the map. Why?

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.

Key facts. Cited by 1 paper in the map (data through 7 Oct 2026). Builds on 46 of its 61 references that are in the map. First posted on arXiv 3 Jun 2026.

Builds on

46 of 61 references.

References · 46

  1. 2018Neural Tangent Kernel: Convergence and Generalization in Neural Networkscited by 2,597
  2. 2016Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Networkcited by 1,898
  3. 2018On Lazy Training in Differentiable Programmingcited by 671
  4. 2018Gradient Descent Provably Optimizes Over-parameterized Neural Networkscited by 655
  5. 2018Gradient Descent Finds Global Minima of Deep Neural Networkscited by 572
  6. 2019Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networkscited by 562
  7. 2016Deep Convolutional Neural Network for Inverse Problems in Imagingcited by 431
  8. 2017How to Escape Saddle Points Efficientlycited by 284
  9. 2017Non-convex learning via Stochastic Gradient Langevin Dynamics: a nonasymptotic analysiscited by 242
  10. 2015Escaping From Saddle Points --- Online Stochastic Gradient for Tensor Decompositioncited by 241
  11. 2020End-to-End Variational Networks for Accelerated MRI Reconstructioncited by 201
  12. 2018Gradient descent aligns the layers of deep linear networkscited by 182
34 more
  1. 2019Towards moderate overparameterization: global convergence guarantees for training shallow neural networkscited by 176
  2. 2022Neural Networks can Learn Representations with Gradient Descentcited by 161
  3. 2020When Do Neural Networks Outperform Kernel Methods?cited by 150
  4. 2022High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representationcited by 146
  5. 2016Toward Deeper Understanding of Neural Networks: The Power of Initialization and a Dual View on Expressivitycited by 138
  6. 2017Theoretical insights into the optimization landscape of over-parameterized shallow neural networkscited by 128
  7. 2016Matrix Completion has No Spurious Local Minimumcited by 124
  8. 2014Phase Retrieval via Wirtinger Flow: Theory and Algorithmscited by 123
  9. 2019A Function Space View of Bounded Norm Infinite Width ReLU Nets: The Multivariate Casecited by 113
  10. 2015Low-rank Solutions of Linear Matrix Equations via Procrustes Flowcited by 97
  11. 2017Recovery Guarantees for One-hidden-layer Neural Networkscited by 94
  12. 2016A Geometric Analysis of Phase Retrievalcited by 92
  13. 2018Overparameterized Nonlinear Learning: Gradient Descent Takes the Shortest Path?cited by 87
  14. 2018Image Inpainting via Generative Multi-column Convolutional Neural Networkscited by 81
  15. 2021Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstructioncited by 77
  16. 2018Phase Retrieval Under a Generative Priorcited by 73
  17. 2017A Hitting Time Analysis of Stochastic Gradient Langevin Dynamicscited by 70
  18. 2017Implicit Regularization in Nonconvex Statistical Estimation: Gradient Descent Converges Linearly for Phase Retrieval, Matrix Completion, and Blind Deconvolutioncited by 63
  19. 2017Globally Optimal Gradient Descent for a ConvNet with Gaussian Inputscited by 61
  20. 2017Learning ReLUs via Gradient Descentcited by 54
  21. 2015Solving Random Quadratic Systems of Equations Is Nearly as Easy as Solving Linear Systemscited by 50
  22. 2018Learning One-hidden-layer ReLU Networks via Gradient Descentcited by 50
  23. 2020Learning a Single Neuron with Gradient Methodscited by 50
  24. 2024Neural network learns low-dimensional polynomials with SGD near the information-theoretic limitcited by 45
  25. 2022HUMUS-Net: Hybrid unrolled multi-scale network architecture for accelerated MRI reconstructioncited by 30
  26. 2025Emergence and scaling laws in SGD learning of shallow neural networkscited by 29
  27. 2015When Are Nonconvex Problems Not Scary?cited by 27
  28. 2019Sharp Restricted Isometry Bounds for the Inexistence of Spurious Local Minima in Nonconvex Matrix Recoverycited by 26
  29. 2023Over-Parameterization Exponentially Slows Down Gradient Descent for Learning a Single Neuroncited by 25
  30. 2023Good regularity creates large learning rate implicit biases: edge of stability, balancing, and catapultcited by 23
  31. 2021Learning a Single Neuron with Bias Using Gradient Descentcited by 17
  32. 2023Learning a Neuron by a Shallow ReLU Network: Dynamics and Implicit Bias for Correlated Inputscited by 16
  33. 2024Simplicity bias and optimization threshold in two-layer ReLU networkscited by 11
  34. 2024Feature Averaging: An Implicit Bias of Gradient Descent Leading to Non-Robustness in Neural Networkscited by 3

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 1.

1 paper.

  1. 2026Removing spurious minima for planar features by skip connectionscited by 0