Provable Acceleration of Nesterov's Accelerated Gradient for Asymmetric Matrix Factorization and Linear Neural Networks

Authors: zhenghao  xu
yuqing  wang
tuo  zhao
rachel  ward
molei  tao
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 03:27:21

Abstract: We study the convergence rate of first-order methods for rectangular matrix factorization, which is a canonical nonconvex optimization problem. Specifically, given a rank-$r$ matrix $\mathbf{A}\in\mathbb{R}^{m\times n}$, we prove that gradient descent (GD) can find a pair of $\epsilon$-optimal solutions $\mathbf{X}_T\in\mathbb{R}^{m\times d}$ and $\mathbf{Y}_T\in\mathbb{R}^{n\times d}$, where $d\geq r$, satisfying $\lVert\mathbf{X}_T\mathbf{Y}_T^\top-\mathbf{A}\rVert_F\leq\epsilon\lVert\mathbf{A}\rVert_F$ in $T=O(\kappa^2\log\frac{1}{\epsilon})$ iterations with high probability, where $\kappa$ denotes the condition number of $\mathbf{A}$. Furthermore, we prove that Nesterov's accelerated gradient (NAG) attains an iteration complexity of $O(\kappa\log\frac{1}{\epsilon})$, which is the best-known bound of first-order methods for rectangular matrix factorization. Different from small balanced random initialization in the existing literature, we adopt an unbalanced initialization, where $\mathbf{X}_0$ is large and $\mathbf{Y}_0$ is $0$. Moreover, our initialization and analysis can be further extended to linear neural networks, where we prove that NAG can also attain an accelerated linear convergence rate. In particular, we only require the width of the network to be greater than or equal to the rank of the output label matrix. In contrast, previous results achieving the same rate require excessive widths that additionally depend on the condition number and the rank of the input data matrix.

Main paper: 
 introduction: 
 Nonconvex optimization is pervasive in the training of modern machine learning models. Despite the success of first-order methods in practice, theoretical understanding of their convergence properties is limited even for simple nonconvex problems. Take the rectangular low-rank matrix factorization problem as an example, which is a canonical nonconvex problem:EQUATIONwhere we solve for two small matrices X ∈ R m×d and Y ∈ R n×d to approximate a big rank-r target matrix A ∈ R m×n with r ≪ min(m, n) and m, n not necessarily equal. Specifically, we consider the over-parameterized regime where d ≥ r, so that the global minimum of (1) is zero. While various direct methods exist for solving (1), we focus on understanding the global convergence behaviors of first-order methods applied to such a nonconvex problem, with the motivation of gathering insight into the training dynamics of neural networks.Most existing results study the simplest first-order method, gradient descent (GD), under different initialization schemes. Note that the initialization scheme matters to convergence analysis1 , due to the fact that (1) is a nonconvex and nonsmooth 2 optimization problem. Thus, proper initialization is important for the fast convergence rates of first-order methods. Ye and Du [2021] show that with small Gaussian random initialization, GD can find X T and Y T such that f (X T , Y T ) ≤ ϵ in T = O(d 4 (m + n) 2 κ 4 log 1 ϵ ) iterations with high probability, where κ denotes the condition number. Jiang et al. [2023] improve this result to O(κ 3 log 1 ϵ ) which has no explicit dimensional dependence on m and n. These analyses rely on balanced initialization where entries of X 0 and Y 0 have the same variance so that the iterates are guaranteed to stay in a smooth region.Moreover, we remark that to the best of our knowledge, we are not aware of any existing theoretical results on rectangular matrix factorization analyzing the global convergence rate of more advanced first-order methods such as Nesterov's accelerated gradient (NAG), which has been proved to achieve faster rates for smooth convex optimization problems [Nesterov, 2013] .Recently, Ward and Kolda [2023] showed that by using an unbalanced random initialization where X 0 is larger than Y 0 , alternating gradient descent (AltGD) that alternatingly optimizes X t and Y t via gradient steps can achieve O(d 2 (d -r + 1) -2 κ 2 log 1 ϵ ) iteration complexity. However, their analysis is specifically designed for AltGD and not applicable to GD, let alone more advanced methods such as NAG which are nevertheless widely used in machine learning practice. Two questions naturally arise here: Q1: Can GD achieve the same convergence rate as AltGD for (1)? Q2: Can more advanced first-order methods (e.g., NAG) achieve faster convergence rate for (1)?• Main Results. We answer the two questions above affirmatively by developing a new theory on first-order methods for (1). Specifically, we consider an unbalanced initialization scheme X 0 = cAΦ and Y 0 = 0, where c > 0 is a large constant and Φ is a Gaussian random matrix. Note that our initialization of X 0 is the same as that in Ward and Kolda [2023] , but they initialized Y 0 using a small Gaussian random matrix. This modification is mainly for simpler analysis and makes little difference in practice. Under our new initialization scheme, we first prove an O(d 2 (d -r + 1) -2 κ 2 log 1 ϵ ) iteration complexity for GD (Theorem 1), matching that of AltGD in Ward and Kolda [2023] . Our analysis is based on a new theoretical framework different from Ward and Kolda [2023] and can be further extended to analyzing NAG. We then show that NAG can attain a provable acceleration with an O(d(d -r + 1) -1 κ log 1 ϵ ) iteration complexity (Theorem 2). We discuss the tightness of our results (Remark 2) and conduct numerical experiments for validation (Section 5). Empirically, we observe that NAG exhibits a much faster rate than GD and our bounds are quite tight.Our analysis technique can also be applied to linear neural networks. We consider unbalanced initialization similar to the one for (1). We show that NAG can achieve an accelerated convergence rate for each overparameterization level (Corollaries 1 to 3), under the commonly adopted interpolation assumption (Assumption 1, see e.g. Du and Hu 2019) . In particular, we only require the network width to be greater than the rank of the output matrix.• Additional Related Work. For matrix factorization, there is a large body of works focusing on the symmetric case, where A is positive semidefinite and A = XX ⊤ [Bhojanapalli et al., 2016 , Li et al., 2018] . However, these analyses are difficult to generalize to the rectangular case (1) due to the additional unbalanced scaling issue 3 . To overcome this, additional balancing regularization is often required [Tu et al., 2016 , Park et al., 2017 , Zhang et al., 2021 , Bi et al., 2022] , which changes the objective function in (1). Du et al. [2018] show that GD can automatically balance the two factors hence explicit regularization is not necessary, but they only establish linear convergence rate for rank-1 matrix and cannot generalize to rank-r case. Some other works remove this regularization for the general matrix sensing problem and show linear convergence rate for general ranks [Ma et al., 2021 , Tong et al., 2021a,b] . These results do not directly apply to our setting as they require singular value decomposition (SVD) at initialization, which consumes roughly the same amount of computation as solving (1). Moreover, these works only consider exact parameterization (d = r), leaving out the overparameterization regime (d > r). Overparameterization may heavily slow down convergence due to the possible singularity of iterates [Stöger and Soltanolkotabi, 2021] , thus some works consider preconditioning for acceleration [Zhang et al., 2023 , Xu et al., 2023] . These preconditioned methods are specifically tailored to symmetric factorization and are not directly comparable with the first-order methods we consider, as their algorithms not only use the gradient. For algorithmic acceleration, Zhou et al. [2020] first propose a computationally tractable modified Nesterov's method for general loss function f that is L-smooth µ-strongly convex to the product XY ⊤ . However, their method still requires balancing regularization when applied to rectangular matrices, and SVD-based initialization that dominates the computation. Moreover, their acceleration pertains to the condition number L/µ of the loss function rather than κ of the target matrix, on which their dependence is O(κ 2 ).For linear neural networks, Du and Hu [2019] and Hu et al. [2020] show linear convergence of GD with Gaussian and orthogonal initialization respectively, and Min et al. [2021] studies convergence rate of gradient flow (GF) for unbalanced initialization. Wang et al. [2021] show that Polyak's heavy ball (HB) method [Polyak, 1964] attains accelerated convergence rate with orthogonal initialization. Liu et al. [2022] further investigate NAG and show a similar accelerated rate for Gaussian initialization. All these previous works consider sufficiently wide networks that depend on the output dimension, the rank, and the condition number of input. The results are summarized in Table 1 . , r = rank(D), r = rank(L), r = ∥D∥ 2 F / ∥D∥ 2 , κ = cond 2 (D), κ 1 = O(κ 2 ), κ 2 = O(κ).GD [Du and Hu, 2019] Gaussian [Hu et al., 2020] Orthogonal Ω rκ 2 (d out + log r δ )Ω rκ 3 (d out + log r δ ) (1 -3 4κ ) t GD(1 -1 4κ ) t HB [Wang et al., 2021] Orthogonal [Liu et al., 2022] GaussianΩ κ 5 ∥D∥ 2 (d out + log r δ ) (1 -1 4 √ κ ) t NAGΩ rκ 5 (d out + log r δ ) (1 -1 2 √ κ ) t NAG (ours, Corollary 1) Unbalanced (12) ≥ r + Ω(log 1 δ ) (1 -1 2 √ κ1 ) t NAG (ours, Corollary 2) Unbalanced+Orth (13) ≥ r (1 -1 2 √ κ ) t NAG (ours, Corollary 3) Unbalanced (14) ≥ d out + Ω(log 1 δ ) (1 -1 2 √ κ2 ) t• Notations. Throughout this paper, ∥•∥ denotes the Euclidean norm of a vector or the spectral norm of a matrix, and ∥•∥ F denotes the Frobenius norm of a matrix. For any matrix, σ i (•) denotes its i-th largest singular value. For a square matrix, λ i (•) denotes its i-th largest eigenvalue. For a nonzero positive semidefinite matrix, λ max (•) and λ min (•) denote its largest and smallest nonzero eigenvalues respectively. For a matrix X, we use col(X) to denote its column space, ker(X) to denote its kernel space and define cond(X) := ∥X∥ X † as its condition number, where X † denotes the pseudoinverse of X. For any positive integer n, I n denotes the identity matrix of size n. We use ⊗ to denote the Kronecker product between matrices, ⊕ to denote the direct sum of vector spaces, and vec(•) to denote the column-first vectorization of a matrix. We use N (µ, σ 2 ) to denote Gaussian distribution with mean µ and variance σ 2 .Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: The main claims made in the abstract and introduction (Section 1) accurately reflect the paper's contributions and scope.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
