PaperlayerView exact revision on arXiv

Source-linked AI summary

Global Exponential Convergence of Two-Layer Linear Network Training

Stephen Y Zhang, Gabriel Peyré

arXiv:2610.09356v1cs.LGmath.OCstat.ML

TL;DR

Two-layer linear network training is nonconvex in its factors, and its time-varying hidden covariance preconditioner may become degenerate. The paper closes the dynamics over covariance and derives global convergence guarantees, including an explicit rate under a positive support gap.

  • Problem

    Two-layer linear networks have nonconvex factors and potentially degenerate, time-varying preconditioning, leaving global convergence to an optimal predictor uncertain.

    Source: §1 · p. 1

  • Method

    The analysis closes rich mean-field factor gradient flow over a finite-dimensional Bures covariance flow and uses conserved spectral margins to bound hidden covariance blocks.

    Source: §1 · pp. 1, 3

  • Results

    With g(Σ0) > 0, the predictor converges linearly at rate κg(Σ0), including for possibly singular initial covariance; for Σ0 = σ^2Id, the loss rate is at least 4σ^2κ.

    Source: §1 · p. 1, §3 · p. 5

  • Takeaways & Limitations

    A positive support gap supplies sufficient hidden-block conditioning for explicit global convergence rates, and includes positive-definite as well as some singular initializations.

    Source: §1 · p. 1

    The deep linear ResNet extension requires a uniform residual-path bound because the loss depends only on the endpoint.

    Source: §4 · p. 9

Abstract

from arXiv · show

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 $\Sigma_0 = \sigma^2 \mathrm{Id}$, the loss converges to the global minimum with linear rate at least $4\sigma^2\kappa$, where $\kappa$ 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.

1 Introduction

The paper studies whether training two-layer linear networks can converge globally despite nonconvex factors and potentially degenerate, time-varying predictor preconditioning. It uses covariance conservation to establish explicit convergence guarantees across several training settings.

1 Introduction

  • A positive covariance support gap yields global linear convergence under smooth PL predictor losses, including for some singular initializations.For Σ0 = σ^2Id, the loss rate is at least 4σ^2κ.Source: §1 · pp. 1, 3
  • Mean-field conservation laws bound hidden covariance blocks, conditioning predictor dynamics and supporting the explicit rate.The rate remains valid under finite-width sampling; factor gradient descent also converges globally under an explicit stepsize bound.Source: §1 · pp. 1, 3
  • Unlike NTK and lazy-training analyses, the approach lets incoming and outgoing weights evolve, making the changing representation essential to the rate.Exact second-moment closure in the linear setting connects training to Bures–Wasserstein covariance geometry and makes a sharp quantitative rate available.Source: §1 · p. 2
  • For deep linear ResNets, the layerwise argument gives linear convergence subject to uniform residual-path control.The paper applies this condition to finite- and continuous-depth networks.Source: §1 · p. 3
  • Heavy-ball dynamics close over the phase covariance, with global exponential convergence under an explicit initial-energy and damping condition.For two-scale white initializations, the admissible damping interval is nonempty at sufficiently large position scales, with fixed initial loss gap and velocity covariance.Source: §1 · pp. 2–3

2 Mean-field conservation laws

The factor dynamics close exactly on a finite-dimensional covariance that follows Bures flow. Conserved spectral quantities yield hidden-block lower bounds, allowing convergence guarantees even when the initial covariance is singular.

2 Mean-field conservation laws

  • The factor flow closes on the covariance as a Bures equation, preserving the characteristic polynomial of JΣ and the conserved quantities Hk.The covariance equation is the Bures gradient flow of F(Σ) = ℓ(Σvu).Source: §2 · p. 4
  • The conserved support margins give uniform lower bounds on the hidden blocks, and a positive support gap can persist even for a degenerate covariance.The margins are preserved along the flow; they need not equal the initial hidden-block minimum eigenvalues.Source: §2 · p. 4
  • Special conserved-spectrum conditions can eliminate hidden blocks and yield lower-dimensional intrinsic dynamics.These reductions are described in Appendix E.Source: §2 · p. 5
  • In the scalar-input/output illustration, equal diagonal blocks follow invariant hyperbolas, while unequal blocks follow fixed-b sections of determinant leaves; both displayed trajectories converge to r = β.For equal blocks, the trajectories converge to (r, s) = (β, √(β^2 + α^2)) when α > 0.Source: Figure 1 · p. 5, §2 · p. 5

3 Linear convergence of training dynamics

A positive initial support gap guarantees global linear convergence of the Bures flow, including for some singular covariances. The section also establishes rate preservation under finite-width sampling and global convergence of factor gradient descent under a sufficiently small stepsize.

