Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruction
Year: 2021

Abstract:
Recently there has been significant theoretical progress on understanding the convergence and generalization of gradient-based methods on nonconvex losses with overparameterized models. Nevertheless, many aspects of optimization and generalization and in particular the critical role of small random initialization are not fully understood. In this paper, we take a step towards demystifying this role by proving that small random initialization followed by a few iterations of gradient descent behaves akin to popular spectral methods. We also show that this implicit spectral bias from small random initialization, which is provably more prominent for overparameterized models, also puts the gradient descent iterations on a particular trajectory towards solutions that are not only globally optimal but also generalize well. Concretely, we focus on the problem of reconstructing a low-rank matrix from a few measurements via a natural nonconvex formulation. In this setting, we show that the trajectory of the gradient descent iterations from small random initialization can be approximately decomposed into three phases: (I) a spectral or alignment phase where we show that that the iterates have an implicit spectral bias akin to spectral initialization allowing us to show that at the end of this phase the column space of the iterates and the underlying low-rank matrix are sufficiently aligned, (II) a saddle avoidance/refinement phase where we show that the trajectory of the gradient iterates moves away from certain degenerate saddle points, and (III) a local refinement phase where we show that after avoiding the saddles the iterates converge quickly to the underlying low-rank matrix. Underlying our analysis are insights for the analysis of overparameterized nonconvex optimization schemes that may have implications for computational problems beyond low-rank reconstruction.

