A Local Polyak-Łojasiewicz and Descent Lemma of Gradient Descent For Overparametrized Linear Models
Authors: Ziqing Xu
Hancheng Min
Salma Tarmoun
Enrique Mallada
Rene Vidal
Venue: Accepted by TMLR ()

Abstract: Most prior work on the convergence of gradient descent (GD) for overparameterized neural networks relies on strong assumptions on the step size (infinitesimal), the hidden-layer width (infinite), or the initialization (large, spectral, balanced). Recent efforts to relax these assumptions focus on two-layer linear networks trained with the squared loss. 
In this work, we derive a linear convergence rate for training two-layer linear neural networks with GD for general losses and under relaxed assumptions on the step size, width, and initialization. A key challenge in deriving this result is that classical ingredients for deriving convergence rates for nonconvex problems, such as the Polyak-Łojasiewicz (PL) condition and Descent Lemma, do not hold globally for overparameterized neural networks. Here, we prove that these two conditions hold locally with local constants that depend on the weights. Then, we provide bounds on these local constants, which depend on the initialization of the weights, the current loss, and the global PL and smoothness constants of the non-overparameterized model. Based on these bounds, we derive a linear convergence rate for GD. Our convergence analysis not only improves upon prior results but also suggests a better choice for the step size, as verified through our numerical experiments.

Published in Transactions on Machine Learning Research (04/2025)
A Local Polyak-Lojasiewicz and Descent Lemma of Gradient Descent For Overparametrized Linear Models
Ziqing Xu
Syinnga@@uuartton.uueen..c
Center for Innovation in Data Engineering and Science (IDEAS) Deportment of Statistics and Data Science, the Wharton School University of Pennsyloania
Hancheng Min
hanchamnn &scas upenn. edu
Center for Innovation in Data Engineering and Science (IDEAS) University of Pennsylvania
Salma Tarmoun
starmoun@sas.upppe edu
Center for Innovation Data Engincering and Science (IDEAS) Graduate Group in Applied Mathematics and Computational Science University of Pennsylvania
Enrique Mallada
mallada edu
Mathematical Institute for Data Science (MINDS) Deportment of Electrical and Computer Engineering Johns Hopkins University
René Vidal
vidalr@upenn.com edu
Center for Innovation in Dala Engineering and Science (IDEAS)
Department of Electrical and Systems Engineering; Department of Radiology University of Pennsyloaana
Reviewed OpenReview: https: // openreview. net/ str/frrm? id= VPL3T43Hab
Abstract
Most prior work on the convergence of gradient descent (GD) for overparameterized neural workorks relies on strong assumptions on the step size (infinitesimal), the hidden-layer width (infinite), or the initialization (large, spectral, balanced). Recent efforts to relax these assumptions focus two-layer linear networks trained with the squared loss. In this work, we derive linear convergence rate for training two-layer linear neural networks with GD for general losses and under relaxed assumptions on the step size, width, and initialization. A key challenge in deriving this result is that classical ingredients for deriving convergence rates for nonconvex problems, such as the Lemma, do not hold globally for overparameterrized neural networks. Here, we prove that these two conditions hold locally with local constants that depend on the weights. Then, we provide bounds on these local constants, which depend on the initialization of the weights, the current loss, and the global PL and smoothness constants of the von-overparameterize model. Based on these bounds, we derive lincar convergence rate for GD. Our convergence analysis not only improves upon prior results but also suggests a better choice for the step size, as verified through our numerical experiments.
(PL) condition and Descent
1
Introduction
Neural networks have shown great empirical success in many real-world applications, such computer vision He al. (2016) and natural language processing Vaaswani al. 2018). However, our theoretical understanding

---

