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

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 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 . 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 Section 4, 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 . 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) . Then specialized truncation or saddle escaping algorithms such as trust region, cubic regularization , or noisy (stochastic) gradient-based methods 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 on problems of the form (II), these results cannot rule out exponentially slow cases in the worst-case . Indeed, it has been argued that in general a more granular analysis of the trajectory of gradient descent beyond the landscape may be necessary . 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 .

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) . 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) . However, these existing analyses contain many idealized and non-realistic 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 .

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. ).

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 . Here, our goal is to recover a low-rank matrix of the form $XX^{T}$ 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 .

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 $XX^{T}$. In the third phase, the local refinement phase, we show that the iterates approximately converge towards the underlying low-rank matrix $XX^{T}$ 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 , 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. For a detailed comparison we refer to Section 3.

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

with $r\geq r_{\star}$. More compactly one can rewrite the optimization problem above in the form

In order to solve the minimization problem (2) we run gradient descent iterations starting from (often small) random initialization. More specifically,

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\geq m$. In this case, there are infinitely many $\bar{U}$ such that $f(\bar{U})=0$, but $\|\bar{U}\bar{U}^{T}-XX^{T}\|_{F}$ is arbitrarily large (see, e.g., [39, Proposition 1]). That is, even if gradient descent converges to a global optimum, i.e. $f\left(\bar{U}\right)=0$, it is a priori not clear whether it has found the low-rank solution $XX^{T}$ (see also Figure 7).

Main results

In this section, we present our main results. Stating these results requires a couple of simple definitions. The first definition concerns the measurement operator $\mathcal{A}$.

We note that for a Gaussian measurement operator $\mathcal{A}$ By that, we mean that all the entries of the (symmetric) measurement matrices $\left\{A_{i}\right\}_{i=1}^{m}$ are drawn i.i.d. with distribution $\mathcal{N}\left(0,1\right)$ on the off-diagonal and distribution $\mathcal{N}\left(0,1/\sqrt{2}\right)$ on the diagonal., RIP of rank $r$ and constant $\delta>0$ holds with high probability, if the number of observations satisfies $m\gtrsim nr/\delta^{2}$ .

The second definition concerns the condition number of the factor $X$.

where $\sigma_{r_{\star}}\left(X\right)$ denotes $r_{\star}$-th largest singular value of $X$.

With these definitions in place we are now ready to state our main results.

We begin by stating our first main result.

Under the assumption that $r\geq 2r_{\star}$ and that the scale of initialization fulfills

Assume that $r_{\star}<r<2r_{\star}$ and that the scale of initialization fulfills

