Overparameterized Nonlinear Learning: Gradient Descent Takes the Shortest Path?

Introduction

In this paper we mostly focus on nonlinear least-squares problems. In Section 5 we discuss results that apply to a broader class of loss functions $\mathcal{L}(\bm{\theta})$.

Classical statistical estimation/learning theory postulates that to find a reliable model that avoids overfitting, the size of the training data must exceed the intrinsic dimensionSome common notions of intrinsic dimension include Vapnik–Chervonenkis (VC) Dimension , Rademacher/Gaussian complexity , as well as naive parameter counting. of the model class $f(\cdot;\bm{\theta})$ used for empirical risk minimization (1.1). For many models such notions of intrinsic dimension are at least as large as the number of parameters in the model $p$, so that this literature requires the size of the training data to exceed the number of parameters in the model i.e. $n>p$. Contrary to this classical literature, modern machine learning models such as deep neural networks are often trained via first-order methods in an over-parameterized regime where the number of parameters in the model exceed the size of the training data (i.e. $n<p$). Statistical learning in this over-parameterized regime poses new challenges: Given the nonconvex nature of the training loss (1.1) can first-order methods converge to a globally optimal model that perfectly interpolate the training data? If so, which of the global optima do they converge to? What are the statistical properties of this model and how does this model vary as a function of the initial parameter used to start the iterative updates? What is the trajectory that iterative methods such as (stochastic) gradient descent take to reach this point? Why does a model trained using this approach generalize to new data and avoid overfitting to the training data?

In this paper we take a step towards addressing such challenges. We demonstrate that in many cases first-order methods do indeed converge to a globally optimal model that perfectly fits the training data. Furthermore, we show that among all globally optimal parameters of the training loss these algorithms tend to converge to one which has a near minimal distance to the parameter used for initialization. Additionally, the path that these algorithms take to reach such a global optima is rather short, with these algorithms following a near direct trajectory from initialization to the global optima. We believe these key features of first-order methods may help demystify why models trained using these simple algorithms can achieve reliable learning in modern over-parametrized regimes without over-fitting to the training data.

2 Insights from Linear Regression

Therefore, using a step size of $\eta\leq\frac{1}{\left\|\bm{X}\right\|^{2}}$ the residual iterates converge at a geometric rate to zero. This yields the first key property of gradient methods for over-parametrized learning:

Key property I: Gradient descent iterates converge at a geometric rate to a global optima.

Let $\bm{\theta}^{*}$ denote the global minima we converge to and $\Pi_{\mathcal{R}}$ and $\Pi_{\mathcal{N}}$ denote the projections onto the row space and null space of $\bm{X}$, respectively. Since the gradients lie on the row space of ${\bm{X}}$ and ${\bm{X}}$ is full row rank, denoting the unique pseudo-inverse solution by $\bm{\theta}^{\dagger}$, we have

The equalities above imply that $\bm{\theta}^{*}$ is the closest global minima to $\bm{\theta}_{0}$; which highlights the second property:

Key property II: Gradient descent converges to the closest global optima to initialization.

Key property III: Gradient descent takes a near direct trajectory to reach the closest global optima.

In this paper we show that similar properties continue to hold for a broad class of nonlinear over-parameterized learning problems.

3 Contributions

Our main technical contributions can be summarized as follows:

We provide a general convergence result for overparameterized learning via gradient descent, that comes with matching upper and lower bounds, showing that under appropriate assumptions over a small neighborhood of the initialization, gradient descent (1) finds a globally optimal model, (2) among all possible globally optimal parameters it finds one which is approximately the closest to initialization and (3) it follows a nearly direct trajectory to find this global optima.

We show that SGD exhibits the same behavior as gradient descent and converges linearly without ever leaving a small neighborhood of the initialization even with rather large learning rates.

We demonstrate the utility of our general results in the context of three overparameterized learning problems: generalized linear models, low-rank matrix regression, and shallow neural network training.

Convergence Analysis for Gradient Descent

The nonlinear least-squares problem in (1.1) can be written in the more compact form

A natural approach to optimizing (2.1) is to use gradient descent updates of the form

starting from some initial parameter $\bm{\theta}_{0}$. For the nonlinear least-squares formulation (2.1) above the gradient takes the form

The particular form of the gradient in (2.2) suggests that the eigenvalues of the Jacobian matrix may significantly impact the convergence of gradient descent. Our main technical assumption in this paper is that the spectrum of the Jacobian matrix is bounded from below and above in a local neighborhood of the initialization.

Here, $\sigma_{\min}(\cdot)$ and $\left\|\cdot\right\|$ denote the minimum singular value and the spectral norm respectively.

Our second technical assumption ensures that the Jacobian matrix is not too sensitive to changes in the parameters of the nonlinear mapping. Specifically we require the Jacobian to have either bounded or smooth variations as detailed next.

holds for some $0\leq\lambda\leq 1$. Here, $\alpha$ and $\beta$ are the bounds on the Jacobian spectrum over $\mathcal{D}$ per Assumption 1. (b) Smooth deviation: For all $\bm{\theta}_{1},\bm{\theta}_{2}\in\mathcal{D}$

With these assumptions in place we are now ready to state our main result.

Consider a nonlinear least-squares optimization problem of the form

Assumption 2 (a) holds over $\mathcal{D}$ with $\lambda=1/2$ and set $\eta\leq\frac{1}{2\beta^{2}}$.

Then, running gradient descent updates of the form $\bm{\theta}_{\tau+1}=\bm{\theta}_{\tau}-\eta\nabla\mathcal{L}(\bm{\theta}_{\tau})$ starting from $\bm{\theta}_{0}$, all iterates obey.

Furthermore, the total gradient path is bounded. That is,

A trivial consequence of the above theorem is the following corollary.

Consider the setting and assumptions of Theorem 2.1 above. Let $\bm{\theta}^{*}$ denote the global optima of the loss $\mathcal{L}(\bm{\theta})$ with smallest Euclidean distance to the initial parameter $\bm{\theta}_{0}$. Then, the gradient descent iterates $\bm{\theta}_{\tau}$ obey