Published in Transactions on Machine Learning Research (04/2025)
of why neural networks work well is still scarce. One unsolved question is why neural networks trained via vanilla gradient descent (GD) enjoy fast convergence although their loss landscape is non-convex. This question has motivated recent work which focuses on deriving convergence rates for overparameterized neural networks. However, most prior work on the linear convergence of GD for overparametrized neural networks requires strong assumptions on the step size (infinitesimal), width (infinitely large), initialization (large, spectral), or restrictive choices of the loss function (squared loss) (See Table 1 for details).
Derivation of convergence of nonlinear networks requires restrictive assumptions. The work of 2022 studies the convergence of GD for neural networks in the neural langent kernel (NTK) regime, which requires the network to have large or infinite width and large initialization. However, Chizat et al. (201 );
(Du et al., 2018b; Lee et al. 2019 Liu et al.
Chen et al. (2022) show that the NTK regime limits feature
learning, and the generalization performance of neural networks in this regime degrades substantially. To go beyond the NTK regime, Mei et al. 20188 ; Chizat & Bach (2018); Sirignano & Spiliopoulos/ (2020); Ding et al.
(2022) study convergence of neural networks in the mean-field regime under the assumption of infinite width and infinitesinal step sizes. However, while such analysis can guarantee convergence a global optimum for wider range of initializations, it still imposes strong assumptions on the network width (infinite) and step size (infinitesimal).
Suboptimal convergence rates of overparametrized linear networks. To relax the assumptions on the width, step size, or initialization, some recent work focused on deriving convergence rates of GD for neural networks with a linear activation function in the context of matrix factorization (Arora et al. 2018 Du et al. 2018b) and linear regression (Du & Hu 2019 Xu al. 2023). In these settings, instead of learning a matrix W directly, one learns an overparametrized version of W defined as the product of L matrices Wi..-We WL. For example, in the case L 2, which is the one will analyze in this paper, this leads to the following no1-overparamettized and overparametrized problems defined, respectively, by
1 min ELS (W)= = 2 -Y--WW - -xW1 W
(uorooverparammtriizee)))
and
min LLS (W1, W2) = 1 - -V W..W2
(overparametrizee)
where X,Y Y are data matrices, W,W, Wi, W2 are weight matrices, and LS stands for least-square objective. Here, overparametrized indicates that, although the function classes represented by XW and XW, W2 coincide when the sizes of W1 and W2 are large enough, XW, W2 introduces additional matrices to represent the same function. Therefore, while the non-overparametrizee loss above defines standard lincar regression problem, the overparametrized loss can be seen as using a two-layer linear nelwork to solve the regression problem.
A classical approach to derive a linear convergence rate for non-convex problems relies on the PL condition and the smoothness inequality (see 22.1 for details). While these two conditions hold globally for non- overparametrized models, they do not hold globally for overparametrized linear works (see 2.2 for details). This is because overparametrization skews the gradient through weight-dependent linear operator T that acts on VlTs (derived based on chain rule)
Vis W2) Wwi,w, W. Vhis = T(Vers) W W2W.
To circumvent this challenge, existing studies (Arora et al. 2018 Du Hu 2019 Xu et al. 2023 use the smoothness inequality of the oonooverparamttrized model, which holds globally, as a substitute for the one of the overparametrized model. In addition, prior works use diffeerent methods and assumptions to derive (local)
PL conditions (See Appendix C for details). For example, Arora et al. (0018) impose unrealistic assumptions
on the initialization (large margir and small imbalance" while Du & Hu 2019) assume large initialization
and large width. Xu et al. 2023 show the overparametrized models satisfy local PL conditions and control
1PL condition has been widely used to derive convergence for non-convex problems Karimi al. 2020 Arora al. 2018
Min al. 2021 Pridovich-Keil ct al. 2023 Xu al. 2023
TTh maryin quantity that measure how close inttialization is the global minimum.
TThe imbalance is quantity that. measures the difference between the weights of two adjacent layers.
2

---

Published in Transactions on Machine Learning Research (04/2025)
Table 1: Comparison with prior work
Work
Loss
Step Size
Width
Initialiationn
Du et xl.
2018bb
Loc
6t
al.
2015
Jacot
Squared
LLu
Finite
Large
Large
loss
Nonlinear networks
2022
Nguyen &
Mou-
dclli
Me
ml.
2018
Chizat
Bach
Squared
Strigaano & Spilooon
Infinitesimal
Large
General
loss
Tos 2020
Ding et al.
2022
Saxe
al.
2013
Squared
Gidel et all
19
Infinitesimal
Finite
Spectral
loss
Linear networks
et al.
2021
Arorg 6f 8L
2018
Du Squared loss
Finite
Finite
Large margin and
sl. 2018a
small imbalance
Squared
Xu al.
2023
loss
Finite
Finite
General
This work
General
Finite
Finite
General
these constants using the weights at initialization through careful choices of the step sizes. However, how the initialization scale and width of the network affect these constants is missing.
While the aforementioned analyscs successfully derive linear convergence rates for overparamotrized lincar networks, all of them only consider squared loss and do not generalize other types of loss functions. Moreover, the analysis techniques rely on the smoothness inequality of the non-overparametrized model. Therefore, they lack insight into the optimization geometry of overparameterized problems. Numerically, the actual convergence rate of the loss under the step sizes proposed in all the above work is slow (See $..2 for details). Hence, an analysis needed that establishes the convergence of neural networks trained using GD with a general loss under more relaxed assumptions. Furthermore, this analysis should offer more accurate predictions of the actual rates convergence. This paper aims to bridge some of these gaps.
1.1 Main Contribution
In this work, we derive linear convergence rates for GD with possibly adaptive step sizes on overparameterized two-layer linear networks with a general loss, finite width, finite step size, and general initialization. Specifically, we make the following contributions:
We analyze the Hessian of two-layer lincar networks and show that the optimization problem satisfies a local PL condition and local Descent Lemma, where we characterize the local PL, constant and local smoothness constant along the descent direction at GD iterates singular values of the weight matrices (See Thcorem 3.1.
[ by their corresponding loss values and the
We show that when the step size satisfies certain constraints (not infinitesimal), the imbalance remains close to its initial value. Based on this property, we prove the local PL and smoothness constants can be bounded along the trajectory GD (See Theorem 3.2 Building on these results, we design an adaptive step size scheduler that yields linear convergence rate for GD. Morcover, our results cover GD with decreasing, constant, and increasing step sizes while prior work (Arora al. 2018 Du & Hu 2019 Xu
et al. 2023 only covers GD with decreasing and constant step sizes.
We show that, under our step size scheduler, the local smoothness constant decreases along the GD trajectory, indicating that the optimization landscape becomes more benign as training proceeds. Building on this observation, we demonstrate that our step size scheduler accelerates convergence (see $4.2
Our analysis allows us to show that when GD iterates are around global minimum, the difference between the local rate convergence of the overparametrized model and the rate the non-overparamettized model is up one factor of the condition number of To (see 2.2 for definitions) which can be made arbitrarily
4In the paper, we adopt the local smoothness constant convenient shorthand refer the smoothness constant. along the descent direction at GD iterates.
3

---

Published in Transactions on Machine Learning Research (04/2025)
close to one by proper initialization. In contrast, the (asymptotic) convergence rates derived in prior work (Arora et al., 2018 Du & Hu 2019 Xu et al. 2023) are adversely affected by the square of the condition number of To, leading to slower rates compared to our results.
1.2 Related work
We now provide detailed description of prior work in addition to the discussion above.
One line of work (Du et al. 2018b Lee et al., 2019 Liu et al. 2022 Nguyen & Mondelli, 2020) studies the convergence of GD with constant step size for squared loss under the assumption that the width and initialization neural networks are sufficiently large, which is also known as the neural tangent kernel (NTK) regime. Under these assumptions, the Training trajectories of a neural network are governed by a kernel determined at initialization and the network weights stay close their initial values. Such properties help them derive a lincar convergence rate of GD. However, the convergence rates derived in (Du et al. 2018b Lec
et al. 2019 Liu et al. 2022 Nguegnang et al.
2021 are inversely proportional to the number of samples the initial loss and can be arbitrarily close to one when the number of samples or the initial loss is sufficiently
large. Moreover, Chizat et al. (2019); Chen et al. (022) show that the NTK regime prohibits feature learning, and the performance of neural networks in this regime degrades substantially.
To relax assumptions on width, step size, and initialization, numerous studies have explored the convergence rates of gradient-based algorithms for neural networks with linear activations. These studies are motivated by observations that linear networks exhibit similar nonlinear learning phenomena to those seen in simulations
of nonlinear networks (Saxe et al.) 2013). For instance, Du & Hu (2019;; Arora et al. (2018) Xu et al. 2023
establish lincar convergence rates for linear networks with squared loss optimized via GD. Their analysos rely on the smoothness inequality of non-overparammelrizer models and do not provide an exact characterization of the optimization landscape in overparametrized settings. In contrast, we deliver a tighter analysis by characterizing the local PL condition and Descont Lemma in overparametrized models, allowing us to derive faster convergence rates and design adaptive step sizes that adjust the evolving local optimization landscape along the GD trajectory (see Appendix C for details). Moreover, Arora al. (2018) and Du & Hu (2019) examine GD with constant step sizes; however, Arora et al. (2018) requires unrealistic initialization conditions (large margin and low imbalance), while Du & Hu restrictive assumptions lead lo convergence rates That substantially differ from those of non-overparametrizel models. To the best of our knowledge, Xu et al. (2023 is the only work deriving a linear convergence rate for linear networks trained via adaptive step-size GD, but it imposes restrictive step-size constraints, resulting in slow convergence when the initial loss is large.
2019 assumes large initialization and width. These
Notation. We use lower case letters to denote scalar, and capital letters A and AT to denote matrix and its transpose. We use Tmax (A) and (m)n(A) to denote the largest and smallest singular values of A, | and ||A]|2 2 to denote its Frobenius and spectral norms, and A[i to denote its (i,j)-th element. For a function f(Z), we use Vf(Z) = 2z fZ) to denote its gradient.
2 Preliminaries
In this paper, we consider using the GD algorithm to solve the following optimization problem and its overparametrized version
min e(W),
(Problem 1)
WEReXm
min WieRook KiERwok.WWaet Xo.
L(W,WW> e(WAW,))
(Problem 2)
We are mostly interested in solving Problem 2, which covers many problems, such as matrix factorization Koren et al. 2009). matrix sensing (Chen Chil 20013). training linear neural networks (Arora et al. 2018 Du et al.
201 8a
Xu et al. 2023). In particular, when ((W) Y XWI?, where X, Y are data matrices, Problom corresponds to training two-layer linear neural network with n inputs, hidden neurons, m outputs, and weight, matrices W, and W2 using the squared loss.
4

---

Published in Transactions on Machine Learning Research (04/2025)
2.1 Convergence rate of GD for Problem
In this section, review the analysis for deriving the convergence rate of GD for Problem We seek to derive the convergence rate of GD for Problem 1 with the following iterations,
W(++1) = ((t) T/ (t)),
(2)
where we will use {(t), Vē(t) as a shorthand for ((((())))(((() respectively.
Throughout the paper, we make the following assumptions.
Assumption 2.1. The loss e(W) is twice differentiable, K K-smooth, and p-strongly convex. Assumption 2.2. minw ((W) =0
Assumption 2. ensures the solution to Problem is unique. Morcover, commonly used loss functions such as the squared loss and the logistic loss with & regularization both satisfy Assumption 2.1 Assumption 2.2
is for the purpose of convenience and brevity of theorems in this work. This assumption can be relaxed (to have arbitrary E*) without affecting the signíficance of our results. Morcover, one have the following inequalities based on the above assumptions for arbitrary W, V E Raxm
K
((V) S S(())++ + (Ve(W), -- W) + 2 ||V
W
Smoothness inequality,
(3)
2 (e((W) 2 (((W)
PL incquality.
(4)
Since strong convexily implies PL condition, equation holds under Assumption 2.1 In $3 we derive the convergence rate of Problem 2 based on the argument of the local PL condition. To be consistent, we highlight the rolc of the PL condition here. Morcover, the analysis in 2. remains applicable when p-strong convexity is relaxed to μ-PL condition.
In (Polyak 1963 Boyd & Vandenberghe 2004], it was shown that whenever 0< t. the GD iteration equation achieves linear convergence. The derivation is based two ingredients: Descent lemma and the PL inequality where Descent lemma is derived from the smoothness inequality.
Descent lemma. Starting from the smoothness inequality in equation 3 GD iterates (W(++1), (()) to derive Descent lemma, i.c.,
one can substitute (V, W) with the
K n
((++))=(+++ - [[(++))-((0]3=(
2 VE(t)
Basod on the PL incquality in equation and Descont lemma above, one can see there strict decrcase the loss at each GD step
Km l(t+1) 1) S {(t) (n [e(t)]] < (1 2pm HKni)e(t).
(5)
where the fact that 0 < n/ < 2, implies 0 2pm. + μK < Moreover, the minimum descent rate in equation 5 is achieved when nt leading to the following linear convergence rate:
11 ((++1) (1 K
E(T) [00))
/4 K
E(0).
(6)
Tightness of the analysis. The previous analysis guarantees lincar convergence rate for any arbitrary non-convex function that is K-smooth and satisfies the p-PL condition. Moreover, one can show that the rate equation optimal the sense that there exists a function that K-smooth and satisfies the μ-PL condition for which the bound on equation is met with equality. Therefore, one would be tempted to apply such an analysis to Problem 2, We will next show that overparameterizatoo introduces several challenges that prevent this analysis from being readily applied.
5

---

Published in Transactions on Machine Learning Research (04/2025)
2.2 Challenges in the Analysis of Convergence of Problem 2
optimized via GD
In this section, first introduce GD with adaptive step size solve Problem 2/ Then, discuss the main challenges in deriving the convergence rate for Problem 2 based the analysis in 2.1
Overparamctrized GD. We consider using GD with adaptive step size de to solve Problem 2
Wi(++1) (i] (2(++1) [2((t] 7. (2)(
where VL(Wi, W2) is computed via the chain rule:
ve(w),2 VL(Wi, W2) = T(Ve(W); WT, W2) = Ve(w) Wi]
(8)
Here T R + ( is a weight-dependenn linear operator that acts on V&(W). Thus, the gradient of L in equation can be viewed as a "skowed/scaled gradient' of that depends on W1, W2. It is this dependence on the weights Wi, W2 that makes it impossible to globally guarantee that equation 3 and equation hold, as shown next.
Proposition 2.1 (Non-existence of global PT. constant and smoothness constant). Under mild assumptions, the PL incquality and smoothness incquality con only hold globally with constants Pover = and K Kover 0c for L(W1, W2).
The proof of the above proposition can be found in Appendix B Moreover, we present a simple example in Appendix to help the readers visually understand Proposition 2.1
The nxn-exiseence of global PL and smoothness constants in the over- wee-paaaaeetrieed models prevents us from using the same proof technique $2.1 to derive the linear convergence of GD. In $ we show that although these constants do not exist globally, choices of the step size of GD, the PL and smoothness constants can be controlled for all iterates of GD. Thus, the linear convergence of GD can be derived.
can characterize them along itcrates of GD. Morcover, under proper
3 Convergence of GD for Problem 2
To deal with the challenges presented in 2.2 in 3..1 we propose a novel PL inequality and Descent Lemma evaluated on the iterates of GD for Problem 2 and show that the local rate of decrease per iteration for Problem 2 is worsened by the condition number of T compared with the convergence rate of Problem I Next, based the results in $3.1 in $3.2 we show that the condition number of T during the training be controlled by its initial value, which helps us deriving convergence rate for GD that depends condition number of T at initialization, the step size, K, and μ.. Morcover, we present a sketch of the proof of Theorem 3.2 to highlight the technical novelty and implications of the theorem in $3.3 Finally, in $3.4
the
and 3.5 we discuss how initialization and width influence the convergence rates derived in Theorem 3.2
Throughout the paper, we assume that the width satisfics h2 min m). This assumption ensures =" L* where L minw,
w, w, L(W,, W2), and thus solving Problem 2 yields the solution to Problem 1 When h min m), Problem 2 enforces rank constraint on the product. Thus, minw E(W) may not be equal to minw,, W, L(W1, W2). We are therefore interested in studying Problem 2 under the assumption h min(n, mn
which is the same setting in (Arora et al. 2018 Du et al. 2018a Xu et al. 2023).
3.1 Local PL Inequality and Descent Lemma for Over-parametrized GD
In 2.2 we saw that there does not exist a global PL constant or global smoothness constant for Problem 2 However, to prove that GD converges lincarly to a global minimum of [Problem 2 is sufficient. for Descent lemma and PL inequality to hold for iterates of GD. The following theorem formally characterizes the local PT. inequality and Descent lemma for Problem

---

Published in Transactions on Machine Learning Research (04/2025)
Theorem 3.1 (Local Descent Lemma and PL condition for GD), At the t-th iteration of GD applied to the Problem 2. Descent lemma and PL inequality hold with local smoothness constant Kt and PI, constant Pt, i.e.,
Kinh (VL()) 2
L((+1) - L(t) - (m 2
IIVOOI 2 IL(())
(9)
Moreover, if the step size n, satisfies n7 > 0 and n,K, < 2, then the following inequality holds
I(++1) (1) 2pmn + peK L(D)(mn)))
(10)
where
Ht = po (Tt)
(11)
K₁ = Ko (T)) / L(t+++66K mmax
max (
(Ti) v L(tme
(12)
and we use L(t) and Te as shorthands for L(Wi(1), W2(()) and T( ((- W2(()), resp.
The proof of the above theorem can be found in Appendix D Notice that Me, Kr are not actually constants since they vary w.r.t. the iteration index t. In this work, we adopt the convention to call them local PL and smoothness constants to be consistent with the terminology in $2
In 2.1 we showed that as long as one chooses = 17, with 0 O< K, GD in equation for Problem achieves linear convergence, with an optimal rate (1 - k) given when = However, we argue T Theorem 3.1 does not imply linear convergence of overparametrized GD even though there always exists sufficiently small N7 > 0 such that ne < The difference is due to the fact that er and K, are changing w.r.t. the iterations. Specifically, if lim it K; O, has lim inf 0<m=<2 p(n, () 1. Thus, equation 10 does not necessarily imply t-x00 that the product of the per-iterate descent [ 1) goes to zero.
Towards linear convergence. Nevertheless, if there exists The 0 that can simultaneously satisfy the constraint n, K, 2 and the uniformly bound 1 - 2 + n. Kenk I < 1, for all f, one can expect the linear convergence
S L(t+1) S p(n,, (t) & p L())
(13)
Guaranteeing a uniform bound as in equation 13 requires one to keep track and control the evolution of (t)), To, ne and L(t). In the next section, we will address these issues. For the time being, we focus next on how the pe; K, in Theorem 3.1 depend on the μr, K, L(t), n, and the current weights.
Characterization of μt. K₁. Theorem and smoothness constant, i.e, Mt: Ke, via a lime-varying linear operator T,. Specifically, the PL constant in equation 11 is the PL constant of E(W), i.e., μ,, scaled by ouin (Tt). Moreover, the smoothness constant in equation 12 consists of two parts. The first one is KTaaas)TT which represents the smoothness constant of E(W), i... K, scaled by 2 max
shows how overparanctrization affects the local PL constant
(Tt). The rest of the terms decrease to zero the loss L(t) approaches zero.
Effect of overparametrizztion decrease per iteration of Problem 2 trained via GD. Around global minimum, p(n; t) takes simplified form, i.., p(nt, f) =)=1] 22p0 (Te)n minn (Tz) (T)) By minimizing the p(T,, t) over one can obtain the optimal local rate decrease per iteration
optimization. Equation 10 in Theorem 3.1 characterizes the rate of
Jmin (Tx)
1 K
min O < KG
t) =1=
(14)
K o max (T,)
K( (TT)
where we use x(T)) to denote the condition number of Te. Compared with the optimal convergence rate of non-overparameetrieed model in $2.1 i.e., 1- # the local rate of decrease of overparametrized model is worsened by x(T)). Moreover, the local rate of decrease becomes faster as T, is well-conditioned.
In the next section, we will show that proper choice of initialization and step sizes ni does indeed lead linear convergence of overparametrized GD, and sufficient condition for Te to be well-conditioned.
7

---

Published in Transactions on Machine Learning Research (04/2025)
3.2 Linear Convergence of Problem 2 with GD
In this section, we first state theorem which shows that GD in equation 7 converges linearly to global minimum of Problemm 2 (See Theorem 3.2) under certain constraints on and the initialization. Then, based
on the convergence rate in Thcorem 3.2 we propose an adaptive step size scheduler that accelerates the convergence. We refer the reader to Table for the definition of varions quantities appearing in this section.
Table 2: Notation
Symbol
Definition
Description
D(t) wi (t) (t) W (2(t)
Imbalance: (almost) training-invariant invariant quantity
max Amax (-D(0)), 0)
X
max( Amaxl (D()), 0)
Quantities related eigenvalues of imbalance. They used define
max An (D(O)). + max
01,02. 91,
A+
x max (D(
A_
A_ max( Am ( -D(0)),
01
A A
(A 2 489 VIA + 2 489
Lower bound of (Ti)
4 X
432
A_
x 48
Upper bound on a max (Tt)
2
2
31
max (0, mmin )
K (0 1'
Lower bound on amin mmin((2 (1)
8
O mmxx (W*) +
K [W)( (0) W'
Upper bound Tmax (W)))WWi())
We now present our main result on the linear convergence of GD for Problem 2
Theorem 3.2 (Lincar convergence of GD for Problem 2). Suppose satisfies Assumption 2.1 and Assump-
lion 2.2 and given & 2 min(m, n), we assume GD in equation is inilialized so that on 0. Then there ewists 7max > O such that Vnoo 17, that satisfies
n70- Mmax and
no S n S min((1 + no: K.
+
(15)
one can derive the following bound for each iteration: pS ur S K2 S KG, and
L(t+1) < L(t)p()), 1) L(t)p(m, 0).
(16)
Morcover, based on equation 16 GD algorithm in equation
converges linearly
L(I+1) < L(0)p(no.0)
(17)
where
2TH + Tu ū= 14 (1 + 202 exxxx/ A = (1+1))) p(n,1) 1) = K, / 22K 0f 3) (2) A' exl.xxpp((poo) [13 2KL(())Amm] (
The proof of the above thcorem is presented in Appendix convergence for [Problem 2 under the assumptions that following remarks:
The above theorem states GD enjoys lincar
0 and certain constraints the We make the
Conditions on the initialization for linear convergence. From Theorem 3.2 we see that if the initialization (0), (0)) satisfics a₁ > then GD converges lincarly with an appropriate choice of the step size. The constraints on No ensure that f > 0 and < p(n.. 0) The assumptions on on has been studied in Min et al. (2022) where the authors show that an 0 when there is either 1) sufficient.
8

---

Published in Transactions on Machine Learning Research (04/2025)
imbalance A > 0 or 2) sufficient margin 31 > O, where A $1 is defined in Table 2 In $3.4 we present two conditions that ensure > 0 which covers commonly used Gaussian initialization, Xavier initialization, and He initialization. Please see $3.4 for a detailed proof and discussions.
Evolution of smoothness constant. One unique feature in our Theorem is the time-varying upper bound K, on the local smoothness constant K, along GD iterates. The constraints on no ensures that 0 < <p(no. 0), A < 1. Thus, K, monotonically decrease to K exp(v/o)aa w.r.t. t. The fact that K, is monotonically decreasing w.r.t. t suggests that the local optimization landscape gets more benign as the training procoeds. Thus in order to achieve fast rate of convergence, there is need for a time-varying choice of step size that adapts to the changes in the local smoothness constant K2 (because theoretically, the optimal choice is T7--- = K. based on equation 10
Requirement the step size. We have mentioned in the previous remark that time-varying step size could be beneficial for convergence. However, prior analyses (Arora et al. 2018 Du & Hu| 2019 Xu
et al. 2023) are all restricted to a constant decaying step size. The main reason is that one requires a uniform spectral bound on Te and (t) throughout the entire GD trajectory to establish linear convergence and such a uniform bound has only been shown under a constant or decaying step size. In our analysis, we show a similar spectral bound can be obtained even with a growing step size (See Appendix E as long as Th S (1 + 7 4n but not too much 7) S *े (ensures sufficient decrease in the loss every iteration). The first bound diverges to infinity exponentially fast, and the second bound has growing lower bound 1 which K,
1 monotonically increases to K eexp(v10)az Thus, initially, the step size is restricted to [Mo. (1 + +1)) & As the training goes on, the binding constraint becomes T) S R.. suggesting that GD can take a step that achieves the theoretically largest descent in the loss. In Davis et al.) (2024), the authors study GD with carefully designed adaptive step sizes (GDPolyak) applied to loss functions that exhibit quartic growth away from the solution set. Their analysis covers mom--oorrrrrmmttriizet matrix sensing problems and student-teacher setups involving a single neuron. GDPolyak differs from our method in how the step sizes are determined. In our approach, the step size at each iteration follows a classic optimization strategy: we estimate the next itcrate's loss based on the current loss and weights (sec equation 16), then choose the step size by minimizing this upper bound (see equation 18]. As a result, our adaplive schedule requires accurate characterization of the local PT, constant and local smoothness constant. By contrast, GDPolyak alternates between constant. step sizes and Polyak step sizes, relying on the latter to adapt to local features of the loss landscape.
Local rate of convergence In Theorem 3.2 we show 1(++1) I()+))) f) and AS when TBe satisfies certain constraints. When i is sufficiently large, or equivalently around any global minimum of Problem 2. the optimal rate of convergence is achieved as (via a proper choice of nc)
1 an +1 (1-exp(vmm)
min t) = 1- P 14
51
K K.
K
exp(
Notice the optimal local rate of convergence compared with the optimal convergence rate of the non-overparameetrized model, the optimal local rate of convergence of the overparametrized model is worsened by : which is an upper bound on <(T)). In $3.5 we provide a sufficient condition for To to be well-conditioned, i.e. en to be close to one, and numerically verfity this condition in Appendix[]
be arbitrarily close to 1 2 01 as Too decreases. Moreover,
as
Comparison in local rate of convergence with SOTA. We compare our results with prior works studying the same problem Arora et al. (2018); Du et al. (2018a); Xu et al. (2023) (See Table 3). Moreover, present detailed discussion the difference between proof techniques used in this work and prior work, and how it leads to different convergence rates. Please see Appendix for details.
Choices of the step size. Recall that for non-overparamttrizee GD, we have E(++1) S (1-2um-tuKn2)ee there exists an optimal choice of ni. = 11 Theorem 3.1 and Theorem 3.2] we show 1(t+1) < h(n,,t)((t) under certain conditions on T/e where t(),,t) E {p(n, t). p(n,,)) t)}. It is natural to use similar approach to select step size at each iteration. To achieve the optimal step size, it suffices minimize the upper bound on L(t+1) to achieve the most decrease
K that minimize the theoretical upper bound on é(t +1). In
9

---

Published in Transactions on Machine Learning Research (04/2025)
Table Comparison of convergence rates between prior work and our work.
local rate of convergence
loss
step size
initialization
Arora et
2018
squared loss
constant
D(0) R 0, 0,
1 Q
p Kag
Du et al.
2018a)
squared loss
decreasing
D(0) R 0, >
no explicit rate
Xu al.
(023))
squared loss
constant
@n
>
1 Q
paj Ka;
our work
general
adaptive
> @1
1 2 Kag
pad
at each iteration. The difference is that we have time-varying upper bound on I(t++) thus the minimizer Th depends on time, and our choice of Ty must respect our constraint on step size in equation [5 to the following choice for rf::
This leads
Th = arg min (1+ n
h(m, t).
(18)
Since p(n, f), p(ne,)) 1) are quadratic in terms of Tj, so the takes the following closed-form solutions depending on which upper bound to use:
min (11 n 1/2 no. K.
if h((n)) = p(T,, t),
(19)
min (11 + 76 / no: k.)
if h(me)) p(n.tt.
The above choices of the adaptive step sizes satisfy the constraints in Theorem 3.2 so they both guarantee linear convergence for over-parametrized GD. Moreover, such choices of give us the following theoretical bound on L(t+1), i.e.,
L(++1) S L() II h(mk, k).
(20)
In 44.1 we provide numerical verification of the close aligument between the theoretical bounds stated above and the actual convergence rate. We also observe an accelerated convergence when employing the step sizes
specified in equation [8 compared with the one proposed in Xu et al. (0223) and Backtracking line search. We refer the readers to 4 for simulation results.
3.3 Proof Sketch
In this section, we present proof sketch of Theorem 3.2 that highlights the key technical contributions. In Theorem 3.1 we established local PL inequality and local Descent Lemma. Moreover, one can show the upper bound
L(L+1)+L(() (1-2 n1 te K2 K.n() = L(1) p(n,,))
(21)
As discussed in 33.1 this inequality alone does not guarantce lincar convergence. For instance, lim; K, 0, then lim4-+o Plt !) 1. To show there exists 0 < p such that p(m,, () < أ holds for all t, the following two-step approach in similar way as it done in Xu et al. (202))
Step one: uniform spectral bounds for To and W. First, we show when 7/t is controlled, one can have the uniform spectral bounds on Te and (t).. The following lemma characterizes this property formally. (Uniform spectral bounds on Te, W().... Under the same assumption and constraints in Theorem 3.2 one has the following uniform spectral bounds on Tr, W(t)
Lemma 3.1
a1 + 2a (1 exp( TT)) < 2 (T,) < 2 m (T:) S a2 exp( (To)
(22)
81 < min (W(1) < mmax (W(I) < 82.
(23)
10

---

Published in Transactions on Machine Learning Research (04/2025)
The above lemma shows the uniform spectral bounds on T. and W, depending a1, ag S1, B2 and no- Moreover, the bounds on the singular values of Te can be arbitrarily close to 01,02 yoo approaches zero.
Similar results have been derived in Xu et al. (2023) where the authors show uniform spectral bounds of Te. W2 for constant step size GD. Our proof strategy is similar to Xu et al. (2023) which relies on the fact that when the GD enjoys linear convergence, the change of imbalance during the training small. For constant step size GD, we can characterize the change using the step size n and the convergence rate p L.e. [1) - D( r S O ( In this work, we discover when we allow step size to grow but not. too fast, i.c. 7) (1 + n we still control the change of imbalance, i.c. [ D(O)Ilr (1-2)) n where A = (1 + (vm, 0). This observation helps us derive the uniform spectral bounds for Tr, ((I) while allowing the step size to grow. The bound on the change of imbalance is not restricted to scenarios when loss converges linearly. In Ghosh et al. (0025), the authors demonstrate that the imbalance gap decreases at a linear rate, even in the rdge-of-stabilitg regime. We mention this work here an additional reference for readers interested in such phenomena.
Step two. Second, we employ an induction-based argument to show that based on Lemma 3.1 show μ Z μ. K; Kr and L(t) converges linearly with the rate Po-
one can
Lemma 3.2 (Induction step to show Kt is bounded and L(t) converges linearly.). Under the same assumption and constraints in Theorem 3.2 assume L(t) enjoys linear convergence with rate p(no, 0) until ileration k; then the following holds for ileration k+1
Pea 2 p. KEt1 S KK-1, with p, Kit1 defined in Theorem 3.2. Moreover,
(24)
show
Pl k+1) p(men k+1) S p(no, 0).
(25)
Equation 24 is direct consequence of Lemma 3.1 and the induction that L(t) enjoys linear convergence until
iteration k. We can lower bound μ4 by subsituting miin (Tx) with the bound in equation 22 and upper bound
K₂ by subsituting maa (T&), max (k)), [(k) with the bound equation 22 equation 23 and ((0) 0)'
respectively. Basod on these results, one can derive the following upper bound on L(k+1 under the same
constraints on ne in Theorem 3.2
KAnR
L(k+1) S L(k) - (TT-
Local Descent lemma in Theorem 3.1
(26)
Kinh Lkk)
S L(k)- 21(m 2
Local PL inequality with p
(27)
= (1 27Tm + TTKK ni) [(kk)
Use K₁ & Kg < Ro Use constraints on no and equation 15
(28)
< (1 27T + TT (.k)))
(29)
S p(n. (k)
(30)
In summary, we show GD in equation achicves lincar convergence with the rate P(1. 0) under the constraints
on N7 and the assumption a1 > in this section. Moreover, we show the local rate of convergence depends on
an In the next section, we theoretically show that commonly used random initialization leads to an > O,
and large width and large initialization scale leads to faster local rate of convergence, i.e., en 02
3.4 Proper initialization ensures 01 >
In $3.2. we see that under the assumption ai 0, Problem 2 trained via GD in equation 7 converges linearly to global minimum under certain constraints the step sizes. In this section, present two conditions on the initialization that ensure 0 > O.
We first show two conditions on the initialization that ensure 0i > 0. Lemma 3.3 (Mild overparamctrization ensures aj > 0). Let (i) W2 are initialized entry-wise i. d. from continuous distribution P. When h 2 m + n, a1 > holds almost surely over random initialization with Wi (0), W2(
11

---

Published in Transactions on Machine Learning Research (04/2025)
Lemma 3.4 (Lemma 1 in (Min et al. 2021 ). Let (O). W2(0) are initialized entry-wise i.i.d. from N(0, kāy 2- For 48 > 0 and h poly(n, m, }), with probability - over random initialization with W1(0), (O) the following holds 2 h1-2p
with 3 S S }
We refer the readers to Appendix F for detailed proof. Both lemmas presented above ensure ai > 0 under different conditions on the width. Compared with Lemma 3.4 Lemma 3.3 considers wider range of distributions that include Gaussian distribution and uniform distribution. Thus, commonly used random initialization schemos, such Xavier initialization (Glorot & Bengio, 2010) and He initialization (He et al., 2015), lead to ai > 0. Moreover, the requirement of overparametrization in Lemma 3.3 is mild compared with the one in Lemma 3.4 i.e., h 2 m + n versus h 2 poly(n, m. }). As result, Lemma 3.3 can be applied to more general overparametrization. On the other hand, the conclusion of Lemma 3.3 is weaker than Lemma 3.4 in the sense that Lemma 3.3 only proves a1 > 0 but does not characterize its magnitude while Lemma 3.4 characterizes the lower bound on on will increase as h increases.
3.5 Large width and proper choices of the variance lead to well-conditioned To
In this section, we show that a proper choice of iniiaalitiion and width can lead to well-conditioned To, ie., 1. Based on the results in $3.2 we show that the local convergence rate of overparametrized model trained via GD can match the rate of the non-overparameetrieed model.
Theorem 3.3. Let (O),), (2) initialized entry-wit wise i.i.d. from (O, RT5) with 1/ < V8 e (0, 1), poly(m, }), with probubilily 1 - 6 over random inilialization Wi (0), W2(O), the following holds
Omin (To)
an 21 ((h2p-1).
(31)
max (To)
a2
We refer the readers Appendix F.1 for detailed proof. The above theorem states the condition number of To approaches one when the width increases to infinity under suitable choices of the variance of Gaussian initialization. Therefore, increasing the width can lead to fast convergence rate.
In Appendix G we will numerically show that with proper choices of the variance of the initialization and the adaptive step size proposed in $3.2 a large width will lead to well-conditioned T, and the convergence rate of the overparametrized model can asymptotically match the rate the non-overparametrizee model.
4 Experiments
In this section, we first present empirical evidence that Theorem actual convergence rale under different initialization in $4.1 Then, in 4.2 we compare the convergence rate of GD using the adaptive step size proposed in equation backtracking line search. Throughout the experiments, we consider Problem 2 with squared loss
3.2
provides a good characterization of the
19 in Section 3.3 of Xu et al. (2023) and
1 -111 -XW,W, L(W1, W2) = 2
(32)
where X, Y 100x10 are data matrices and Wi, W22 E 100h are the weights. This can be viewed a two-layer linear network with input and output dimensions 10 and the width of the hidden layer to be h. Throughout the simulations we choose h E (500, (000,4000) We choose c 0.5, d = 1.01 in Theorem F. The initialization of the weights and generation of data matrices are as follows: (1))), W2(0) E IR 1000 and have entry-wise i..dd samples drawn from (0, 1). We generate X as a random orthogonal matrix, and Y = XW) + a&e where E IR 10x10 10 and are entry-wise i.id. samples drawn from N/ 1). When o² is large, the initial loss is large, thus the margin small. Morcover, we experimentally observe that the initial imbalance grows w.r.t. h. The choices of h and allow us to test our results in different regimes.
4.1 Evaluation of the Tightness of the Theoretical Bound on the Convergence Rate
Figure compares the actual convergence rate of L(t) versus the theoretical upper bound in $3.2 for different choices of and h, and dissimilar 0n In all cases, the theoretical upper bound follows the actual loss well. ag
12

---

Published in Transactions on Machine Learning Research (04/2025)
Stop
with Equattenl99) Equabion220y with Atrp) - pint, B
plm.
in Equation15) High- pla.
Equationn18)
Equatlen(20) with hiai 201.
=0.5549 h=500. o²=00.01
m:0.5545, 5000 0²=0.1
mm05550 5=500, o'=1.0
lopio Ln
Jopa LEO
Japha LiD
Iterations
iteraiions
Iterations
=000744,, 000000
-.6674 1166666h-1100 1-00 0'=0.1
0.66725 of
logio Lie
logio LEBI
log LEB
ierations
teraaions S -008829,, 4-4000 o²-0.1
Inerations 0-88286, /-4000 oP-1.0
& --0.8201, A-0000 000000000000
logno Up
logn LRO
logha Lit)
lterations
Rerations
lterations
Figure Tightness of the theoretical upper bound versus reconstruction error L(t) for different choices of step size in $3.2 shown in different colors. We run the simulations for nine different settings of initialization and data generation. For each setting, repeat the simulation thirty times. The triangle lines represent the thcoretical upper bound on the training loss in equation 17 and equation 20, The solid lines represent the mean of the log of the reconstruction error L(l). The shaded area is the mean of log10 L(l) plus and minus one standard deviation.
Moreover, we observe for each adaptive step size scheme, the theoretical bounds and the actual rate of convergence become slower decreases. This is because our bounds the local rate of convergence depend and the smaller az?
ai the slower the convergence rate.
4.2 Comparison with Prior Work and Backtracking Line Search
In this subsection, we compare the adaptive step sizes proposed in Xu al. (2023), backtracking line search with the step sizes proposed in equation 19 with h(n.) p(n., t). We set the hyperparameters of the adaptive step size scheme proposed in Xu et al. (2023) to be C1 0.5, 02 1.5, which is the same setting in their sinulations. We refer the readers to Appendix G for detailed descriptions of backtracking line scarch.
19 achieves the fastest convergence compared with Xu et al. (2023) and backtracking line search in different settings. This because for the adaptive step size schoduler this work, the step size at each iteration has closed form (See equation 19), thus the time for picking the optimal step size per iteration is negligible. The only time-consuming part is to find yo since one needs to solve equation 107 and equation 102 get Tmax- For the step size proposed in Xu
Figure 2 shows that the step size choice proposed in equation
13

---

Published in Transactions on Machine Learning Research (04/2025)
tackrrccing
Stop by
Equatoor( Nigh (lpp))
=0.5509 h=500, o'=1.0
8:00.5599 h=500. o²=1.0
mm0.5509, A=500, g²=1.0
log10 MN
logho LN
loghn UN
Seconds
Seconds
Seconds
0-.6725 h=1000, df=1.0
- 100000 o'=10
0.06725 h-1000, o'=1.0
logio LLD
Joggo UN
logro Lt
Seconds
Seconds -00.8266, /-4000, o²=1.
Seconds
3) S-0.8266, 11-000, g²-1.0
E-0.8286, /-4000, 02=1.0
loglo Lif)
logia Lit)
logno Lit)
Seconds
Seconds
Seconds
Figure 2: Evolution of the loss and of the step size for different choices of the step size schedule under different initialization and data generation. We run the simulations thirty times. For each setting, repeat. the simulation thirty times. The solid lines represent the mean of log10 of the reconstruction error L(t). The shaded area is the mean of logio L(t) plus and minus one standard deviation.
et al. 2023 the algorithm consists of solving third-ordes polynomial at each iteration, which results in larger computational time. Moreover, the adaptive scheduler proposed our work follows a sharper characterization of the local convergence rate than Xu et al. (2023), and the adaptive step size scheduler in this work theoretically converges of order an faster than the one proposed Xu et al. (2023). For the backtracking line search algorithm, since the algorithm iteratively searches for the step size at each iteration. Therefore, the time cost for cach itcration is high well.
as
5 Conclusion
This paper studies the convergence of GD for optimizing two-layer linear networks with general loss functions. Specifically, we derive a lincar convergence rate for finite-width networks initialized outside the NTK regime. We use common framework for studying the convergence of GD for the non-convex optimization problem, iee PL condition and Descent lemma. Although the loss landscape neural networks does not satisfy the PI. condition and Descent lemma with global constants, we show that when the step size small, both conditions satisfy locally with constants depending on the singular value the weights, the current loss, and the singular value of the products. Furthermore, we prove that the local PL and smoothness
14

---

Published in Transactions on Machine Learning Research (04/2025)
constants can be uniformly bounded by the initial imbalance, margin, PL constant, and smoothness constant of the non-overparamellerizel model. Finally, we provide an explicit convergence rate dependent on margin, imbalance, and the condition number of the non-overparrmeeteriee model. Based on this rate, we propose an adaptive step size scheme that accelerates convergence compared to a constant step size.
Acknowledgements
The authors acknowledge the support of the Office of Naval Research (grant 503405-78051), the National Science Foundation (grants 2031985, 2212457 and 2330450), and the Simons Foundation (grant 814201).
References
Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. arXiv preprint arXiv:1810. 02281, 2018.
Stephen P Boyd and Lieven Vandenberghe. Conver optimization. Cambridge university press, 2004.
Yuxin Chen and Yuejie Chi. Spectral compressed sensing via structured matrix completion. In International Conference on Machine Learning, pp. 414-422. PMLR, 2013.
Zhengdao Chen, Eric Vanden-Eijnden, and Joan Bruna. On feature learning in neural networks with global convergence guarantees. arXiv preprint orXiv:2204.10782 2022.
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over-paraamterizeed models using optimal transport. Advances in neural information processing systems, 31, 2018.
Lenaic Chizal, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in Neural Information Processing Systems, 32, 2019.
Damek Davis, Dmitriy Drusvyatskiy, and Liwei Jiang. Gradient descent with adaptive slepsize converges (nearly) linearly under fourth-order growth. arXiv preprint arXiv:2409.1971, 2024.
Zhiyan Ding, Shi Chen, Qin Li, and Stephen J Wright. Overparameterrzation of deep resnet: Zero loss and mean-field analysis. J. Mach. Learn. Res., 23:48 1, 2022.
Simon Du and Wei Hu. Width provably matters in optimization for deep linear neural networks. In International Conference on Machine Learning, pp. 1655 1664, 2019.
Simon S Du, Wei Hu, and Jason D Lec. Algorithmic regularization in learning deep homogencous models:
Layers are automatically balanced. Advances in Newral Information Processing Systems, 31, 2018a.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over- parameterized neural networks. arXiv preprinl arXiv: X10.00000.0204,, 2018b.
Sara Fridovich-Keil, Fabrizio Valdivia, Gordon Wetzstein, Benjamin Recht, and Mahdi Soltanolkotabi. Gradient descent provably solves nonlinear tomographic reconstruction. ArXin, pp. arXiv 2310, 2023.
Avrajit Ghosh, Soo Min Kwon, Rongrong Wang, Saiprasad Ravishankar, and Qing Qu. Learning dynamics of deep matrix factorization beyond the edge of stability. In The Thirleenth International Conference on Learning Representations, 2025.
Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien. Implicit regularization of discrete gradient dynamics in lincar neural networks. Advances in Neural Information Processing Systems, 32, 2019.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference artificial intelligence and statistics, pp. 249-256. JMLR Workshop and Conference Proceedings, 2010.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification, 2015.
15

---

Published in Transactions on Machine Learning Research (04/2025)
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEFE Conference on Computer Vision and Pollern Recognition (CVPR), June 2016.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximal-gradient methods under the polyak-tojasievicz condition, 2020.
Yehuda Koren, Robert Bell, and Chris Volinsky, Matrix factorization techniques for recommender systems. Computer, 42(8):30-37, 2009.
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as lincar models under gradient descent. Advances in neural information processing systems, 32, 2019.
Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Loss landscapes and optimization in over-parameterried non-linear systems and neural networks. Applied and Computational Harmonic Analysis, 59:85 116, 2022.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, E16(33))E7665--7777 2018.
Hancheng Min, Salma Tarmou,, René Vidal, and Enrique Mallada. On the explicit role of initialization the convergence and implicit bias of overparametrized linear networks. In International Conference on Machine Learning, pp. 7760-7768. PMLR, 2021.
Hancheng Min, Salma Tarmoun, René Vidal, and Enrique Mallada. Convergence and implicit bias of gradient flow overparametrized linear networks. arXiv preprint ar arXic:2105.0683. 2022.
Hancheng Min, René Vidal, and Enrique Mallada. On the convergence of gradient flow on multi-layer linear models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.). Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. (1850-24887.0 PMLR, 23-29 Jul 2023. URL https://proceeddingg mlr wttp:/ngs........s/////22222. wwt.:pessssv/mmmmeesee mtmmddmmmmmmmmmmmmmmm
Gabin Maxime Nguegnang, Holger Rauhut, and Ulrich Terstiege. Convergence of gradient descent for learning linear neural networks. arXiu preprint arxiv:2108. 100.00000000000 2021.
Quynh N Nguyen and Marco Mondelli. Global convergence of deep networks with one wide layer followed by pyramidal topology. Advances in Neural Information Processing Systems, 33:1 13:11961-1197, 2020.
Boris Teodorovich Polyak. Gradient methods for the minimisation of functionals. USSR Computational Mathematics and Mathematical Physics, 344):864-878, 1963.
Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. ar preprint arXiv:1312. 61200, 2013.
Justin Sirignano and Konstantinos Spiliopoulos. Mcan field analysis of noural networks: law of large numbers. SIAM Journal on Applied Mathematics, 80(2):725-755 2020.
Salma Tarmoun, Guilherme Franca, Benjamin D Haeffele, and René Vidal. Understanding the dynamics of gradient flow in overparameterized linear models. In International Conference on Machine Learning, pp. 10153-10161. PMLR, 2021.
Ashish Vaswani, Samy Bengio, Eugene Brevdo, Francois Chollet, Aidan N Gomez, Stephan Gouws, Llion Jones, Lukasz Kaiser, Nal Kalchbrenner, Niki Parmar, et al. Tensor2tensor for neural machine translation. arXiv preprint 1803 2018.
Roman Vershynin. digh-dimmessonal probability; An introduction with applications in data science, volume 47. Cambridge university press, 2018.
16

---

Published in Transactions on Machine Learning Research (04/2025)
Ziqing Xu, Hancheng Min, Salma Tarmoun, Enrique Mallada, and René Vidal. Linear convergence of gradient descent for finite width over-parametrized linear networks with general initialization In International Conference on Artificial Intelligence and Statistics, pp. 2262 2284, PMLR, 2023.
17

---

Published in Transactions on Machine Learning Research (04/2025)
A
Preliminary lemmas
In this section, we present preliminary lemma which will be used in the following sections.
Lemma A.1 (Inequality on the Frobenius norm). For matrix A, B, D, have
(A, |||| P |B]|F
(33)
2|| < |||| 3 IB
(34)
|AB + CDI < [o (A) + ( [1 + [D]
(35)
|| P + 1B - 2|| + B 1
(36)
min (A) < |AB P
2
o max (A)
(37)
min (B) < ||B|i S _ (B)
(38)
Lemma A.1 has been derived and used multiple times in prior work. We refer the readers to Appendix C in
Xu et al. (023) for detailed proof.
Lemma A.2 (Singular values of T). The largest and smallest singular values of TT are given as
min (7) muin (W1) + o - (W2),
max (T) = o max (W1) + 2 max (W2).
(39)
Proof. First, one can see
T o T(U:W,, W2) UWWW, + W.W,U,
(40)
where T is the adjoint of T. Then, we use Min-max theorem to show
Amin (T 7) = iin (Wi) + oi (W2), Amax o oT) = o m + max (W2).
(41)
Let the singular value decompositions of W1, W2 be
12
W1 UNEVV
Σ 1
Σ U2
(42)
where r1 rank(Wi), rank(W2), and (or fo2.1 are of descending order. Then, one has the following
Amin (T'o)) min (U,UW2W W.W,Tu
|| Fr=l
= min (U, (U,UWWW,)) + min (U, (U.WW.TTU |01r=1
|| r=
2 m (W1) + o (W2).
(43)
On the other hand, if one choose U Ui 2 the following equation holds
(U,WWWW + W.W,TU)
(11,,, 2 vi 22 W2W,
WWW, WT U1 )
=(1mm M2 11 U2
T
Σ
o "2,
(u1m u2 Σ of "1, (.
22
Σ
Vin Unn u u2
+
o tr(ue, un Ul
1=1
= + 00,22
(44)
18

---

Published in Transactions on Machine Learning Research (04/2025)
U2,r2 = holds for all i #12. Therefore, where the last line is based on the fact that 01, =0, u2,r2
based on equation 43 and equation 44 one has
Amin (T T) = min (Wi) + a omin (W2).
(45)
Similarly, we can show
Amax (T* TT p omax (W2) + o ( (W2).
(46)
Lemma A.3 (Singular values of random matrix). Given m,n E N with S Let A E guxm be random matric with i..d. standard normal entries. For any 8 > O, with probability at least 1 - 2 exp(5?), one has
Vn m ö Tmin (A) S maxx (A) S vn vm
(47)
The proof of this lemma can be found in Vershynin (0018).
B Non-existence of Global PL Constant and Smoothness Constant for Problem 2
In this section, we show that under mild assumptions, the PL inequality and smoothness inequality can only hold with constants Pover and Kover = oo for Problem 2 We then make the following assumption on Problem Assumption B.1. (W1, W:2) = 0) is not a global minimizer of Problem
Based on the above assumption, one has the following proposition.
Proposition B.1 (Non-existence of global PL constant and smoothness constant). Under Assumption B. the PL inequality and smoothness inequality can only hold with constants Hovvr = and Kover = oc for T.W1, W2).
Proof. We first show Mover = 0. The gradient L. is given as follows
VLL W2 = [/ + wall2
(48)
Notice when W1, W22 are zero matrices, the RHS of the above equation is zero. Therefore, we have
[[V,WW, 32 =0. 0. On the other hand, under Assumption B.1 since (W1, W2) (0,0) not a min-
imizer of Problem 2 have L(W,, W2) 40. Thus, the PL inequality can only hold globally with Pover 0,
I/V,(W, W2 [/ee(W +[V((W)) Will P LWW, W2).
(49)
Then, we show Kover for Problem 2
We consider the smoothness incquality evaluated arbitrary
(W1, Wo) and the minimizer (Wi. W2) of Problem 2
+ Kover Z" - 2
L(W,, W2) < L(W W - (VL(WT, W3), 2 -
(50)
2
where we use Z, z* in short for (W,, W2), (Wi, W2). Since (W*, W minimizes Problem 2 VL(Wr. Wr) = mj+ and L(W W3) = 0. Thus, equation 50 is equivalent to the following
we have
2L(W,, W2)
28WW,
Kover
(51)
1|Zw Zw- 11? Zzw Zw- 11
On the other hand, since ((W) is p-strongly convex w.r.t. W, the following inequality holds for arbitrary U,V
E(U) > ((V) + (VE(V), vy 4
(52)
19

---

Published in Transactions on Machine Learning Research (04/2025)
Loss Velue of tix)
PL Constant of (is,.)) xx)
600
Xg) - xyagg
Smoothness Constant of Lixg-xs)
3
9
X
Figure 3: Plot of a toy example illustrating the loss function ((x) and its overparametrized version L(x1, x2); along with the corresponding local PL constant Mover (T1, and smoothness constant definitions of ((x), (21,,T)), Mover (T1, 12), and Kover (U1, are given in equation 57 equation 58
Kover (#1, T2). The
and
equation 59
We substitute U,Y V with WWW, TT in equation 52
we have
e(W W,
2/ [WW - (W 11
(53)
Finally, we combine equation 51 and equation
53
and derive the following lower bound on Kover
28(W, BT
Kover 2
Based on equation 51
Vw - Zw- 11
w - -W((W)) 1/:
Apply equation 53 to W()
IW, W 11 W
MlW,,,," WWWWW-WW WW +WWW, (W T
[| - [Wi-W/ll + 112- - W2
M][(W] W +W; WG) TI
[W, - Wi + [ WW
[[(1] WT ||
W2) T>
IV
Apply Lemma A.1
2| Wi + 1W2 W 1
min (Wa) [W Wr +o min (Wi) W W IP:
IV
NF
Apply Lemma A.1
(54)
W1 Wr + Il22 W2
Similarly, also derive the following lower bound on Kover
/ Giin (W2) /
W1 W 2- + 2 (W1 W W3 PP
Kaver
(55)
2
[[1 W + W22 W
We take the average of the lower bound on Kover in equation 54 and equation 55
(innn + oii (W2) W1 Wi | + (o (W1) + o 2ii
(W))
W2
Kover 2
[ Wr + W W
1V
min oi (W1) + o (WT), 02 (W2) + 02 oi (W?))
min o (W1) 02 min (W₂)
(56)
Due to the arbitrary choices of Wi, W2, can let miin minn(W) and amin (W2) to be arbitrarily large, thus the smoothness incquality for Problem 2 can only hold globally with Kover oc.
To illustrate Proposition 1. we examine how the PI. constant and smoothness constant evolve in a simple
one-dimensional setting. Specifically, we compare the non-overparumetrizel loss
((x) (x-1)
(57)
20

---

Published in Transactions on Machine Learning Research (04/2025)
with the overpanametrized loss
(2)) (2).r2 1
(58)
In this simple example, one can exactly compute local smoothness constants and PL constants as follows
Hover (2 (2) (VL(I1,32))
2L(x1, = 2 2²2
Kover (2), max (2 +s+ [ -
40f (2r132 1)²
(59)
2
Since the local smoothness constants and PL constants of (2)1 now depend on the input, we use Pover (x1, (2), (2) to denote them.
In Figure 3 we plot the loss landscapes for both {(x) and L(x). We observe that while {(x) is strongly convex and smooth, with global PL and smoothness constants p= K= = 1, the overparameterized loss L(x) is non-convex. Moreover, its local PL constant Mover (21, vanishes when 21 =2 =0, and its local smoothness constant Kover (211 x2) diverges as + ||>|| + Consequently, if one attempts to enforce the PL inequality and Descent Lemma with global constants in the overparameterized setting, the only possibilities are
Hover min Mover (21, and Kover max Kover Kover 22) (2), 21.22
=00
(60)
21,32
C A Detailed Comparison with Prior Work
In this section, we present a detailed comparison to Arora et al. (2018) highlight the difference technical details and improvement on the convergence rate.
| Du et al. (2018a):
Xu et al. 2023 to
Summary of the strategy of proof in (Arora et al., 2018 Du et al., 2018a; Based on GD update in equation 7 one derive the following update on the product W(t)
Xu et al. 2023
W(++1)= = W(t) - Tr T((VE()) + (()
(61)
where T is the adjoint of T. Then, substituting equation 61 into the smoothness inequality of the non- overparametrized model in equation 3. we can derive the following upper bound on the loss at itcration t+1 using the loss at iteration (Also see Lemmaa31 in Xu el al. (2023
Lemma C.1. If at the t-th iteration of GD applied the over-porumetrized loss L, the step size m satisfies
KTπ 1o max (T:) + ne |M/
min (T)) -/V
(62)
then the following inequality holds
I((+1) S ((1.))
(63)
where
p(n,t) 1- 2mp (T)) K pm á max (T:) + 2nT max (W())I + 2n 2222nkKo o (anaaxTTTT)) (Tx) (F)) + t 2Ko max (
(64)
Improvement of the local rate of decrease. First, one can see the local rates of decrease in both work are polynomials of degree four and depend on n), Vl(() and singular values of Te, W,. Moreover, around any global minimum, i.c., L(t)= 0, [[V[(t)]lr = O, we have the following local rate of decrease per iteration
1 - 2 min (Ti) + 1/ K uo max (Te)
local rate of decrease in prior work,
1 2 min (Ti) + 7/ K no 2 (T)) 02 max (Tt)
local rate of decrease in this work,
(65)
21

---

Published in Transactions on Machine Learning Research (04/2025)
and the optimal local rates of decrease regardless of the constraints the n
o min
(Tt)
1
optimal local rate of decrease in prior work,
K
(T))
max
μ
min (Tt)
optimal local rate of decrease in this work.
(66)
K
max
(Tt)
Thus, one can see our characterization of local Descent lemma and PL inequality leads to faster local rates of decrease compared with prior results by
(T)) Nevertheless, equation 66 does not imply linear convergence
(Ta)
since if lim munn (T)) 10000
mmax (Tc) = 0, one would not expect sufficient decrease per iteration. In order to show linear
convergence, one needs to provide uniform lower bound T (7t) max (Ti)* Vt.
Improvement of the local rate of convergence. In this work, we show when the step sizes satisfy certain constraints (See Theorem 3.2 there exist uniform spectral bounds for the condition number of Te, i.e.,
Gminn (T+) max (7:) < c(1o) ag vt where a1, a2 only depend on the initial weights and c(no) is constant approaching one 17o decreases. Thus, the optimal final rate of convergence derived in this work is
optimal local rate of convergence in this work.
(67)
K 02
In prior work, the rates in (Du et al. 2018a Arora et al. 2018
extremely slow in practice (See Section 4 in (Xu et al. 2023)) ). In (Xu et al., 2023 the authors introduce two auxiliary constants O << < <1,02 1, VL.
( and show that one can uniformly bound the condition number of Tr during training, i.e. (7.) < cza2 Moreover, they enforce problem-dependent assumptions the choices of C1 According to Clain E.1 in [Xu et al. 2023), is at most and can be arbitrarily small when the initial loss is large. Thus, the local rate of convergence in [Xu et al. 2023 is at most, in our notation,
11 af K
1-
optimal local rate of convergence in (u et al. 2023
(68)
c2
When comparing equation 67 and equation 68 one can directly conclude the local rate of convergence derived in this work is much faster than the rate derived in (Xu al.) 2023). Moreover, the optimal local rate of convergence of the overparametrized model in this work is different from the optimal rate of convergence of the non-overparametrrieed model up to a factor of 2i, which shows overparametrization has benign effect if one can control @1 through properly initialization the weights. However, such results are not shown in the ng work of Arora el al. 2018 Du et al. 2018a Xu et al. 2023
Proof of Theorem 3.1
In this section, we present the proof of Theorem 3.1
Theorem D.1 (Restate of Thcorem 3.1 At the t-th iteration of GD applied to the Problem the Descent lemma and PL inequality hold with local smoothness constant K₁ and PL constont pr, i.e.,
Kinh
L(++1) - (n 2
(VL()))
[IL(())
LL).).. (()(t)
(69)
Morcover, if the step size satisfies 1) > O and 1y K2 < 2, then the following inequality holds
L(t+1) S L(t 2pm + pcKmn = L(t)p(n,,t,
(70)
where
14 min (Tr),
(71)
K, = Ko 2 o (W)) + 2//2/KLL 2()
(72)
and we use I((), Te as shorthand for L(Wi (1), (1),,, T( (W), (W))))(())) resp.
22

---

Published in Transactions on Machine Learning Research (04/2025)
Proof. We first show that the local PL incquality holds.
(t()) (e((t)TG(C |
VL(1))
(e(())(()13 + ((e(L) W
2 min (W2 ( miin (Wi (O))I
Apply Lemma A.1
2 min ma (W)()))L(1) 2 (WI())
Apply PL inequality of &
(t).
where the last equality uses the fact that min (T:) o min (Wi(t)) o (Wa(t)) Then, we show that Descent lemma holds with local smoothness constant Kr. We can view L((+1) using the following second Taylor approximation,
(1
= L(tt) - m 1 [V + ((IVL( f(11 ((1-)(a (1)r)(g,,))) H()gg))
(73)
where we use Z411 Ze short for (Wi+++1,, 1), (2(++1)), (Wi(t), W2())) respectively, and 9e denote the unit vector of the gradient direction. Moreover, the (() is defined as follows,
VL(t) [V
H(r) (C)
(tt)
(74)
Notice the integral in the equation 73 does not have closed-form solution. We the following two-step approach to derive upper bound the RHS of equation 73
Step one. We first show that one can upper bound (gi, H (0)gi) using the singular values of T,, K and L(f).
The following lemma characterizes formally.
Lemma D.1 (Upper bound on (g., (0)g:)) We have the following upper bound
(gt. (0) Ko² max (T)) + 22K((t)
(75)
The proof of Lemma D.I is presented at the end of this section.
Step two. Then, for any E 1), we can show Kg,, (H - H(r)) is bounded, which leads to an upper bound on (gt, H())))) The following lemma characterizes the upper bound on (ge, (T)g.)
Lemma D.2 (Uniform upper bound on (g., H(Tgg:)) For any T E [0, 1), we have
(9, H(T)g))) Ko aa (7))
(T)) 2KL(t)m Kr
max
The proof of Lemma
D.2 is presented at the end of this section.
23

---

Published in Transactions on Machine Learning Research (04/2025)
Based on equation 73 and Lemma
D.2 one can derive Descent lemma
/ (1 (1--)) T) H(()g
-(t+1)= - m (V P (
Equation 73
$ L(t) TL ( P ( /, (1- max kig. H ( (T)g:)/dr
/
$ L(t) TL (V // + V
(1-77 (1-T)_e KdT
Lemma D.2
TiKK
L(t) T (L)
[L(E)
2
oi Kt
[(1) (17-
(76)
Therefore, Descent lemma is proved.
Now we present the proof of Lemma D.1 and Lemma be used in the proof
D.2 We first define the following quantity which will
(s) T1(W((1) ST V w; T(O,, W2(() 8mp Ww, T(t))
(77)
A(s) = (t) -t) (t),2()) +Wi(t)Vw, (w,_L())W(((+++(((Vvvv))
+ L(),
(78)
where A(8) is the product of (1) - 81) Vw; [(0) and (2( 88 Vw, I(I). Moreover, we have M(0)= = L(l),
( and {(As))
Then, we present several lemmas that will be used in the proof of Lemma D.1 and Lemma D.2
Lemma D.3. Given (tt E R W2(t) E Rmxh at t-th iteration, the following holds
w. L(t) w. Wr L()) |F S VV,, L()) (())l Vw, ((t w. [(t) |F S max (t)))
(79)
(80)
Proof. Based on Lemma A.1 one has A| F F <|4 + BBB Thus, let Vw, Vw,L(t), and we complete the proof of equation
For equation 80 has the following
[w,,
S mmax (W (tt))
equation 37 in Lemma A.1
< 2ax (W (t)
K-smooth of t,
(81)
which completes the proof.
Lemma D.4. Given Wt(t)& E Roxh W2(() e Rm at t-th iteration, the following holds
VV, L())W((1) WW(()V T()TTI o² nax (Ti) (K))))
(82)
Proof. We prove this lemma using the results from Lemma A.1 and Lemma
A.2
IV,, =WW) l(t) l Tle |VVCO/:WW(())))==I/Wwv. (1) Ve(() Sa (W2()))lV((0)1}+++ m ( (VEOO)) (T)) (
Proporty of norm
See definition of T.
equation 37 in Lemma A.1
Temma A.2
(T)) 2KL(t)
K-smooth of L.
(83)
max
24

---

Published in Transactions on Machine Learning Research (04/2025)
Lemma D.5. Given Wi (t) E Ruxh W2(t) E Rmxh at t-th iteration, for any (0, the following holds
(V((A(s)) S mik / (2 (T)) + 2n K2 o max (W()
(84)
Proof. Based on Lemma
A.1 and the assumption that E is K-smooth, has the following
(VE(A(s)) VV(AO))
<s]] - A(0)]l
L()
< +Wi(t) - L(t) 1Fe
(W)))L(())((((((()))))))))) L(1) " L0V,,,,
(85)
where the last line is due to the fact that s E 1].
Then, based on Lemma and Lemma D.4 one has the following,
(V((A(s)
{K) - A(0)ll kw, (W (t) 1|r VV L(tVV I(t) 1r Km (T)) /(K(()) + vi. 22KK 2 (W)))(()))))
(86)
which completes the proof.
Lemma D.1 (Upper bound (g., (0)gg:)) We have the following upper bound
(, (0) < Ko² max (Tx) + 22K((().
(87)
Proof. First, we notice (gt, (0)g:)) is the second-order directional derivative of I(f) w.r.t. the gradient direction,
1
d M()
(g,, H()gg)
(88)
VLOD ds
a=0
Moreover, can compute
d dg
(s) as follows
d
ds ((s)
d
dsa L (Wi(t)
- (MV
&
(A((s)
Definition of A(s)
ds
0=0
d (VE(A(()),
d
A(s))
ds
ds
= -(VE(A(s)),
A(s)) _ A(s), vPE(A(s))
ds A(5))
(89)
dzz
25

---

Published in Transactions on Machine Learning Research (04/2025)
Under the assumption that l is K-smooth and Lemma A.1
one can derive the following upper bound on dg (s)
d²
(s)
ds
==0
de
d +(
d A(s))
(V)(A(s))
A())
A(s), ve((A(s))
d52
ds
ds
d
d
S <VE(A(s)) A(s)) K
i is K-smooth
752
18
=
=2(V((()) W. (t) V w. L(t) X
+w)()) 11?
<2 VV W. V(1VV,, ws ((1) 1r
K(o (Wi d (W2(1)) V L(DIIE + Ww, [(0))
Lemma A.1
[i((t)) [1/W: [(t)l [())
+ K(O max (W) + m (W2(t)) VW, L())) + [w,
Lemma
SV <V2KL(() [IV W: [( IVW L(O
K-smooth of E
K(o (1)) max (W2(1)) Vw, T(IIII: + Vw, [())))))
(90)
Finally, we derive the the upper bound on (9:, (0)g.) based on equation 90
dl M(s)
(g.. H()g)))
VLOI) ds2
==
V L(E)IE + IIV
[()) (V2KL(t)) + o (W((t) + d (W2(t))
[V,, [()) T(t))
=V2KL(()) + o m (W) o (W2(t))
= /2KL(1) + mase (Ti),
(91)
where the last line is based on Lemma A.2
Lemma D.2 (Upper bound on (gt. H(r)g))) For any E [0, 1), we have
(g,, (T))() S Kr,
(92)
where
K, Ko² max 2 max (W(t)) ( +3Ko (T)))KKL(t))))
(93)
Proof. First, the same method to compute (ot, H(T)gt) as it done in Lemma D.1
d²
(gc. H(T)g)) = [
ds M(+t))
(94)
==0
Based on similar calculations in equation 89 one has
d²
ds2 M(s+7)
==
d
l (A(s+T))
ds²
==0
ds (Ve(A(())))) A(s+r))
ds
==
((((A++)) d
=(Ve(A((+TT), A(8+7)) ds
A(s+T))
(95)
de²
ds m=0
26

---

Published in Transactions on Machine Learning Research (04/2025)
Under the assumption that E is K-smooth, Lemma A.1 and Lemma D.1
show
d²
M(x M(++r)
ds²
de
A(s ( A(5+7), ((((A(++T) d A(c+T))
(VE(A(s+T))
do
do
==0
d² A(s +K| A(s+T)|1 d5
& (V((A(++7)),
ds
Vw, Vw,
+ 1
2374444)))174440000030400000000) Vw, L(t)
I(O) + (,,0)00060000...00000000011,",,0 w, L(t) 1|3
(VN V/, L(t)
IF
L(D) 11 IVw, WWW,w, (()) [I 22 44
+ ((T, ( (W)
+ - (V) 10, 1F
VV
+
4m, K L(t WWW 11- [w), ||F
(96)
where the last line is derived based on the fact that (0, 1].
Notice in equation
90 we have shown
Vw. ww,.C()))HVVV,OOOWAA... (WW w, L(t) 11? (V2KT(()) + m (T:))
Vw,
[|/ T(1 + (0000)))))))0(T00))))))))
(97)
Morcover, in Lemma
D.3 Lemma
D.4 and Lemma D.5 we have shown
Vw, L() S V/,
VV, 4() < 2Ko (W()))((
(98)
(w,,(())2())) (w 5 2. (Ti) (2T)))
(V((A())) n.K 22K (Te) + 2n2 K2 m (t))
(99)
27

---

Published in Transactions on Machine Learning Research (04/2025)
Thus, one can further upper bound equation 96
follows
dc
do2 ((++T)
2-0
$ (. ww, L(t)
+ - WW 1,,, L(0 + K Vw, (()
+ 4np K Vw, VLw, (Wi()))) 110 Tw,, [( ||
A wi L(U1) 1V 4()113) (V2KL(1) max (Ti)
1VV,,, + 11 V 2KL( (T)) 2n k o max (W())L()) - -I/V L(O1 +IV,,, L(0Il²))
)
+ Anf K2 omax W (t) VIV (())? + [Vw.)(t)]
2n Ko (Tt) (t) Vw, L()) + [(,
(100)
As a result, we can show
d² M(++))
(g,, H(T)g)) = [V dsa
< Ko imaax ( ( + +3Ko (T) 2()))
max (ttt))(tt)r)
E Proof of Theorem 3.2
Then, we provide detailed proof. Theorem E.1 (Tinear convergence of GD for Problem 2). Assume the GD algorithem equation 7 initialized such that on > O. Pick any 0 < c < > Let f1o
In this section, first. introduce the generalized form of Theorem 3.2
(1) be the unique posilive solulion of the following equalion
6X K exp( [1 1+3 1--3//2K /2////KL// ((0))o))
(101)
and To (2) be the smallest positive solution of the following equation
4KL(0) = 111eepp(-njp x A)
(102)
Then, we define Mmaa min(1)" (1) % satisfies
(2)
log + 22)) ) For any no and nn such that 0 < no S Mmax and m
no S n S min + i k.)
(103)
one can derive the following linear convergence rate for GD
T(++1) < L(t)((m < ((1)p(mo, 0) s (1))p(m.0)++)
(104)
where
p(n. t) 1 1-22pm 2pm. juk ū= 4=1101= 20 (1-eex((mmn)) A (1+ n6 p(pu 0).
K, = 22 +6K (,, (
5nn the case when equation
102 does not have positive salution, we 10
28

---

Published in Transactions on Machine Learning Research (04/2025)
Notice Theorem 3.2 in 3.2 can be viewed special case of Theorem E.1 with 222 d = 2.
Before presenting the proof Theorem 1 we first show that the constraints on do not induce an empty set, or equivalently 7max > 0. Alongside, we provide several inequalities that are implied by the constraints on 1/u, which the proof Theorem is relied on.
Existence of yo To show the exislence Yo, it is equivalent show that Mmmax > First, since on, 02 are positive, have log + 201
0. Moreover, one can see when yoo log (1
201 we have j 0.
(1) > O. The LHS of equation 101 increases as vJo increases, and it equals zero as Vio- 0. Second, we show 1
Thus, there exists a unique positive solution of equation 101 which is equivalent to 1l > 0. Notice
Ko = v 22L(0))6K' (2 - (133001/333000000111 (0)
(105)
Therefore, 0< 7/0 S % implies 0 < Kn I which is equivalent to 1/o S 1. to 0 p(nu. 0) < 1 and A > O.
This constraint further leads Ko.
Finally, we show T/= > O. Notice when Ylu = 0, the RHS and LHS of equation 102 both equal zero. Moreover, when 7o > 0, one can rewrite equation 102 as follows
4KL(0) = (1 exp( (xp(-16)) x (1 A)
4KL(0) (1 exp(-116) (2 K.ono d - 0))
4KL(0)) 1-c 1 - cxpp (
x (2 juomm 1o d-1 p(no, 0))-
(106)
exp( -16)
Then, we study the order of both sides of equation 106 in terms of 7o in the regime where 0 < Y/0 S
min log(1 + 20g) Ro
Since < < < 1, and the LHS of equation 106 is of order O(% it decreases
monotonically to zero as vo approaches zero. The RHS of equation 100 the product of two terms, i.e., 1-exp--16) and 2μa - nKono - - - 0). We notice Th approaches zero as to decreases to zero. Thus, expl 76)
1-exp(-10 converges to one as To decreases zero. Moroover, when fo I min [oog(1+ + 20g ) RR we have exp( exp(-16)
# 0 and 0 P(1o, 0) < 1. Therefore, the RHS of equation 106 is of order O(1). As result, when no >> 0 is sufficiently small, one has
4KL(0)) < exp((()))) x A)-
(107)
(2) Moreover, equution 102 has positive roots, and we use Tho to denote its smallest positive root. The following holds for all 0 170 S 1lo
4KL(0) (1 exx(xxx))6)) (1 A)
(108)
If equation 102 does not have positive root, then equation 107 holds for all positive 7/o-
To summarize, we have shown that Tmax
0, and the vo always exists.
Moreover, when no satisfies
0 Nl < Mmax, the following holds
A 0, p(mo, 0), A 1, 4KL(0) < expe -6)) x (1 - A)
(109)
Now we present the proof of Theorem
Proof. We employ an induction-based approach to prove Theorem E.1 by iteratively showing the following properties hold for all iteration & when do and nn satisfy the constraints in Theorem E.
29

---

Published in Transactions on Machine Learning Research (04/2025)
Art): L(t)
L(t (n) 0).
A2(t): : 81 S min (W))) - max (W(1)) Σ 82.
As(t) [t) D(0)lee exp(()))1::
Aa(t): on + 202 (1 exp(116)) min (Tt) S max (Tt) I 02 exp(1)).
Assume (k),, A2(),, As()), (k)) hold at iteration 1.2, ,t, then show they all hold for iteration [+1.
Prove A1(++1) hold.
We first show that under the constraints in Theorem E.1 and the induction assumption, one can lower bound and upper bound ir and K, using t and K, respectively, which characterized by the following lemma.
Lemma E.1. The following lower bound and upper bound on tr and K2 hold respectively
p Me
Kr S K.
(110)
The proof of the above lemma can be found at the end of Appendix E In Theorem 3.1 we have shown that the local PL inequality and Descent lemma hold with local PL constant and local smoothness constant K,
Kink (LL
L(t+1) S L(t) (Tt
(VL()) Z m.((t).
(111)
2
2
Therefore, one has
Koy
I(t+1) S I(()- -(1
[IVI(())
Kith
$ [(I -2p
T(())
Under the constraints 0 < 1 <
2
Ke
Kiyk
$ [(t)) (1c 2
I(())
Lemma
(1 2T + j.Km)) T(())
(1 2pm + ((t) ((t)
p(n, (t)
Lemma b
(112)
Finally, we show p(m,1)) 1) S plm,, 0).
pln, t) <
2im + j.k
1
- 1 - 27 TK,
Use S 11 S K, Use K₁ < Ko-
< 1 - 2im + KKond = 0))
(113)
Therefore, A (1(++) holds.
Prove A2(t+1) hold.
Since we have shown A₁ (1++1) holds, one has L(l+1) S L(0). Moreover, based on the assumption that ((W) is u-strongly convex and K-smooth, one has the following inequality
K
W((+1) w < {((+1) 1((+1) < (W(+1)-W''
(114)
2
30

---

Published in Transactions on Machine Learning Research (04/2025)
Then we can show max (W & B2 follows
max (W (t+
o max (t++1) W + )
S muax (W*) [IW W -
Weyl's inequality
S max (W") + (((++)--W11
2 L(t+1)
S maxx (W*)
V
p
2
S max (W*)
Use I(t+1] S T,(0)
(115)
For B1 S min (L++)) same result has been derived in Min et al. (2023). We refer the readers to Appendix B in Min et al. (2023) for details.
Prove 1y(t+11 hold.
We first present the following lemma that bounds [| (rk)llr for all k.
Lemma E.2. One has the following upper bound on 1
(D(++))-((bIl) S 2K o max L(k).
(116)
The proof of the above lemma can be found at the end of this section.
Based on Lemma E.2
one can show that 1g(++1) holds
(((+1))-(00)))) S
|[ + 11-
k=0
2K 2Kn200m max (Tk)
Lemma E.2
2K 1/
02 max (0)
Use (1)),, Vk 1, ,
2K 1/ exp (x2
Use As(k), Vk = 1,...
k=0
<
2K(1 + n * (0))mo) 0)*
Use S (1 + =
k=
2K exp(1i)ng
A
A (1 18) 0)
k=
2K nnaa exp(16)
A <
(117)
1-A
Prove (t+++) hold.
We first present the following two lemmas which will be used to prove that (t(++1) hold.
Lemma E.3. One can use a1, 02 lower and upper bound the singular values of To
a S 2ii (To) < 2 (To) S ag
(118)
Lemma E.4. One can bound the deviation of the singular values of Ty using the deviation of the imbalance DD)) - D(
2 (Tx) 2 01 - D() D(O) = Ttd max (Tx) 5 02 2)) D(O)]lF T
(119)
(120)
31

---

Published in Transactions on Machine Learning Research (04/2025)
The proof of Lemma E.3 and Lemma E.4 can be in Xu et al [2023] Appendix C.
Notice T + =01 22g Therefore, if one can show
T S exp(16)a2.
(121)
Then, the following holds directly
o max (T++1) S T:41 S exp(16)02,
(122)
min (T++)) 2 T 01 +202 2T a1 + 20z exp(())))
(123)
Therefore, it suffices to show equation 121 holds. We start from Lemma E.4
T a2 22|1 - D(0))
aa2 + 4KL(0)n)a2(0) exp(16)
Use As(t+1)
1
02 exp(16)
S + +1 - -xx((-16) x (1 - A)
Equation 107
1-A
cxp(16))2.
(124)
Now, we present the proof of lemmas used in the proof of Theorem E.1 All lemmas presented below are based on the assumption that A1 A2(k), A. hold for all iterations & = 2,, ... t and the constraints presented Theorem F.T For convenience, we do not state these assumptions and constraints repetitively. Lemma E.1. The following lower bound and upper bound on pr and Kg hold respectively
Me, Kr S K.
(125)
Proof. We start with the lower bound on pe. Due to the assumption that A+(t) hold, one has the following lower bound
Pe uo min (T:) pa1
(126)
For the upper bound on K2, first show that based the assumption that (1k) hold for all t, one has
L(t) s L(t-1))((mo))) 0) s L(0)p(m..0)
(127)
Then, based on equation 127 (t) and the constraint that ve (1 r/8) 5np, we can derive the following upper bound on K,
K, =K m 2K)+)Kk(()++ max (W()) +3K (T)) (TTT22KLL 2KL(t)m
<Ka2 exp(n6) + (2w.L(0))mmm.m...o.. 0)* + 6K 32 0)
+ 3Kaz (otocccpc6662/22177770000
<Kazexp(16) +
(0 + (1
+ +Ka2 +B8roo.enp(ril1222XX000000o0 01 ) 3m
Use mh E (1 + =
(128)
where the last. line follows the definition of A =
+ (1+ 11))7 0).
Lemma E.2. One has the following upper bound on [
[D(k+1)- -()++) 2K 0 max ( T.k))
(129)
32

---

Published in Transactions on Machine Learning Research (04/2025)
Proof. In equation i and equation
- have
() Wi( Nh (K))) Wa(k+1) = W2(k) na Ve(a)) ( W.:):).
(130)
There, we can compute D( - (k) as follows
D(k+1)- (kk) -W(kk+ (k(++1) (
(K))) ( + (
-(Wi(k) - (((()))(())) (Wi ((W(() -n (((K)))2(())
- -W2(A)) Th V((K) (V (W2(1)
nA Ve(A) WG (W(
(i(() ((())))(())-)))) (2(k)
=n (W2(k)) V((k) V((I:)W((k) -
(1) V Ve(k)Wi(k))
(131)
Based on the above equation, one bound 1)
follows
10(*++))D(bb))) (K)) VE(E) V((H))22 (A)) V(()))(((
Property of norm
- ((() VE(A)
Vε(A) ((())
equation 37 S n max (V2(+)) ni max (W)
- (
K-smooth of l s 2K of (Kk)
(132)
F Verification of the assumption o₁ > 0
this section, we provide two conditions that ensure a >
In Min et al. (2021) the authors show the following lemma which guarantees a1
Lemma F.1 (Lemma in (Min et al. 2021)). Let (0))), W2(0) initialized entry-wise i.i.d. from N(, K with 1 S S 2- For 48 > O and h poly(n, m, }), with probability over random inétialization with W1(0), (O), the following holds
in h 1-2p
(133)
The above theorem states when Problem 2 is sufficiently overparametrized, i.c., h 2 poly(n, m, }), Gaussian initialization with proper variance ensures ai has a positive lower bound 112-2p, Moreover, the lower bound increases as increases.
Next, going to show with mild overparamctrization,
ensure 01 0.
Lemma F.2 (Mild overparametrization ensures aı 0). Let W1 (O),), W2(0) are initialized entry- wise i. d. from a continuous distribution P. When h 2 m+n, the following holds almost surely over random initialization with W1(0), (0), W2)
0.
(134)
Compared with Lemma F.1 Lemma considers wider range of distributions that include Gaussian distribution and uniform distribution. Thus, commonly used random initialization schemes, such Xavier initialization (Glorot & Bengio 2010) and He initialization (He et al. 2015), lead to 0. Moreover, the requirement of overparametrizatie
i.e., h 2 m+nn versus 2 poly(n, m, })- As a result, Lemma F can be applied to more general overparamet rization... On
in Lemma F.2 is mild compared with the one in Lemma
the other hand, the conclusion of Lemma F.2 is weaker than Lemma F.I in the sense that Lemma only
proves 0n but do not characterize its magnitude.
Before presenting the proof of Lemma
we first present two lemmas that will be used in the proof.
33

---

Published in Transactions on Machine Learning Research (04/2025)
Lemma F.3. Let A E R h be random matriz with entry- wise drawn i.i.d. from continuous distribulion IP. Then A is of full column rank almost surely.
We refer the readers to (eershynin 2018 for detailed proof. Lemma F.4. A sufficient condition for on > 0 is amtn (D(0) > 0.
The proof of this lemma can be found in (Min et al. 2021
Now we present the proof of Lemma F.2
Proof. Based on Lemma F 14 it suffices show that one almost surely has Jm+n (D())) > 0 over random initialization with Wi(0), W (0). We use proof by contradiction. Assume o m (D(0)) 0, then one has dim(ker ((0) h
+ 1.
On the other hand, Lemma F.3 implies with probability one, [W] (0), W (0)] E glx(ntm) is of full column rank. Our next step is to show dim(ker D(0)) < - m. If this is true, then there is a contradiction. Thus, one directly has om+n(D(0)) 0.
For any ARAh that satisfics D(O)u 0, we can write this equation as follows
(1
D(O) = 0 +> [W] (0,, (0)
== 0
(135)
-W2(0)
Since [W]] (0), W:T (0)) is of full column rank, the above equation is equivalent to
wi (0
(2()) v := 0,
(136)
and dimnker D(0) < h
F.1 Large width and proper choices of the variance lead well-conditioned To
In this section, provide proof for Thcorem
3.3 We first restate Theorem 3.3 here
Theorem F.1 (Restate of Theorem 3.3 Let (O), W2(0) initialized entry-wise i.i.d. from N(0, バeс
with $. 4O (0, 1), wh 2 poly(m, n, J), with probability 1 1- 5 over random initialization Wi (0), W₂(0), the following condition holds
Tmin (To)
a1 ((120-1)
(137)
max
02
Proof. Since a1, a2 are the lower and upper bounds for the singular values of To, it is straightforward to see
min (To)
01
(138)
max (Tp)
Thus, it suffices to show @1 > 1- 1(h=-1)
We provide upper bounds and lower bounds of a1 02 separately to prove Theorem 3.3 In Min et al. (2022). the authors provide the following lower bound on an under the
1).
same setting as Theorem
3.3
Lemma F.5 (Lemma
Min al. (2022) Under the some setting as Theorem 3.3 with probability
over - 5 over random initialization over W, (0). W2(0), the following holds
a1 2 2h + 2B²h h-2p - 4Bh
Wi(0) (0) 2
(139)
where B /m + n + log 2
34

---

Published in Transactions on Machine Learning Research (04/2025)
We refer the readers to Min et al. (0022) for detailed proof.
Then, we derive an upper bound on 02 where the definition is given as (See also Table 2
2+ + X + 4,5g
A_ +
X + 4,35
(140)
2
2
where X_, 2,, B2 is defined as follows
A_ max mmax (-DO),,
== max (=(0))max(((0))0))) (D(0)),),
K (O),, (0)-W*|p
82 = 0 (W*) ++
(141)
Based on Lemma F.5 one can upper bound 82 as follows
K (O),, (0) W'
B2 max (W*)
K + +
S max (W*)
(142)
Then, we provide upper bound for Amax >max (D(0)). Similar analysis can be used to derive an upper bound for ( D(0)-
First, based on Lemma A.3 the following holds with probability at least 1 -
2
Tmax Tnax(h"[W1100 (2))) s vh + v m + n + log
(143)
2 0"
Therefore, max ([W1, W2])= < 4 - Bh-P. Then, based on this upper bound, one can derive the following upper bound on Fmax (D(0))
WT(()) 1
([wi(O), w2(C)]
max (D(0)) = o Tmuxx
Weg (0)
(i)(o) (0) -wT((o)
I max
W7 [1),), W2(0))
) xx
= o (Wi(O), (2(
< 1-2
2p 2
Bp -2p 2p
(144)
Similarly, one can show ( show the following upper bound on 022
+ 2h 2p + B²h 2p, By combining all the results, one can
x + X2 + +482
X_ +
X + 4
2
2
A4 X_ 2,3
0max (D(0)) + Tmax (-D()))) + 2,32
K (IIW* | 2/mBl -20)
< 21: 1-pp 2p + 4h 29 2B2h -p + 20 (W*) +2*)+2+-
22
35

---

Published in Transactions on Machine Learning Research (04/2025)
With the bounds a1 a, derive the following bound on Omin Ta) munx (To)
Mmin (To)
max (To)
02
212-2p + 2 22 4Bh -p
2h 1-11-pp-- 2p 4h 2p B + 2B2h + 20 max (W,) + 2/ 11 2/mBhi 20)
8h 2p B + 20 mux (W)) + 2 (I 2/mBh+ 20)
11-
2h 1- + - 2p 2B2h 2p m (W+) 2 K (IIW-II 2/mBhh -pp) 2p)
12h
(145)
where the last line holds because the dominating term in the numerator and denominator is of order ((1) and 2r) separately.
G Simulation
In this section, we first numerically verify that with proper initialization (See Theorem 3.3 larger width will lead to a well-conditioned To. Then, we present experiments showing that the overparametrized model trained with GD following the adaptive step size proposed in $3.2 can almost match the rate of non-overparametrizel model. Throughout the experiments, we consider Problem 2 with squared loss
I(W1,W) = 5 511 W,W
1 - 1
(146)
G.1 Large width leads to well-conditioned To
In this section, we compare the <(To) under different scales of the variance of the initialization, i.e., and difference width of the networks, i.c., h. We choose E (0.275, 0.375, 0.475) and h from [500, 2000]. We generate the data matrix X as a 10 10 orthogonal matrix and Y = xe +0,00, where O E 10x10 is
entry-wise i.idd. sampled from (00 0.1). The weight matrices Wi, W2 are initialized following Theorem
0220
516
18
81 0.000
012
3.040
0.40
0.35
014
600
800
1030
1250 1400 Width
1905
1950
2000
605 930 1000
1250 1400
1505
1930
2000
600
1996 2000 000
1000
3200 1450
16W
Widthh
Wilth
Figure 4: <(o)) under different choices of p and h. We repeat the simulation thirty times and plot the average value of w(Tn).
Figuro shows that for fixed p, larger width leads to well-conditioned <(T)). Morcover, for fixed width, 01 In Theorem 3.3 we show 01 2 1-12(h 55. One can see we decrease as either h p. the lower bound on 0g decreases. Therefore, the simulation results support Theorem 3.3
smaller p will lead smaller ag
36

---

Published in Transactions on Machine Learning Research (04/2025)
G.2 Overparametrized GD can almost match the rate of the non-overparrmetriiee GD
In this section, compare the rale of overparametrized GD following the adaplive step size in $3.2 with the non-overparametrizee GD. We initialize weight matrices W1, W2 E R as entry-wise (0, f. The data matrices are gencrated as ==UEV, = =XW WT +O 0.1) where U.V E R5x5 are random orthogonal matrices and E E R5x5 5 is a diagonal matrix where the diagonal entry is uniformly i..d. drawn from [1.8, 2.3] d (X) and the step size for the overparametrized GD
The step size for the non-overparametriied GD ?) =
follows equation 19 with h(n, t) p(n,, t). Figure shows that as one increases the width of the networks, the
705597 500
m0.8620, 100000
=0 7286, fm 1000
in Cza Drompoeeatr
with c-ns for verpanmeer
Souateen for Onerenementii
Nopna utt
lopaa Lit
Aop Liti
terations
herations
Iterations
Figure 5: Comparison of convergence rate of GD for the non-overparrmmetrized model and overparametrized model. We run the simulations thrity times. The red line represents log {(t) and the blue line reprosents 1og10 L(t). The shaded area represents plus and minus one standard deviation of the reported loss.
overparametrized GD can almost match the rate of the non-overparametriied GD asymptotically. Morcover, as the width increases, 01 increases and the rate of the overparametrized GD is more close to the one of the be
.g non-overparametriied GD. This is because, in $ we show that the optimal local rate of convergence 21, Therefore, one increases the width, o2 approaches one, and this leads to
arbitrarily close to 1 K az the rate of overparametrized GD approaches 1 k
G.2.1 Detailed description of backtracking line search
For backtracking line search, the algorithm is described as follows: In the simulation in 4.2 we chooser 0.1
Algorithm Backtracking Line Search.
Given Data matrices X, X,Y, initialization (O) W2(0), and hyperparameters nhe T, T. Result W Wz that minimize L(W1, W2) = =111 WTT. l
for f 0, L,-- T do = 7et while L(W1(1) 1/ Ww, L(1, (Z( - ne Vws((() > L(I) - do
?) 7T)
end while
Wi (t + 1) (t) n, Vw, L(t)
W2(1 1) (t(1) tI. V w. T(t)
end for
and 7 = 0.9.
Figure ssoows that for fixed p, larger width leads to well-conditioned <(T)). Moreover, for a fixed width, smaller p will lead smaller a1 In Theorem 3.3 we show a1 > - 2). One can see we decrease as
as" either h p, the lower bound on ag decreases. Therefore, the simulation results support Theorem 3.3
37


Reviews:
[{"review_id": "jbpTEPwbPW", "reviewer": "Reviewer_wrn2", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n1. The paper is clearly written and easy to follow. Also, the presentation of the paper is good.\n2. The topic of this paper seems interesting, investigating the the convergence of GD for training two-layer linear neural networks is the base for understanding the convergence behavior for training deep neural networks.\n3. The authors provide rigorous theoretical analysis, analyzing linear convergence across diverse scenarios as well as offering a comprehensive analysis of how factors like step size, network width, and initialization influence the process, which to my knowledge could be an interesting results to the field.\n4. The authors provide some valuable empirical results.\n\n** Weakness **\n\nIn my view, the primary concerns relate to the significance of the contributions in comparison to existing literature. Specifically, when compared to closely related works [1, 2], which also establish linear convergence, particularly [2], which considers a very similar setting except for the step-size constraint, I could not say the contribution is quite significant. A more substantial contribution might involve extending the analysis to a different dimension. At present, a more meaningful results would be to demonstrate the convergence analysis in the context of deep networks or under the suboptimal assumption.\n\n[1]  Arora, Sanjeev, et al. \"A convergence analysis of gradient descent for deep linear neural networks.\" arXiv preprint arXiv:1810.02281 (2018).\n\n[2] Xu, Ziqing, et al. \"Linear convergence of gradient descent for finite width over-parametrized linear networks with general initialization.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2023.\n\nTypos, Page 5, \"such an analysis to Problem Problem 2. \" -> \"such an analysis to Problem 2. \"", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper investigates the convergence of GD for training two-layer linear neural networks, addressing limitations in prior work that relied on stringent assumptions about step size, network width, and initialization. By leveraging the PL condition and the Descent Lemma, tthe authors bound constants dependent on initialization and loss, proving linear convergence for general losses. They propose an adaptive step size scheme, validated empirically, which accelerates convergence compared to fixed step sizes. The analysis extends to matrix factorization, improving singular value tracking for adaptive step size computation.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See weakness.", "broader_impact_concerns": "I have not found any discussions about the limitations and potential negative societal impact. But in my opinion, this may not be a problem, since the work only focuses on the analyzing the convergence of GD. Still, it is highly encouraged to add corresponding discussions.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "B0nZu3nWrJ", "reviewer": "Reviewer_gdTu", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\nS1. Given the complexity of the global convergence analysis of neural networks, as is also discussed in the paper, analysis of the local convergence in two-layer linear networks naturally can help further the understanding of the problem. To my knowledge, such results have not appeared before in the setup introduced in the paper.\n\nS2. The paper uses clear language, and there are in-detail explanations for a reasonable number of questions a reader might have. For instance, there are explanations for why studying global convergence is problematic, when the initialization requirements are satisfied, and why the step size requirements are defined a certain way.\n\n**Weaknesses**\n\nW1. There is no discussion on which commonly used losses would satisfy the loss assumptions. \n\nW2. While I appreciated the thorough explanations, all the small details, sometimes also distributed across the paper, made it overwhelming and hard to follow at times. I understand that, to some extent, it is because of the nature of the results. Nevertheless, I suggest what can be improved in the Requested Changes.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper studies the convergence of the gradient descent in linear networks and derives a linear local convergence rate.\nTo achieve this, the submission proves the local versions of the Polyak-Łojasiewicz condition and Descent lemma in Theorem 3.1 and, based on them, obtains the local convergence results in Theorem 3.2.\nSimilarly to recent related works, the preprint allows for finitely wide networks and common initialization schemes, and it \nimproves on other results by allowing for more general loss functions and adaptive learning rates. The theoretical bounds on the convergence rates are verified experimentally.\nFinally, the work derives an algorithm for the step-size choice based on Theorem 3.2 and demonstrates experimental improvements over the previous methods.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**Crucial**\n\nC1. Add a discussion on which loss functions would satisfy the assumptions. \nWhere possible, rephrase the general loss claim. General loss sounds like any or almost any loss is acceptable.\n\nC2. Enhancing the clarity of the paper:\n\n1. Reference the specific statements in the Main Contribution section.\n    \n2. Reference assumptions in the theorem statements: assumptions on the loss and other assumptions used (if any). The assumption that $h \\geq \\min (n, m)$ is also required to obtain the results, as I understand.\n\n3. Add names for the quantities where possible in Table 2.\n\n4. Add proof intuitions for the theorems or mention the key techniques used.\n\n**Minor**\n\nC1. Specify what an overparametrized network means. Definitions of Problem 1 and Problem 2, as well as definitions of the overparametrized and non-overparametrized problems at the top of page 2, suggest that overparametrization is a property that depends on the network depth.\n\nC2. Specify what LS is on page 2.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "dZ8m7mpLJI", "reviewer": "Reviewer_QAmF", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This is a nice paper with stengths:\n\n1) The problem is important since convergence analysis for overparamterzied networks are important. \n2) The experiments and theory matches well. The difficulty at hand is clearly explained. \n3) Although their exist past works that uses local PL condition to prove convergence in the same setting. The **criteria for TMLR is not novelty** but that claims should support evidence. The later condition holds for this paper, but not necessarily the first one. \n\nHere are some weakness that can be addresable:\n\n\n1) Please discuss how the local PL inequality behaves near a saddle. seems like the coefficient needs to be very small. some landscape diagram can help reader grasp the problem. for example 1d comparisons of (x-1)^2 and (x^2-1)^2 and showing that PL condition is violated near the saddle of the latter, can be a good way to show this.\n\n2) local pseudo convexity (instead of strong convexity) have been previously used in problems such as in https://arxiv.org/pdf/2310.03956. Better to cite such papers where these local PL conditions are previously used in analysis. \n\n3) Some comparisons to work such as https://arxiv.org/pdf/2409.19791 are missing, which also shows linear convergence on adaptive step size GD on such overparamterized non-convex optimization problems. However, the choice of adaptive steps are different from the current work. But a discussion should be helpful. \n\n4) What if a constant step-size is used? then what is the largest step-size that can be used to ensure convergence? Does this step-size depend on the initialization or imbalance? Works such as https://openreview.net/forum?id=J4Dvxv7WnG show that the step-size that can drive overparamterized linear networks is independent of such initlaization or imbalance. More specifically for GD, the imbalance gap monotonically decays with linear convergence rate even beyond the stability threshold. Moreover, even beyond this stability limit, the loss do not diverge but shows stable oscillations. Discussing the observations are helpful since the current work studies the descent lemma.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the convergence of overparamterized linear networks for GD. As the classical ingredients such as PL condition and smoothness condition do not hold globallay for these losses, the authors propose a local analysis. They show that instead it is possible to propose a local PL and smoothness condition that depends on the condition number of the layer factors. This quantity may change say for example to near a saddle or a global minima, still allowing them to prove convergence.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please see the points mentioned in the weakness. Addressing 1-4 and citing the papers in proper context would be beneficial.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors have studied convergence of gradient descent (GD) for two-layer linear networks under a general loss and finite step-size. The key contribution is to show that a local PL inequality and a local smoothness inequality hold along the trajectory of GD leading to local descent lemma for such networks. The authors further leverage on this descent lemma to design  an adaptive step-size, which achieves a linear convergence rate for GD.

The reviewers have initially raised concerns regarding comparison with related work that use local PL inequality to achieve convergence and those obtaining linear convergence rate for linear networks under different step-sizes compared to this paper. The authors have clarified that their analysis explicitly characterizes how the network width and the initialization scheme impact the convergence rate. The reviewers acknowledge rigorous theoretical results with providing sufficient insights and discussion on the key theorems as major strengths of this paper. 


In the revised version, the authors have clarified their notion of overparameterization on page 2 “Here, overparametrized indicates that, although the function classes represented by $XW$ and $XW_1W_2\cdots W_L$ coincide under certain width assumptions, we introduce additional matrices to represent the same function. The overparametrized loss above can be seen as using a deep linear network to learn the target $Y$.” So depth of the network is the key point. On page 4 and the rest of the paper, they introduce (Problem 2) for two-layer networks as an overparameterized problem. This is not consistent with the notion earlier discussed on page 2. I’d recommend that the paper is accepted after addressing this comment.
Decision: Accept (poster)