heading: 1 Introduction
text: Many contemporary problems in machine learning and signal estimation spanning deep learning to low-rank matrix reconstruction involve fitting nonlinear models to training data. Despite tremendous empirical progress, theoretical understanding of these problems poses two fundamental challenges. First, from an optimization perspective, fitting these models often requires solving highly nonconvex optimization problems and except for a few special cases, it is not known how to provably find globally or approximately optimal solutions. Yet simple heuristics such as running (stochastic) gradient descent from (typically) small random initialization is surprisingly effective at finding globally optimal solutions. A second generalization challenge is that many modern learning models including neural network architectures are trained in an overparameterized regime where the parameters of the model exceed the size of the training dataset. It is well understood that in this overparameterized
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
regime, these large models are highly expressive and have the capacity to (over)fit arbitrary training datasets including pure noise. Mysteriously however overparameterized models trained via simple algorithms such as (stochastic) gradient descent when initialized at random continue to predict well or generalize on yet unseen test data. In particular, it has been noted in a number of works that for many modern machine learning architectures, the scale of initialization is important for the generalization/test behavior [1, 2]. It has been noted that stronger generalization performance is typically observed for a smaller scale initialization. Indeed, small random initialization followed by (stochastic) gradient descent iterative updates is arguably the most widely used learning algorithm in modern machine learning and signal estimation.
There has been a large number of exciting results aimed at demystifying both the optimization and generalization aspects over the past few years. We will elaborate on these results in detail in the supplementary, however, we would like to briefly mention the common techniques and their existing limitations. On the optimization front a large body of work has emerged on providing guarantees for nonconvex optimization which can roughly be put into two categories: (I) smart initialization+local convergence and (II) landscape analysis+saddle escaping algorithms. Approaches in (I) focus on showing local convergence of local search techniques from carefully designed spectral initializations [3, 4, 5, 6, 7, 8, 9, 10]. Approaches in (II) focus on showing that in some cases the optimization landscape is benign in the sense that all local minima are global (no spurious local minima) and the saddle points have a direction of strict negative curvature (strict saddle) [11]. Then specialized truncation or saddle escaping algorithms such as trust region, cubic regularization [12, 13], or noisy (stochastic) gradient-based methods [14, 15, 16, 17] are deployed to provably find a global optimum. Both approaches fail to fully explain the typical behavior of local search techniques in practice. Indeed, for many nonconvex problems local search techniques or simple variants, when initialized at random, quickly converge to globally optimal solutions without getting stuck in local optima/saddles without the need for sophisticated initialization or saddle escaping heuristics. We note that while for differentiable losses eventual convergence to local minimizers is known from a random initialization [18] on problems of the form (II), these results cannot rule out exponentially slow cases in the worst-case [19]. Indeed, it has been argued that in general a more granular analysis of the trajectory of gradient descent beyond the landscape may be necessary [20]. For example, some recent advances has been made by analysing the trajectory of gradient descent using a leave-one-out analysis for the phase retrieval problem [21].
Similarly, there has been a lot of exciting progress on the generalization front, especially for neural networks. Specific to generalization capabilities of gradient-based approaches these results broadly fall into two categories: (a) the first category is based on a linearization principle which characterizes the performance of nonlinear models such as neural networks by comparing it to a linearized kernel problem around the initialization (a.k.a. Neural Tangent Kernels) [22, 23, 24, 25, 26, 27, 28]. This has often been referred to as ”lazy training". (b) the second category is based on a continuous limit analysis in the limit of width going to infinity and learning rate going to zero (mean-field analysis) [29, 30, 31, 32, 33]. However, these existing analyses contain many idealized and nonrealistic assumptions (e.g. requiring large, random initialization in (a), which typically leads to worse generalization than what is observed in practice, or unrealistically large widths in (b)) and therefore cannot fully explain the success of overparameterized models or serve as a guiding principle for practitioners [34].
Despite the aforementioned exciting recent theoretical progress many aspects of optimization and generalization and in particular the role of random initialization remains mysterious. This leads us to the main challenge of this paper
Why is small random initialization combined with gradient descent updates so effective at finding globally optimal models that generalize well despite the nonconvex nature of the optimization landscape or model overparameterization?
In this paper we wish to take a step towards addressing the above challenge by demystifying the critical role of small random initialization in gradient-based approaches. Specifically we show that
Small random initialization followed by a few iterations of gradient descent behaves akin to spectral initialization.
By that, we mean more precisely, that if the initialization is chosen small enough, then in the initial stage of the training, gradient descent implicitly behaves like spectral initialization techniques such as
those commonly used in techniques based on the method of moments. This implicit spectral bias of gradient descent from random initialization puts the gradient descent iterations on a particular trajectory towards solutions that are not only globally optimal but also generalize well for overparameterized models. We also show that with small random initialization this implicit spectral bias phenomenon is more prominent for more overparameterized models in the sense that it materializes after fewer iterations. This intriguing phenomenon is depicted in Figure 1 in the context of a low-rank reconstruction problem. This figure clearly demonstrates that the first few iterations of gradient descent starting from a small random initialization are virtually identical to that of running power iterations (a popular algorithm to find the spectral initialization, see, e.g. [35]).
Concretely we focus on the problem of low-rank matrix recovery, which appears in many different application areas such as recommendation systems, phase retrieval, and quantum tomography [36]. Here, our goal is to recover a low-rank matrix of the form XXT from a few linear measurements. We consider a natural, non-convex approach based on matrix factorization, where we minimize the loss function via gradient descent. In this paper, we show that, regardless of the amount of overparameterization used, for small random initialization vanilla gradient descent will always converge towards the low-rank solution. This holds as long as the measurement operator obeys a popular restricted isometry property [37].
Our analysis consists of three phases. The first phase is the aforementioned spectral or alignment phase where we show gradient descent from small random initialization behaves akin to spectral initialization, which is a key insight of this paper. Indeed, we show that the first few gradient descent iterates can be accurately approximated by power method iterates. Next, we show that after this first spectral or alignment phase, gradient descent enters a second phase, which we refer to as saddle avoidance phase. In this phase, we show that the trajectory of the gradient iterates moves away from degenerate saddle points, while the iterates maintain almost the same effective rank as XXT . In the third phase, the local refinement phase, we show that the iterates approximately converge towards the underlying low-rank matrix XXT with a geometric rate up to a certain error floor which depends on the initialization scale. In particular, by decreasing the scale of initialization this error threshold can
be made arbitrarily small. While in this paper our main focus is on low-rank matrix reconstruction, we believe that our analysis holds more generally for a variety of contemporary machine learning and signal estimation tasks including neural networks.
Finally we note that while a similar setting has already been studied in [38], our analysis goes beyond it in many important ways. For example, our result holds for any amount of overparameterization and allows for arbitrarily small initialization. Maybe most importantly, we study the spectral phase phenomenon at initialization.