The theorem and corollary above show that if the Jacobian of the nonlinear mapping is well-conditioned (Assumption 1) and has bounded/smooth deviations (Assumptions 2) in a ball of radius $R$ around the initial point, then gradient descent enjoys three intriguing properties.

Thus the distance between the global optima GD converges to and the initial parameter $\bm{\theta}_{0}$ is within a factor $4\frac{\beta}{\alpha}$ of the distance between the closest global optima to $\bm{\theta}_{0}$ and the initialization. This shows that among all global optima of the loss, the GD iterates converge to one with a near minimal distance to the initialization. In particular, (2.4) shows that for all iterates the weighted sum of the distance to the initialization and the misfit error remains bounded so that as the loss decreases the distance to the initialization only moderately increases.

Gradient descent follows a short path: Another interesting aspect of the above results is that the total length of the path taken by gradient descent remains bounded. Indeed, based on (2.7) the length of the path taken by GD is within a factor of the distance between the closest global optima and the initialization. This implies that GD follows a near direct route from the initialization to a global optima!

We would like to note that Theorem 2.1 and Corollary 2.2 are special instances of a more general result stated in the proofs (Theorem 9.3 stated in Section 9.2).Theorem 2.1 and Corollary 2.2 above are a special case of this theorem with $\lambda=1/2$ and $\rho=1$. This more general result requires Assumptions 1 and 2 to hold in a smaller neighborhood and improves the approximation ratios. Specifically, this more general result allows the radius $R$ to be chosen as small as

Also the approximation ratios in Corollary 2.2 can be improved to

However, this requires a smaller learning rate and hence leads to a slower converge guarantee.

More samples leads to a slower convergence rate by degrading the condition number of the Jacobian,

The required convergence radius $R$ increases proportional to $\sqrt{n}$ and we need Jacobian to be well-behaved over a larger neighborhood for fast convergence.

A natural question about the results discussed so far is whether the size of the local neighborhood for which we require our assumptions to hold is optimal. In particular, one may hope to be able to show that a significantly smaller neighborhood is sufficient. We now state a lower bound showing that this is not possible.

Consider a nonlinear least-squares optimization problem of the form

holds for all $\bm{\theta}$. Also, for any $\alpha$ and $\beta$ obeying $\alpha,\beta\geq 0$ and $\beta\geq\alpha$, there also exists a linear regression problem where running gradient descent updates of the form $\bm{\theta}_{\tau+1}=\bm{\theta}_{\tau}-\eta\nabla\mathcal{L}(\bm{\theta}_{\tau})$ starting from $\bm{\theta}_{0}=0$ with a sufficiently small learning rate $\eta$, all iterates $\bm{\theta}_{\tau}$ obey

Convergence Analysis for Stochastic Gradient Descent

Arguably the most widely used algorithm in modern learning is Stochastic Gradient Descent (SGD). For learning nonlinear least-squares problems of the form (2.1) a natural implementation of SGD is to sample a data point at random and use that data point for the gradient updates. Specifically, let $\{\gamma_{\tau}\}_{\tau=0}^{\infty}$ be an i.i.d. sequence of integers chosen uniformly from $\{1,2,\ldots,n\}$, the SGD iterates take the form

Here, $G(\bm{\theta}_{\tau};\gamma_{\tau})$ is the gradient on the $\gamma_{\tau}$th training sample. We are interested in understanding the trajectory of SGD for over-parameterized learning. In particular, whether the three intriguing properties discussed in the previous section for GD continues to hold for SGD. Our next theorem addresses this challenge.

Furthermore, suppose one of the following statements is valid.

Assumption 2 (a) holds over $\mathcal{D}$ and set $\eta\leq\frac{\alpha^{2}}{\nu\beta^{2}B^{2}}$.

Furthermore, on this event the SGD iterates never leave the local neighborhood ${\cal{D}}$.

Case studies

In this section we specialize and further develop our general convergence analysis in the context of three fundamental problems: fitting a generalized linear model, low-rank regression, and neural network training.

A natural approach for fitting such GLMs is via minimizing the nonlinear least-squares misfit of the form

The above theorem demonstrates that when fitting GLMs in the over-parameterized regime, gradient descent converges at a linear to a globally optimal model. Furthermore, this convergence is to the closest global optima to the initialization parameter. Also, we can deduce from (4.2) that the total gradient path length when using a step size on the order of $\frac{1}{\left\|\bm{X}\right\|^{2}\Gamma^{2}}$ is bounded by

2 Low-rank regression

This approach, originally proposed by Burer and Monteiro , shifts the search space from a large low-rank positive semidefinite matrix $\bm{\Theta}\bm{\Theta}^{T}$ to its factor $\bm{\Theta}$. In this section we study the behavior of GD and SGD on this problem in the over-parameterized regime where $n<dr$.

This theorem shows that with modest over-parametrization $dr\gtrsim n$, GD linearly converges to a globally optimal model and achieves zero loss. Note that degrees of freedom of $d\times r$ matrices is $dr$ hence as soon as $n>dr$, gradient descent can no longer perfectly fit arbitrary labels highlighting a phase transition from zero loss to non-zero as sample size increases. Furthermore, our result holds despite the nonconvex nature of the Burer-Monteiro approach.

3 Training shallow neural networks

We can now rewrite our input-output model in the more succinct form

The theorem below provides geometric global convergence guarantees for one-hidden layer neural networks in a simple over-parametrized regime.

Beyond nonlinear least-squares

Our first result shows that when the PL inequality holds around a minimally small neighborhood of the initialization, the intriguing properties of gradient descent discussed in Theorem 2.1 and Corollary 2.2 continue to hold beyond nonlinear least-squares problems.

with $\eta\leq 1/L$, all iterates $\bm{\theta}_{\tau}$ obey the following inequalities