3 Linear convergence of training dynamics

  • A positive initial support gap g(Σ0) guarantees global linear convergence, and positive definiteness is sufficient but not necessary.The trajectory preserves covariance rank, and its predictor converges to a global minimizer.Source: §3 · pp. 5–6
  • For a full-rank isotropic initialization Σ0 = σ^2 Id_m, a+(Σ0) = a−(Σ0) = σ^2, giving a rate of 4σ^2λ+.Source: §3 · p. 6
  • When Q has a nontrivial kernel, the predictor converges to a global minimizer in Σ_vu⋆ + ker Q; if Q is positive definite, it converges to the teacher predictor.Source: §3 · p. 6
  • If the empirical covariance is within relative error ε of a population covariance with positive support gap, the sampled flow has rate exponent 2κ(1−ε)g(eΣ0).The empirical initial loss gap is the prefactor and need not equal the population loss gap.Source: §3 · p. 7
  • Factor gradient descent converges globally for sufficiently small stepsizes that control descent and accumulated conservation error.The factors and covariance converge geometrically, with the limiting predictor globally optimal; sufficiently small stepsizes can certify exponents arbitrarily close to the gradient-flow exponent.Source: §3 · pp. 7–8

4 Deep linear residual networks

Deep linear ResNets converge globally under layer-margin and residual-path conditions. The result extends to continuous depth when the covariance profile meets the stated assumptions and the training trajectory is well defined.

Finite-depth setting

  • The convergence guarantee requires the residual predictors to have an operator-norm bound independent of depth; layerwise conservation alone does not control forward and backward propagation.Source: §4 · pp. 8–9
  • With positive-definite initial layer covariances and residual-path control, finite-depth ResNets converge globally: every layer covariance and factor converges, and the endpoint is globally optimal.Source: §4 · pp. 8–9
  • e^-4B lower-bounds the product-conditioning factor when S ≥ 2B; for fixed B, a positive average layer support gap suffices for a rate uniform in depth.Source: §4 · p. 9

Continuous-depth limit

  • The continuous-depth formulation assumes a well-defined global training trajectory and uniformly conditioned residual propagators.Source: §4 · p. 9
  • For continuous depth, profiles starting with positive-definite covariance at almost every depth converge in L1 to a profile whose network attains the global minimum.The smallest initial eigenvalue need not be bounded away from zero uniformly over depth.Source: §4 · p. 9

5 Heavy-ball momentum convergence

With heavy-ball momentum, two-layer linear network dynamics close over the joint covariance of positions and velocities, not the position covariance alone. Conformal conservation laws and energy dissipation yield a global convergence certificate under explicit initialization and damping conditions.

5 Heavy-ball momentum convergence

  • Inertial dynamics close on the joint position-velocity covariance because adding velocity as a state variable prevents closure on position covariance alone.The phase covariance tracks position, mixed, and velocity blocks; the predictor remains the vu block of the position covariance.Source: §5 · pp. 9–10

Appendix C.7 gives the particle equations, measure formulation, and blockwise expansion.

  • Conformal phase conservation laws preserve the spectrum and trace powers of the exponentially rescaled phase covariance, while mechanical excess energy dissipates at rate 2γK_t.Without damping, the phase generator is symplectic; damping rescales the associated bilinear form exponentially.Source: Appendix C.7 · p. 10
  • Theorem 5.2 certifies exponential decay of mechanical excess energy when the initialization and damping satisfy a phase-energy condition that keeps the position covariance within a spectral tube.Under that condition, E_t ≤ 3E_0e^(-λ_s(γ)t), and the factors converge exponentially to limits with an optimal predictor.Source: Appendix C.7 · p. 10
  • For two-scale white initialization, sufficiently large position scales make the certified damping condition hold for fixed deformation budget, loss gap, and velocity covariance.The resulting admissible damping interval is nonempty; the certificate provides explicit bounds on the interval.Source: Appendix C.7 · p. 11
  • Figure 3 maps certified damping windows through the interiors of the G_s(γ)=0 curves, while Figure 4 compares empirical losses with certified energy bounds.The Figure 4 upper bounds use explicit, approximately optimal rates rather than rates fitted to trajectories; displayed losses stay below the bounds, although predictor loss may oscillate as energy decreases.Source: Appendix C.7 · p. 11, Figure 3 · p. 11, Figure 4 · p. 11

6 Conclusion

  • The paper extends covariance conditioning as a common convergence mechanism to deep linear ResNets with bounded residual paths and to heavy-ball dynamics with an explicit damping certificate.The momentum certificate includes two-scale white initializations without requiring a phase-covariance identity.Source: §6 · p. 12