heading: 2 Low-rank matrix recovery via non-convex optimization
text: As mentioned earlier in this paper we focus on reconstructing a (possibly overparameterized) Positive Semidefinite (PSD) low rank matrix from a few measurements. In this problem, given m observations of the form
yi = ⟨Ai,XXT ⟩ = Tr (AiXXT ) i = 1, . . . ,m, (1) we wish to reconstruct the unknown matrix XXT . Here, X ∈ Rn×r⋆ with 1 ≤ r⋆ ≤ n is a factor of the unknown matrix and {Ai}mi=1 are known symmetric measurement matrices. A common approach to solving this problem is via minimizing the loss function
min Ū∈Rn×r f(Ū) ∶= minŪ∈Rn×r 14m m∑ i=1 (yi − ⟨Ai, Ū ŪT ⟩) 2 ,
with r ≥ r⋆. More compactly one can rewrite the optimization problem above in the form min Ū∈Rn×r f(Ū) ∶= minŪ∈Rn×r 14 ∥A (Ū ŪT −XXT )∥2`2 , (2) where A ∶ Rn×n Ð→ Rm is the measurement operator defined by [A (Z)]i ∶= 1√m ⟨Ai, Z⟩. In order to solve the minimization problem (2) we run gradient descent iterations starting from (often small) random initialization. More specifically, Ut+1 = Ut − µ∇f (Ut) = Ut + µA∗ [y −A (UtUTt )]Ut= Ut + µ [(A∗A) (XXT −UtUTt )]Ut. where we have set U0 = αU is the initialization matrix, A∗ denotes the adjoint operator of A and y = (yi)mi=1 ∈ Rm denotes the measurement vector. Here, U ∈ Rn×r is a typically random matrix which represents the form of the initialization and α > 0 is a scaling parameter. There are two challenges associated with analyzing such randomly initialized gradient descent updates. The first is an optimization challenge. Since f is non-convex it is a priori not clear whether gradient descent converges to a global optimum or whether it gets stuck in a local minima and/or saddle. The second challenge is that of generalization. This is particularly pronounced in the overparameterized scenario where the number of parameters are larger than the number of data points i.e. rn ≥m. In this case, there are infinitely many Ū such that f(Ū) = 0, but ∥Ū ŪT −XXT ∥F is arbitrarily large (see, e.g., [39, Proposition 1]). That is, even if gradient descent converges to a global optimum, i.e. f (Ū) = 0, it is a priori not clear whether it has found the low-rank solution XXT (see also Figure 5).

heading: 3 Main results
text: In this section, we present our main results. Stating these results requires a couple of simple definitions. The first definition concerns the measurement operator A. Definition 3.1 (Restricted Isometry Property (RIP)). The measurement operator A ∶ Rn×n Ð→ Rm satisfies RIP of rank r with constant δ > 0, if it holds for all matrices Z of rank at most r (1 − δ) ∥Z∥2F ≤ ∥A (Z)∥2`2 ≤ (1 + δ) ∥Z∥2F . (3) We note that for a Gaussian measurement operator A 1, RIP of rank r and constant δ > 0 holds with high probability, if the number of observations satisfies m ≳ nr/δ2 [37, 40].
1By that, we mean that all the entries of the (symmetric) measurement matrices {Ai}mi=1 are drawn i.i.d. with distributionN (0,1) on the off-diagonal and distributionN (0,1/ √ 2) on the diagonal.
The second definition concerns the condition number of the factor X . Definition 3.2 (condition number). We denote the condition number of X ∈ Rn×r⋆ by κ ∶= ∥X∥ σr⋆(X) ,
where σr⋆ (X) denotes r⋆-th largest singular value of X . With these definitions in place we are now ready to state our main results. Due to space limitations in the main paper we focus on the case of r ≥ 2r⋆. With refer the reader to the supplementary for results covering all r ≥ r⋆ including two special cases: (1) the fully overparameterized case, i.e., r = n along with comparisons with existing work, in this case [38], and (2) the scenario that U has the same number of parameters as X , i.e., r = r⋆. Theorem 3.3. Let X ∈ Rn×r∗ and assume we have m measurements of the low rank matrix XXT of the form y = A (XXT ) with A the measurement operator. We assume A satisfies the restricted isometry property for all matrices of rank at most 2r⋆+1 with constant δ ≤ cκ−4r⋆−1/2. To reconstruct XXT from the measurements we fit a model of the form Ū ↦ A (Ū ŪT ) with Ū ∈ Rn×r via running gradient descent iterations of the form Ut+1 = Ut − µ∇f (Ut) on the objective (2) with a step size obeying µ ≤ cκ−4∥X∥−2. Here, the initialization is given by U0 = αU , where U ∈ Rn×r has i.i.d. entries distributed as N (0,1/√r). Furthermore, we assume r ≥ 2r⋆ and that the scale of initialization fulfills
α ≲ min⎧⎪⎪⎨⎪⎪⎩ (min{r;n})1/4 κ1/2n3/4 (2κ2 √
n min{r;n}) −6κ2 ; 1 κ7n ⎫⎪⎪⎬⎪⎪⎭∥X∥. (4) Then, after
t̂ ≲ 1 µσmin (X)2 ln( C1nκ min{r;n} ⋅max{1; κr⋆min{r;n} − r⋆} ⋅ ∥X∥α ) iterations we have that
∥Ut̂UTt̂ −XXT ∥F∥X∥2 ≲ n 2κ81/16r1/8⋆(min{r;n})15/16 ⋅ α21/16∥X∥21/16 , (5) holds with probability at least 1 −Ce−c̃r. Here, c, c̃,C,C1 > 0 are fixed numerical constants. Note that the test error ∥Ut̂UTt̂ −XXT ∥2F can be made arbitrarily small by choosing the scale of initialization α small enough. In particular, the dependence of the test error on α is polynomial and the dependence of the number of iterations on α is logarithmic, which means that reducing the test error by scaling down α introduces only modest additional computational cost. Hence, as long as the rank at most 2r⋆ + 1 RIP with constant δ ≤ cκ−4r⋆−1/2 holds, gradient descent converges to a point in the proximity of the low-rank solution, whenever the initialization is chosen small enough regardless of the choice of r. This holds even when the model is overparameterized i.e. rn ≫ m and the optimization problem has many global optima many of which do not obey UUT ≈ XXT . This result thus further demonstrates that when initialized with a small random initialization gradient descent has an implicit bias towards solutions of low-rank or small nuclear norm. This is in sharp contrast to Neural Tangent Kernel (NTK)-based theory for low-rank matrix recovery (see [23, Section 4.2]) which will not approximately recover the ground truth matrix XXT due to the larger scale of initialization required when using that technique.
As discussed in Section 2, the restricted isometry property holds with high probability for a sample complexity m ≳ nr2⋆κ8 for Gaussian measurement matrices. Up to constants, this sample complexity is optimal in n, while it is sub-optimal in r⋆ and κ compared to approaches based on nuclear-norm minimization (see, e.g., [37]). While there is numerical evidence that the true scaling of m in r⋆ should also be linear in the non-convex case [41], we note that the optimal dependence of the sample complexity on r⋆ is a major open problem in the field, as the sample complexities in all theoretical results for non-convex approaches in the literature scale at least quadratically in r⋆. Interpretation: Recall from Section 1 that our convergence analysis can be divided into three phases: the spectral phase, the saddle avoidance phase, and the local refinement phase. As it will become clear from the proofs in the supplementary when r ≥ 2r⋆ the bound on the number of iterations can
be decomposed as follows
t̂ ≲ 1 µσmin (X)2 ⎡⎢⎢⎢⎢⎣ ln(2κ2 √ n min{r;n} )
´¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¸¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¶ Phase I: spectral/alignment phase
+ ln(σmin (X) α
) ´¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¸¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¶
Phase II: saddle avoidance phase
+ ln(max{1; κr⋆ min{r;n} − r⋆ } ∥X∥ α )
´¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¸¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¶ Phase III: local refinement phase
⎤⎥⎥⎥⎥⎦ .
(6)
First, we note that the duration of all three phases scales inversely with σmin (X)2. This is due to the fact that in all three phases the dynamics associated the smallest singular value of X is the slowest one and hence needs the most time to complete.
In the spectral phase, the eigenvectors corresponding to the leading r⋆ eigenvalues of UtUTt become aligned with the eigenvectors corresponding to the leading r⋆ eigenvalues of A∗A (XXT ). We observe in (6) that in the spectral phase increasing r, i.e. the amount of parameters, decreases the number of iterations in this phase. As we will explain in the supplementary, the reason is that increasing r decreases the angle between the column space of the initialization U0 and the span of the eigenvectors corresponding to the leading r⋆ eigenvalues of A∗A (XXT ) used in spectral initialization. As a consequence, gradient descent needs fewer iterations to align these two subspaces. In the saddle avoidance phase (Phase II), σr⋆ (Ut) , the r⋆th largest singular value of Ut, grows geometrically until it is on the order of σmin (X). Hence, this duration depends on the ratio between the σmin (X) and the the scale of initialization α. This is clearly reflected in the upper bound on the number of needed iterations in equation (6).
In Phase III, the local refinement phase, the matrix UtUTt converges towards XX T . In particular, at iteration t̂ the test error obeys (5). We observe that a smaller α allows for a smaller test error in (5) but per (6) this higher accuracy is achieved with a modest increase in the required iterations.