results for matrix factorization: 
 We start with formalizing our initialization scheme for matrix factorization problem (1). Let Φ ∈ R n×d be a Gaussian random matrix with i.i.d. entries [Φ] i,j ∼ N (0, 1/d). We initializeEQUATIONwhere c > 0 is a constant to be specified later. Typically, we require c to be larger than a certain threshold, which depends on the dimensions, the extreme singular values of A, and possibly the condition number of X 0 . We note that changing c would not affect cond(X 0 ), hence there is no recursive definition. As we mentioned, ( 2) is a modified version of the initialization in Ward and Kolda [2023] , where we replace the small random Gaussian matrix Y 0 by 0 and choose c independently of the step size. We set Y 0 = 0 mainly for simplicity, and our analysis can be extended to the case where Y 0 is a sufficiently small Gaussian random matrix. While the initialization of X 0 differs from standard Gaussian initialization, it has the following interpretation: Suppose we start from t = -1 and let X -1 = c ′ Φ ′ and Y -1 = c ′′ Φ for some 0 < c ′ ≪ c ′′ ≪ 1 and Gaussian random matrix Φ ′ , then by taking a gradient step with step size c/c ′′ we get X 0 ≈ cAΦ and Y 0 ≈ 0. This initialization of X 0 also coincides with the first step of randomized singular value decomposition, which is also referred to as sketching (see e.g. [Halko et al., 2011] ).With initialization (1), we can analyze the global convergence rates of various first-order methods.Consider gradient descent (GD) first. The gradient of the squared Frobenius error in (1) is given by∇ X f (X, Y) = (XY ⊤ -A)Y, ∇ Y f (X, Y) = (XY ⊤ -A) ⊤ X.For t ≥ 0, the GD update with constant step size η > 0 is written asEQUATION) Let R t := X t Y ⊤ t -A denote the residual, then f (X t , Y t ) = 1 2 ∥R t ∥ 2 F .We have the following convergence rate for GD.Theorem 1 (GD convergence rate). For 0 < τ < c 1 , denote δ = 3e -(d-r+1)•min{log 1 c 1 τ ,c2, 1 2 } , where c 1 and c 2 are universal constants. Denote L = σ 2 1 (X 0 ), µ = σ 2 r (X 0 ). Let η = 2 L+µ , c ≥ c := √ dσr(A) 12τ ( √ d- √ r-1) cond 4 (X0)∥A∥ Fcond 2 (X0)-1 be a sufficiently large constant. Then with c plugged in initialization (2), GD returns X t and Y t with probability at least 1 -δ such that∥R t ∥ F ≤ 3c 2 σ 2 1 (A) 64 1 - µ L t .In particular, if c = c, then GD finds∥R T ∥ F ≤ ϵ ∥A∥ F in T = O d 2 κ 2 τ 2 (d -r + 1) 2 • log C ϵ iterations, where C = 27τ 2 (d-r+1) 2 16d 2 cond 4 (X0)κ 2 cond 2 (X0)-1 .Theorem 1 shows that GD converges in O(d 2 (d -r + 1) -2 κ 2 log 1 ϵ ) iterations with initialization (2), and the constant prefactor does not have dependence on the ambient dimension m and n. This matches the convergence rate for AltGD derived in Ward and Kolda [2023] . The step size 2L+µ is commonly used in optimization literature and leads to optimal convergence rate [Nesterov, 2013] . While the bound on ∥R t ∥ F in Theorem 1 does not explicitly depend on ∥A∥ F , the norm still affects the convergence rate through the choice of c defined in (2). When c = c = O( ∥A∥ F ), the bound linearly depends on ∥A∥ F . When c > c, the bound linearly depends on c 2 , which dominates ∥A∥ F .We then consider Nesterov's accelerated gradient (NAG) method [Nesterov, 2013] applied to (1). We take the form of NAG that is originally designed for smooth strongly convex loss function ℓ:z t+1 = zt -η∇ℓ(z t ), zt+1 = z t+1 + β(z t+1 -z t ),where η is the step size, β is the momentum parameter, and z or z in our case consists of both X and Y. If we focus on the {z t } sequence with zt = (X t , Y t ) and plug in the objective function in (1), then with X -1 = X 0 and Y -1 = Y 0 , the NAG update is given byEQUATION)We have the following convergence rate for NAG.Theorem 2 (NAG convergence rate). For 0 < τ < c 1 , define δ as in Theorem 1. DenoteL = σ 2 1 (X 0 ), µ = σ 2 r (X 0 ). Let η = 1 L , β = √ L- √ µ √ L+ √ µ , c ≥ c := 29 d(2 √ d+ √ r)∥A∥ F •κ τ 3 ( √ d- √ r-1) 3 σ 2r (A) be a constant. Then with c plugged in initialization (2), NAG returns X t and Y t with probability at least 1 -δ such that∥R t ∥ F ≤ c 2 σ 2 1 (A) 64 cond(X 0 ) 1 - √ µ 2 √ L t .In particular, if c = c then NAG finds∥R T ∥ F ≤ ϵ ∥A∥ F in T = O dκ τ (d -r + 1) • log C ϵ iterations, where C = 841d(2 √ d+ √ r) 64τ 3 ( √ d- √ r-1) 3 • κ 3 cond(X0) .Theorem 2 shows that NAG can achieve O(d(d -r + 1) -1 κ log 1 ϵ ) iteration complexity with high probability. The dependence on the condition number κ is improved from being quadratic to linear. Moreover, the dependence on the dimension is also improved. As shown in Theorem 1, the GD iteration number has an O(d 2 ) dependence in the worst case (d = r). Here, NAG has at most O(d) dependence. The level of overparameterization d will affect both the convergence rate and the probability of success. To ensure a small fail probability δ, it requires d = r -1 + Ω(log 1 δ ). Again, the step size 1 L and momentum√ L- √ µ √ L+√ µ are commonly used in the literature [Nesterov, 2013] .Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We state all settings and assumptions required for our results and discuss limitations (e.g. exact rank-r A, Y 0 = 0, etc.) in Sections 1, 2, 4 and 6.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