Furthermore, the total path length of gradient descent is bounded via

Similar to Corollary 2.2 a trivial consequence of the above theorem is the following corollary.

Consider the setting and assumptions of Theorem 5.2 above. Let $\bm{\theta}^{*}$ denote the global optima of the loss $\mathcal{L}(\bm{\theta})$ with smallest Euclidean distance to the initial parameter $\bm{\theta}_{0}$. Then, the gradient descent iterates $\bm{\theta}_{\tau}$ obey

We end this section by discussing a simple lower bound which demonstrates that the required radius over which the Local PL result must hold per Theorem 5.2 is optimal up to a factor of two.

Numerical Experiments

To verify our theoretical claims, we conducted experiments on MNIST classification and low-rank matrix regression. To illustrate the tradeoffs between the loss function and the distance to the initial point, we define normalized misfit and normalized distance as follows.

We consider MNIST digit classification task and use a standard LeNet model from Tensorflow https://github.com/tensorflow/models/blob/master/research/slim/nets/lenet.py. This model has two convolutional layers followed by two fully-connected layers. Instead of cross-entropy loss, we use least-squares loss, without softmax layer, which falls within our nonlinear least-squares framework. We conducted two set of experiments with $n=500$ and $n=5000$. Both experiments use Adam with learning rate $0.001$ and batch size $100$ for $1000$ iterations. At each iteration, we record the normalized misfit and distance to obtain a misfit-distance trajectory similar to Figure 9.13. We repeat the training $20$ times (with independent initialization and dataset selection) to obtain the typical behavior.

In Figure 2(b) and 3(b) we increase the sample size to $n=5000$. Similar to the first case, during the initial phase ($\text{misfit}\leq 0.4$) the loss-distance curve is a straight line and levels off later on. Compared to $n=500$, leveling off occurs earlier and is more visible. For instance, at $\text{misfit}=0.2$, output layer FC2 has distance of $0.5$ for $n=5000$ and $0.25$ for $n=500$. This is consistent with Theorem 2.1 which predicts (i) more samples imply a Jacobian with worse condition number and (ii) the global minimizer lies further away from the initialization and it is less-likely that the Jacobian will be well-behaved over this larger neighborhood.

2 Low-rank regression

We consider a synthetic low-rank regression setup to test the predictions of Theorem 4.2. We generate input matrices with i.i.d. standard normal entries and labels with i.i.d. Rademacher entries. We set $r=4$ and $d=100$ and initialize $\bm{\Theta}_{0}$ according to Theorem 4.2. We vary the sample size to be $n\in\{25,50,100,200\}=\{dr/16,dr/8,dr/4,dr/2\}$ and run gradient descent for $200$ iterations with a constant learning rate per Theorem 4.2. We observe a linear tradeoff in terms of misfit-distance to initialization with a narrow confidence interval consistent with our theoretical predictions in Figure 9.13. In the large sample size ($n=dr/2$), the problem is less over-parameterized and the confidence intervals become notably wider especially when the misfit is close to zero (i.e. by the time we reach a global minima). As predicted by our main theorem, the distance to initialization $\bm{\Theta}_{0}$ increases gracefully as the number of labels $n$ increases.

Prior Art

Implicit regularization: There is a growing interest in understanding properties of overparameterized problems. An interesting body of work investigate the implicit regularization capabilities of (stochastic) gradient descent for separable classification problems including . These results show that gradient descent does not converge to an arbitrary solution, for instance, it has a tendency to converge to the solution with the max margin or minimal norm. Some of this literature apply to regression problems as well (such as low-rank regression). However, for regression problems based on a least-squares formulation the implicit bias/minimal norm property is proven under the assumption that gradient descent converges to a globally optimal solution which is not rigorously proven in these papers.

Overparameterized neural networks: A few recent papers study the benefits of overparameterization for training neural networks and related optimization problems. Very recent works show that overparameterized neural networks can fit the data with random initialization if the number of hidden nodes are polynomially large in the size of the dataset. Similar to us, these works argue that there is a global minima around the random initialization. However these works are specialized towards neural nets and similar to us the bounds on the network size to achieve global optimality appear to be suboptimal.We note that while both our results and these papers are suboptimal for one-hidden layer neural networks, they are not directly comparable with each other. We assume $n\leq d$ where as these papers assume $poly(n)\lesssim k$. Also the assumptions on the activations are different from each other. In contrast, we focus on general nonlinearities and also focus on the gradient descent trajectory showing that among all the global optima, gradient descent converges to one with near minimal distance to the initialization. We would also like to note that the importance of the Jacobian for overparameterized neural network analysis has also been noted by other papers including and also which investigate the optimization landscape and properties of SGD for training neural networks. An equally important question to understanding the convergence behavior of optimization algorithms for overparameterized models is understanding their generalization capabilities this is the subject of a few interesting recent papers . While our results do not directly address generalization, by characterizing the properties of the global optima that (stochastic) gradient descent converges to it may help demystify the generalization capabilities of overparametrized models trained via first order methods. Rigorous understanding of this relationship is an interesting and important subject for future research.

Stochastic methods: SGD performance guarantees are typically in expectation rather than in probability. Martingale-based methods have been utilized to give probabilistic guarantees . The main challenge in nonconvex analysis of SGD, is to ensure SGD iterates stay within a region where nonconvex analysis can apply even when using rather large learning rates. While a few papers show that SGD stays in a specific region with high probability in specific instances, these results require using very small learning rates (which translates into very small variance) to ensure standard concentration arguments apply. In contrast, our approach allows for much larger learning rates by using martingale stopping time arguments. Our approach is in part inspired by which studies SGD for nonconvex phase retrieval but involves different assumptions on the loss.