Note that the test error $\|U_{\hat{t}}U_{\hat{t}}^{T}-XX^{T}\|_{F}^{2}$ can be made arbitrarily small by choosing the scale of initialization $\alpha$ small enough. In particular, the dependence of the test error on $\alpha$ is polynomial and the dependence of the number of iterations on $\alpha$ is logarithmic, which means that reducing the test error by scaling down $\alpha$ introduces only modest additional computational cost. Hence, as long as the rank at most $2r_{\star}+1$ RIP with constant $\delta\leq c\kappa^{-4}{r_{\star}}^{-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\gg m$ and the optimization problem has many global optima many of which do not obey $UU^{T}\approx XX^{T}$. 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 $XX^{T}$ 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\gtrsim nr_{\star}^{2}\kappa^{8}$ for Gaussian measurement matrices. Up to constants, this sample complexity is optimal in $n$, while it is sub-optimal in $r_{\star}$ and $\kappa$ compared to approaches based on nuclear-norm minimization (see, e.g., ). While there is numerical evidence that the true scaling of $m$ in $r_{\star}$ should also be linear in the non-convex case , we note that the optimal dependence of the sample complexity on $r_{\star}$ 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_{\star}$.

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, when $r\geq 2r_{\star}$ the bound on the needed number of iterations (see inequality (5)) can be decomposed as follows

First, we note that the duration of all three phases scales inversely with $\sigma_{\min}\left(X\right)^{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_{\star}$ eigenvalues of $U_{t}U_{t}^{T}$ become aligned with the eigenvectors corresponding to the leading $r_{\star}$ eigenvalues of $\mathcal{A}^{*}\mathcal{A}\left(XX^{T}\right)$. We observe in (9) 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 Section 5.2, the reason is that increasing $r$ decreases the angle between the column space of the initialization $U_{0}$ and the span of the eigenvectors corresponding to the leading $r_{\star}$ eigenvalues of $\mathcal{A}^{*}\mathcal{A}\left(XX^{T}\right)$ used in spectral initialization. As a consequence, gradient descent needs fewer iterations to align these two subspaces.

In the saddle avoidance phase (Phase II), $\sigma_{r_{\star}}\left(U_{t}\right)$ , the $r_{\star}$th largest singular value of $U_{t}$, grows geometrically until it is on the order of $\sigma_{\min}\left(X\right)$. Hence, this duration depends on the ratio between the $\sigma_{\min}\left(X\right)$ and the the scale of initialization $\alpha$. This is clearly reflected in the upper bound on the number of needed iterations in equation (9).

In Phase III, the local refinement phase, the matrix $U_{t}U_{t}^{T}$ converges towards $XX^{T}$. In particular, at iteration $\hat{t}$ the test error obeys (6). We observe that a smaller $\alpha$ allows for a smaller test error in (6) but per (9) this higher accuracy is achieved with a modest increase in the required iterations.

The following result deals with the scenario $r=r_{\star}$, that is, the iterates $U_{t}$ have as many parameters as the ground truth matrix $X$.

for some $0<\varepsilon<1$. Then with probability at least $1-C\varepsilon+\exp\left(-cr_{\star}\right)$ after

Here $C,c>0$ are fixed numerical constants.

Note that by choosing $\alpha$ small enough we can make the test error in (10) arbitrarily small. In particular, this means that then well-known local convergence results can be applied showing that $U_{t}U_{t}^{T}$ converges linearly to $XX^{T}$ (see, e.g., ) . Thus, this result implies that if the measurement operator fulfills the restricted isometry property, gradient descent with small, random initialization will converge to the ground truth matrix $X$ in polynomial time. It is known that under the RIP assumption the loss landscape is benign in the sense that there are no local optima that are not global and all saddles have a direction of negative curvature. However, such results do not imply that vanilla gradient descent converges quickly (i.e. in polynomial time) to a global optimum, as gradient descent may take exponential time to escape from saddle points.

To the best of our knowledge, this is the first in the non-overparameterized setting $r=r_{\star}$ result which shows the convergence of vanilla gradient descent to the ground truth from a random initialization using only the restricted isometry property in polynomial time. The only other paperin the low-rank matrix recovery literature, which shows fast convergence of vanilla gradient descent to the ground truth from a random initialization, is . In this work, the problem of phase retrieval has been studied, which can be formulated as a low-rank matrix recovery problem with $r=r_{\star}=1$. The paper shows that gradient descent converges from a random initialization to the ground truth with a near-optimal number of iterations. However, the proof in this paper leverages the rotation-invariance of the Gaussian measurements vectors via carefully constructed auxiliary sequences. In contrast, Theorem 3.4 above relies only on the restricted isometry property and no further assumptions on $\mathcal{A}$ are needed.

3 Special case: r=n𝑟𝑛r=n with orthonormal initialization

Here $c>0$ is a fixed numerical constant.

Note that this result improves over [38, Theorem 4.1] in several aspects. First of all, in it is assumed that the measurement operator $\mathcal{A}$ has the rank-$4r$ restricted isometry property with constant $\delta\lesssim\kappa^{-6}r_{\star}^{-1/2}\log^{-2}\frac{n}{\alpha}$. In particular, this suggests that this result cannot handle the scenario that the scale of initialization $\alpha$ becomes arbitrarily small, as this would also require that the restricted isometry constant $\delta$ becomes arbitrarily small as well. This in turn would require an arbitarily large sample size. Moreover, requires a step size of at most $\mu\lesssim\kappa^{-6}r_{\star}^{-1/2}\log^{2}n\|X\|^{-2}$, whereas the above theorem only needs the weaker assumption $\mu\lesssim\kappa^{-2}\|X\|^{-2}$. These improvements aside the main difference between our result and this prior work is that we can handle any $r$ by formalizing an intriguing connection between small random initialization and spectral learning.

Related work

Global convergence guarantees for nonconvex low-rank matrix recovery: As mentioned earlier in Section 1, there is a large body of work on developing global convergence guarantees for nonconvex problems. In the context of low-rank matrix recovery, several papers have demonstrated that low-rank reconstruction problems in a variety of domains can be solved via nonconvex gradient descent starting from spectral initialization. More precisely, this has been shown for phase retrieval , matrix sensing , blind deconvolution , and matrix completion . However, in practice often random initialization is used in lieu of specialized spectral initialization techniques. To remedy this issue, more recent literature , focusses on studying the loss landscape of such problems. These papers show that despite their non-convexity under certain assumptions these loss landscapes are benign in the sense that there are no spurious local minima, (i.e. all minimizers are global minima) and saddles points have a strict direction of negative curvature (a.k.a. strict saddle) . Then specialized truncation or saddle escaping algorithms such as trust region, cubic regularization or noisy (stochastic) gradient-based methods are deployed to provably find a global optimum. These papers however do not directly develop global convergence for gradient descent (without any additional modification) from a random initialization. For differentiable losses eventual convergence to local minimizers is known from random initialization but these results do not provide convergence rates and only guarantee eventual convergence. Indeed, gradient descent may converge exponentially slowly in the worst-case . In contrast to the above literature our result in Theorem 3.4 (in the case of $r=r_{\star}$) shows that gradient descent from a small random initialization converges rather quickly to the global optima. As mentioned earlier, we are able to establish this result by demonstrating that in the initial phase gradient descent iterates are intimately connected to the spectral initialization techniques discussed above. Furthermore, the above spectral initialization followed by local convergence or landscape analysis techniques cannot be directly applied in in the overparameterized case ($r>r_{\star}$) whereas our analysis works regardless of model overparameterization.

We would like to mention that even more recently the paper proves the convergence of gradient descent starting from a random initialization for low-rank recovery problems via an interesting leave-one-out analysis. To the best of our knowledge, this is the only existing result, which provides convergence guarantees for vanilla gradient descent from random initialization for low-rank matrix recovery problems in the non-overparameterized setting $r=r_{\star}$. However, the leave-one-out analysis heavily relies on the independence and the rotation invariance of the measurements. Also similar to the above this analysis does not seem to easily lend to generalization in the overparameterized regime. In contrast, our proof techniques rely on standard restricted isometry assumptions without requiring the independence of the measurements and does provide generalization guarantees with model overparameterization ($r>r_{\star}$). Moreover, in it has been shown that Riemannian gradient descent converges with nearly linear rate to the true solution from a random initialization in the population loss scenario.

Overparameterization in low-rank matrix recovery: In the influential work it has been conjectured and in the special case that the measurement matrices commute proven that gradient descent on overparameterized matrix factorization converges to the solution with the minimal nuclear norm. This phenomenon is now often referred to as implicit regularization. In , evidence is provided that adding depth even increases the tendency of gradient descent to converge to low-rank solution. In it has been shown that there are certain scenarios where the conjecture in does not hold. In theoretical and empirical evidence has been provided that gradient flow with infinitesimal initialization is equivalent to a certain rank-minimization heuristic.

In this paper, we shed further light on the implicit regularization of gradient descent. In particular, we provide a precise analysis of the initial stage and relate it to the power method and our analysis explains how overparameterization is beneficial in the initial stages. Closest to our work is the paper , which studies a special case of the problem analysed in this paper. More precisely, this paper considers the special case $r=n$ with orthonormal initialization. We also applied our theory to this exact same setting, see Section 3.3, where we include a detailed comparison for this special scenario. Most importantly our theory is able to handel the case $\alpha\rightarrow 0$, which the result in seems not to be able to. Moreover, analysing the full range of possible choise of $r$ requires a careful analysis of the spectral phase, which is one key novely of this paper compared to .

In it has been shown that in certain scenarios, where the measurement matrices $A_{i}$ are positive semidefinite (PSD), the equation $y=\mathcal{A}\left(UU^{T}\right)$ has a unique low-rank solution. This means that in these scenarios the PSD constraint by itself might lead to a low-rank matrix recovery, which makes implicit regularization by gradient descent meaningless in this setting. However, note that these results not apply to the scenario studied in this paper, as we assume the measurement matrices $A_{i}$ to be Gaussian, which, in particular, means that they are not positive semidefinite. In particular, in our setting it can be shown that there are infinitely many solutions to the equation $y=\mathcal{A}\left(UU^{T}\right)$ with arbitrarily large test error .

Gradient-based generalization guarantees for overparameterized tensors and neural networks: A recent line of work is concerned with connecting the analysis of neural network training with the so-called neural tangent kernel (NTK) . The key idea is that for a large enough initialization, it suffices to consider a linearization of the neural network around the origin. This allows connecting the analysis of neural networks with the well-studied theory of kernel methods. This is also sometimes referred to as lazy training, as with such an initialization the parameters of the neural networks stay close to the parameters at initialization. However, there is a line of work, which suggests that NTK-analysis might not be sufficient to completely explain the success of neural networks in practice. The paper provides empirical evidence that by choosing a smaller initialization the test error of the neural network decreases. A similar performance gap between the performance of the NTK and neural networks has been observed in , where it has been shown that the performance gap is larger if the covariance matrix is isotropic.

There is also a line of work , which is concerned with the mean-field analysis of neural networks. The insight is that for sufficiently large width the training dynamics of the neural network can be coupled with the evolution of a probability distribution described by a PDE. These papers use a smaller initialization than in the NTK-regime and, hence, the parameters can move away from the initialization. However, these results do not provide explicit convergence rates and require an unrealistically large width of the neural network.

For the problem of tensor decomposition it has also been shown that gradient descent with small initialization is able to leverage low-rank structure . This is relevant to neural network analysis, since in a relationship between tensor decomposition and training neural networks has been established. In it has been shown that neural networks with ReLU function and trained by SGD can outperform any kernel method. One crucial element in their analysis is that the early stage of the training is connected with learning the first and second moment of the data.

While in this paper we do not study overparameterized tensor or neural network models we note that the NTK-theory can also be applied to low-rank matrix recovery (see [23, Section 4.2]). This means that if the scale of initialization is chosen large enough and the number of parameters is larger than the number of measurements, i.e. $nr\gtrsim m$, then gradient descent will converge linearly to a global minimizer with zero loss. However, since for this approach the parameters will stay close to the initialization, this approach will not recover the ground truth matrix $XX^{T}$. Hence, an NTK analysis will not yield good generalization. In contrast in this paper we have seen that choosing a small initialization is a remedy for low-rank matrix recovery. So in this sense our result can be viewed as going beyond the lazy training in NTK theory. In fact we believe that similar analysis to the one developed in this paper for low-rank recovery can be used to analyze a much broader class of overparameterized models including the analysis of neural networks. We defer this to a future paper.

Linear neural networks: In the convergence of gradient flow and gradient descent is studied for (deep) linear neural networks of the form

Overview and key ideas of the proof

In this section, we briefly discuss the key ideas and techniques in our proof. We begin by discussing a simple decomposition, which is utilized throughout our proofs. Next, in Sections 5.2 and 5.3 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 $U_{t}$ 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 $U_{t}U_{t}^{T}$ 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. Figure 2 depicts these three phases.

Let us remark that the proof in the related work decomposes the convergence analysis into two phases, which roughly correspond to Phase II and Phase III in our proof. However, the proof details are quite different since we use a different decomposition into signal and noise term, see Section 5.1.

This decomposition has the following two simple properties, which will be useful throughout our proofs.

The column space of the noise term is orthogonal to the column span of $X$, i.e. $V_{X}^{T}U_{t}W_{t,\perp}=0$.

When $V_{X}^{T}U_{t}$ is full rank, then the signal term has rank $r_{\star}$ and the noise term has rank at most $r-r_{\star}$.

The first statement follows directly from the observation $V_{X}^{T}U_{t}W_{t,\perp}W_{t,\perp}^{T}=V_{X}^{T}U_{t}\left(\text{Id}-W_{t}W_{t}^{T}\right)=0$. The second statement is a direct consequence of the definition of $W_{t}$. ∎

We would like to note that decomposing $U_{t}$ into two terms has appeared in prior work such as as well as in earlier work in the compressive sensing literature. However, uses a different decomposition. A key advantage of our decomposition is that it only depends on $U_{t}$ and $X$, whereas the decomposition in depends on all previous iterates $U_{0},U_{1},\ldots,U_{t-1}$.

2 The spectral/alignment phase

In this section we turn our attention to giving an overview of the key ideas and proofs of the spectral/alignment phase. More specifically, we will argue that in the first few iterations gradient descent implicitly performs a form of spectral initialization. By that, we mean that after the first few iterations the column span of the signal term $U_{t}W_{t}W_{t}^{T}$ is aligned with the column span of $X$ and that $\|U_{t}W_{t,\perp}\|$ is relatively small compared to $\sigma_{\min}\left(U_{t}W_{t}\right)$, meaning that the signal term dominates the noise term.

We now provide the main intuition behind the analysis in our spectral/alignment phase. Our starting point is the observation that for the gradient at the initialization $U_{0}=\alpha U$ it holds that

In particular, we observe that for $\alpha>0$ sufficiently small the second term is negligible. Hence, we have that

In the first few iterations (i.e. small $t$) we expect the matrix $U_{t}$ to be small and continue to scale commensurately with $\alpha$ and we expect that a similar approximation holds for the first iterations. Hence, for $\alpha$ sufficiently small we can approximate $U_{t}$ by

Figure 3 clearly illustrates that the first few iterations of gradient descent behave essentially identical to (11) confirming our intuition and proofs.

To be more precise about the aforementioned alignment with $X$, let the singular value decomposition of $\mathcal{A}^{*}\mathcal{A}\left(XX^{T}\right)$ be given by $\mathcal{A}^{*}\mathcal{A}\left(XX^{T}\right)=\sum_{i=1}^{n}\lambda_{i}v_{i}v_{i}^{T}$. It follows that

It is well-known that when the operator $\mathcal{A}$ obeys the restricted isometry property we have

that $\lambda_{r_{\star}}\left(Z_{t}\right)/\lambda_{r_{\star}+1}\left(Z_{t}\right)$ grows exponentially. In particular, this means that

Since $U_{0}$ is a random Gaussian matrix, for an appropriate choice of $t$, we will be able to show that the matrix $U_{t}$ has the following two properties with high probability, where $L=\text{span}\left\{v_{1};\ldots;v_{r_{\star}}\right\}$ and $L_{t}$ is the projection of $U_{t}$ onto its best rank-$r_{\star}$ approximation:

There is a sufficiently large gap between $\sigma_{r_{\star}}\left(U_{t}\right)$ and $\sigma_{r_{\star}+1}\left(U_{t}\right)$, i.e., $\frac{\sigma_{r_{\star}}\left(U_{t}\right)}{\sigma_{r_{\star}+1}\left(U_{t}\right)}\geq\Delta>1$, where $\Delta$ is an appropriately chosen constant.

We have that $\|V_{L^{\perp}}^{T}V_{L_{t}}\|$ is small. Since the column space of $\mathcal{A}^{*}\mathcal{A}\left(XX^{T}\right)$ is aligned with the column space of $X$, this also implies that $\|V_{X^{\perp}}^{T}V_{L_{t}}\|$ is small.

This confirms that in the first few iterations, gradient descent indeed implicitly performs akin to spectral initialization with the column space of $U_{t}$ aligned with the column space of $X$. However, this does not yet fully complete our analysis for the spectral/alignment phase, since critical to the analysis of second phase we need certain properties to hold for the signal and noise terms $U_{t}W_{t}$ and $U_{t}W_{t,\perp}$ (see Section 5.1) rather than the singular value decomposition of $U_{t}$. However, using the properties of the SVD of $U_{t}$, which are listed above, we will establish the following properties of $U_{t}W_{t}$ and $U_{t}W_{t,\perp}$ .

The column space of $U_{t}W_{t}$ is aligned with the column space of $X$: $\|V_{X^{\perp}}^{T}V_{U_{t}W_{t}}\|\leq c\kappa^{-2}$.

The spectral norm of the noise term is not too large compared to the minimum singular value of the signal term, i.e., $2\sigma_{\min}\left(U_{t}W_{t}\right)\geq\|U_{t}W_{t,\perp}\|$.

The spectral norm of the noise term is bounded from above in the sense that i.e., $\|U_{t}W_{t,\perp}\|\ll\sigma_{\min}\left(X\right)$.

The spectral norm of $U_{t}$ is bounded, i.e., $\|U_{t}\|\leq 3\|X\|$.

3 The saddle avoidance phase and the refinement phase

In the next two phases, we will show that the signal term $U_{t}W_{t}W_{t}^{T}U_{t}^{T}$ converges towards $XX^{T}$, whereas the spectral norm of the noise term, i.e. $\|U_{t}W_{t,\perp}\|$, stays small. For that, we show that throughout this process the columns of the matrices $X$ and $U_{t}W_{t}$ stay approximately aligned, i.e., the angle $\|V_{X^{\perp}}^{T}V_{U_{t}W_{t}}\|$ stays small. This latter property also ensures that after the spectral phase the iterates are not too close to well known saddle points of the optimization landscape (it is known that this problem may have degenerate saddle points at a point $\bar{U}$ obeying rank$(\bar{U})<r_{*}$ ). See Figure 4 for a depiction of the gradient flows of the landscape when $r_{*}=1$.

Next we sketch the proofs of Phase II and Phase III in more detail.

Phase II: In this phase, we will show that the minimal singular value of the signal term, $\sigma_{\min}\left(U_{t}W_{t}\right)$ grows exponentially, until it holds that $\sigma_{\min}\left(U_{t}W_{t}\right)\geq\frac{\sigma_{\min}\left(X\right)}{\sqrt{10}}$. To this aim, we show that

holds under suitable assumptions (see Lemma 9.1). In order to show that the spectral norm of the noise term $\|U_{t}W_{t,\perp}\|$ grows much slower than $\sigma_{\min}\left(U_{t+1}W_{t+1}\right)$, we establish the inequality

(see Lemma 9.2). The next inequality (see Lemma 9.3) shows that $\|V^{T}_{X^{\perp}}V_{U_{t}W_{t}}\|$ stays sufficiently small

As mentioned above, this implies in particular, that $U_{t}$ stays sufficiently far away from saddle points $\bar{U}$, which are rank-deficient, e.g., $\text{rank}\left(\bar{U}\right)<r_{\star}$.

Phase III: After we have shown that $\sigma_{\min}\left(U_{t}W_{t}\right)\geq\frac{\sigma_{\min}\left(X\right)}{\sqrt{10}}$ holds for some $t$, we enter the local refinement phase. We start by observing that the error $\|XX^{T}-U_{t}U_{t}^{T}\|_{F}$ can be decomposed into two summands, i.e.

(see Lemma B.4). We will bound the second summand by using inequality (13), which is also valid for the third phase. We will show that the first summand decreases at a linear rate. For that, we establish the inequality

Hence, by using inequality (14) we will be able to show that $\|U_{t}U_{t}^{T}-XX^{T}\|_{F}$ is decreasing, as long as the spectral norm of the noise term stays sufficiently small.

Numerical experiments

In this section, we perform several numerical experiments to corroborate our theoretical results.

Spectral phase and alignment under different levels of overparameterization. First, we examine how the angle between these two subspaces (i.e. $\|V_{L^{\perp}}^{T}V_{L_{t}}\|$) changes in the first few iterations. We depict the results for different $r$ in Figure 5(a). 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 $U_{t}$ 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_{L^{\perp}}^{T}V_{L_{t}}\|$ becomes small with fewer iterations. This is in accordance with our theory for $r_{\star}\leq r\leq n$ (see, for example, the first summand on the right-hand side of equation (9)), 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^{T}_{L^{\perp}}V_{L_{t}}\|$. This holds even though in this case the rank of $U_{0}$ is still not larger than $n$. One potential explanation for this phenomenon might be that for such a choice of $r$ the matrix $U_{0}$ is better conditioned.

Growth of $\sigma_{r_{\star}}\left(U_{t}\right)$ and saddle avoidance. In Figure 5(b) we depict how $\sigma_{r_{\star}}\left(U_{t}\right)$ 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 5(a), Phase I, the spectral phase takes longer to complete.

Evolution of the test error and the refinement phase. Similarly, in Figure 5(c) we depict how the (normalized) test error $\|U_{t}U_{t}^{T}-XX^{T}\|_{F}/\|XX^{T}\|_{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 (9)).

Test error under different scales of initialization. In the next experiment, we focus on understanding how the scale of initialization $\alpha$ affects the generalization error $\|U_{t}U_{t}^{T}-XX^{T}\|_{F}^{2}$. For that, we set $r=180$ and run gradient descent with for different choices of $\alpha$. We stop as soon as the training error becomes small ($f\left(U_{t}\right)\leq 0.5\cdot 10^{-9}$). We depict the results in Figure 6. We see that the test error decreases as $\alpha$ decreases. In particular, this figure indicates that the test error depends polynomially on the scale of initialization $\alpha$. This is in line with our theory, where we also show that the test error decreases at least with the rate $\alpha^{21/16}$ (see inequality (6) 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 $\|U_{t}U_{t}^{T}-XX^{T}\|_{F}^{2}$ and the train error $f\left(U_{t}\right)$ changes throughout training and, in particular, how this depends on the scale of initialization. To this aim, we run gradient descent with $4\cdot 10^{5}$ iterations. We see that for a small scale of initialization, $\alpha=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 $U_{t}$ is ill-conditioned in the sense that $\sigma_{r_{\star}}\left(U_{t}W_{t}\right)$ is much larger than $\|U_{t}W_{t,\perp}\|$. It is an interesting future research direction to extend our theory to this part of the training.

For large scale of initialization $\alpha=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 , where the parameters stay close to the initialization during the training. We depict the results in Figure 7.

Number of iterations until convergence: In the last experiment, we set $\alpha=10^{-3}$ and examine how many iterations are needed until the test error $\|U_{t}U_{t}^{T}-XX^{T}\|_{F}^{2}$ falls below a certain threshold of $10^{-4}$ for different values of $r$ obeying $5\leq r\leq 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 8. 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 5, 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.

Preliminaries

Before we are going into the details of the proof, we are collecting some useful definitions.

2 Restricted isometry property and related properties

As discussed in Section 2, we are going to assume that the measurement operator $\mathcal{A}$ satisfies the restricted isometry property. However, as it turns out, the following two slightly weaker properties will suffice for our proof.

The following lemma shows that these two properties are induced by the standard restricted isometry property (Definition 3.1).

Suppose that $\mathcal{A}$ has the restricted isometry property as in (3) for all matrices of rank $r+1$ with constant $\delta_{1}<1$. Then $\mathcal{A}$ has the spectral-to-spectral restricted isometry property of rank $r$ with constant $\sqrt{r}\delta_{1}$.

Suppose that $\mathcal{A}$ has the restricted isometry property as in (3) for all matrices of rank $2$ with constant $\delta_{2}<1$. Then $\mathcal{A}$ has the spectral-to-nuclear restricted isometry property with constant $\delta_{2}$.

From it follows that if $\mathcal{A}$ has the restricted isometry property of rank $r+r^{\prime}$ with constant $\delta_{1}<0$, then it holds for all matrices $Z$ with rank at most $r$ and all matrices $Y$ with rank at most $r^{\prime}$ that

In order to show the second claim consider the eigenvalue decomposition $Z=\sum_{i=1}^{n}\lambda_{i}v_{i}v_{i}^{T}$. We compute that

where in the second inequality we used the spectral-to-spectral restricted isometry property (with $r=1$), which holds due to the first part of this proof. This finishes the proof of the second statement.

where the last line is due to inequality (16). This finishes the proof of Lemma 7.3. ∎

Analysis of the spectral phase

In the following we will provide an analysis of the spectral phase, where the proofs of the technical lemmas are deferred to Appendix A. Our first goal is to show that in the first few iterations $U_{t}$ can be approximated by

The next lemma shows how well $U_{t}$ can be approximated by $\widetilde{U_{t}}$ for $t\leq t^{\star}$. To formulate it, we set $E_{t}=U_{t}-\widetilde{U_{t}}$.

Suppose that $\mathcal{A}$ satisfies the rank-$1$ RIP with constant $\delta_{1}$. For all integers $t$ such that $1\leq t\leq t^{\star}$ it holds that

The next lemma gives a lower bound for $t^{\star}$. In particular, this shows how long the approximation in Lemma 8.1 is valid.

Let $\widetilde{U_{t}}$ be as defined before and consider the eigenvalue decomposition $\mathcal{A}^{*}\mathcal{A}\left(XX^{T}\right)=\sum_{i=1}^{n}\lambda_{i}v_{i}v_{i}^{T}$. Then we have that

and denote by $L$ the subspace spanned by the eigenvectors, which correspond to the largest $r_{\star}$ eigenvalues of the matrix $M=\mathcal{A}^{*}\mathcal{A}\left(XX^{T}\right)$. Note that $L$ is also the subspace spanned by the eigenvectors corresponding to the largest $r_{\star}$ eigenvalues of the matrix $Z_{t}$. Denote by $L_{t}$ the subspace spanned by the left-singular vectors of $U_{t}=Z_{t}U_{0}+E_{t}$, which corresponds to the largest $r_{\star}$ singular values.

Since $Z_{t}$ is computed via a power method we expect for $t$ large enough that $\lambda_{r_{\star}}\left(Z_{t}\right)\gg\lambda_{r_{\star}+1}\left(Z_{t}\right)$. Moreover, if, in addition, $\|E_{t}\|$ is sufficiently small, we expect in this case that the subspace $L$ is aligned with the subspace $L_{t}$. This is made precise by the following lemma.

Then the following three inequalities hold.

Recall that we are interested in bounds for the quantities $\sigma_{r_{\star}}\left(U_{t}W_{t}\right)$, $\|V_{X^{\perp}}^{T}V_{U_{t}W_{t}}\|$, and $\|U_{t}W_{t,\perp}\|$, i.e., properties of the signal and noise term. However, in the lemma above, we have obtained instead bounds for $\sigma_{r_{\star}}\left(U_{t}\right)$, $\|V_{X^{\perp}}^{T}V_{L_{t}}\|$, and $\sigma_{r_{\star}+1}\left(U_{t}\right)$, i.e. for the singular value decomposition of $U_{t}$. However, if $\|V_{X^{\perp}}^{T}V_{L_{t}}\|$ is small, these quantities are closely related to each other, as the next lemma shows.

Assume that $\|V_{X^{\perp}}^{T}V_{L_{t}}\|\leq\frac{1}{8}$ for some $t\geq 1$. Then it holds that

By combining Lemma 8.3 and Lemma 8.4, we obtain the following technical result.

Let $XX^{T}$ be a low-rank matrix of rank $r_{\star}$. Assume that

where $\widetilde{c}_{2}>0$ is a sufficiently small, absolute constant. Then it holds that

In order to utilize Lemma 8.5, we need to insert bounds for the approximation error $\|E_{t}\|$, which we have derived in the Lemmas 8.1 and 8.2. This yields the following lemma.

where $v_{1}$ denotes the eigenvector corresponding to a leading eigenvalue of the matrix $\mathcal{A}^{*}\mathcal{A}\left(XX^{T}\right)$. Assume that the step size satisfies $\mu\leq c_{2}\kappa^{-2}\|X\|^{-2}$. Then after

Here $c_{1},c_{2},c_{3}>0$ are absolute constants only depending on the choice of $c$.

Note that the result above holds for any initlialization $U$. To complete the proof we are going to utilize the fact that $U$ is a random matrix with Gaussian entries. This yields the following lemma, which is the main result of this section.

Assume that the step size satisfies $\mu\leq c_{2}\kappa^{-2}\|X\|^{2}$. Then with probability at least $1-p$, where

The proof of Lemma 8.7 requires the following theorem, which gives a non-asymptotic lower bound for the smallest singular value of a Gaussian matrix.

With this theorem in place we can prove Lemma 8.7.

In order to proceed we are going to distinguish the following two cases.

Case 1: $r\geq 2r_{\star}$ Note that by choosing $\varepsilon>0$ appropriately, we obtain from Theorem 8.8 combined with inequality (41) that with probability at least $1-O\left(\exp\left(-cr\right)\right)$ it holds that

By combining the inequalities (39), (40), and (42) with Lemma 8.6 the claim follows in the case that $r\geq 2r_{\star}$.

Case 2: $r_{\star}\leq r\leq 2r_{\star}$ Similar to the first case, we note that by choosing $\varepsilon>0$ appropriately, we obtain by applying Theorem 8.8 combined with inequality (41) that with probability at least $1-\left(C\varepsilon\right)^{r-r_{\star}+1}-\exp\left(-cr\right)$ it holds that

By combining the inequalities (39), (40), and (43) with Lemma 8.6 the claim follows. ∎

Analysis of the saddle avoidance and refinement phases

Before stating and proving the main result of this section, Theorem 9.6, we will first collect some useful lemmas. Their proofs are deferred to Appendix B.

In Phase II we will show that $\sigma_{\min}\left(U_{t}W_{t}\right)$ grows until it reaches $\sigma_{\min}\left(U_{t}W_{t}\right)\geq\frac{\sigma_{\min}\left(X\right)}{\sqrt{10}}$. For that, we note

where $(a)$ and $(b)$ follow from the definition of $W_{t}$. Hence, in order to show that $\sigma_{\min}\left(U_{t}W_{t}\right)\geq\frac{\sigma_{\min}\left(X\right)}{\sqrt{10}}$ it suffices to show that $\sigma_{\min}\left(V_{X}^{T}U_{t}\right)\geq\frac{\sigma_{\min}\left(X\right)}{\sqrt{10}}$. For that, we will use the next lemma, which shows that $\sigma_{\min}\left(V_{X}^{T}U_{t}\right)$ grows exponentially.

Assume that $\mu\leq c\|X\|^{-2}\kappa^{-2}$, $\|U_{t}\|\leq 3\|X\|$, and that $\|V_{X^{\perp}}^{T}V_{U_{t}W_{t}}\|\leq c\kappa^{-1}$. Moreover, suppose that

Furthermore, assume that $V_{X}^{T}U_{t}$ has full rank. Then it holds that

Here $c>0$ is constant, which is chosen small enough.

The next lemma will allow us to show that the noise term $\|U_{t}W_{t,\perp}\|$ is growing slower than $\sigma_{\min}\left(V_{X}^{T}U_{t+1}\right)$.

Assume that $\mu\leq c\min\left\{\|X\|^{-2};\|\left(\mathcal{A}^{*}\mathcal{A}-\text{Id}\right)\left(XX^{T}-U_{t}U_{t}^{T}\right)\|^{-1}\right\}$ and that $\|U_{t}\|\leq 3\|X\|$. Moreover, suppose that $V_{X}^{T}U_{t+1}W_{t}$ has full rank and that $\|V_{X^{\perp}}^{T}V_{U_{t}W_{t}}\|\leq c\kappa^{-1}$. Then it holds that

Here, $c>0$ is an absolute constant chosen small enough.

The next lemma shows that the angle between the column space of the signal term $U_{t}W_{t}$ and column space of $X$ stays sufficiently small.

Assume that $\|U_{t}W_{t,\perp}\|\leq 2\sigma_{\min}\left(U_{t}W_{t}\right)$ and $\|U_{t}\|\leq 3\|X\|$ holds. Moreover, assume that

where $c>0$ is a small enough absolute constant. Then it holds that

The next lemma will show that we have $\|U_{t}\|\leq 3\|X\|$ for all $t$, a technical assumption which is needed in the above lemmas.

Assume that $\|U_{t}\|\leq 3\|X\|$, $\mu\leq\frac{1}{27\|X\|^{2}}$, and

Then it also holds that $\|U_{t+1}\|\leq 3\|X\|$.

With these lemmas in place, we will be able to show that $\sigma_{\min}\left(U_{t}W_{t}\right)\geq\frac{\sigma_{\min}\left(X\right)}{\sqrt{10}}$ holds after sufficiently many iterations. Hence, we can enter Phase III, the local refinement phase.

The next lemma is concerned with this third phase. It shows that $U_{t}W_{t}W_{t}^{T}U_{t}^{T}$ converges towards $XX^{T}$, when projected onto the column space of $X$. We are going to provide a somewhat more general version of the lemma than what is needed in the proofs of our main results, since it may be of independent interest. For that, let ${\left|\kern-1.07639pt\left|\kern-1.07639pt\left|\cdot\right|\kern-1.07639pt\right|\kern-1.07639pt\right|}$ be a matrix norm, which satisfies ${\left|\kern-1.07639pt\left|\kern-1.07639pt\left|ABC\right|\kern-1.07639pt\right|\kern-1.07639pt\right|}\leq\|A\|{\left|\kern-1.07639pt\left|\kern-1.07639pt\left|B\right|\kern-1.07639pt\right|\kern-1.07639pt\right|}\|C\|$ for all matrices $A,B,C$. Furthermore, we assume that ${\left|\kern-1.07639pt\left|\kern-1.07639pt\left|A\right|\kern-1.07639pt\right|\kern-1.07639pt\right|}={\left|\kern-1.07639pt\left|\kern-1.07639pt\left|A^{T}\right|\kern-1.07639pt\right|\kern-1.07639pt\right|}$ for all matrices $A$. For example, this property is fulfilled by all Schatten-$p$ norms.

Assume that $\|U_{t}\|\leq 3\|X\|$ and that $\sigma_{\min}\left(U_{t}W_{t}\right)\geq\frac{1}{\sqrt{10}}\sigma_{\min}\left(X\right)$. Moreover, assume that $\mu\leq c\kappa^{-2}\|X\|^{-2}$, $\|V_{X^{\perp}}^{T}V_{U_{t}W_{t}}\|\leq c\kappa^{-2}$, and

where $\Delta_{t}:=XX^{T}-U_{t}U_{t}^{T}$. Then it holds that

Here, $c>0$ is an absolute constant chosen small enough.

When applying this lemma in our proof, we are going to set ${\left|\kern-1.07639pt\left|\kern-1.07639pt\left|\cdot\right|\kern-1.07639pt\right|\kern-1.07639pt\right|}=\|\cdot\|_{F}$. However, we believe that this lemma might be of independent interest, as it shows that $U_{t}U_{t}^{T}$ converges linearly towards $XX^{T}$ with respect to several different norms.

Having collected all the necessary ingredients, we can state and prove the main theorem of this section.

where $c_{2}>0$ is a small enough absolute constant. Then after

The proof of Theorem 9.6 shows that the number of iterations needed to complete Phase II is smaller than

and that the number of iterations needed to complete Phase III is smaller than

Phase II: In this phase, we will prove that $\sigma_{\min}\left(V_{X}^{T}U_{t}\right)$ is growing exponentially until it is at larger than $\frac{\sigma_{\min}\left(X\right)}{\sqrt{10}}$, while $\|U_{t}W_{t,\perp}\|$ grows at a much slower rate. For that, set

We will prove by induction that for $t_{\star}\leq t\leq t_{1}$ the following inequalities hold:

Note that when the inequalities above hold, then from the definition of $t_{1}$ above and inequality (55) we can derive that we must have that

For $t=t_{\star}$, we first note that inequalities (56), (57), and (58) follow directly from our assumptions. In order to prove inequality (55) we note that

where inequality $(a)$ is a consequence of assumption (53) and inequality $(b)$ follows from the definition of $\gamma$. Assume now that we have shown these four inequalities for some $t<t_{1}$. In order to prove them for $t+1$ we note first that

In inequality $(a)$ we applied the triangle inequality and for inequality $(b)$ we used the restricted isometry property as well as Lemma 7.3. Inequality $(c)$ is due to the induction assumption (57). In inequality $(d)$ we used the assumption $\delta\leq\frac{c_{1}}{\kappa^{4}\sqrt{r_{\star}}}$ as well as the induction assumption (56). For inequality $(e)$ we used $t\leq t_{1}$ as well as (59) and for inequality $(f)$ we used (52).

Next, we observe that by Lemma 9.1 we have that

In $(a)$ we have used that $\sigma_{\min}\left(V_{X}^{T}U_{t}\right)\leq\frac{\sigma_{\min}\left(X\right)}{\sqrt{10}}$, which follows from $t<t_{1}$. Using the induction assumption, this implies inequality (55). Moreover, the inequality chain above shows that $V_{X}^{T}U_{t+1}W_{t+1}$ has full rank. This allows us to apply Lemma 9.2, which implies that

where in inequality $(a)$ we used (58) as well as (60) and that the constant $c_{1}$ is chosen sufficiently small. This shows inequality (56). Next, due to inequality (60), our induction assumptions, and Lemma 9.4 we obtain that $\|U_{t+1}\|\leq 3\|X\|$, which shows inequality (57). Next, we note that by Lemma 9.3 we have that

where in inequality $(a)$ we used the induction hypothesis (57) as well as (60). Inequality $(b)$ follows from inequality (60) and our assumption on the step size $\mu$. In inequality $(c)$ we used the induction assumption $\|V_{X^{\perp}}^{T}V_{U_{t}W_{t}}\|\leq c_{2}\kappa^{-2}$. By choosing the constant $c_{1}>0$ small enough, this implies inequality (58) and, hence, finishes the induction step.

Note that from the definition of $t_{1}$ and from inequality (55) we obtain inequality (59). Hence, we obtain that

where inequality $(a)$ follows from inequality (56) and inequality $(b)$ follows from (59). Inequality $(c)$ follows from choosing the absolute constant $c_{2}>0$ small enough. Inequality $(d)$ follows from the fact that $\gamma\leq\sigma_{\min}\left(X\right)$. This finishes the proof of the second phase.

Phase III: In the third phase, we analyse the refinement of the signal $U_{t}$. For that, we define

Similar as in Phase II, we are going to show inductively that the following inequalities are fulfilled for $t_{1}\leq t\leq\hat{t}$:

For $t=t_{1}$ we note the inequalities (66), (68), and (69) follow from the results in Phase 1. Inequality (67) follows directly from setting $t=t_{1}$. Moreover, for $t=t_{1}$, inequality (70) follows from the observation that

where we have used that $\|U_{t_{1}}W_{t_{1}}\|\leq\|U_{t_{1}}\|\leq 3\|X\|$ by induction assumption (68).

For the induction step from $t$ to $t+1$ (with $t<\hat{t}$), we note first that with similar arguments as in Phase 1 we can show that

where inequality $(a)$ follows from (67). Inequality $(b)$ follows from (65) as well as the elementary inequality $\ln\left(1+x\right)\leq x$. Inequality $(c)$ follows from the assumptions $\gamma\leq c_{2}\frac{\sigma_{\min}\left(X\right)}{\min\left\{r;n\right\}\kappa^{2}}$ and $\delta\leq\frac{c_{1}}{\kappa^{4}\sqrt{r_{\star}}}$. This puts us in a position to apply our technical lemmas. We note that by Lemma 9.1 we have that

Note that for $\sigma_{\min}\left(V_{X}^{T}U_{t}\right)\leq\frac{1}{2}\sigma_{\min}\left(X\right)$ it holds that $(\ast)\geq 1$ and thus it follows that (66) holds for $t+1$ in this case. In the case of $\frac{1}{2}\sigma_{\min}\left(X\right)\leq\sigma_{\min}\left(V_{X}^{T}U_{t}\right)$ we obtain that

where in inequality $(a)$ we used the induction hypothesis (57) and in inequality $(b)$ we used the assumption $\mu\leq c_{1}\kappa^{-2}\|X\|^{-2}$. Hence, we have shown that also in this case the inequality (66) holds for $t+1$. Note that the previous inequality chain also implies that $V_{X}^{T}U_{t+1}W_{t}$ is invertible. Hence, in a similar way as in Phase II for inequality (56) we can verify that (67) holds for $t+1$. Moreover, note that from Lemma 9.4, induction assumption (57), and the assumption on the step size $\mu$ it follows that $\|U_{t+1}\|\leq 3\|X\|$. Moreover, inequality (69) can be shown analogously as in Phase $1$.

Next, we want to apply Lemma 9.5. For that, we compute that

Now choose symmetric matrices $Z_{1},Z_{2}$ of rank at most $r_{\star}$ such that $XX^{T}-U_{t}W_{t}W_{t}^{T}=Z_{1}+Z_{2}$ and $\langle Z_{1},Z_{2}\rangle=0$. It follows that

where inequality $(a)$ is a consequence of the restricted isometry property of order $2r_{\star}+1$ (see inequality (15) in Lemma 7.3).Note that if we would have assumed that $\mathcal{A}$ satisfies the restricted isometry property of order $3r_{\star}$ the decomposition of $XX^{T}-U_{t}W_{t}W_{t}^{T}U_{t}^{T}=Z_{1}+Z_{2}$ would not have been necessary. Instead, we could have directly applied inequality (15) in Lemma 7.3 with $Z=XX^{T}-U_{t}W_{t}W_{t}^{T}U_{t}^{T}$. Inequality $(b)$ follows from $\langle Z_{1},Z_{2}\rangle=0$. Next, denote by $U_{t}W_{t,\perp}W_{t,\perp}^{T}U_{t}^{T}=\sum_{i=1}^{n}\lambda_{i}v_{i}v_{i}^{T}$ the eigendecomposition of $U_{t}W_{t,\perp}W_{t,\perp}^{T}U_{t}^{T}$. Again, using inequality (15) in Lemma 7.3 it follows that

Combining the last three inequality chains, we obtain that

Inequality $(a)$ follows from the restricted isometry property. Inequality $(b)$ follows from the fact that $U_{t}W_{t,\perp}$ has rank at most $r-r_{\star}$ and inequality $(c)$ follows from $t<\hat{t}_{1}$ (see (64)). In inequality $(d)$ we used the assumption that $\delta\leq\frac{c_{1}}{\sqrt{r_{\star}}\kappa^{4}}\leq\frac{c}{2\kappa^{2}}$ with $c>0$ being the constant in Lemma 9.5. In an analogous way we can establish the inequalities

This shows that inequality (49) is fulfilled (with ${\left|\kern-1.07639pt\left|\kern-1.07639pt\left|\cdot\right|\kern-1.07639pt\right|\kern-1.07639pt\right|}$ being the Frobenius norm $\|\cdot\|_{F}$). Hence, we can apply Lemma 9.5 and we obtain that

where in the last inequality we used the induction assumption (70). We note that this shows that (70) also holds for $t+1$, if we can show that

where in the last inequality we used (61) and (67). Hence, for $c_{2}>0$ small enough, inequality (71) is implied by

By rearranging terms and using the elementary inequality $\ln\left(1+x\right)\geq\frac{x}{1-x}$, we see that this in turn is implied by

Hence, (65) shows (71), which shows inequality (70) for $t+1$. This finishes the induction step.

Conclusion: In order to finish the proof we distinguish two cases. Namely, in the first case we stop at $\hat{t}=\hat{t}_{1}$ and in the second case we stop at $\hat{t}=\hat{t}_{2}$. First, we consider the case $\hat{t}=\hat{t}_{1}$. Then we obtain that

where inequality $(a)$ follows from Lemma B.4. Inequality $(b)$ follows from (70) and inequality (71). Inequality $(c)$ is due to $\hat{t}=\hat{t}_{1}$ and the definition of $\hat{t}_{1}$. This proves inequality (54) in the case that $\hat{t}=\hat{t}_{1}$. Next, we consider the second case, $\hat{t}=\hat{t}_{2}$. We calculate that

Inequality $(a)$ follows from Lemma B.4 and inequality $(b)$ is due to the fact $\hat{t}=\hat{t}_{2}$ (see (65)). Inequality $(c)$ is due to the fact that $U_{\hat{t}}W_{\hat{t}}$ has rank at most $r-r_{\star}$. For inequality $(d)$ we used inequality (67) combined with inequality (61) and inequality $(e)$ follows from the fact that $\hat{t}\leq\hat{t}_{1}$ (see (63)). Inequality $(f)$ is due to the assumption $\gamma\leq c_{2}\frac{\sigma_{\min}\left(X\right)}{\min\left\{r;n\right\}\kappa^{2}}$. Inequality $(g)$ follows from the fact that $\gamma\leq\|X\|$. This proves inequality (54) in the case that $\hat{t}=\hat{t}_{2}$.

In order to finish the proof we need to show (54) follows from our definition of $\hat{t}$. For that we note that we obtain from $\hat{t}\leq\hat{t}_{1}$ and the definition of $\hat{t}_{1}$ (see equation (65)) that

Combining this estimate with inequality (59) shows (54). ∎

Proof of the main results

In order to finish the proof we will distinguish two cases:

Case $r\geq 2r_{\star}$: Due to (4) and (72) we can apply Lemma 8.7. Hence, with probability at least $1-O\left(\exp\left(-cr\right)\right)$ after

with $1\lesssim\beta\lesssim\frac{n\kappa^{4}}{\min\left\{r;n\right\}}$. Our goal is to apply Theorem 9.6 with $\gamma=\frac{\alpha\beta}{4}$. For that we need to check that

condition (77) is fulfilled, when the constant in (4) is chosen sufficiently small. Hence, by Theorem 9.6 after

Note that for the total amount of iterations we have that

where in inequality $(b)$ we have used $\beta\gtrsim 1$ and chosen the constant $C_{1}>0$ large enough. This finishes the proof of the first part.

Case $r_{\star}<r<2r_{\star}$: As in the first case, we can apply Lemma 8.7. Hence, with probability at least $1-\left(C\varepsilon\right)^{r-r_{\star}+1}+O\left(\exp\left(-cr\right)\right)$ after

iterations the inequalities (35), (36), (37), and (38) hold with $\frac{\varepsilon}{r}\lesssim\beta\lesssim\frac{\kappa^{4}n}{\varepsilon}$. Again, we want to apply Theorem 9.6 with $\gamma=\frac{\alpha\beta}{4}$. For that we need to check that

condition (78) holds true, when the constant in (7) is chosen sufficiently small. Hence, by Theorem 9.6 after

Note that for the total amount of iterations we have that

where in equality $(a)$ we have used $\frac{r_{\star}}{r-r_{\star}}\geq 1$, which follows from $r_{\star}<r\leq 2r_{\star}$. Inequality $(b)$ follows from $\frac{\varepsilon}{r}\lesssim\beta$ as well as from choosing $C_{2}>0$ large enough. This finishes the proof.

2 Proof of Theorem 3.4

As in the proof of the second part of Theorem 3.3 we can show that with probability at least $1-C\varepsilon+O\left(\exp\left(-cr_{\star}\right)\right)$ after

iterations the inequalities (35), (36), (37), and (38) hold with $\frac{\varepsilon}{r_{\star}}\lesssim\beta\lesssim\frac{\kappa^{4}n}{\varepsilon}$. Now define the matrix $\widehat{U_{t_{\star}}}$by adding a zero column to $U_{t_{\star}}$, i.e.,

Clearly, we can run gradient descent on $\widehat{U}_{t_{\star}}$ instead of $U_{t_{\star}}$ with the same step size, which gives us a sequence $\widehat{U}_{t_{\star}},\widehat{U}_{t_{\star}+1},\widehat{U}_{t_{\star}+2},\ldots$ to which we can apply Theorem 9.6 with $\gamma=\frac{\alpha\beta}{4}$ and $r=r_{\star}+1$. However, note that the last column always stays zero, which means that the results of this theorem also apply to $U_{t_{\star}},U_{t_{\star}+1},U_{t_{\star}+2},\ldots$. Hence, after

Note that for the total amount of iterations we have that

3 Proof of Theorem 3.5

We start by noting that in the special case $r=n$, the required assumptions for Theorem 9.6 are already fulfilled at the initialization $t_{0}=0$. This means that in this special case we do not need to analyze the spectral phase. This is shown by the following lemma.

It follows that $\|V_{X^{\perp}}^{T}V_{U_{0}W_{0}}\|=0$, which verifies the first equality. In order to see that the second equality holds we note that

The third equality follows directly from the definition of $U_{0}$. This finishes the proof. ∎

Now we are in a position to give a proof of Theorem 3.5.

By Lemma 10.1 we have that $\|V_{X^{\perp}}^{T}V_{U_{0}W_{0}}\|=0$, $\sigma_{\min}\left(U_{0}W_{0}\right)=\alpha$, and $\|U_{0}\|=\alpha$. This allows us to apply Theorem 9.6 with $t_{0}=0$ and $\gamma=\alpha$, which yields that after

Conclusion

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 or in blind deconvolution . Moreover, while our main results, e.g. Theorem 3.3 do require early stopping, our simulations (e.g. Figure 7(a)) 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_{\star}$ 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.

Acknowledgements

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.

References

Appendix A Proofs for the spectral phase

Claim: Set $\hat{E}_{i}:=\mu\mathcal{A}^{*}\mathcal{A}\left(U_{i-1}U_{i-1}^{T}\right)U_{i-1}$. Then, for $t\geq 1$ it holds that

Proof of the claim: We will prove the claim by induction. For $t=1$ we note that

which proves the claim for $t=1$. Now suppose that the claim holds for some $t$. We obtain that

where the last line follows from the definition of $\hat{E}_{t+1}$. By using the induction hypthesis we obtain that

where in the first line we used the submultiplicativity of the operator norm and in the second line we used the triangle inequality. In the third line we used that $\|\mathcal{A}^{*}\mathcal{A}\left(XX^{T}\right)\|=\lambda_{1}\left(M\right)$. Hence, we have shown that

In order to proceed assume now $t\leq t^{\star}$. Then by inequality (80) and the previous inequality we obtain that

A.2 Proof of Lemma 8.2

In order to finsh the proof, we are going to derive a lower bound for $t^{\star}$. First we note that by the definition of $t^{*}$ followed by elementary algebraic manipulations for $t<t^{*}$ we have

A.3 Proof of Lemma 8.3

Proof of inequality (18): Due to Weyl’s inequality we have that

where the second inequality follows from the Courant-Fisher minimax theorem (see, e.g., [67, Appendix A]). Now we note that

Proof of inequality (19): From Weyl’s inequality it follows that

Denote by $U=V_{U}\Sigma_{U}W_{U}^{T}$ the singular value decomposition of $U$. Then we can compute that

The first line is due to the Courant-Fisher minimax theorem and $U_{0}=\alpha U$. The last line follows again from the Courant-Fisher minimax theorem. Together with inequality (82) this implies the third claim.

Proof of inequality (20): First, we note that

Note that since $V_{L}^{T}V_{U}$ has rank $r_{\star}$, the matrix $Z_{t}V_{L}V_{L}^{T}U$ must have rank $r_{\star}$ as well. In particular, since $Z_{t}V_{L}V_{L}^{T}U=V_{L}V_{L}^{T}Z_{t}V_{L}V_{L}^{T}U$ this means that $L$ is the subspace spanned by the left-singular vectors of $Z_{t}V_{L}V_{L}^{T}U$ corresponding to the largest $r_{\star}$ singular values. Due to Wedin’s sin $\theta$ theorem we obtain that

where in the last line we also used (19). (Note that the assumption (17) guarantees that the denominator is positive, which is a necessary condition for an application of Wedin’s sin $\theta$ theorem.) Now we observe that

Together with the previous inequality chain, this shows inequality (20). ∎

A.4 Proof of Lemma 8.4

Before proving Lemma 8.4, we are going to introduce some notation. Let $U_{t}=\sum_{i=1}^{r}\sigma_{i}u_{i}v^{T}_{i}$ be the singular value decomposition of $U_{t}$. Define $L_{t}:=\sum_{i=1}^{r_{\star}}\sigma_{i}u_{i}v^{T}_{i}$ and $N_{t}:=\sum_{i=r_{\star}+1}^{r}\sigma_{i}u_{i}v^{T}_{i}$. Denote by $L_{t}=V_{L_{t}}\Sigma_{L_{t}}W_{L_{t}}^{T}$ and $N_{t}=V_{N_{t}}\Sigma_{N_{t}}W_{N_{t}}^{T}$ the singular value decomposition of those two matrices.

We start by proving the following technical lemma. It says that if the subpace spanned by the columns of $X$ and $L_{t}$ are aligned, then also the subspaces given by $W_{t}$ and $W_{L_{t}^{\perp}}$ will be closely aligned.

Assume that $\|V_{X^{\perp}}^{T}V_{L_{t}}\|\leq 1/2$. Then it holds that

In order to control the denominator we note that

In the last line we have used the assumption $\|V_{X^{\perp}}^{T}V_{L_{t}}\|\leq 1/2$. Hence, we have shown that

Now we are in a position to prove Lemma 8.4.

Proof of inequality $\eqref{ineq:initspec2}$: First, we observe that due to Lemma A.1 and the assumption $\|V_{X^{\perp}}^{T}V_{L_{t}}\|\leq\frac{1}{8}$ we have that