heading: 4 A glimpse of our analysis
text: In our proofs, we show that the trajectory of the gradient descent iterations can be approximately decomposed into three phases: (I) a spectral or alignment phase where we show that gradient descent from random initialization behaves akin to spectral initialization allowing us to show that at the end of this phase the column spaces of the iterates Ut and the ground truth matrix X are sufficiently aligned, (II) a saddle avoidance phase, where we show that the trajectory of the gradient iterates move away from certain degenerate saddle points , and (III) a refinement phase, where the product of the gradient descent iterates UtUTt converges quickly to the underlying low-rank matrix XX
T . The latter result holds up to a small error that is commensurate with the scale of the initialization and tends to zero as the scale of the initialization goes to zero.
To formalize the above, we use L and Lt to denote the subspaces spanned by the eigenvectors corresponding to the r⋆ largest eigenvalues of the matrix A∗A (XXT ), and UtUTt , respectively. Moreover, for a subspace L of dimension r⋆ we use VL ∈ Rn×r⋆ to denote an orthonormal matrix whose columns span the subspace L. Note that L is the subspace, which is obtained by commonly used spectral methods. Using this notation, Figure 2a depicts the three phases described above.
In the spectral phase we will prove that we can approximate the iterate Ut by
Ut ≈ (Id + µA∗A (XXT ))t´¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¸¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¶=∶Zt U0 = Zt1U0 ∶= Ũt. (7)
We note that the matrix Z1 = Id + µA∗A (XXT ) is the basis for the commonly used spectral initialization, where typically a factorization of the rank r∗ approximation of this matrix is used as the initialization [6, 5, 42]. Therefore, the approximation (7) suggests that gradient descent iterates modulo the normalization are akin to running power method on Z1. Hence, we expect that at the end of the spectral phase the subspace Lt to be closely aligned with the subspace L, i.e. the subspace obtained by commonly used spectral initialization techniques. In particular, this also implies that Lt is also aligned with the subspace X . Figure 2b clearly illustrates that the first few iterations of gradient descent behave essentially identical to the power method, confirming our intuition.
The description of the second and third phase is more elaborate and technical in nature and we defer to the supplementary for a more detailed and intuitive explanation. However, to give a brief
description, in these two phases we will decompose the iterates Ut into the sum of two matrices, a "signal" matrix of rank r⋆ and a "noise" matrix of rank at most r − r⋆. In Phase (II) we will prove that the smallest singular value of the signal term, which is approximately the same as σr⋆ (Ut) grows, whereas the spectral norm of the noise matrix grows at a much slower rate. We will also show that in this phase the columns of the signal term stay approximately aligned with the span of the matrix X . As soon as the smallest singular value of the signal term of Ut is approximately at the same order than the smallest singular value of X we enter Phase (III). In that Phase we provide a local convergence argument, which shows that the signal term of Ut converges towards X (up to a rotation), whereas the noise term stays small.