proof sketch for convergence rates: 
 We now provide the proof sketch for Theorems 1 and 2. Our proof is based on induction. We start with the assumptions that X t and Y t are not too far from X 0 and Y 0 respectively and the initial residual is bounded by some constant, which are guaranteed at time t = 0. Given the induction assumptions, we then track the dynamics of residual R t and decompose it into linear and higher-order parts. We can show that the linear part is contracted and the higher-order part shrinks exponentially, together implying that ∥R t+1 ∥ F = O(θ t ) for some θ ∈ (0, 1) and X t+1 and Y t+1 is still within a bounded region around initialization. This shows the induction assumptions for the next iterate, thus by invoking the induction we complete the proof.The key to our proof is to show the contraction and its rate. Firstly, the linear part of the dynamics is not a contraction over the whole space, thus we need to identify in which subspace it is a contraction. Secondly, we need to quantify the rate of contraction to get global convergence rates. These necessitate the following proposition about the properties of X 0 with initialization (2).Proposition 1. For any τ, c > 0, A ∈ R m×n being a rank-r matrix with condition number κ := cond(A), Φ ∈ R n×d being a random matrix with i.i.d. entries from N (0, 1/d), the following holds for X 0 = cAΦ with probability at least 1 -δ:τ ( √ d - √ r -1) √ d c • σ r (A) ≤ σ r (X 0 ) ≤ σ 1 (X 0 ) ≤ 2 √ d + √ r √ d c • σ 1 (A),where δ = 3e -min{(d-r+1) log 1 c 1 τ ,c2d, d 2 } , c 1 and c 2 are universal constants. When it holds, the condition number of X 0 is bounded:cond(X 0 ) ≤ 2 √ d + √ r τ ( √ d - √ r -1) • κ ≤ 6d τ (d -r + 1) • κ.By Proposition 1, the top singular value of X 0 is bounded from above by σ 1 (A), and the r-th singular value of X 0 is bounded from below by σ r (A), hence we have cond(X 0 ) = O(κ). Moreover, X 0 has rank r with probability 1 and thus it preserves the column space of A, i.e., col(X 0 ) = col(A). This subspace preservation property will be passed to subsequent iterations of first-order methods and is critical to our analysis. In particular, we will show this space corresponds to the contraction subspace.As mentioned, we track the dynamics of R t for GD to prove Theorem 1. Let r t = vec(R t ) denote the vectorized residual, then the GD update (3) corresponds to the following dynamics: Proposition 2 (GD dynamics). Let P t = X t+1 -X t and Q t = Y t+1 -Y t denote the update steps for t ≥ 0. Then GD (3) admits the following dynamics:EQUATIONwhereH t = (Y t Y ⊤ t ) ⊗ I m + I n ⊗ (X t X ⊤ t ) and ξ t = η(H 0 -H t )r t + vec(P t Q ⊤ t ).The linear part at time t is (I mn -ηH t )r t , which is approximately (I mn -ηH 0 )r t when X t and Y t are close to their initialization. The approximation error along with the higher-order term vec(P t Q ⊤ t ) is contained in ξ t . It follows immediately from Proposition 2 thatr t+1 = (I mn -ηH 0 ) t+1 r 0 + t s=0 (I mn -ηH 0 ) t-s ξ s .If T GD := I mn -ηH 0 is a contraction map, i.e., it has all eigenvalues bounded |λ i (T GD )| ≤ ρ for some ρ ∈ [0, 1), and the nonlinear error ξ t shrinks exponentially at rate θ ∈ (ρ, 1), then we have ∥r t ∥ = O(θ t ). However, for d < min(m, n)/2, T GD cannot be a contraction map for any η, as the rank of H 0 is at most (m + n)d < mn. In fact, if X 0 is initialized as in (2), then rank(H 0 ) = nr < mn regardless of the choice of d. As H 0 has no full rank, T GD must have a non-trivial eigensubspace corresponding to eigenvalue 1. In the following lemma, we show that r t and ξ t are not in this "bad" subspace but rather in a contracted subspace as desired. Lemma 1 (Eigensubspace). Let H ⊆ R mn denote the linear subspace containing all eigenvectors of H 0 with positive eigenvalues. If X 0 is initialized as in (2), then we haveH = (col(A)) n and {r t , ξ t } t≥0 ⊂ H,where H 0 , r t and ξ t are defined as in Proposition 2.Given that r t and ξ t are in the contracted subspace H throughout all iterations, the convergence rate is determined by the contractivity of T GD over this subspace, which corresponds to the condition number of X 0 with initialization (2). Lemma 2 (GD contractivity). Let L = σ 2 1 (X 0 ), µ = σ 2 r (X 0 ), and H be defined as in Lemma 1. Let η ∈ (0, 2 L ), then for any v ∈ H, ∥T GD v∥ ≤ max{|1 -ηL| , |1 -ηµ|} ∥v∥ .In particular, if η = 2 L+µ , then ∥T GD v∥ ≤ L-µ L+µ ∥v∥.By Lemmas 1 and 2, the linear part of GD dynamics contracts r t and ξ t , and the rate of contraction is ρ = max{|1 -ηL| , |1 -ηµ|}. To complete the proof, it remains to bound the magnitude of error ξ t and show induction conditions for the next iteration. This is guaranteed by the following lemma. Lemma 3 (Nonlinear error). If there exist θ ∈ (0, 1) and some constants C 1 and C 2 such that for any s ≤ t, the GD dynamics (5) yields∥r s ∥ ≤ C 1 θ s ∥r 0 ∥, ∥X s -X 0 ∥ F ≤ C 2 , ∥Y s -Y 0 ∥ F ≤ C 2 , then we have vec(P s Q ⊤ s ) ≤ C 3 θ 2s ∥r 0 ∥ 2 and ∥η(H 0 -H s )r s ∥ ≤ C 4 θ s ∥r 0 ∥for some constants C 3 and C 4 depending on C 1 and C 2 . Moreover, if C 1 and C 2 satisfyEQUATION)then we have ∥X t+1 -X 0 ∥ F ≤ C 2 and ∥Y t+1 -Y 0 ∥ F ≤ C 2 .Lemma 3 shows that ∥ξ t ∥ = O(θ t ) if the residual shrinks exponentially and the iterates are not too far from initialization, which in turn implies that X t+1 and Y t+1 are also within the C 2 -balls around their initialization. It turns out that there is a set of valid coefficients for the induction to go through as long as the c in ( 2) is sufficiently large. Therefore, by choosing c properly and plugging in ρ = L-µ L+µ and θ = 1 -µ L , we prove the convergence rate for GD, and the iteration complexity follows immediately from Proposition 1. The complete proof is provided in Appendix B.6.Remark 1. In our proof, the unbalanced initialization guarantees the existence of induction constants in Lemma 3. The amount of unbalance affects the constant factors but will not affect the convergence rate (1 -µ L ) t . To be explicit, suppose we initializeX 0 = c 1 AΦ 1 ∈ R m×d , Y 0 = c 2 Φ 2 ∈ R n×d , where [Φ 1 ] i,j ∼ N (0, 1/d) and [Φ 2 ] i,j ∼ N (0, 1/n), then by replacing H 0 in Proposition 2 with H ′ 0 = I n ⊗ (X 0 X ⊤ 0 ), we can generalize the proof and obtain the same convergence rate when c 1 is sufficiently large and c 1 c 2 = O(1). Meanwhile, we have ∥R 0 ∥ F ≤ (1 + O(c 1 c 2 )) ∥A∥ F with high probability [Ward and Kolda, 2023] . Therefore, when c 1 is fixed, a smaller c 2 yields a smaller initial loss, resulting in a smaller constant factor. Meanwhile, the convergence rate remains the same as the condition number of H ′ 0 is not affected, and the shift H t -H ′ 0 is controlled for small c 2 .3.2 Proof Sketch for NAG Convergence Rate (Theorem 2)We now turn to prove Theorem 2. Similar to GD, we track the residual dynamics of NAG.Proposition 3 (NAG dynamics). Let P t = X t+1 -X t and Q t = Y t+1 -Y t denote the update steps for t ≥ 0. Then NAG (4) admits the following dynamics:EQUATIONwhereH t = (Y t Y ⊤ t ) ⊗ I m + I n ⊗ (X t X ⊤ t ), ξ t = ζ t + ι t , ζ t = vec(P t Q ⊤ t ) + β vec(P t-1 Q ⊤ t-1 ) + βη vec(R t-1 Y t-1 Q ⊤ t-1 + P t-1 X ⊤ t-1 R t-1 ), ι t = (1 + β)η(H 0 -H t )r t -βη(H 0 -H t-1 )r t-1 .As Proposition 3 shows, NAG dynamics (7) has additional momentum terms involving P t and Q t . When β = 0, it reduces to the GD dynamics (5). The introduction of momentum terms allows the linear part in (7) to contract r t and ξ t faster. To be more explicit, letEQUATIONdenote the linear part of the system. The next lemma shows NAG improves the rate of contraction.Lemma 4 (NAG contractivity). Let η = 1 L , β = √ L- √ µ √ L+ √ µ , then for all (u, v) ∈ H × H, T NAG u v ≤ 1 - µ L u v .The price to pay for the faster rate of contraction is the additional perturbations. The ι t term characterizes dynamics shift, which can be controlled as GD in Lemma 3. The ζ t term characterizes higher-order terms in the dynamics (7), which can be controlled by the updates P t and Q t . In GD, these terms correspond to the gradient so that they can be bounded if R t shrinks and X t and Y t are not too far away from X 0 and Y 0 . In NAG, we haveP t = ηR t Y t + η t s=1 β t-s+1 R s Y s ,and a similar equation holds for Q t . If R t shrinks at rate θ > θ 2 ≥ β, then we have an O(θ t ) upper bound for ∥P t ∥ F and ∥Q t ∥ F . We formalize the argument in the following induction lemma. Lemma 5. Suppose 0 < β ≤ θ 2 < θ < 1. If there exist some constants C 1 and C 2 such that for any s ≤ t, the NAG dynamics (7) yieldsr s r s-1 ≤ C 1 θ s r 0 r -1 , ∥X s -X 0 ∥ F ≤ C 2 ,and∥Y s -Y 0 ∥ F ≤ C 2 , then we have ∥ζ t ∥ ≤ C 3 θ 2t r 0 r -1 2 , and ∥ι t ∥ ≤ C 4 θ t r 0 r -1for some constants C 3 and C 4 depending on C 1 and C 2 . Moreover, if C 1 and C 2 satisfyEQUATION)then we have ∥X t+1 -X 0 ∥ F ≤ C 2 and ∥Y t+1 -Y 0 ∥ F ≤ C 2 .Lemma 5 is similar to Lemma 3. Again by choosing a sufficiently large c to initialize X 0 , we can find a set of feasible coefficients for the induction. In particular, we plug in ρ = 1 -√ µ √ L , θ = 1 - √ µ 2 √ L and β = √ L- √ µ √ L+√ µ , then c defined in Theorem 2 ensures the success of induction, hence the accelerated convergence rate of NAG is proved. The complete proof is provided in Appendix C.4. Remark 2. Our analysis differs from that of Ward and Kolda [2023] . Their analysis is based on the Polyak-Łojasiewicz (PL) inequality [Łojasiewicz, 1963] : f (X t , Y) is approximately µ-PL and L-smooth in Y, and the unbalanced initialization (large X 0 small Y 0 ) ensures that only Y matters to the convergence rate, as X is not changing by much. Since the objective function in (1) is quadratic in X, the problem has condition number κ := L µ = O(κ 2 ). With these notations, the complexity in Ward and Kolda [2023] reads as O(κ log 1 ϵ ), which is standard for PL functions. However, PL inequality cannot fully capture the properties of (1), and the analysis in Ward and Kolda [2023] does not apply to the case where X t and Y t are updated simultaneously rather than alternatingly. In fact, if we fix X ≡ X 0 and optimize Y only, then our initialization (2) makes the problem quasi-strongly convex (QSC), which is strictly stronger than PL [Necoara et al., 2019] . For QSC functions, NAG can achieve O( √ κ log 1 ϵ ) convergence rate Necoara et al. [2019] , while for PL functions the rate can only be Ω(κ log 1 ϵ ) [Yue et al., 2023] . We note that simultaneously optimizing X and Y causes the nonconvexity issue and hence (1) does not fit in the framework for QSC functions as it requires convexity. Our results in Theorems 1 and 2 match the ones for QSC functions and Theorem 2 further matches the lower bound for general smooth strongly convex functions [Nemirovski and Yudin, 1983] , which generally exhibit more favorable properties than nonconvex optimization problems to which (1) belongs. Hence, we conjecture that our rate bounds are tight for both GD and NAG. However, rigorous theory is yet to be constructed to solidify our conjecture.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: We provide anonymized code in the zip file for experiments in Section 5 as supplement materials.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