B Measure formulation and Wasserstein conservation laws

The measure formulation represents neurons by their distribution and training by a Wasserstein flow. For objectives depending on second moments, this flow closes exactly as a Bures flow of the covariance, enabling width-independent conservation laws and spectral bounds.

B.1 Neuron distributions and Wasserstein flow

  • An n-neuron network is represented by its empirical neuron law, whose second moment is WW^T/n and whose loss is the predictor loss evaluated at the covariance's vu block.Any positive semidefinite covariance can also be realized by a centered Gaussian law.Source: §B.1 · p. 17
  • The Wasserstein gradient flow moves each neuron along the negative spatial gradient of the first variation, and its PDE formulation includes empirical measures.The empirical characteristics recover the finite-network factor dynamics.Source: §B.1 · p. 17
  • A covariance trajectory determines a Wasserstein gradient flow when an initial law is transported by the associated characteristic flow, but matching second moments alone is insufficient.The transported law and the given covariance trajectory have the same second-moment evolution and initial data.Source: §B.1 · p. 17

B.2 Covariance-dependent objectives

  • For any sufficiently regular objective depending on second moments, the Wasserstein flow closes exactly on the Bures gradient flow of those moments.The covariance evolves autonomously as Σ̇_t = −2(D_tΣ_t + Σ_tD_t).Source: §B.2 · p. 18

B.3 Width consistency and covariance-level laws

  • A mean-field conservation law is constant along every admissible trajectory and initialization in its flow family; it is covariance-level when it depends only on the second moment.A single such law induces a conserved law at every finite width.Source: §B.3 · p. 18

B.4 Gram-function laws and their spectral restriction

  • The fundamental mean-field laws include conserved spectral laws of JΣ and, for particles, a conserved Gram imbalance whose nonzero spectrum matches that of −Γ/n.The spectral laws remain conserved for any initial law with finite second moment.Source: §B.4 · p. 19

C.1 Loss estimates

  • Under the smooth PL assumptions, PL implies quadratic growth without requiring convexity, though the best PL and quadratic-growth constants need not coincide.A known larger quadratic-growth constant improves distance bounds but does not change the loss rates based on κ.Source: §C.1 · p. 20

C.2 Covariance closure and conservation laws

  • The signed spectral support margins are both positive exactly for positive-definite covariances, but a positive sum suffices for convergence and therefore also covers some singular covariances.Each positive margin lower-bounds its corresponding hidden block; the support gap can remain positive even when the full covariance is degenerate.Source: §C.2 · p. 21

C.3 Linear convergence

The section establishes global convergence and trajectory bounds for covariance flows, then examines when hidden-block bounds can be obtained by alternative criteria. It shows that a vanishing support gap can permit arbitrarily slow convergence, even at fixed trace and initial loss.

C.3 Linear convergence

  • A positive conserved support gap g(Σ_0) guarantees global linear convergence of the Bures flow, including for some singular initial covariances.The covariance limit has an optimal predictor, and the trajectory bounds apply to particle factors after normalization by √n.Source: §C.3 · p. 21
  • Coercivity can provide an alternative hidden-block bound from determinant conservation, but this criterion is generally coarser than the spectral-margin rate.For positive-definite initial covariance, the stated bound keeps both hidden blocks above cI.Source: §C.3 · pp. 22–23
  • At Σ_0 = αI, the support-gap rate is 4ακ, whereas the determinant-based rate can be extremely small in dimension.For general initial covariances, the margin-based rate depends on the signed spectrum without an additional dimension factor.Source: §C.3 · p. 23
  • On a degenerate boundary with zero support gap, some initial predictors converge to a nonoptimal stationary point, while the zero-target case can converge algebraically.These boundary cases fall outside the positive-support-gap theorem.Source: §C.3 · p. 23
  • When the support gap vanishes, trace and initial loss alone cannot ensure a uniform positive exponential rate, even for strongly convex loss with a unique minimizer.In the fixed-trace, fixed-loss-gap example, the initial relative decay rate is 4δ and the time to reach r=1/2 diverges as δ approaches zero.Source: §C.3 · p. 23

C.4 Finite-width convergence

Finite-width sampling can preserve a positive support gap and the associated convergence guarantee under a relative covariance bound, but small-width isotropic sampling may lose that gap. Factor gradient descent also converges globally under an explicit stepsize bound, with adjustable rates approaching the flow exponent.

C.5 Discrete-time convergence

  • A positive initial support gap guarantees global factor-gradient-descent convergence for sufficiently small steps, with a bound depending on the initial covariance, loss gap, and loss parameters.The stepsize condition controls both descent and accumulated conservation error.Source: §C.4 · p. 24