Nonconvex optimization: A key idea for solving nonconvex optimization problems is ensuring that optimization landscape has desirable properties. These properties include Polyak-Lojasiewicz (PL) condition and the regularity condition (e.g. local strong convexity) . PL condition is particularly suited for analyzing overparameterized problems and has been utilized by several recent papers . Unlike these works, we show that overparameterized gradient descent trajectory stays in a small neighborhood and we only need properties such as PL to hold over this region. There is also a large body of work that study the applications discussed in this paper in the over determined regime $p\leq n$. For instance, Low-rank regression and generalized linear models have been considered by various works including in such an overdetermined setting. More recently, provable first order methods for learning neural networks have been investigated by multiple papers including in the overdetermined setting.

Discussion and future directions

This work provides new insights and theory for overparameterized learning with nonlinear models. We first provided a general convergence result for gradient descent and matching upper and lower bounds showing that if the Jacobian of the nonlinear mapping is well-behaved in a minimally small neighborhood, gradient descent finds a global minimizer which has a nearly minimal distance to the initialization. Second, we extend the results to SGD to show that SGD exhibits the same behavior and converges linearly without ever leaving a minimally small neighborhood of initializtion. Finally, we specialize our general theory to provide new results for overparameterized learning with generalized linear models, low-rank regression and shallow neural network training. A key tool in our results is that we introduce a potential function that captures the tradeoff between the model misfit and the distance to the initial point: the decrease in loss is proportional to the distance from the initialization. Our numerical experiments on real and synthetic data further corroborate this intuition on the loss-distance tradeoff.

In this work we address important challenges surrounding the optimization of nonlinear over-parametrized learning and some of its key features. The fact that gradient descent finds a nearby solution is a desirable property that hints as to why generalization to new data instances may be possible. However, we emphasize that this is only suggestive of the generalization capabilities of such algorithms to new data. Indeed, developing a clear understanding of the generalization capabilities of first order methods when solving over-parameterized nonlinear problems is an important future direction. Making progress towards this generalization puzzle requires merging insights gained from optimization with more intricate tools from statistical learning and is an interesting topic for future research.

Proofs

We introduce the following matrix and vector which play a crucial role in the convergence analysis of our algorithms

2 Gradient descent convergence proofs (Theorem 2.1 and Corollary 2.2)

Theorem 2.1 and Corollary 2.2 are a special case of a more general result stated below. Theorem 2.1 and Corollary 2.2 then follows by setting $\lambda=1/2$ and $\rho=1$.

Consider a nonlinear least-squares optimization problem of the form

Assumption 2 (a) holds over $\mathcal{D}$ and set $\eta\leq\frac{\lambda}{\beta^{2}}$.

Then, running gradient descent updates of the form $\bm{\theta}_{\tau+1}=\bm{\theta}_{\tau}-\eta\nabla\mathcal{L}(\bm{\theta}_{\tau})$ starting from $\bm{\theta}_{0}$, all iterates obey.

Furthermore, the total gradient path is bounded. That is,

Let $\bm{\theta}^{*}$ denote the global optima of the loss $\mathcal{L}(\bm{\theta})$ with smallest Euclidean distance to the initial parameter $\bm{\theta}_{0}$. Then, the gradient descent iterates $\bm{\theta}_{\tau}$ also obey

Proof Sketch. To prove the above theorem we begin by noting that the residual $\bm{r}_{\tau}$ satisfies the recursion

where ${\bm{C}}(\bm{\theta}_{\tau}):={\cal{J}}(\bm{\theta}_{\tau+1},\bm{\theta}_{\tau}){\cal{J}}(\bm{\theta}_{\tau})^{T}$. Here, (a) follows from fundamental rule of calculus and (b) from the gradient identity $\nabla\mathcal{L}(\bm{\theta}_{\tau})=\mathcal{J}^{T}(\bm{\theta}_{\tau})\bm{r}_{\tau}$. If ${\bm{I}}-\eta{\bm{C}}(\bm{\theta}_{\tau})$ has spectral norm less than $1$, the the residual verctors will converge linearly. We build on this observation and show that one only needs this requirement over a minimally small neighborhood of $\bm{\theta}_{0}$. To this aim, we first introduce a potential set which contains the space of parameters that can be reached by gradient descent.

Note that $\mathcal{P}(\bm{\theta}_{0},R_{\zeta})\subseteq{\mathcal{B}}(\bm{\theta}_{0},R_{\zeta})$. Our first lemma shows that, if an iterate $\bm{\theta}_{\tau}\in{\cal{P}}:=\mathcal{P}(\bm{\theta}_{0},R_{\zeta})$, then the next iterate $\bm{\theta}_{\tau+1}$ stays in the set ${\cal{D}}:=\mathcal{B}(\bm{\theta}_{0},R_{\zeta})$.

In the above, (a) follows from the upper bound on the Jacobian over $\mathcal{D}$ per Assumption 1, (b) from the fact that $\eta\leq\frac{1}{\beta^{2}}$, (c) from $\alpha\leq\beta$, and (d) from $\zeta\leq\alpha$. The latter combined with the triangular inequality yields

concluding the proof of $\bm{\theta}^{+}\in\mathcal{D}$.

The next lemma establishes the convergence to a global minima that lies in a minimally small local neighborhood under a Jacobian condition (9.10). The proof of this lemma is deferred to Section 9.2.1.

Furthermore, the total gradient path is bounded. That is,

The next lemma shows that (9.10) indeed holds. We defer the proof of this lemma to Section 9.2.2.

Assumption 2(a) holds over $\mathcal{D}$ and $\eta\leq\frac{1}{\beta^{2}}$

With these lemmas in place we are now ready to prove Theorem 9.3. Proof of Theorem 9.3: Set $\zeta=(\lambda-\eta\beta^{2}/2)\alpha$ and observe that

Based on the above, the assumptions of Theorem 9.3 also subsume those of Lemma 9.6. Thus (9.11), (9.12), and (9.13) hold for all $\tau$.

This completes the bounds (9.2), (9.3), and (9.4) of Theorem 9.3. The proofs of (9.5) and (9.6) follow immediately from (9.3) and (9.4) by noting that for any global optima (including the closest global optima to $\bm{\theta}_{0}$ denoted by $\bm{\theta}^{*}$) we have

