Stochastic gradient descent, weighted sampling, and the randomized Kaczmarz algorithm

Abstract

We obtain an improved finite-sample guarantee on the linear convergence of stochastic gradient descent for smooth and strongly convex objectives, improving from a quadratic dependence on the conditioning $(L/\mu )^2$ (where $L$ is a bound on the smoothness and $\mu$ on the strong convexity) to a linear dependence on $L/\mu$. Furthermore, we show how reweighting the sampling distribution (i.e. importance sampling) is necessary in order to further improve convergence, and obtain a linear dependence in the average smoothness, dominating previous results. We also discuss importance sampling for SGD more broadly and show how it can improve convergence also in other scenarios. Our results are based on a connection we make between SGD and the randomized Kaczmarz algorithm, which allows us to transfer ideas between the separate bodies of literature studying each of the two methods. In particular, we recast the randomized Kaczmarz algorithm as an instance of SGD, and apply our results to prove its exponential convergence, but to the solution of a weighted least squares problem rather than the original least squares problem. We then present a modified Kaczmarz algorithm with partially biased sampling which does converge to the original least squares solution with the same exponential convergence rate.

Appendix: Proofs

Our main results utilize an elementary fact about smooth functions with Lipschitz continuous gradient, called the co-coercivity of the gradient. We state the lemma and recall its proof for completeness.

1.1 The co-coercivity Lemma

Lemma 8.1

(Co-coercivity) For a smooth function $f$ whose gradient has Lipschitz constant $L$,

$$\begin{aligned} ||{\nabla f(\varvec{x}) - \nabla f(\varvec{y})} ||_2^2 \le L\left\langle \varvec{x}-\varvec{y}, \nabla f(\varvec{x}) - \nabla f(\varvec{y})\right\rangle . \end{aligned}$$

Proof

Since $\nabla f$ has Lipschitz constant $L$, if ${\varvec{x}}_{\star }$ is the minimizer of $f$, then (see e.g. [32], page 26)

(8.1)

Now define the convex functions

$$\begin{aligned} G(\varvec{z}) = f(\varvec{z}) - \left\langle \nabla f(\varvec{x}), \varvec{z}\right\rangle , \quad \text {and}\quad H(\varvec{z}) = f(\varvec{z}) - \left\langle \nabla f(\varvec{y}), \varvec{z}\right\rangle , \end{aligned}$$

and observe that both have Lipschitz constants $L$ and minimizers $\varvec{x}$ and $\varvec{y}$, respectively. Applying (8.1) to these functions therefore gives that

$$\begin{aligned} G(\varvec{x}) \le G(\varvec{y}) - \frac{1}{2L}||{\nabla G(\varvec{y})} ||_2^2, \quad \text {and} \quad H(\varvec{y}) \le H(\varvec{x}) - \frac{1}{2L}||{\nabla H(\varvec{y})} ||_2^2. \end{aligned}$$

By their definitions, this implies that

Adding these two inequalities and canceling terms yields the desired result. $\square$

1.2 Proof of Theorem 2.1

With the notation of Theorem 2.1, and where $i$ is the random index chosen at iteration $k$, and $w=w_{\lambda }$, we have

where we have employed Jensen’s inequality in the first inequality and the co-coercivity Lemma 8.1 in the final line. We next take an expectation with respect to the choice of $i$. By assumption, $i\sim {\mathcal {D}}$ such that $F(\varvec{x}) = \mathbb {E} f_i(\varvec{x})$ and $\sigma ^2 = {\mathbb {E}}\Vert \nabla f_i({\varvec{x}}_{\star })\Vert ^2$. Then ${\mathbb {E}}\nabla f_i(\varvec{x})=\nabla {\mathbf {F}}(\varvec{x})$, and we obtain:

$$\begin{aligned} {\mathbb {E}}||{{\varvec{x}}_{k+1}-{\varvec{x}}_{\star }} ||_2^2&\le ||{{\varvec{x}}_k- {\varvec{x}}_{\star }} ||_2^2 - 2\gamma \left\langle {\varvec{x}}_k-{\varvec{x}}_{\star }, \nabla {\mathbf {F}}({\varvec{x}}_k)\right\rangle \\&+ 2\gamma ^2 {\mathbb {E}}\left[ L_i \left\langle {\varvec{x}}_k-{\varvec{x}}_{\star }, \nabla f_i({\varvec{x}}_k) - \nabla f_i({\varvec{x}}_{\star })\right\rangle \right] + 2\gamma ^2 {\mathbb {E}}||{\nabla f_i({\varvec{x}}_{\star })} ||_2^2 \\&\le ||{{\varvec{x}}_k- {\varvec{x}}_{\star }} ||_2^2 - 2\gamma \left\langle {\varvec{x}}_k-{\varvec{x}}_{\star }, \nabla {\mathbf {F}}({\varvec{x}}_k)\right\rangle \\&\quad +\, 2\gamma ^2 \sup _i L_i {\mathbb {E}}\left\langle {\varvec{x}}_k-{\varvec{x}}_{\star }, \nabla f_i({\varvec{x}}_k) - \nabla f_i({\varvec{x}}_{\star })\right\rangle + 2\gamma ^2 {\mathbb {E}}||{\nabla f_i({\varvec{x}}_{\star })} ||_2^2 \\&= ||{{\varvec{x}}_k- {\varvec{x}}_{\star }} ||_2^2 - 2\gamma \left\langle {\varvec{x}}_k-{\varvec{x}}_{\star }, \nabla {\mathbf {F}}({\varvec{x}}_k)\right\rangle \\&\quad +\, 2\gamma ^2 \sup L \left\langle {\varvec{x}}_k-{\varvec{x}}_{\star }, \nabla {\mathbf {F}}({\varvec{x}}_k) - \nabla {\mathbf {F}}({\varvec{x}}_{\star })\right\rangle + 2\gamma ^2 \sigma ^2 \end{aligned}$$

We now utilize the strong convexity of $F(\varvec{x})$ and obtain that

$$\begin{aligned}&\le ||{{\varvec{x}}_k- {\varvec{x}}_{\star }} ||_2^2 - 2\gamma \mu (1 - \gamma \sup L) ||{{\varvec{x}}_k-{\varvec{x}}_{\star }} ||_2^2 + 2\gamma ^2 \sigma ^2 \\&= (1 - 2\gamma \mu (1 - \gamma \sup L)) ||{{\varvec{x}}_k-{\varvec{x}}_{\star }} ||_2^2 + 2\gamma ^2 \sigma ^2 \end{aligned}$$

when $\gamma \le \frac{1}{ \sup L}$. Recursively applying this bound over the first $k$ iterations yields the desired result,

$$\begin{aligned} {\mathbb {E}}||{{\varvec{x}}_k-{\varvec{x}}_{\star }} ||_2^2&\le \Big (1 - 2\gamma \mu (1 - \gamma \sup L)\Big )\Big )^k||{{\varvec{x}}_0- {\varvec{x}}_{\star }} ||_2^2\\&\quad +\, 2\sum _{j=0}^{k-1} \Big (1 - 2\gamma \mu (1 - \gamma \sup L)\Big )\Big )^j \gamma ^2\sigma ^2 \\&\le \Big (1 - 2\gamma \mu (1 - \gamma \sup L)\Big )\Big )^k ||{{\varvec{x}}_0- {\varvec{x}}_{\star }} ||_2^2 + \frac{\gamma \sigma ^2}{\mu \big ( 1 - \gamma \sup L \big )}. \end{aligned}$$