extension to linear neural network: 
 Our analysis can be extended to the squared loss training of two-layer linear neural networks, which is equivalent to the following optimization problem:EQUATIONHere, D ∈ R n×N corresponds to all input data concatenated together, L ∈ R m×N denotes the labels, N is the total number of training data samples, and d is the network width. We make the following interpolation assumption, which is commonly adopted in the study of the convergence rate of linear neural networks [Du and Hu, 2019 , Hu et al., 2020 , Wang et al., 2021] .There is A with cond(A) = O(1) such that L = AD, rank(L) = r.Under Assumption 1, we can establish a linear convergence rate for NAG when the initialization is sufficiently unbalanced and X 0 contains the column space of L.EQUATION)where p = √ μ 144 √ L does not depend on the scaling of X 0 . If we choose η = 1 L and β = √ L- √ μ √ L+ √ μ , then the t-th iterate of NAG (X t and Y t ) will correspond to residual R t = X t Y ⊤ t D -L satisfying ∥R t ∥ F ≤ σ 2 r (X 0 )σ min (D) 576 1 - √ μ 2 L t . Equivalently, let C = σ 2 r (X0)σmin(D) 576∥LD ⊤ ∥ F, then the iteration complexity for ϵ relative error isT = O σ 1 (X 0 ) λ max (DD ⊤ ) σ r (X 0 ) λ min (DD ⊤ ) log C ϵ .As Theorem 3 shows, if our initialization guarantees the column space of X 0 contains columns of L, then the residual shrinks at a linear rate. In the worst case, the columns of L span the whole space of R m , hence d should be at least m. However, when the data exhibits some low-dimensional properties, e.g., D is low-rank, then r can be much smaller than m and N . In this case, an initialization similar to (2) can meet the requirement of Theorem 3. Moreover, note that the convergence rate depends on both D and X 0 , hence by orthonormalization we can make cond(X 0 ) = 1 for a faster rate. When r ≤ d ≪ min(m, N ), such orthonormalization is affordable as it takes O(md 2 ) time rather than O(mN 2 ) in the worst case. We summarize these initialization options:EQUATIONHere, Orth(•) denotes the orthonormalization result whose columns are orthonormal. By applying singular value bounds and invoking Theorem 3, we obtain the following corollaries.Corollary 1. Suppose initialization (12) is applied with some sufficiently large c. For any0 < τ < c 1 , 0 < δ < 1, if d ≥ r -1 + Ω(log 1 δ ), then with probability at least 1 -δ, NAG finds X T and Y T such that f (X T , Y T ) ≤ ϵ LD ⊤ 2 F where T = O d • cond(L) τ (d -r + 1) λ max (DD ⊤ ) λ min (DD ⊤ ) log 1 ϵ .Corollary 2. Suppose initialization (13) is applied with some sufficiently large c.If d ≥ r, then with probability 1, NAG finds X T and Y T such that f (X T , Y T ) ≤ ϵ LD ⊤ 2 F where T = O λ max (DD ⊤ ) λ min (DD ⊤ ) log 1 ϵ .Corollary 3. Suppose initialization (14) is applied with some sufficiently large c. For any0 < τ < c 1 , 0 < δ < 1, if d ≥ m -1 + Ω(log 1 δ ), then with probability at least 1 -δ, NAG finds X T and Y T such that f (X T , Y T ) ≤ ϵ LD ⊤ 2 F where T = O d τ (d -m + 1) λ max (DD ⊤ ) λ min (DD ⊤ ) log 1 ϵ .Remark 3. While we only consider NAG in this section, our analysis can be directly applied to GD and obtain Oσ 2 1 (X0)λmax(DD ⊤ ) σ 2r (X0)λmin(DD ⊤ ) log 1 ϵ convergence rate with initializations (12) to (14).Corollaries 2 and 3 show accelerated convergence rate of NAG, as their dependence on the condition numberκ := λmax(DD ⊤ ) λmin(DD ⊤ ) = cond 2 (D) is O( √ κ) rather than O(κ), matching the results in Wang et al. [2021] for HB and Liu et al. [2022] for NAG. Meanwhile, Corollary 1 has an additional dependence on cond(L). Under Assumption 1, cond(L) = O( √ κ) and hence the overall dependence is O(κ). Although this is slower than NAG with initialization (13) or ( 14), it still outperforms GD with initialization ( 12), which has O(κ 2 ) dependence. Compared to previous results listed in Table 1 , we only require the network width to be Ω(r +log 1 δ ) or Ω(m+log 1 δ ) depending on the initialization and there is no additional dependence on the input rank or condition number. When the data is low-rank, NAG with initialization (12) enables the sublinear-width (w.r.t. output dimension and sample size) network to converge linearly. It can be further accelerated if orthonormalization is adopted (13), which echos the orthogonal initialization in Hu et al. [2020] , Wang et al. [2021] . In the general case, our analysis still provides a tighter result, as ( 14) only requires the width to be Ω(m + log 1 δ ). 
numerical experiment: 
 We validate our results via numerical experiments. For matrix factorization (1), we constructA = UΣV ⊤ ∈ R 100×80, where Σ ∈ R 5×5 is diagonal with σ 1 (Σ) = 1 and σ 5 (Σ) = 0.2, and U and V are orthonormal matrices. We set different levels of overparameterization (d ≥ 5) and initialize X 0 and Y 0 according to (2) with c = 50 √ d. For linear neural network (10), we construct the input data matrix D = UΣV ⊤ ∈ R 80×120 , where Σ ∈ R 5×5 is diagonal with σ 1 (Σ) = 1 and σ 5 (Σ) = 0.5, U is orthonormal and V is Gaussian. We use a Gaussian matrix A ∈ R 100×80 to construct the label matrix L = AD. We keep c = 50 √ d and initialize X 0 and Y 0 according to (12). We run all experiments with 10 different initialization seeds and take the average. The left plot is for (1), and the right plot is for (10).We first compare GD and AltGD. For matrix factorization, We use the same initialization and the same step size η = 2/(L + µ), where L and µ are computed as defined in Theorems 1 and 2. For linear neural networks, L and µ are replaced by L and μ in Theorem 3. As shown in Figure 1 , they perform very similarly and the loss curves are overlapped. To better illustrate, we additionally use η = 1/L for GD, and it performs differently from GD/AltGD with η = 2/(L + µ). To further illustrate the tightness of our theory, we compare our theoretical predictions with the actual loss in matrix factorization, as shown in Figure 3 . We set c = 200 √ d and σ 5 (Σ) ∈ {0.1, 0.01}, keeping other settings unchanged. The theoretical prediction at step t is computed as(1 -µ/L) 2t • f (X 0 , Y 0 ) for GD and (1 - √ µ/(2 √ L)) 2t • f (X 0 , Y 0 )for NAG. We observe that the slope of the predicted loss closely matches the actual loss, supporting the tightness of our theory, especially for GD. Additional experiments are provided in Appendix E. 
