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

1 Introduction . 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 can not 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 can not 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 . 2 Low-rank matrix recovery via non-convex optimization . 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 ) .

Summary: 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.