This concludes the proof of Theorem 9.3. All that remains is to prove Lemmas 9.6 and 9.7 which are the subject of the two sections below.

We will prove this lemma by induction. Assume the claim holds until iteration $\tau$. First, since (9.12) holds, applying Lemma 9.5 and using the facts that $\eta\leq 1/\beta^{2}$ and $\zeta\leq\alpha$, we can conclude that $\bm{\theta}_{\tau+1}\in{\cal{D}}$.

For the norm of the residual using the fact that ${\bm{C}}(\bm{\theta})\succeq\lambda{\cal{J}}(\bm{\theta}){\cal{J}}(\bm{\theta})^{T}$ (per assumption (9.10)) we have

Here, (a) follows from (9.7), (b) from (9.10) and the upper bound on the spectral norm of the Jacobian, (c) and from merging the terms on the right hand side. Combining (9.2.1) with ${\sigma_{\min}(\mathcal{J}(\bm{\theta}_{\tau}))}\geq\alpha$, and using $\eta\leq\lambda/\beta^{2}$, we conclude that

completing the proof of (9.11). For the remainder of discussion, denote ${\gamma}=(\lambda-\eta\beta^{2}/2)\eta$. ${\gamma}$ is nonnegative due to upper bound on $\eta$ and we have

We now turn our attention to proving (9.12). To this aim we start from (9.2.1) and complete the square to conclude that

Also note that using the upper bound on spectrum of $\mathcal{J}$ and ${\gamma}\leq\lambda\eta\leq\frac{1}{\beta^{2}}$ we have

Thus, taking square root from both sides of (9.2.1) we reach the following identity for changes in the norm of residual

To combine the identities (9.14) and (9.17) in such a way to yield our theorem we proceed by defining the potential/Lyapunov function below with $\zeta=\alpha{\gamma}/\eta$.

A unique feature of the $\mathcal{V}_{\tau}$ potential is that it is non-increasing. To see this note that using (9.17) we have

Finally using the definition of $\mathcal{V}_{\tau}$ and its non-increasing property (9.2.1) we have

concluding the proof of (9.13) and Lemma 9.6 when we substitute $\zeta=(\lambda-\eta\beta^{2}/2)\alpha$.

2.2 Proof of Lemma 9.7

we consider the two cases related to Assumption 2 separately.

If Assumption 2(a) holds then for any $\bm{\theta}_{1},\bm{\theta}_{2}\in\mathcal{D}$ we have

Thus for $\bm{\theta},\bm{\theta}^{+}\in\mathcal{D}$ we have

Next, suppose Assumption 2(b) holds. Then, for any $\bm{\theta}_{1},\bm{\theta}_{2}\in\mathcal{D}$ we have

Repeating the previous argument (with Assumption 2(a)), we again conclude with (9.21).

3 Lower bounds proofs (Theorem 2.3)

We begin by proving (2.12). To show this we first use the upper bound on the Jacobian matrix to prove that the nonlinear mapping is Lipschitz. To this aim note that

completing the proof of the Lipschitz property. This Lipschitz property combined with the triangular inequality allows us to conclude

For any $\alpha,\beta\geq 0$ obeying $\alpha\leq\beta$ and any $\bm{\theta}$, we have

If $\eta\leq 1/\beta^{2}$, all iterations satisfy $0\leq\theta_{\tau}\leq\theta^{\star}=\gamma$. On the other hand, the misfit in each iteration obeys

4 SGD proofs (Proof of Theorem 3.1)

We begin our SGD analysis by writing the SGD iterates in terms of the Jacobian matrix. To this aim define the matrix ${\cal{J}}(\bm{\theta}_{\tau};\gamma_{\tau})$ which keeps the $\gamma_{\tau}$-th row of ${\cal{J}}(\bm{\theta}_{\tau})$ and sets the remaining rows to zero. We note that

Similar to the GD proof we begin by noting that the residual $\bm{r}_{\tau}$ satisfies the recursion

Here, (a) follows from the fundamental rule of calculus, (b) from the stochastic update rule, and (c) from combining the form of the stochastic gradient in (9.23) with the definition of ${\bm{C}}(\bm{\theta}_{\tau};\gamma_{\tau})$.

It is completely unclear if SGD stays inside a neighborhood around the initial model to ensure the on average convergence argument discussed above is useful. We will develop a novel martingale-based argument to show that SGD does indeed stay in this local neighborhood. We briefly discuss the intuition behind this approach here. Since SGD is inherently random, ideally, we would like to show that, a variant of (2.4) holds. Specifically, define

Figure 4 provides a pictorial illustration of this potential function.

4.2 Decrease of the expected misfit

In this section we will show that under the assumption that SGD iterates always remain close to the initialization, the expected value of the norm of the residual will decrease in each iteration. Concretely, in this section we prove the following lemma.

Assume $\bm{\theta}\in\mathcal{D}^{\prime}:=\mathcal{B}(\nu/2)$ with $\nu$ a scalar obeying $\nu\geq 3$. Also assume the Jacobian associated with $f$ obeys Assumption 1 over the set $\mathcal{D}:=\mathcal{B}(\nu)$ and the rows of the Jacobian have bounded Euclidean norm over this set, that is

Assumption 2(a) holds over ${\cal{D}}$ and $\eta\leq\frac{\alpha^{2}}{2\beta^{2}B^{2}}$.

For simplicity of exposition of the proof of this lemma we define $r(\bm{\theta})=f(\bm{\theta})-\bm{y}$ and $r(\bm{\theta}^{+})=f(\bm{\theta}^{+})-\bm{y}$. We prove the lemma in three steps.

Step I: We show that as long as $\eta\leq\frac{1}{\beta B}$, then $\bm{\theta}^{+}\in\mathcal{D}$.