conclusion and future work: 
 We establish the convergence rate of GD and NAG for rectangular matrix factorization (1) under an unbalanced initialization and show the provable acceleration of NAG. We further extend our analysis to linear neural networks (10) and show the acceleration of NAG without excessive width requirements in previous work. Numerical experiments are provided to support our theory.We believe our analysis can be extended to initialization where X 0 ≈ cAΦ and Y 0 ≈ 0 rather than exact equalities. Relaxing the exact rank-r condition to approximately rank-r is also a possible generalization. The linear neural network model considered in this paper cannot fully capture the practical settings. We leave the extension to nonlinear activations for future work.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: We specify all important experiment details in Section 5.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [No] Justification: Our experiments do not require error bars.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [NA] Justification: There is no societal impact of the work performed. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [NA] Justification: The paper does not use existing assets.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: A.1 Singular Value Bounds for Random Matrix Proposition 4 (Rudelson and Vershynin [2009] ). Let A be an N × n random matrix, N ≥ n, whose elements are i.i.d. zero mean sub-Gaussian random variables with unit variance. Then for τ ≥ 0, we haveP σ n (A) ≤ τ ( √ N - √ n -1) ≤ (c 1 τ ) N -n+1 + e -c2Nwhere c 1 , c 2 > 0 depend (polynomially) only on the sub-Gaussian moment. Proposition 5 (Vershynin [2010] ). Let A be an N × n random matrix, N ≥ n, whose elements are i.i.d. zero mean Gaussian random variables with unit variance. Then for t ≥ 0, we haveP σ 1 (A) ≥ √ N + √ n + t ≤ e -t 2 2 .A.2 Proof of Proposition 1Proof of Proposition 1. Singular value decompose A asA = UΣV ⊤ , then X 0 = cUΣV ⊤ Φ. Since V ⊤ V = I r , the columns of V ⊤ Φ ∈ R r×d are independent Gaussian vectors with distribution N (0, 1 d V ⊤ V) = N (0, 1 d I r ). By Proposition 4 in Appendix A, we haveP σ r (V ⊤ Φ) ≤ τ 1 - √ r -1 √ d ≤ e -(d-r+1) log 1 c 1 τ + e -c2dfor some universal constants c 1 and c 2 and any τ ≥ 0. On the other hand, by Proposition 5 in Appendix A, we haveP σ 1 (V ⊤ Φ) ≥ √ d + √ r + √ s √ d ≤ e -s 2 .Plugging in s = d and applying the union bound yieldP τ ( √ d - √ r -1) √ d ≤ σ r (V ⊤ Φ) ≤ σ 1 (V ⊤ Φ) ≤ 2 √ d + √ r √ d ≥ 1 -δ,where δ = 3e -min{(d-r+1) log 1 c 1 τ ,c2d, d 2 } . The proposition follows immediately from the fact thatc • σ r (V ⊤ Φ)σ r (A) ≤ σ r (X 0 ) ≤ σ 1 (X 0 ) ≤ c • σ 1 (V ⊤ Φ)σ 1 (A).B Missing Proofs for GD B.1 Auxiliary Lemma Lemma 6. Suppose {a t } t≥0 and {b t } t≥0 are two non-negative sequences satisfyinga t+1 ≤ ρ • a t + b t , b t ≤ θ t • c 0 , where 0 ≤ ρ < θ < 1, c 0 ≥ 0,then the following holds for all t ≥ 0:a t ≤ θ t • a 0 + c 0 θ -ρ .Proof. The inequality holds trivially for t = 0. For t ≥ 0, we havea t+1 = ρ t+1 • a 0 + t s=0 ρ t-s θ s • c 0 = ρ t+1 • a 0 + θ t+1 -ρ t+1 θ -ρ • c 0 = θ t+1 • a 0 + 1 θ -ρ • c 0 .Proof of Proposition 2. According to (3), we haveR t+1 = X t+1 Y ⊤ t+1 -A = (X t + P t )(Y t + Q t ) ⊤ -A = R t -η R t Y t Y ⊤ t + X t X ⊤ t R t + P t Q ⊤ t .Applying vectorization on both sides yields r t+1 = r t -ηH t r t + β(r t -r t-1 ) + vec(P t Q ⊤ t ) = (I mn -ηH t )r t + vec(P t Q ⊤ t ). Hence we have the result.Proof of Lemma 1. By Proposition 1, the symmetric matrix H 0 = I n ⊗ (X 0 X ⊤ 0 ) has nr positive eigenvalues, and the eigensubspace of these positive eigenvalues isH = n i=1 col(X 0 ) = n i=1 col(A).According to the GD update (3),col(X t+1 ) ⊆ col(X t ) + col(X t Y ⊤ t Y t ) + col(AY t ) ⊆ col(X t ) + col(A), hence by induction we conclude col(X t ) ⊆ col(A) for all t ≥ 0. As a result, we haver t = vec(X t Y ⊤ t -A) ∈ H. For ξ t , notice that col(R t Y t Y ⊤ t + X t X ⊤ t R t ) ⊆ col(R t ) + col(X t ) ⊆ col(A) and col(P t Q ⊤ t ) = col((X t+1 -X t )(Y t+1 -Y t ) ⊤ ) ⊆ col(X t+1 ) + col(X t ) ⊆ col(A), thus we have ξ t = η • vec(R t Y 0 Y ⊤ 0 + X 0 X ⊤ 0 R t -R t Y t Y ⊤ t -X t X ⊤ t R t ) + vec(P t Q ⊤ t ) ∈ H.Proof of Lemma 2. Since I mn commutes with symmetric matrix H 0 , we can simultaneously diagonalize the two matrices and getλ i (T GD ) = 1 -ηλ mn-i (H 0 ), ∀i = 1, 2, . . . , mn. When η ∈ (0, 2 L ), λ i (T GD ) = 1 for i = 1, 2, . . . , (m -r)n. Let {v i } mn i=1 be orthonormal eigenvec- tors, v i corresponds to λ i (T GD ), then we have Span({v i } (m-r)n i=1 ) = ker(H 0 ) ⊥ H. Consequently, ∥T GD v∥ = T GD mn i=1 ⟨v, v i ⟩ v i = mn i=(m-r)n+1 ⟨v, v i ⟩ 2 λ 2 i (T GD ) ≤ max (m-r)n+1≤i≤mn |λ i (T GD )| ∥v∥ = max{|1 -ηL| , |1 -ηµ|} ∥v∥ ,where the second identity is from v ∈ H ⊥ ker(H 0 ). Plugging in the step size yields the second result.Proof of Lemma 3. For all s ≤ t, by assumption we have∥P s ∥ F = η ∥R s Y s ∥ F ≤ η ∥Y s ∥ ∥R s ∥ F ≤ η(∥Y 0 ∥ + ∥Y s -Y 0 ∥) ∥R s ∥ F ≤ η(∥Y 0 ∥ + ∥Y s -Y 0 ∥ F ) ∥R s ∥ F ≤ η(∥Y 0 ∥ + C 2 ) ∥R s ∥ F ≤ η(∥Y 0 ∥ + C 2 )C 1 θ s ∥r 0 ∥ .Similarly, we have∥Q s ∥ F ≤ η(∥X 0 ∥ + C 2 )C 1 θ s ∥r 0 ∥ .Combining the two bounds yieldsvec(P s Q ⊤ s ) = P s Q ⊤ s F ≤ ∥P s ∥ F ∥Q s ∥ F ≤ C 3 θ 2t ∥r 0 ∥ 2 ,whereC 3 = η 2 C 2 1 (∥X 0 ∥ + C 2 )(∥Y 0 ∥ + C 2 ).For the second part, we have∥(H 0 -H s )r s ∥ = R s (Y 0 Y ⊤ 0 -Y s Y ⊤ s ) + (X 0 X ⊤ 0 -X s X ⊤ s )R s F ≤ R s (Y 0 Y ⊤ 0 -Y s Y ⊤ s ) F + (X 0 X ⊤ 0 -X s X ⊤ s )R s F ≤ Y 0 Y ⊤ 0 -Y s Y ⊤ s ∥R s ∥ F + X 0 X ⊤ 0 -X s X ⊤ s ∥R s ∥ F ≤ (2 ∥Y 0 ∥ + ∥Y s -Y 0 ∥ F ) ∥Y s -Y 0 ∥ F ∥R s ∥ F + (2 ∥X 0 ∥ + ∥X s -X 0 ∥ F ) ∥X s -X 0 ∥ F ∥R s ∥ F ≤ 2(∥X 0 ∥ + ∥Y 0 ∥ + C 2 )C 2 ∥R s ∥ F ≤ C 4 θ s ∥r 0 ∥ , where C 4 = 2η(∥X 0 ∥ + ∥Y 0 ∥ + C 2 )C 1 C 2 .Finally, when (6) holds, we have∥X t+1 -X 0 ∥ F ≤ t s=0 ∥P s ∥ F ≤ η(∥Y 0 ∥ + C 2 )C 1 1 -θ ∥r 0 ∥ ≤ C 2 .Similarly, we have ∥Y t+1 -Y 0 ∥ F ≤ C 2 .Proof of Theorem 1. Let C 1 to C 4 be constants defined in Lemma 3. Define ρ = L-µ L+µ , θ = 1 -µ L , a t = C 1 ∥r t ∥, and b t = C 1 ∥ξ t ∥ for t ≥ 0. By Proposition 2 and lemmas 1 and 2 we havea t+1 ≤ ρ • a t + b t for all t ≥ 0. It remains to show that b t ≤ θ t • c 0 . By initialization (2), a 0 = C 1 ∥r 0 ∥ = C 1 ∥A∥ F , b 0 = 0. Let C 1 = µ(L+µ)p 2∥A∥ F L(1+p) and C 2 = p √ L where p = µ(L-µ) 24L 2 ∈ (0, 1). Plugging η = 2 L+µ , ∥X 0 ∥ = √ L and ∥Y 0 ∥ = 0 into C 3 and C 4 yields C 3 = µ 2 p 3 ∥A∥ 2 F L(1 + p) , C 4 = 2µp 2 ∥A∥ F . Let c 0 = C 1 (C 3 ∥r 0 ∥ + C 4 ) ∥r 0 ∥ ,then we can show the following relations:EQUATIONIndeed, by Proposition 1, with probability at least 1 -δ, our choice of c guaranteesEQUATION)Our goal is to showa 0 + c 0 θ -ρ = C 1 ∥A∥ F + C 1 (C 3 ∥A∥ F + C 4 ) ∥A∥ F • L(L + µ) µ(L -µ) ≤ C 2 1 ∥A∥ F ,which is equivalent to∥A∥ F + µp 3 L(1 + p) + 2p 2 • L(L + µ) L -µ ≤ µ(L + µ)p 2L(1 + p) .The above inequality holds when:EQUATIONEQUATIONEQUATIONLet p = µ(L-µ) 24L 2 , then we have p < 1, pL < µ andp 2 L -µ ≤ p L -µ = µ 24L 2 ≤ 1 6L , 2pL L -µ ≤ µ 12L ≤ µ 6L(1 + p) ,thus ( 18) and ( 19) hold. Finally, (17) holds in view of (16):µ(L + µ)p 6L(1 + p) ≥ µp 6 = µ 2 (L -µ) 144L 2 ≥ ∥A∥ F .Combining the results proves the (15). Now we can proceed with the induction in Lemma 3. Firstly, ∥r 0 ∥ ≤ C 1 ∥r 0 ∥ as C 1 ≥ 1 by ( 15), and ∥X 0 -X 0 ∥ F = ∥Y 0 -Y 0 ∥ F = 0 ≤ C 2 .Suppose the induction conditions in Lemma 2 holds for s ≤ t, then we haveb s = C 1 ∥ξ s ∥ ≤ C 1 (C 3 θ 2s ∥r 0 ∥ 2 + C 4 θ s ∥r 0 ∥) ≤ c 0 • θ s .Consequently, by Lemma 6 and (15) we havea t+1 ≤ θ t+1 • a 0 + c 0 θ -ρ ≤ C 2 1 • θ t+1 ∥A∥ F ,thus ∥r t+1 ∥ ≤ C 1 θ t+1 ∥r 0 ∥. Moreover, by our construction of C 1 and C 2 , (6) always holds, thus we also have∥X t+1 -X 0 ∥ F ≤ C 2 and ∥Y t+1 -Y 0 ∥ F ≤ C 2 .All conditions for the t + 1 step are satisfied, hence the proof is completed by induction. Plugging in C 1 and the choice of c yields the results.Proof of Proposition 3. According to the NAG update rule, we haveR t+1 = X t+1 Y ⊤ t+1 -A = (X t + P t )(Y t + Q t ) ⊤ -A = R t + P t Y ⊤ t + X t Q ⊤ t + P t Q ⊤ t = R t + (β(X t -X t-1 ) -(1 + β)ηR t Y t + βηR t-1 Y t-1 ) Y ⊤ t + X t β(Y ⊤ t -Y ⊤ t-1 ) -(1 + β)ηX ⊤ t R t + βηX ⊤ t-1 R t-1 + P t Q ⊤ t = R t -(1 + β)η R t Y t Y ⊤ t + X t X ⊤ t R t + β(X t Y ⊤ t -X t-1 Y ⊤ t-1 ) + βη R t-1 Y t-1 Y ⊤ t-1 + X t-1 X ⊤ t-1 R t-1 + β(X t Y ⊤ t + X t-1 Y ⊤ t-1 ) -β X t-1 Y ⊤ t + X t Y ⊤ t-1 + βη R t-1 Y t-1 Y ⊤ t + X t X ⊤ t-1 R t-1 -R t-1 Y t-1 Y ⊤ t-1 -X t-1 X ⊤ t-1 R t-1 + P t Q ⊤ t = R t -(1 + β)η R t Y t Y ⊤ t + X t X ⊤ t R t + β(R t -R t-1 ) + βη R t-1 Y t-1 Y ⊤ t-1 + X t-1 X ⊤ t-1 R t-1 + β(X t Y ⊤ t + X t-1 Y ⊤ t-1 -X t-1 Y ⊤ t -X t Y ⊤ t-1 ) + βη R t-1 Y t-1 Y ⊤ t + X t X ⊤ t-1 R t-1 -R t-1 Y t-1 Y ⊤ t-1 -X t-1 X ⊤ t-1 R t-1 + P t Q ⊤ t .r t+1 = r t -(1 + β)ηH t r t + β(r t -r t-1 ) + βηH t-1 r t-1 + β vec(X t Y ⊤ t + X t-1 Y ⊤ t-1 -X t-1 Y ⊤ t -X t Y ⊤ t-1 ) + βη vec(R t-1 Y t-1 Y ⊤ t + X t X ⊤ t-1 R t-1 -R t-1 Y t-1 Y ⊤ t-1 -X t-1 X ⊤ t-1 R t-1 ) + vec(P t Q ⊤ t ) = (1 + β)(I mn -ηH t )r t -β(I mn -ηH t-1 )r t-1 + ψ t + ϕ t .Hence we have r t+1 r t = (1 + β)(I mn -ηH 0 ) -β(I mn -ηH 0 ) I mn 0r t r t-1 + ξ t 0 .Proof of Lemma 4. Suppose λ is an eigenvalue of T NAG , then we havedet(T NAG -λI 2mn ) = det((β + λ 2 -(1 + β)λ)I mn + (η(1 + β)λ -ηβ)H 0 ).Since H 0 is symmetric, it can be simultaneously diagonalized with I, hence the above equation becomesλ 2 -(1 + β)λ + β + η(1 + β)λ i (H 0 )λ -ηβλ i (H 0 ) = 0for some 1 ≤ i ≤ mn. Solving the equation yieldsλ = 1 2 (1 + β)(1 -ηλ i (H 0 )) ± (1 -ηλ i (H 0 )) (-4β + (1 + β) 2 (1 -ηλ i (H 0 ))) .For i > nr, λ i (H 0 ) = 0, hence λ = 1 or λ = β. The corresponding eigen subspaces areH 1 = (u ⊤ , v ⊤ ) ⊤ | u = v ∈ ker(H 0 ) , H β = (u ⊤ , v ⊤ ) ⊤ | u = βv ∈ ker(H 0 ) . The dimensions are dim(H 1 ) = dim(H β ) = (m -r)n.It is easy to verify that whenever 0 < β < 1,H 1 ⊕ H β = ker(H 0 ) × ker(H 0 ).The complement space of H 1 ⊕ H β corresponds to the eigen subspace for non-trivial eigenvalues. By checking the dimension and orthogonality, we have(H 1 ⊕ H β ) ⊥ = H × H.For i ≤ nr, the subspace is H × H and the contraction condition requires0 < η < 2(1 + β) (1 + 2β)σ 2 1 (X 0 ) = 2(1 + β) (1 + 2β)L .By checking the monotonicity of |λ| with respect to 1-ηλ i (H 0 ) ∈ [1 -ηL, 1 -ηµ], we have |λ| ≤ max 1 2 (1 + β)(1 -ηµ) + (1 -ηµ) (-4β + (1 + β) 2 (1 -ηµ) , 1 2 -(1 + β)(1 -ηL) + (1 -ηL) (-4β + (1 + β) 2 (1 -ηL) . If we choose step size η = 1 L , momentum β = √ L- √ µ √ L+ √ µ , then we have |λ| ≤ 1 -µ L .Proof of Lemma 5. According to Lemma 3,ξ t = ζ t + ι t , ζ t = vec(P t Q ⊤ t ) + β vec(ηR t-1 Y t-1 Q ⊤ t-1 + ηP t-1 X ⊤ t-1 R t-1 + P t-1 Q ⊤ t-1 ) ι t = (1 + β)η(H 0 -H t )r t -βη(H 0 -H t-1 )r t-1 .We first bound ∥P t ∥ F and ∥Q t ∥ F . For every 0 ≤ s ≤ t, we have∥R s Y s ∥ F ≤ ∥Y s ∥ ∥R s ∥ F ≤ (∥Y 0 ∥ + ∥Y s -Y 0 ∥) ∥R s ∥ F ≤ (∥Y 0 ∥ + ∥Y s -Y 0 ∥ F ) ∥R s ∥ F ≤ (∥Y 0 ∥ + C 2 ) ∥R s ∥ F .Similarly,R ⊤ s X s F ≤ (∥X 0 ∥ + C 2 ) ∥R s ∥ F .By assumption, we have∥R s ∥ F ≤ r s r s-1 ≤ C 1 θ s r 0 r -1 .As a result, the momentum terms can be bounded:EQUATIONandEQUATIONwhere we use β ≤ θ 2 < θ in the last steps.Next, we bound ∥ζ t ∥. Using the triangle inequality, we get∥ζ t ∥ ≤ P t Q ⊤ t F + β ηR t-1 Y t-1 Q ⊤ t-1 + ηP t-1 X ⊤ t-1 R t-1 + P t-1 Q ⊤ t-1 F .For the first term, we haveP t Q ⊤ t F ≤ ∥P t ∥ F ∥Q t ∥ F ≤ η 2 C 2 1 (∥X 0 ∥ + C 2 )(∥Y 0 ∥ + C 2 ) (1 -θ) 2 θ 2t r 0 r -1 2 .For the second term, we haveβ ηR t-1 Y t-1 Q ⊤ t-1 + ηP t-1 X ⊤ t-1 R t-1 + P t-1 Q ⊤ t-1 F ≤β (η ∥R t-1 ∥ F (∥Y t-1 ∥ ∥Q t-1 ∥ F + ∥X t-1 ∥ ∥P t-1 ∥ F ) + ∥P t-1 ∥ F ∥Q t-1 ∥ F ) ≤ η 2 C 2 1 (∥X 0 ∥ + C 2 )(∥Y 0 ∥ + C 2 )(3 -2θ) (1 -θ) 2 θ 2t r 0 r -1 2 .As a result, we have∥ζ t ∥ ≤ C 3 θ 2t r 0 r -1 2 ,whereC 3 = η 2 C 2 1 (∥X0∥+C2)(∥Y0∥+C2)(4-2θ) (1-θ) 2.We then show upper bound for ∥ι t ∥. Using the triangle inequality, we getEQUATION)For any s ≤ t, we have∥(H 0 -H s )r s ∥ = R s (Y 0 Y ⊤ 0 -Y s Y ⊤ s ) + (X 0 X ⊤ 0 -X s X ⊤ s )R s F ≤ R s (Y 0 Y ⊤ 0 -Y s Y ⊤ s ) F + (X 0 X ⊤ 0 -X s X ⊤ s )R s F ≤ Y 0 Y ⊤ 0 -Y s Y ⊤ s ∥R s ∥ F + X 0 X ⊤ 0 -X s X ⊤ s ∥R s ∥ F ≤ (2 ∥Y 0 ∥ + ∥Y s -Y 0 ∥ F ) ∥Y s -Y 0 ∥ F ∥R s ∥ F + (2 ∥X 0 ∥ + ∥X s -X 0 ∥ F ) ∥X s -X 0 ∥ F ∥R s ∥ F ≤ 2(∥X 0 ∥ + ∥Y 0 ∥ + C 2 )C 2 ∥R s ∥ F ≤ 2(∥X 0 ∥ + ∥Y 0 ∥ + C 2 )C 1 C 2 θ s r 0 r -1 .Plugging it into (22) yields∥ι t ∥ ≤ 2(∥X 0 ∥ + ∥Y 0 ∥ + C 2 )C 1 C 2 ((1 + β)ηθ t + βηθ t-1 ) r 0 r -1 ≤ C 4 θ t r 0 r -1 ,whereC 4 = 2η(∥X 0 ∥ + ∥Y 0 ∥ + C 2 )C 1 C 2 (1 + 2θ).Finally, given ( 9) and ( 20), we have∥X t+1 -X 0 ∥ F ≤ t s=0 ∥P s ∥ F ≤ ηC 1 (∥Y 0 ∥ + C 2 ) (1 -θ) 2 r 0 r -1 ≤ C 2 ,where the last inequality is from our assumption on C 2 . Similarly, by ( 21), we have∥Y t+1 -Y 0 ∥ F ≤ t s=0 ∥Q s ∥ F ≤ ηC 1 (∥X 0 ∥ + C 2 ) (1 -θ) 2 r 0 r -1 ≤ C 2 .Proof of Theorem 2. By initialization, we have∥r 0 ∥ = ∥r -1 ∥ = ∥A∥ F . Let C 1 to C 4 be constants defined in Lemma 5. Define ρ = 1 - √ µ √ L , θ = 1 - √ µ 2 √ L , a t = C 1 (r ⊤ t , r ⊤ t-1) , and b t = C 1 ∥ξ t ∥ for t ≥ 0. It is easy to verify that β ≤ θ 2 < θ < 1 and ρ < θ < 1. By Proposition 3 and lemmas 1 and 4 we havea t+1 ≤ ρ • a t + b t for all t ≥ 0. It remains to show that b t ≤ θ t • c 0 . For the initial step, a 0 = √ 2C 1 ∥A∥ F , b 0 = 0. Let C 1 = µp 4 √ 2∥A∥ F (1+p) and C 2 = p √ L where p = √ µ 144 √ L ≤ 1 144 < 1, then we have C 3 = µp 3 (2 + µ L ) 8 ∥A∥ 2 F (1 + p) , C 4 = µp 2 (3 -µ L ) 2 √ 2 ∥A∥ F . Let c 0 = √ 2C 1 ( √ 2C 3 ∥A∥ F + C 4 )∥A∥ F , then we can show the following relations:EQUATIONIndeed, by Proposition 1, with probability at least 1 -δ, our choice of c guaranteesEQUATIONthus C 1 ≥ 1. Here, we use the bound p ≤ 1 144 < 1 to verify the numerical constant. It remains to showa 0 + c 0 θ -ρ ≤ √ 2C 2 1 ∥A∥ F ,which is equivalent to∥A∥ F + p 3 √ µL(2 + µ L ) 2 √ 2(1 + p) + p 2 √ µL(3 -µ L ) √ 2 ≤ µp 4 √ 2(1 + p) , Since we set p = √ µ 144 √L < 1, each one of the three terms on the left hand side is upper bounded by µp 12 √ 2(1+p) , hence the inequality holds. The relations (23) guarantee the induction conditions in Lemma 5, thus we have∥r t+1 ∥ ≤ √ 2C 1 θ t+1 ∥A∥ F ≤ c 2 σ 2 1 (A) 64 ∥A∥ F cond(X 0 ) θ t+1 ∥A∥ F ,where the last inequality uses p > 0 and Proposition 1.Let rt = vec( Rt ), then we have the following dynamics. Lemma 7.Let P t = X t+1 -X t and Q t = Y t+1 -Y t denote the momentum. Let R t = X t Y ⊤ t D-L denote the residual, Rt = X t Y ⊤ t DD ⊤ -LD ⊤ denote the projected residual, rt = vec( Rt ) ∈ R mn .Then NAG has the following dynamics:EQUATIONwhereH t = (DD ⊤ Y t Y ⊤ t ) ⊗ I m + (DD ⊤ ) ⊗ (X t X ⊤ t ), ξ t = ζ t + ι t , ζ t = vec(P t Q ⊤ t DD ⊤ ) + β vec(P t-1 Q ⊤ t-1 DD ⊤ ) + βη vec(( Rt-1 Y t-1 Q ⊤ t-1 + P t-1 X ⊤ t-1 Rt-1 )DD ⊤ ), ι t = (1 + β)η(H 0 -H t )r t -βη(H 0 -H t-1 )r t-1 .Proof of Lemma 7. We denote R t = X t Y ⊤ t D -L as the residual, Rt = R t D ⊤ as the projected residual, then the NAG update for (10) can be written asEQUATIONThe result follows from ( 26) by direct computation.Lemma 8. Let H ⊆ R mn denote the linear subspace containing all eigenvectors of H 0 = (DD ⊤ ) ⊗ (X 0 X ⊤ 0 ) with positive eigenvalues. If col(X 0 ) = col(L) and Y 0 = 0, then we haveH = col(D ⊗ L) and {r t , ξ t } t≥0 ⊂ H,where H 0 , rt and ξ t are defined as in Lemma 7.Proof. By Theorem 4.2.15 in Horn and Johnson [1994] , we have the following eigenvalue decomposition for Kronecker product:H 0 = (U D ⊗ U 0 )(Σ 2 D ⊗ Σ 2 0 )(U D ⊗ U 0 ) ⊤ , where D = U D Σ D V ⊤D and X 0 = U 0 Σ 0 V ⊤ 0 are singular value decompositions of D and X 0 . Therefore, we haveH = col(U D ⊗ U 0 ) = col(D ⊗ X 0 ) = col(D ⊗ L).In particular, the eigenvalues (not ordered) areλ (i-1)m+j (H 0 ) = λ i (DD ⊤ )λ j (X 0 X ⊤ 0 ) = σ 2 i (D)σ 2 j (X 0 ), i ∈ [n], j ∈ [m], where σ j (X 0 ) > 0 for 1 ≤ j ≤ r, σ j (X 0 ) = 0 for r + 1 ≤ j ≤ d. By Assumption 1, L = AD, thus we have vec(LD ⊤ ) = vec(LI N D ⊤ ) = (D ⊗ L)I N ∈ col(D ⊗ L) = H. Meanwhile, vec(X t Y ⊤ t DD ⊤ ) = (D ⊗ X t ) vec(Y ⊤ t D) ∈ col(D ⊗ X t ) ⊆ col(D ⊗ X 0 ) = H, thus we have rt ∈ H. Similarly, we have ξ t ∈ H. Lemma 9 (NAG contraction). If we choose step size η = 1 L and momentum β = √ L- √ μ √ L+ √ μ where L = σ 2 1 (X 0 ) • λ max (DD ⊤ ), μ = σ 2 r (X 0 ) • λ min (DD ⊤ ), then for all (u, v) ∈ H × H, H defined in Lemma 8, T NAG u v ≤ 1 - μ L u v .Proof. Following the same line of proof for Lemma 4 in Appendix C.2 and substituting the eigenvalues in Lemma 8, we obtain the result.Lemma 10. Suppose 0 < β ≤ θ 2 < θ < 1. If there exist some constants C 1 and C 2 such that for any s ≤ t, the NAG dynamics (7) yields rs rs-1EQUATION∥Y s -Y 0 ∥ F ≤ C 2 , then we have ∥ζ t ∥ ≤ C 3 θ 2t r0 r-1 2 , and ∥ι t ∥ ≤ C 4 θ t r0 r-1for some constants C 3 and C 4 depending on C 1 and C 2 . Moreover, if C 1 and C 2 satisfy(max(∥X 0 ∥ , ∥Y 0 ∥) + C 2 ) ηC 1 r0 r-1 ≤ (1 -θ) 2 C 2 ,then we have∥X t+1 -X 0 ∥ F ≤ C 2 , ∥Y t+1 -Y 0 ∥ F ≤ C 2 .Proof of Lemma 10. Following the same of proof for Lemma 5 in Appendix C.3, we haveEQUATIONandEQUATION)As a result, we haveEQUATIONβ (η Rt-1 Y t-1 Q ⊤ t-1 + ηP t-1 X ⊤ t-1 Rt-1 + P t-1 Q ⊤ t-1 )DD ⊤ F ≤βλ 1 (DD ⊤ ) η Rt-1 F (∥Y t-1 ∥ ∥Q t-1 ∥ F + ∥X t-1 ∥ ∥P t-1 ∥ F ) + ∥P t-1 ∥ F ∥Q t-1 ∥ F ≤ η 2 C 2 1 (∥X 0 ∥ + C 2 )(∥Y 0 ∥ + C 2 )(3 -2θ)λ 1 (DD ⊤ ) (1 -θ) 2 θ 2t r0 r-1 2 ,Combining the inequalities, we get∥ζ t ∥ ≤ C 3 θ 2t r0 r-1 2 ,whereC 3 = η 2 C 2 1 (∥X0∥+C2)(∥Y0∥+C2)(4-2θ)λ1(DD ⊤ ) (1-θ) 2.Similarly, we have∥ι t ∥ ≤ 2(∥X 0 ∥ + ∥Y 0 ∥ + C 2 )C 1 C 2 λ 1 (DD ⊤ )((1 + β)ηθ t + βηθ t-1 ) r0 r-1 ≤ C 4 θ t r0 r-1 , where C 4 = 2η(∥X 0 ∥ + ∥Y 0 ∥ + C 2 )C 1 C 2 (1 + 2θ)λ 1 (DD ⊤ ).Finally, by ( 27), we have∥X t+1 -X 0 ∥ F ≤ t s=0 ∥P s ∥ F ≤ ηC 1 (∥Y 0 ∥ + C 2 ) (1 -θ) 2 r0 r-1 ≤ C 2 ,where the last inequality is from our assumption on C 2 . Similarly, by (28), we have∥Y t+1 -Y 0 ∥ F ≤ t s=0 ∥Q s ∥ F ≤ ηC 1 (∥X 0 ∥ + C 2 ) (1 -θ) 2 r0 r-1 ≤ C 2 .Proof of Theorem 3. By initialization, we have∥r 0 ∥ = ∥r -1 ∥ = LD ⊤ F . Let C 1 to C 4 be constants defined in Lemma 10. Define ρ = 1 - √ μ √ L , θ = 1 - √ μ 2 √ L , a t = C 1 Rt R-1 F, and b t = C 1 ∥ξ t ∥ for t ≥ 0. It is easy to verify that β ≤ θ 2 < θ < 1 and ρ < θ < 1. By Lemmas 7 to 9 we havea t+1 ≤ ρ • a t + b t for all t ≥ 0. It remains show that b t ≤ θ t • c 0 . For the initial step, a 0 = √ 2C 1 LD ⊤ F , b 0 = 0. Let C 1 = μp 4 √ 2∥LD ⊤ ∥ F (1+p) and C 2 = p √ L where p = √ μ 144 √ L ≤ 1 144 < 1, then we have C 3 = μp 3 8 ∥LD ⊤ ∥ 2 F (1 + p) 2 + μ L , C 4 = μp 2 2 √ 2 ∥LD ⊤ ∥ F 3 - μ L . Let c 0 = √ 2C 1 ( √ 2C 3 LD ⊤ F + C 4 ) LD ⊤F , then we can show the following relations: Given our choice of constants, there holdEQUATIONIndeed, by ( 11), we have C 1 ≥ 1. It remains to showa 0 + c 0 θ -ρ ≤ √ 2C 2 1 LD ⊤ F ,which is equivalent to LD ⊤ F + μ Lp 3 2 √ 2(1 + p) 2 + μ L + μ Lp 2 √ 2 3 - μ L ≤ μp 4 √ 2(1 + p) .∥r t+1 ∥ ≤ √ 2C 1 θ t+1 LD ⊤ F ≤ μ 576 ∥LD ⊤ ∥ F 1 - √ μ 2 L t+1 LD ⊤ F .By Assumption 1, we have row(L) ∈ row(D) = col(D ⊤ ), thus we have ∥R t ∥ F = X t Y ⊤ t D -L F ≤ σ -1 min (D) (X t Y ⊤ t D -L)D ⊤ F ≤ σ 2 r (X 0 )σ min (D) 576 1 - σ r (X 0 ) λ min (DD ⊤ ) 2σ 1 (X 0 ) λ max (DD ⊤ ) t .This section provides additional experiments. Firstly, we investigate larger-sized problems by setting (m, n) = (1200, 1000) for matrix factorization and (m, n, N ) = (500, 400, 600) for linear neural networks. We keep other settings the same as for Figure 2 and compare the performances of GD and NAG. The results are provided in Figure 4 . As illustrated, the conclusion that NAG performs better than GD and overparameterization accelerates convergence remains valid for large matrices. Secondly, we conduct additional experiments on GD and NAG with different values of c and plot the results in Figure 5 . We set d = 20, while other settings remain the same as in Figure 2 . As illustrated, when c is sufficiently large, further increasing c has little effect on the convergence rate, which is consistent with our theory.We also investigate general unbalanced initializationX 0 = c 1 AΦ 1 ∈ R m×d , Y 0 = c 2 Φ 2 ∈ R n×d , where [Φ 1 ] i,j ∼ N (0, 1/d) and [Φ 2 ] i,j ∼ N (0, 1/n).We set d = 20, while other settings remain the same as in Figure 2 . We keep c 1 = 50 and set different values of c 2 . The results are plotted in Figure 6 . As illustrated, changing c 2 within a range only results in different initial losses (intercept), while the convergence rates (slope) are not significantly affected. This supports our claim in Remark 1.There are some works[Wang et al.,2022, 2023] proving convergence of GD for general initialization under large learning rate and similar objective functions, but nonasymptotic convergence analysis is very challenging and highly dependent on initialization.2 Here, the nonsmoothness refers to the lack of uniform Lipschitz constant for the gradient in the full domain.3 In the symmetric case, the solution's uniqueness is up to rotation, whereas in (1) it is also up to scaling. 