C.4 Finite-width convergence

  • When eΣ0 = σ2IdD+1 and 1 ≤ n ≤ D/16, the empirical support gap is zero with probability at least 1 −4e−cD, despite population gap 2σ2.The example does not determine whether the sampled trajectories converge.Source: §C.4 · p. 25
  • With probability at least 1 −δ, the stated sampling conditions yield an empirical covariance with positive support gap and globally convergent factor flow.The factors converge to a predictor in S⋆; the sampling width condition ensures the relative covariance bound with probability at least 1 −δ.Source: §C.4 · pp. 24–25
  • For balanced dimensions and fixed σ, width proportional to m preserves the dimension-free exponent but allows expected initial sampling error of order m.Keeping that error bounded in expectation requires width of order dudv when the teacher norm is fixed.Source: §C.4 · p. 26

C.5 Discrete-time convergence

  • With a sufficiently small step, choosing c = (1 −ϵ)g(Σ0) certifies exponent 2κ(1 −ϵ)g(Σ0), arbitrarily close to the flow exponent.A deterministic stepsize chosen before sampling also gives high-probability convergence to an optimal predictor under the stated sampling conditions.Source: §C.5 · pp. 28–29

C.6 Deep linear residual networks

The deep-network extension obtains linear convergence under residual-path control, while examples show that positive layer margins alone do not ensure that control. For heavy-ball momentum, the closed phase-covariance dynamics admit conditional energy decay; a separate decomposition identifies the minimum-energy lift of a covariance velocity.

C.6 Deep linear residual networks

  • Under a uniform residual-path bound, the finite- and continuous-depth arguments yield convergence to an optimal endpoint; without that bound, positive conserved margins need not suffice.The counterexample has unbounded residuals and loss tending to 1/2 above the global minimum, while the zero-target variant also violates path bounds despite loss tending to zero.Source: §C.6 · pp. 30–31

C.7 Heavy-ball momentum

  • The heavy-ball state closes over the joint phase covariance of positions and velocities, and its full phase flow is global for positive-semidefinite initial covariance.Global existence alone does not imply uniform position ellipticity or convergence to a minimizer.Source: §C.7 · p. 32
  • For heavy-ball dynamics, if the position covariance stays between aId and bId, the phase energy decays at an explicit rate depending on a, κ, γ, and the energy and smoothness bounds.The conditional estimate is Et ≤ 3E0e−λt on the interval where the covariance bounds hold.Source: §C.7 · p. 33

C.8 An exact decomposition of hidden kinetic energy

  • For a fixed positive-definite position covariance and prescribed covariance derivative, the graph lift D = C, Q = 0 minimizes kinetic energy.The additional positive-semidefinite Schur complement contributes exactly tr Q/2.Source: §C.8 · pp. 34–35

C.9 Tightened conditional momentum rates

The tightened momentum analysis refines conditional energy-decay bounds and their prefactors, then uses the corrected energy estimate to certify a global covariance tube. Its optimized exponent is optimal within the stated Lyapunov bounds, not necessarily for the training dynamics.

C.9 Tightened conditional momentum rates

  • The corrected energy Vt = Et + ε ˙∆t bounds the loss gap by a decaying exponential with its own prefactor.When ˙∆0 = 0, the first energy prefactor is E0/(1 −q).Source: §C.9 · p. 35
  • The unique maximizing correction at fixed η is admissible and gives the optimal rate within the Lyapunov bounds.The approximate parameter choice used in the figures yields an exact bound, but its exponent is no larger than the optimized bound and is not claimed to be the underlying dynamics' optimal rate.Source: §C.9 · pp. 36–37
  • A global bootstrap holds when the corrected energy satisfies the stated deformation-budget condition, with the bound retaining its energy prefactor.A sufficient condition using only initial position bounds and total energy replaces C0 by (1 + q)/(1 −q).Source: §C.9 · p. 37

C.10 Large-position-scale certificate asymptotics

The appendix analyzes damping and deformation budgets for the heavy-ball convergence certificate, including its large-position-scale behavior. It also characterizes the split Bures directions and their generated Lie algebra, which underpin covariance-level conservation laws.