Step II: We prove that the matrix $\bm{C}(\bm{\theta}):={\cal{J}}(\bm{\theta}^{+},\bm{\theta}){\cal{J}}^{T}(\bm{\theta})$ obeys

Step III: We use Step I and II to show the inequalities (9.27) and (9.28) which are equivalent to

Using this inequality we can conclude that

Here, (a) follows from (9.32), (b) from the fact that $\bm{\theta}\in\mathcal{D}^{\prime}:=\mathcal{B}(\nu/2)$, and (c) from $\eta\leq\frac{1}{\beta B}$. Furthermore, the simple fact that $\alpha\leq\beta$ implies that

Here, (a) follows from (9.33) and the fact that $\left\|{\cal{J}}(\bm{\theta}^{+},\bm{\theta})\right\|\leq\beta$ and (b) follows from the fact that $\bm{\theta}\in\mathcal{D}^{\prime}:=\mathcal{B}(\nu/2)$. Combining (9.34) and (9.4.2) we conclude that $\bm{\theta}^{+}\in\mathcal{D}:=\mathcal{B}(\nu)$.

The proof of (9.29) is very similar to the proof of Lemma 9.7 with $\lambda=1/2$. In particular, under Assumption 2(a) the exact same argument yields (9.29). To show the result under Assumption 2(b) we combine (9.2.2) from the proof of Lemma 9.7, (9.32), and $\bm{\theta}\in\mathcal{D}^{\prime}=\mathcal{B}(\nu/2)$ to conclude that

From the arguments of Steps I and II we know that

$\|{\bm{C}}(\bm{\theta})\|\leq\beta^{2}$ and $\|{\cal{J}}(\bm{\theta}^{+},\bm{\theta})\|\leq\beta$,

${\bm{C}}(\bm{\theta})\succeq\frac{1}{2}{\cal{J}}(\bm{\theta}){\cal{J}}(\bm{\theta})^{T}$.

Using (ii) ${\cal{J}}(\bm{\theta}^{+},\bm{\theta}){\cal{J}}(\bm{\theta}^{+},\bm{\theta})^{T}\preceq\beta^{2}{\bm{I}}_{n},$ so that

Furthermore, ${\cal{J}}(\bm{\theta};\gamma){\cal{J}}(\bm{\theta};\gamma)^{T}$ is a diagonal matrix with a single nonzero entry which is bounded by $B^{2}$. Thus,

Using the latter two inequalities allows us to conclude

Here, (a) follows from (9.36), (b) from the fact that the step size obeys $\eta\leq\frac{\alpha^{2}}{2B^{2}\beta^{2}}$, (c) from ${\sigma_{\min}(\mathcal{J}(\bm{\theta}))}\geq\alpha$, and (d) from (9.37). These inequalities allow us to conclude

Here, (a) follows from the calculation in (9.4.1) applied to $\bm{r}(\bm{\theta})$ and $\bm{r}(\bm{\theta}^{+})$, (b) from (9.4.2), (c) from (9.37), and (d) from completing the square. Finally, note that using the upper bound on the spectrum of the Jacobian and the fact that $\eta\leq\frac{\alpha^{2}}{2\beta^{2}B^{2}}\leq\frac{1}{2\beta^{2}}$Note that $\alpha\leq B$. we have

so that the term inside the parentheses of right-hand sided of (9.39) is positive. Consequently, combining Jensen’s inequality with the square root of both sides of (9.39) yields

concluding the proof (9.30). To prove (9.31) we use the penultimate inequality from (9.39) together with the fact that $\sigma_{\min}\left({\cal{J}}(\bm{\theta})\right)\geq\alpha$ to conclude that

4.3 Bounding the increase of expected average distance to anchor points

In this section we will show that under the assumption that SGD iterates always remain close to the initialization, the expected value of the average distance to the anchor points will not significantly increase in each iteration. Specifically, the anchor points we pick are an $\epsilon$ cover of the neighborhood of the initialization denoted by $\mathcal{P}=\{\bm{p}_{1},\bm{p}_{2},\ldots,\bm{p}_{K}\}$. and we monitor the following average distance

Concretely, in this section we prove the following lemma.

Using (9.42) and $\|{\cal{J}}^{T}(\bm{\theta})\|\leq\beta$, we also have

We also prove the following simple lemma.

Proof Using the assumption $\bm{\theta}\in\mathcal{B}(\nu/2)$, we have

Combining (9.44) and (9.46), we conclude that

Using the latter two identities we conclude that

Dividing both sides by $K$ completes the proof of (9.41).

4.4 Shortest path potential is a supermartingale

In this section we show that the shortest path potential

is a supermartingale. Specifically we prove the following lemma.

Also assume the rows of the Jacobian have bounded Euclidean norm over this ball, that is

Furthermore, suppose one of the following statements is valid.

Assumption 2 (a) holds over $\mathcal{D}$ and set $\eta\leq\frac{\alpha^{2}}{\nu\beta^{2}B^{2}}$.

Turning our attention to the supermartingale property, define $\bm{r}_{\tau}=f(\bm{\theta}_{\tau})-\bm{y}$ and note that when $\bm{\theta}_{\tau}\in\mathcal{B}(\nu/2)$, by Lemmas 9.9 and 9.10 we have

Summing these two identities with a scaling of the first inequality by $12$ and the second one by $\alpha$, we obtain

4.5 SGD remains in the local neighborhood

In this section we show that SGD iterates remain close to the initialization. Specifically we prove the following lemma.

Consider the setup of Lemma 9.11 and the potential function $\mathcal{V}$ from (9.48). Also define the stopping time $T=\min\{\tau:\bm{\theta}_{\tau}\not\in\mathcal{B}(\nu/2)\}$. Under the stated assumptions,

The $\mathcal{V}_{\tau\wedge\widetilde{T}}$ term is measurable with respect to filteration $\mathcal{F}_{\tau}$, hence

Now that we established $\mathcal{U}_{\tau}$ is a supermartingale, Martingale maximal inequality implies that