Decision: The paper studies the problem of factorization of rectangular matrices. The authors prove the convergence rate of gradient descent with quadratic dependence on the condition number of the data matrix, and the convergence rate of Nesterov's accelerated gradient descent with linear dependence on the condition number. They also extend their analysis to linear neural networks. Three reviewers are unanimously positive, and one review with borderline reject is short and less informative. I decide to accept the paper.

Reviews: review: paper_summary: paper_summary: In this papers the authors analyze the convergence of  Nesterov Accelerated Gradient algorithm for a) rectangular matrix factorisation and b) linear neural networks. By using imbalanced initialisation,  the authors come up with linear rates of convergence impoving upon state-of-the-art regarding dependence of the rates condition numbers of sought matrices.
strength_weakness: strength_weakness: - The authors didn't mention relevant works studying imbalance effect on the similar  problems e.g.  [1].
- The results address the linear neural network case but it's not obvious how the analysis could be extended to account for non-linearities or non-smooth activation functions. 

[1] Min, Hancheng, et al. "On the explicit role of initialization on the convergence and implicit bias of overparametrized linear networks." International Conference on Machine Learning. PMLR, 2021.
questions: questions: - In [1] the authors provided rates of convergence of gradient flows in the imbalance initialisation regime. Could these results provide some insights on how to generalise the derived rate for different amount imbalance of imbalance?
- Could the authors provide some insights on whether the theoretical results on linear neural network  could be extended to non-linear and possibly non-smooth activation functions?