heading: 5 Numerical experiments
text: In this section, we perform several numerical experiments to corroborate our theoretical results. Experimental setup. For the experiments we set the ground truth matrix X ∈ Rn×r⋆ to be a random orthogonal matrix with n = 200 and r⋆ = 5. Moreover, we use m = 10nr⋆ = 50n random Gaussian measurements. The initialization U is chosen as in Theorem 3.3 and we use a step size of µ = 1/4 which is consistent with these theorems. We note that while all experimental depictions are based on a single trial, in line with the NeurIPS guidelines we have drawn these curves multiple times (not depicted) and the behavior of the plots do not change.
Depiction of the three phases and the role of overparameterization. In our first experiment, we want to examine how increasing the number of parameters via increasing the number of the columns r of the matrix Ut ∈ Rn×r, affects the spectral phase. To this aim we set the scale of initialization to α = 1/ (70n2). Let L denote the subspace spanned by the eigenvectors corresponding to the leading r⋆ singular values of A∗A(XXT ) and Lt denotes the subspace spanned by the left-singular vectors corresponding to the largest r⋆ singular values of Ut.
Spectral phase and alignment under different levels of overparameterization. First, we examine how the angle between these two subspaces (i.e. ∥V TL⊥VLt∥) changes in the first few iterations. We depict the results for different r in Figure 3a. We see that in the first few iterations, i.e. in the spectral phase, this angle converges towards zero. This confirms the main conclusion of this paper that the first few iterations of gradient descent from small random initialization indeed behaves akin to running power method for spectral initialization. This experiment also shows that changing the number of columns r of Ut has an interesting effect on the spectral phase. In particular, increasing r allows the gradient descent algorithm to learn the subspace L with fewer iterations, i.e. ∥V TL⊥VLt∥ becomes small with fewer iterations. This is in accordance with our theory for r⋆ ≤ r ≤ n (see, for example, the first summand on the right-hand side of equation (6)), where we show that more overparameterization allows gradient descent to leave the spectral phase earlier. Interestingly, this improvement continues to hold even when increasing r beyond n allowing for even faster convergence of ∥V TL⊥VLt∥. This holds even though in this case the rank of U0 is still not larger than n. One potential explanation for this phenomenon might be that for such a choice of r the matrix U0 is better conditioned. Growth of σr⋆ (Ut) and saddle avoidance. In Figure 3b we depict how σr⋆ (Ut) grows during the training for different choices of r. We see that the curves look similar, although for smaller r the growth phase sets in at a slightly later time. This is due to the fact that for smaller r, as we have seen in Figure 3a, Phase I, the spectral phase takes longer to complete.
Evolution of the test error and the refinement phase. Similarly, in Figure 3c we depict how the (normalized) test error ∥UtUTt −XXT ∥F /∥XXT ∥F evolves during the training for different choices of r. We observe that for smaller r the third phase sets in slightly later. Again, this is due to the fact that for smaller r the spectral phase takes slightly longer to complete (see inequality (6)).
Test error under different scales of initialization. In the next experiment, we focus on understanding how the scale of initialization α affects the generalization error ∥UtUTt −XXT ∥2F .
For that, we set r = 180 and run gradient descent with for different choices of α. We stop as soon as the training error becomes small (f (Ut) ≤ 0.5 ⋅ 10−9). We depict the results in Figure 4. We see that the test error decreases as α decreases. In particular, this figure indicates that the test error depends polynomially on the scale of initialization α. This is in line with our theory, where we also show that the test error decreases at least with the rate α21/16 (see inequality (5) in Theorem 3.3).
Change of test and train error during training. In the next experiment, we set r = 180 and examine how the test error ∥UtUTt −XXT ∥2F and the train error f (Ut) changes throughout training and, in particular, how this depends on the scale of initialization. To this aim, we run gradient descent with 4 ⋅ 105 iterations. We see that for a small scale of initialization, α = 10−3, which is the scenario studied in this paper, both test error and train error decrease throughout training.
We observe that in the beginning, as described our theory, both test and train error decrease rapidly. After that the decrease of both test and train error slows down significantly. Moreover, the train error converges towards zero, in contrast to the test error. One reason for the slow convergence in this phase might be that Ut is ill-conditioned in the sense that σr⋆ (UtWt) is much larger than ∥UtWt,⊥∥. It is an interesting future research direction to extend our theory to this part of the training. For large scale of initialization α = 0.5, we observe a very different behaviour. We see that the train error converges with linear rate until machine precision is reached. However, the test error barely changes throughout the training. This scale of initialization corresponds to the lazy training regime [34], where the parameters stay close to the initialization during the training. We depict the results in Figure 5. Number of iterations until convergence: In the last experiment, we set α = 10−3 and examine how many iterations are needed until the test error ∥UtUTt −XXT ∥2F falls below a certain threshold of 10−4 for different values of r obeying 5 ≤ r ≤ 30. For each choice of r we run the experiment ten times and then average the number of iterations for each choice of r. The results are depicted in Figure 6. We observe that increasing the number of columns r from 5 to 10, i.e., a small amount of overparameterization, decreases the number of iterations needed. After that the number of iterations
needed stays roughly constant. This observation is in line with Figure 3, where we have seen that overparameterization leads to fast decrease of the test error in the spectral phase (with diminishing speedup as r becomes larger and larger) without affecting the other two phases.