4.6 Putting everything together (completing the proof of Theorem 3.1)

With this recursion established, we take conditional expectations to obtain

5 GLM proofs (Proof of Theorem 4.1)

Let $\bm{\theta}^{*}=\Pi_{\mathcal{N}}(\bm{\theta}_{0})+\bm{\theta}^{\dagger}\in{\cal{G}}$. By construction $\bm{\theta}^{*}$ is the closest global minima to $\bm{\theta}_{0}$ as the null space projections match. We will argue that the gradient descent iterations linearly converge to $\bm{\theta}^{*}$.

Towards this goal, note that $\bm{y}=\phi\left(\bm{X}\bm{\theta}^{*}\right)$ and note that the gradient descent iterations are given by

Now, for two vectors $\bm{a}$ and $\bm{b}$ obeying $\bm{a}\neq\bm{b}$ define $\phi^{\prime}(\bm{a},\bm{b})=\frac{\phi(\bm{a})-\phi(\bm{b})}{\bm{a}-\bm{b}}$ (with the devision interpreted as entry by entry) and note that by the mean value theorem $\phi^{\prime}(\bm{a},\bm{b})\geq\gamma$. Also note that, we can write $\phi({\bm{X}}\bm{\theta}_{\tau})-\phi({\bm{X}}\bm{\theta}^{\star})=\text{diag}(\phi^{\prime}({\bm{X}}\bm{\theta}_{\tau},{\bm{X}}\bm{\theta}^{\star})){\bm{X}}(\bm{\theta}_{\tau}-\bm{\theta}^{\star})$. Consequently, setting $\bm{h}_{\tau}=\bm{\theta}_{\tau}-\bm{\theta}^{\star}$ and ${\bm{D}}_{\tau}=\text{diag}(\phi^{\prime}({\bm{X}}\bm{\theta}_{\tau}))\text{diag}(\phi^{\prime}({\bm{X}}\bm{\theta}_{\tau},{\bm{X}}\bm{\theta}^{\star}))$, we have

To continue further, we use the fact that ${\bm{D}}_{\tau}$ is diagonal with entries between $\gamma^{2}$ and $\Gamma^{2}$. This combined with the fact that the matrices $\widetilde{\bm{X}}$ and $\bm{X}$ have the same eigenvalues allow us to conclude that $\gamma^{2}\sigma_{\min}^{2}\left(\bm{X}\right)\bm{I}\preceq\widetilde{{\bm{X}}}^{T}{\bm{D}}_{\tau}\widetilde{{\bm{X}}}\preceq\Gamma^{2}\left\|\bm{X}\right\|^{2}\bm{I}$. Thus, for $\eta\leq\frac{1}{\Gamma^{2}\left\|\bm{X}\right\|^{2}}$

completing the proof of (4.2). Furthermore, note that

6 Low-rank recovery proofs (Proof of Theorem 4.2)

To specialize Theorem 2.1 we begin by calculating the Jacobian ${\cal{J}}(\bm{\Theta}):={\cal{J}}({\text{vect}(\bm{\Theta})})$ which is given by an $n\times dr$ matrix of the form

In order to verify the assumptions of Theorem 2.1, in this section we gather some key lemmas related to the Jacobian matrix that building on top of each other play a crucial role in our proofs. We defer the proofs to Appendix A. The first key lemma which will play a crucial role in our proofs is that the nuclear norm $\|{\text{mat}\left({\cal{J}}(\bm{\Theta})^{T}\bm{v}\right)}\|_{\star}$ is uniformly bounded for all $\bm{v}$ and $\bm{\Theta}$ with unit Frobenius/Euclidean norms.

holds with probability at least $1-e^{-2dr}$.

The next lemma concerns the average of the nuclear norm of a Gaussian matrix multiplied by a diagonal matrix.

Furthermore, assume $r\leq d$ and $n\leq cdr$ with $c$ a fixed numerical constant. Then,

Next we bound the spectrum of the Jacobian matrix in a ball around the initialization $\bm{\Theta}_{0}$.

6.2 Completing the proof of Theorem 4.2

We will prove this theorem by a direct application of Theorem 2.1. To this aim we need to calculate the various parameters in this theorem.

Hence, with probability at least $1-e^{-\frac{n}{2}}$, the following holds

Now that Theorem 2.1 applies, all that remains is to upper bound these quantities in the upper bound on the learning rate. Per Theorem 2.1 we need to ensure

and use $\min(1,\sqrt{\frac{d}{n}})\geq\frac{1}{\sqrt{r}}$. Proceeding, we use this naive bound to simplify the final expressions. This yields the step size requirement of

Observing $\alpha^{2}/\beta^{2}=1/r$ and substituting $\eta$ and convergence rate $1-\eta\alpha^{2}/2$ concludes the proof.

7 Neural net proofs (Proof of Theorem 4.3)

We begin by noting that the Jacobian matrix in this case is equal to

To prove this theorem we use Theorem 2.1 with $R=\infty$. We just need to calculate the various parameters and verify that the assumptions hold.

Bounding the spectrum of $\mathcal{J}$. We begin by calculating $\alpha$ and $\beta$. To this aim note

Thus, using the bounds on $\phi^{\prime}$

Bounding the Lipschitz parameter of $\mathcal{J}$. To calculate $L$ note that

In the above (a) follows from the fact the square of the spectral norm of concatenation of matrices is bounded by sum of squares of the spectral norms of the individual matrices. Thus we can use

The proof is complete by applying Theorem 2.1.

8 PL proofs

Suppose (5.1) and (5.2) hold until step $\tau$. This implies $\bm{\theta}_{\tau}\in{\cal{D}}$ and local PL is applicable. If ${\cal{L}}(\bm{\theta}_{\tau})=0$, then $\bm{\theta}_{\tau}$ is global minimizer and since ${\cal{L}}$ is differentiable ${\nabla{\cal{L}}(\bm{\theta}_{\tau})}=0$ which in turn implies that $\bm{\theta}_{\tau+1}=\bm{\theta}_{\tau}$ and thus (5.2) holds for $\bm{\theta}_{\tau+1}$. Otherwise, ${\cal{L}}(\bm{\theta}_{\tau})>0$ and using the triangular inequality we can conclude that