Minor: 
Eq. under line 132: Should $x_t$  be replaced by $z_t$ (??
limitations: limitations: Yes
paper_summary: paper_summary: The paper calculates convergence rates of the gradient descent and the Nesterov's accelerated gradient descent algorithms for factorization of rectangular matrices- a nonconvex optimization problem. Their analysis is for algorithms when the factor matrices are initialized as follows:  one matrix is initialized as the original matrix multiplied with Gaussian random matrix with appropriate scaling and the other is a zero matrix. They show convergence rates with quadratic dependence on condition number of the data matrix for GD and linear dependence for NAG. They also extend their analysis to linear neural networks and back their theoretical findings with empirical results.
strength_weakness: strength_weakness: 1) The experimentation is done on very small matrices. In practice much larger rectangular matrices are often encountered. It would give more insights if the experiments were also performed with moderate and large size matrices.
questions: questions: See Weaknesses
limitations: limitations: See Weaknesses
paper_summary: paper_summary: This paper considers the convergence of the first-order optimization method, which includes the gradient descent and the nesterov's acccelerated gradient 
for the marix factorization and the linear neural network. 

In Section 2.1, they analyze the gradient descent algorithm on the matrix factorization (c.f. Thm. 1). 
In Section 2.2, they analyze the NAG (c.f. Thm 2).  Section 3 illustrates the proof strategy. 
Section 4 extends the analysis to the linear neural network (c.f. Thm. 3) and Section 5 presents the numerical experiments.
strength_weakness: strength_weakness: The major concern is the potential impact of this paper.  Generally speaking, this paper studies a well-studied problem with quite standard technique. 

