Global Exponential Convergence of Two-Layer Linear Network Training

Stephen Y Zhang; Gabriel Peyré

cs.LGOct 7, 2026

Global Exponential Convergence of Two-Layer Linear Network Training

Authors: Stephen Y Zhang, Gabriel Peyré

Organizations: Flatiron Institute & NYU · CNRS, ENS, PSL Univ.

Abstract

We prove global exponential (linear) convergence with an explicit rate in the rich scaling for wide two-layer linear networks trained with smooth Polyak-Lojasiewicz predictor losses. Gradient flow in the factors closes exactly in terms of a finite-dimensional Bures flow of the neuron law covariance, in which the predictor dynamics are preconditioned by hidden covariance blocks. Mean-field conservation laws provide uniform spectral lower bounds on the hidden preconditioning blocks when the initial covariance satisfies a spectral support gap condition. This condition encompasses positive definiteness while still allowing for singular initializations. For an initial covariance $Σ_0 = σ^2 \mathrm{Id}$ , the loss converges to the global minimum with linear rate at least $4σ^2κ$ , where $κ$ is the PL constant. We establish stability of this rate under finite-width sampling, as well as global convergence of factor gradient descent for an explicit stepsize interval depending on smoothness, the initial loss, and conserved spectral margins. Our argument extends layerwise to deep linear ResNets, subject to a residual-path bound. In the case of heavy-ball momentum, training dynamics close instead over positions and velocities in terms of a lifted phase covariance. Linear convergence holds under an explicit condition on the energy and damping, specifying a window of admissible dampings. For two-scale white initializations, this interval is nonempty for sufficiently large position scales, with a fixed initial loss gap and velocity covariance. Numerical experiments illustrate the covariance geometry and compare the predicted and observed rates.

Figures & tables

Panel of Figure 4 — Figure 4: Empirical predictor-loss trajectories (solid) and the corresponding upper bounds on mechanical excess energy (dotted) for two-scale white initializations with varying $\sigma_{vu}$ and damping $\gamma$ : (a) $\sigma_{vu}=200$ , (b) $\sigma_{vu}=500$ , and (c) $\sigma_{vu}=1000$ . Each theoretical curve uses the explicit, approximately optimal rate on a tube certified by Theorem 5.2 .

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Notation	Meaning	First reference
$d_{u},d_{v},n,m$	Input/output dimensions, width, and $m:=d_{u}+d_{v}$	Section 1
$U,V,W:=\binom{U}{V},\mathcal{L}$	Factor matrices, stacked positions, and factor loss	Section 1
$\Sigma,\Sigma^{u},\Sigma^{v},\Sigma^{vu}$	Uncentered moment, hidden blocks, and predictor	Section 2
$\ell,\ell_{\star},\mathcal{S}_{\star}$	Predictor loss, infimum, and minimizer set	Assumption 1.1 ; Section 3
$L,\kappa,\lambda$	Smoothness, PL, and QG constants	Assumption 1.1 ; Eq. ( 23 )
$\mathcal{Q},C_{z},\Sigma^{vu}_{\star}$	Quadratic Hessian, input second moment, teacher	Section 1

Appendix

Table 2

Explore similar work

May 29, 2026cs.LG

Balancing Learning Rates Across Layers: Exact Two-Step Dynamics and Optimal Scaling in Linear Neural Networks

We study optimal learning-rate selection in two-layer and three-layer linear neural networks trained to learn linear target functions. In particular, we derive the exact closed-form expressions for the gradients and test loss after one and two steps of gradient descent, enabling a precise characterization of early training dynamics. We characterize how learning rates should scale under the gradient approximation in the first two steps, and prove that performing updates with this approximation yields a tractable surrogate loss with a tight, small approximation error. This formulation enables the theoretical analysis of layer-wise learning rates and reveals a distinct early-training regime: test loss can be minimized by unequal learning rates at the initial step, while equal learning rates become optimal in subsequent steps. Our numerical experiments validate the theory and demonstrate the importance of balancing layer-wise learning rates early during training. The code is available at: https://github.com/TDCSZ327/Layer-Balancing.

Tianyu Pang, Vignesh Kothapalli, Shenyang Deng +3

Dartmouth College · Stanford University · Virginia Tech

Batch Two-Layer Neural Networks Gradient Descent Sequential Scaling

Ask PDF

Jun 3, 2026cs.LG

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

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.

Berk Tinaz, Changzhi Xie, Mahdi Soltanolkotabi

Department of Electrical and Computer Engineering University of Southern California Los Angeles, CA, USA · Department of Computer Science University of Southern California Los Angeles, CA, USA

Rectified Linear Unit Networks Gradient Descent Stationary Point Neurons

Ask PDF

May 11, 2026math.OC

On the global convergence of gradient descent for wide shallow models with bounded nonlinearities

A surprising phenomenon in the training of neural networks is the ability of gradient descent to find global minimizers of the training loss despite its non-convexity. Following earlier works, we investigate this behavior for wide shallow networks. Existing results essentially cover the case of ReLU activations and the case of sigmoid activations with scalar output weights. We study a large class of models that includes multi-head attention layers and two-layer sigmoid networks with vector output weights. Building upon [Chizat and Bach, 2018], we prove that all non-global minimizers of the training loss are unstable under gradient descent dynamics. Thus, when the initial distribution of the parameters has full support (which includes the popular Gaussian case), and in the many hidden neurons or attention heads limit, continuous-time gradient descent can only converge to global minimizers. Establishing the instability of non-global minimizers corresponds to the construction of an ``escaping active set'' -- we complete the proof of [Chizat and Bach, 2018] to construct this set for models with bounded nonlinearities and scalar output weights. We also extend this construction to new cases for models with vector output weights. Finally, we show the well-posedness and the stability with respect to discretization of the mean field training dynamic for sub-Gaussian initializations.

Romain Petit, Clarice Poon, Gabriel Peyré

CNRS and ENS, PSL Universit´e · Mathematics Institute, University of Warwick

Gradient Descent Flat Minima Descent Model Weights Convergence Nonlinearity+1

Ask PDF