heading: 6 Conclusion and Broader Impact
text: In this paper we focused on demystifying the role of initialization when training overparameterized models by showing that small random initialization followed by a few iterations of gradient descent behaves akin to popular spectral methods. We also show that this implicit spectral bias from small random initialization, which is provably more prominent for overparameterized models, also puts the gradient descent iterations on a particular trajectory towards solutions that are not only globally optimal but also generalize well.
We think that our results give rise to a number of interesting future research directions. For example, one could extend our results to scenarios where the measurement matrices are more structured such as in matrix completion [43] or in blind deconvolution [44]. Moreover, while our main results, e.g. Theorem 3.3 do require early stopping, our simulations (e.g. Figure 5a) indicate that early stopping is not needed. It would be interesting to examine whether we can remove the early stopping requirement. It is also an interesting future avenue to examine whether the quadratic dependence of the sample complexity m on r⋆ in our results is really needed. Moreover, while in this paper our main focus was on low-rank matrix reconstruction, we believe that our analysis holds more generally for a variety of contemporary overparameterized machine learning and signal estimation tasks including neural network training. This is a tantalizing future research direction.
Despite being theoretical/foundational in nature our results have potential for broader practical impact. In particular, low rank reconstruction problems are an important component of many recommender engines and our insights may guide better algorithm and systems designs for such engines. More broadly, training overparameterized models using stochastic GD starting from small random initialization is the work-horse of modern learning ncluding deep learning and our insights may in the long term help enable more efficient/reliable training with a smaller carbon footprint and improved test accuracy. As with other technologies such insights may potentially also be used nefariously.
Acknowledgments and Disclosure of Funding
M.S. is supported by the Packard Fellowship in Science and Engineering, a Sloan Research Fellowship in Mathematics, an NSF-CAREER under award #1846369, the Air Force Office of Scientific Research Young Investigator Program (AFOSR-YIP) under award #FA9550-18-1-0078, DARPA Learning with Less Labels (LwLL) and FastNICS programs, and NSF-CIF awards #1813877 and #2008443.