Also, there are some small typos that are quite annoying. For example, the $\|A\|_F$  in Thm. 1and Thm. 2 can cancel out. Thm 2 (line 14), the GD should be NAG. Please do another round of proof-reading.
questions: questions: 1. Why the convergence rate is independent of $\|A\|_F$?
2. Can you explain the line 807? Why the eigenvalues of $\lambda_i(T_{GD})$ ($1\leq i \leq (m-r)n$ can be ignored? The previous line suggest these eigenvalues are one.
limitations: limitations: NA
paper_summary: paper_summary: In this study, the authors established the convergence rate of the gradient descent and Nesterov's accelerated gradient descent methods for the asymmetric matrix factorization and the 2-layer linear network. The authors proved that an unbalanced initialization can lead to linear convergence of both methods, and the Nesterov's acceleration result in a faster convergence rate. Numerical experiments are implemented to support the theory.
strength_weakness: strength_weakness: I do not see major weaknesses. The authors could consider discussing more related literature and include more intermediate steps in the main manuscript. Also, I think the authors can consider illustrating the effect of an unbalanced initialization in numerical experiments.
questions: questions: - I think the following two papers also discussed the asymmetric low-rank matrix optimization problem. It would be better if the authors could discuss them and the references therein:

[1] Zhang, H., Bi, Y., & Lavaei, J. (2021). General low-rank matrix optimization: Geometric analysis and sharper bounds. Advances in Neural Information Processing Systems, 34, 27369-27380.

[2] Bi, Y., Zhang, H., & Lavaei, J. (2022, June). Local and global linear convergence of general low-rank matrix recovery problems. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 36, No. 9, pp. 10129-10137).


- Following the above comment, I wonder if the authors can briefly discuss the potential and challenges of extending the results to more general low-rank matrix optimization problems.

- Lines 79-80: I am not sure why the method in [Stöger and Soltanolkotabi, 2021] is considered a preconditioned method.

- Theorem 1: In my understanding, the value $\|R_t\|_F$ should be proportional to the scale of $A$. However, in the bound of $\|R_t\|_F$, the right hand-side grows with the size of $A$. Maybe there is a typo?

- The above comment also applies to Theorem 2.

- It would be better to mention in Theorems 1-2 that the bound $T$ is derived using the bound in Proposition 1?

- Line 182: I think the shrinkage rate $\theta \in (0, \rho]$?

- I wonder if the authors can compare the performance of GD/NAG with different values of c?
limitations: limitations: See my comments in the Questions section.

Comments: title: Thank you for the comment
comment: Thank you for the response. I will stick to my original score. 

I appreciate that the author agrees with my suggestions on the presentation improvement and promise to make the revisions accordingly.
Dear Reviewer,

Thank you for your time in reviewing our manuscript and offering valuable comments and suggestions. We've addressed the questions raised in your review with our author rebuttals. We'd like to know if these responses have sufficiently addressed your concerns. If any points require further clarification or discussion, please let us know.

Thank you again for your effort.

Best wishes, 

Authors
Dear Reviewer, 

Thank you for your time in reviewing our manuscript and offering valuable comments. We've addressed the weaknesses and questions raised in your review with our author rebuttals. We'd like to know if these responses have sufficiently addressed your concerns. If any points require further clarification or discussion, please let us know.

Thank you again for your effort. 

Best wishes,

Authors
Thank you for your rebuttal to my comments and for the additional experiments you conducted to explore how imbalanced initialization affects the rate of convergence. I will keep my score as it is.
I have read the other reviews and the rebuttal. I will keep my score.