so that $\epsilon_{\tau}(\eta)>0$ for $\eta>0$. Using this definition in (9.57) together with the PL condition for $\bm{\theta}_{\tau}\in{\cal{D}}$, we arrive at

completing the proof of (5.1). To conclude with the result on the shortest path, we add (9.58) from $\tau=0$ to $\infty$ to conclude that

8.2 PL lower bound proof (Proof of Theorem 5.4)

Proof Suppose there exists $\bm{\theta}\in{\cal{D}}$ satisfying ${\cal{L}}(\bm{\theta})=0$. Since ${\cal{L}}$ is differentiable and minimized at $\bm{\theta}$ the gradient must vanish, i.e. ${\nabla{\cal{L}}(\bm{\theta})}=0$. From smoothness of the loss we conclude that

Next, observe that (i) ${\cal{L}}(0)=\gamma^{2}\mu/2$ and (ii) any global minimizer $\bm{\theta}$ satisfies $\bm{y}={\bm{X}}\bm{\theta}^{\star}={\bm{X}}\bm{\theta}$ hence we have that

Acknowledgements

M. Soltanolkotabi is supported by the Packard Fellowship in Science and Engineering, an NSF-CAREER under award #1846369, the Air Force Office of Scientific Research Young Investigator Program (AFOSR-YIP) under award #FA9550-18-1-0078, an NSF-CIF award #1813877, and a Google faculty research award.

References

Appendix A Proof of key lemmas for low-rank recovery

Given the random nature of the matrices $\bm{X}_{i}$, ${\text{mat}\left({\cal{J}}(\bm{\Theta})^{T}\bm{v}\right)}$ defines a random process $\Gamma_{\bm{v},\bm{\Theta}}$ indexed by $\bm{\Theta}$ and $\bm{v}$ that can be rewritten in the form

Next, we also show that $\left\|\Gamma_{\bm{v},\bm{\Theta}}\right\|_{*}$ concentrates well around this expectation. To show this we use the fact stated above that $\left\|\Gamma_{\bm{v},\bm{\Theta}}\right\|_{*}=\left\|{\bm{Y}}{\bm{{\Sigma}}}\right\|_{*}$ is a function of a Gaussian matrix $\bm{Y}$. Furthermore, $\left\|{\bm{Y}}{\bm{{\Sigma}}}\right\|_{*}$ is Lipschitz as for any two matrices $\bm{Y}_{1},\bm{Y}_{2}$ we have

Here ${\bm{V}}$ follows from dual representation of the nuclear norm and is a matrix with spectral norm bounded by $1$ maximizing $\langle\bm{V},\left(\bm{Y}_{2}-\bm{Y}_{1}\right)\bm{\Sigma}\rangle$. Thus for fixed $\bm{v}$ and $\bm{\Theta}$, $\left\|\Gamma_{\bm{v},\bm{\Theta}}\right\|_{*}$ is a $1$-Lipschitz function of a Gaussian matrix $\bm{Y}$. Thus utilizing concentration of Gaussian measure combined with (A.1) implies

Using (A.3) with $t=4\sqrt{dr}$ combined with the above covering bound we conclude that for $n\leq dr$

which implies that $\text{OPT}=\left\|\Gamma_{\bm{\Theta}^{\star},\bm{v}^{\star}}\right\|_{*}\leq 12\sqrt{dr}$, completing the proof. In the above (a) follows from the triangular inequality, (b) from the linearity of $\bm{\Gamma}_{\bm{v},\bm{\Theta}}$ with respect to $\bm{v}$ and $\bm{\Theta}$ and the definition of OPT, and (c) from the bound on the cover.

A.2 Proof of Lemma 9.14

Note that for a Gaussian random vector $\bm{g}\sim\mathcal{N}(\bm{0},\bm{I}_{d})$ we have

Here, (a) follows from Holder’s inequality, (b) from (A.5), (c) from the fact that $\left\|{\bm{G}}{\bm{{\Sigma}}}\right\|\leq\left\|{\bm{G}}{\bm{{\Sigma}}}\right\|_{F}$ and $\left\|{\bm{G}}{\bm{{\Sigma}}}\right\|_{*}\leq\sqrt{r}\left\|{\bm{G}}{\bm{{\Sigma}}}\right\|_{F}$, (d) from Cauchy Schwarz, and (e) from (A.2), (f) from (A.5), and (g) from the fact that $\sqrt{3re^{-2r}}\leq\frac{3}{4}$. The above chain of inequalities thus allow us to conclude that

A.3 Proof of Lemma 9.15

holds with probability at least $1-e^{-2dr}$.

We next turn our attention to the lower bound. Given the random nature of the matrices $\bm{X}_{i}$, ${\text{mat}\left({\cal{J}}(\bm{\Theta})^{T}\bm{v}\right)}$ defines a random process $\bm{\Gamma}_{\bm{v}}$ indexed by $\bm{v}$ which can be rewritten in the form

holds with probability at least $1-e^{-\gamma dr}-e^{-2dr}\geq 1-2e^{-\gamma dr}$.

A.4 Proof of Lemma 9.16

holds with probability at least $1-e^{-2dr}$. Using Lemma 9.15,

holds with probability at least $1-2e^{-\gamma dr}$. Combining the latter two bounds, using $\bm{\Theta}\in{\cal{D}}$ and definition of ${\cal{D}}$, and applying the triangle inequality we conclude that

holds with probability at least $1-3e^{-\gamma dr}$. Using the fact that $\frac{\left\|\bm{A}\right\|_{*}}{\sqrt{r}}\leq\left\|\bm{A}\right\|_{F}\leq\left\|\bm{A}\right\|_{*}$ we thus have