Review summaries:
The authors study the nonconvex approach to the positive semidefinite (PSD) matrix sensing problem, a much-studied model problem for tractable nonconvex optimization in the literature. Specifically, they consider the "overparameterized" version of this problem: there is an unknown ground truth factor matrix $X \in \mathbb{R}^{n \times r_{\star}}$ which one observes through linear measurements $\mathcal{A}(X X^T) \in \mathbb{R}^m$; one seeks to find $X$ by minimizing the nonconvex objective $U \mapsto (1/4) \\| \mathcal{A} ( UU^T - XX^T) \\|^2$ ($\ell^2$ norm) using randomly-initialized gradient descent from initialization $\alpha U \in \mathbb{R}^{n \times r}$ with $r \geq r_{\star}$,  where $\alpha > 0$ is a small scale parameter (e.g. inverse polynomial in $n$) and $U$ is an i.i.d. gaussian matrix with entries of suitable variance. The authors analyze this problem in the setting where the measurement operator $\mathcal{A}$ has the RIP (e.g. enough gaussian measurements) and the general setting of $r \geq r_{\star}$; they prove that gradient descent recovers the true factor matrix $X$ up to symmetry to any desired accuracy (measured in frobenius norm), where the precision is improved by decreasing the magnitude of $\alpha$ (at the cost of a logarithmic gain in the number of gradient descent iterations required) and with constants that depend on the condition number of $X$ (Theorem 3.3). The approach to the proof is claimed as novel, with a connection made between overparameterization in this model and the common spectral initialization schemes in use, and synthetic experiments are presented that support the authors' separation of their approach to the convergence proof into distinct phases, and the predictions of their main theorem.   

This paper takes an initial step to explain why small random initialization followed by a few iterations gradient descent on non-convex optimization can often converge to a global minimum point with also good generalization ability. Specifically, this paper focuses on low-rank matrix recovery, a highly non-convex optimization problem with spurious local optimal points, and shows that: (1) This method has similar behavior to the popular spectral methods with three phases—a spectral or alignment phase, a saddle avoidance/refinement phase, and a local refinement phase. (2) Larger parameter space tends to have a larger spectral bias which helps convergence to a global optimal point with also smaller generalization error. 

This paper studies the role of small random initialization followed by gradient descent in the setting of low-rank matrix sensing. Authors show that the trajectory of the gradient descent from small random initialization behaves similarly as spectral initialization in the initial stage, and then moves from degenerate saddle points and converges to the ground truth. Numerical experiments are provided to validate the theory.