C.10 Large-position-scale certificate asymptotics

  • The boundary rate is a limiting supremal certificate: the theorem applies at strictly admissible interior budgets whose rates approach it.The bootstrap condition need not hold at the boundary itself.Source: §C.10 · p. 38
  • At fixed deformation budget, certificate admissibility reduces to a quadratic inequality that gives explicit damping boundaries.The strict condition is 2γ^2 − csγ + ρs < 0.Source: §C.10 · p. 38
  • The unique deformation budget maximizing damping-window nonemptiness lies in (1/4, 1/3), and it does not optimize the upper damping endpoint or decay exponent.The scalar equation for this budget is explicit and uniquely solvable without computing a trajectory.Source: §C.10 · p. 38
  • At fixed damping, increasing the initial position scale makes the certified deformation vanish and the certified rate converge to κγ/(4L).At fixed position scale, the certificate does not extend to arbitrarily large damping; its admissible region eventually disappears.Source: §C.10 · p. 40

D Counting covariance-level mean-field laws

D.1 Split Bures vector fields

  • The appendix treats common covariance laws by considering smooth split losses and their associated split Bures vector fields.The constant-gradient fields arise from linear losses.Source: §D.1 · p. 42

D.3 Completeness of the spectral laws

For split two-layer linear losses, the locally generic common covariance laws are determined by the spectrum of JΣ. The appendix contrasts this finite-dimensional linear-network result with a ReLU imbalance law and describes a special spectral reduction of the covariance dynamics.

D.3 Completeness of the spectral laws

  • The trace powers of JΣ are conserved, and locally and generically they are the only smooth covariance-level laws common to all split losses.Equivalently, such laws are locally smooth functions of the eigenvalues of JΣ; the characteristic polynomial gives a complete local/generic list.Source: §D.3 · pp. 43–44

D.4 A ReLU contrast

  • For one-homogeneous networks such as ReLU, the flow preserves the pushforward distribution of ∥u∥^2 − ∥v∥^2, but this does not yield a closed finite-dimensional Bures covariance equation.No completeness claim for ReLU conservation laws is made.Source: §D.4 · p. 45

E Intrinsic dynamics

  • A special two-point conserved spectrum determines hidden covariance blocks from the predictor, allowing a smaller intrinsic equation.The appendix derives exact reductions from these special conserved spectral values.Source: E Intrinsic dynamics · p. 46

E.1 A two-point signed spectrum: hyperbolic isotropy

On the hyperbolically isotropic spectral leaf, each predictor has a unique compatible covariance and the Bures dynamics close directly on the predictor. A distinct phase-covariance condition yields a reduced description using two symmetric matrices and a known scale.

E.1 A two-point signed spectrum: hyperbolic isotropy

  • On the hyperbolically isotropic leaf, every predictor has a unique compatible covariance, and its Bures dynamics close on the predictor alone.The leaf is characterized by the two-point signed spectrum; the displayed predictor equation uses the loss gradient and square-root preconditioners shifted by α^2Id.Source: §E.1 · p. 46
  • The general convergence theorem gives loss exponent 4ακ on this leaf.Here the conserved spectral levels satisfy a+(Σ0) = a−(Σ0) = α.Source: §E.1 · p. 46
  • Relaxed balance does not generally ensure predictor closure; the stated coincidence requires n = du = dv and both factors invertible.Without rank qualifications, states with the same predictor and balance value can have different predictor velocities.Source: §E.1 · p. 47

E.2 A single phase spectral scale: conformal symplectic isotropy

  • Conformal symplectic isotropy reduces the phase covariance on a positive-position chart to two symmetric matrices, A and C, and the known scale θt.The reduced block identities include D = C = C⊤ and Q = θ^2A^-1.Source: §E.2 · p. 48
  • Even under this phase constraint, the cross blocks Σvu and Σqp do not generally form a closed state; the full position block remains an intrinsic variable.The general convergence theorem controls D by energy without requiring this special reduction.Source: §E.2 · p. 48

F A nonquadratic example: smoothed Lasso

The appendix illustrates the smooth PL loss class with an entrywise smoothed-Lasso penalty on the predictor. Its figure compares covariance initialization and sampling cases, while the nonsmooth limit falls outside the stated loss assumption.

F A nonquadratic example: smoothed Lasso

  • The smoothed entrywise ℓ1 penalty acts on the predictor, so covariance closure and spectral conservation laws apply without modification.The loss is strongly convex and satisfies the paper’s smooth PL assumption.Source: F A nonquadratic example: smoothed Lasso · p. 49
  • Positive smoothing does not generally yield exact sparsity, and the nonsmooth limit δ = 0 is outside Assumption 1.1.Source: F A nonquadratic example: smoothed Lasso · p. 49
  • Figure 6 compares white initial covariances across α values, a scaled input-output diagonal pair, and population versus finite-width covariance flows.The finite-width panel uses widths 32, 128, 512, and 2048, with five trials each.Source: Figure 6 · p. 49

Related papers

Loading 2610.09356v1…