On the Regularization Effect of Stochastic Gradient Descent applied to Least Squares

In this paper, we consider the finite-dimensional linear inverse problem

where $A\in\mathbb{R}^{n\times n}$ is an invertible matrix, $x\in\mathbb{R}^{n}$ is the (unknown) signal of interest and $b$ is a given right-hand side. Throughout this paper, we will use $a_{1},\dots,a_{n}$ to denote the rows of $A$. Equivalently, we will try to solve the problem

Following Needell, Srebro & Ward [30], we can interpret this problem as

The Lipschitz constant of $f_{i}$ is $\|a_{i}\|_{\ell^{2}}$ which motivates the following basic form of stochastic gradient descent: pick one of the $n$ functions with likelihood proportional to the Lipschitz constant and then do a gradient descent on this much simpler function resulting in the update rule

This is also known as the Algebraic Reconstruction Technique (ART) in computer tomography [10, 14, 15, 26], the Projection onto Convex Sets Method [3, 5, 6, 9, 38] and the Randomized Kaczmarz method [2, 7, 8, 11, 13, 12, 20, 21, 22, 24, 25, 27, 28, 29, 30, 31, 33, 36, 39, 40, 41, 42, 43, 44]. Strohmer & Vershynin [41] showed that

where $\|A\|_{F}^{2}$ is the Frobenius norm. In practice, the algorithm often converges a lot faster initially and this was studied in [17, 18, 39]. In particular, [39] obtains an identity in terms of the behavior with regards to the singular values showing that singular vectors associated to large singular values are expected to undergo a more rapid decay. Motivated by this insight, we provide rigorous bounds that quantify this energy cascade from large singular values to small singular values by identifying an interesting inequality for SGD when applied to Least Squares.

We discuss a simple example that exemplifies the phenomenon that we are interested in. Let us take $A\in\mathbb{R}^{100\times 100}$ by picking each entry independently at random from $\mathcal{N}(0,1)$ and then normalizing the rows to $\|a_{i}\|=1$. The right-hand side is $b=(1,1,\dots,1)$ and we initialize with $x_{0}=\textbf{0}\in\mathbb{R}^{n}$.

Here, $\|A\|_{F}$ denotes the Frobenius norm. This shows that we expect $x_{k}-x$ to be indeed mainly a linear combination of singular vectors associated to small singular values since those are the ones undergoing the slowest decay. It also mirrors the bound obtained by Strohmer & Vershynin [41] since

While being interesting in itself, this identity by itself does not fully explain the behavior shown above: it is only in expectation with no control of the variance. Moreover, the inner product does initially undergo some fluctuations. Taking the same type of matrix as above, we see an example of such fluctuations in Fig. 2.

This type of question is well studied, we refer to Ali, Dobridan & Tibshirani [1], Defossez & Bach [4], Jain, Kakade, Kidambi, Netrapalli, Pillutla & Sidford [16], Neu & Rosasco [32], Oymak & Soltanolkotabi [34], Schmidt, Le Roux & Bach [37] and references therein. The connection of Stochastic Gradient Descent applied to Least Square Problems and the Randomized Kaczmarz Method has been pointed out by Needell, Srebro & Ward [30]. We also mention the papers by Jiao, Jin & Lu [17] and Jin & Lu [18] who studied a similar question and noted that there is energy transfer from large singular values to small singular values.

The main goal of this note is to provide a simple explanation for the rapid initial regularization: the expected decay of $\|A(x_{k}-x)\|$ under SGD can be bounded from above by a term involving $\left\|A^{T}A(x_{k}-x)\right\|_{2}$: this is the same term except that the matrix has been applied to the existing quantity one more time. This increases the norm of the underlying vector except when $A(x_{k}-x)$ is mainly the linear combination of singular vectors with small singular values. So while this is not the case, we actually inherit strong decay properties and this leads to the rapid initial regularization. We now make this precise.

The inequality also holds for $A\in\mathbb{R}^{m\times n}$ with $m\geq n$ as long as $Ax=b$ has a unique solution. We note that $\alpha\geq 1$ is usually quite small: for random matrices, we would expect $\alpha\sim 2$. This is also true for matrices discretizing PDEs via finite elements. The main point of the inequality is that the last term has an additional factor of $A^{T}$: we can rewrite it as

This shows that the presence of large singular vectors in $x_{k}-x$ forces large decay on $\|A(x_{k}-x)\|_{2}^{2}$. Conversely, once the algorithm has reached the plateau phase (see Fig. 1), the terms

are nearly comparable. Thus, this forces $x_{k}-x$ to be mostly orthogonal to most singular vectors corresponding to large singular values: that, however, shows that it is mainly comprised of small singular vectors and thus explains why $\|A(x_{k}-x)\|\ll\|x_{k}-x\|$ is possible in cases where $x_{k}$ is far away from $x$. In particular, this suggests why the method could be effective for the problem of finding a vector $x$ so that $Ax\approx b$. One way is to initialize stochastic gradient descent with $x_{0}=0$ and run it for a while – due to the difference in scales, second-order norms regularizing first-order norms, we observe that $Ax_{k}$ converges quite rapidly; whether it converges to something sufficiently close to $b$ for the purpose at hand, is a different question.

An interesting way to illustrate the result is in terms of partial differential equation. Suppose we try to solve $-\Delta u=f$ on some domain $\Omega\subset\mathbb{R}^{n}$. After a suitably discretization, this results in a discrete linear system $Lu=f$, where $L\in\mathbb{R}^{n\times n}$ is a discretization of the Laplacian $-\Delta$. By an abuse of notation, $u$ denotes a decent approximation of the continuous solution and $f$ a discretization of the continuous right-hand side. However, we also have more information: since $L$ discretizes the Laplacian, we expect that

Here, the first term correspond to the size of $u$ in the Sobolev space $\dot{H}^{1}$ while the second term is the size of $u$ in the Sobolev space $\dot{H}^{2}$. In fact, this is a common way to define discretized approximations of Sobolev space, also known as the spectral definition since they are defined in terms of the spectrum of $L$. Suppose now we compute a sequence of approximations $u_{k}$ via the method outlined above. Then Theorem 1 can be rephrased as

What is of great interest here is that the decay of the error in $\dot{H}^{1}$ is driven by decay of the error in $\dot{H}^{2}$ (which is usually larger).

We illustrate this with an example. Choosing $A\in\mathbb{R}^{500\times 500}$ at random (and then, for convenience, normalize the rows to $\|a_{i}\|=1$), we solve $Ax=(1,1,\dots,1)$ starting with a random initial vector $x_{0}$ where each entry is chosen independently from a standardized $\mathcal{N}(0,1)$ distribution. We consider both the evolution of $\|Ax_{k}-b\|_{2}^{2}$ across multiple runs as well as the size of

which is the term from our Theorem quantifying the expected decay at each step.

We see in Fig. 3. that over 3000 periods, the approximation decays roughly by a factor $\sim 800$ (with little variation across multiple runs). The bound in the Theorem implies an expected decay of $-0.23$ per time-step which, across 3000 time steps, leads to a total of roughly $\sim 696$ in decay.

This shows that the same phenomenon does indeed happen at all scales of ‘smoothness’. The applicability of the result is, naturally, depending on the growth of $\alpha_{\ell}$ in $\ell$, though, generically, one would not expect this to be badly behaved: a priori, there is no good reason to expect that the row of a matrix happens in any way to be the linear combination of singular vectors associated to large singular values – though, naturally, this can happen (for example, if $A$ has one very large entry on the diagonal).