Accelerating SGD using flexible variance reduction on large-scale datasets

Abstract

Stochastic gradient descent (SGD) is a popular optimization method widely used in machine learning, while the variance of gradient estimation leads to slow convergence. To accelerate the speed, many variance reduction methods have been proposed. However, most of these methods require additional memory cost or computational burden on full gradient, which results in low efficiency or even unavailable while applied to real-world applications with large-scale datasets. To handle this issue, we propose a new flexible variance reduction method for SGD, named FVR-SGD, which can reduce memory overhead and speed up the convergence using flexible subset without extra computation. We analyze the details of convergence property for our method, and linear convergence rate can be guaranteed while using flexible variance reduction. Some efficient variants for distributed environment and deep neural networks are also proposed in this paper. Several numerical experiments are conducted on a genre of real-world large-scale datasets. The experimental results demonstrated that FVR-SGD outperforms the currently popular algorithms. Specifically, the proposed method can achieve up to 40% reduction in the training time to solve the optimization problem of logistic regression, SVM and neural networks.

Appendix

Theorem 1

Suppose that$w_*$ is optimum and$\nabla F(w_*)=0$, $\phi ^l_m$ denotes the accumulatedw inm-th epoch and $\tilde{G}_m = \sum _{l=1}^{2n} F(\phi _m^l)/2n$, define

$$\begin{aligned} \alpha&> 0,\\ \beta&= 2n\left( 2 \eta + \frac{2\eta \alpha L^2}{\mu } - 4 \eta ^2 L\right) ,\\ \gamma&= (2n-2K)\left( 2 \eta + \frac{2\eta \alpha L^2}{\mu } - 4 \eta ^2 L\right) \\&\quad + 2K\left( 8 \eta ^2 L + \frac{16 \eta ^2 L^2}{\mu } + \frac{2 \eta L^2}{\mu }\right) ,\\ \delta&= \max \left( 1-\mu \eta +\eta , \frac{\gamma }{\beta }\right) , \end{aligned}$$

there exists$\beta > 0, 0<\delta <1$that makes:

$$\begin{aligned}&{\mathbb {E}} \parallel w_{m+2}^{1} - w_* \parallel ^2 + \beta {\mathbb {E}}(\tilde{G}_{m+1}-F(w_*)) \\&\quad \le \delta [{\mathbb {E}} \parallel w_{m+1}^{1} - w_* \parallel ^2 + \beta {\mathbb {E}}(\tilde{G}_{m}-F(w_*)], \end{aligned}$$

which means our method converges linearly.

Proof

Given any i, suppose

$$\begin{aligned} g_i(w)=f_i(w)-f_i(w_*)-\nabla f_i(w_*)^\top (w-w_*), \end{aligned}$$

we know that $g_i(w)$ is convex and $\nabla g_i(w_*)=0$, so $g_i(w_*)=\min \limits _\eta g_i(w)$. Based on the continuity of $g_i(w)$, we can obtain

$$\begin{aligned} 0&=g_i(w_*) \le \min \limits _\eta [g_i(w-\eta \nabla g_i(w))] \\&\le \min \limits _\eta [g_i(w)-\eta \Vert \nabla g_i(w)\Vert _2^2 + 0.5L\eta ^2 \Vert \nabla g_i(w)\Vert _2^2 ]\\&=g_i(w)- \frac{1}{2L}\Vert \nabla g_i(w)\Vert _2^2. \end{aligned}$$

That is,

$$\begin{aligned}&\Vert f_i(w)-f_i(w_*)\Vert _2^2 \le 2L[f_i(w)-f_i(w_*)\\&\quad -\nabla f_i(w_*)^\top (w-w_*)]. \end{aligned}$$

Taking the expectation of i and using the fact that ${\mathbb {E}} [\nabla f_i(w_*)]=0$, we obtain

$$\begin{aligned} {\mathbb {E}} \parallel \nabla f_i(w) - \nabla f_i(w_*) \parallel ^2 \le 2L (F(w)-F(w_*)), \end{aligned}$$

(17)

Inequality (18) can be obtained from the continuity (10) and convexity (11) of F.

$$\begin{aligned} \parallel \nabla f_i(w) - \nabla f_i(w_*) \parallel ^2&\le L^2\parallel w - w_* \parallel ^2 \nonumber \\&\le \frac{2L^2}{\mu } (F(w) - F(w_*)). \end{aligned}$$

(18)

Now we can take expectation of $v^k_{m+1}$ with respect to $i_k$ and obtain

$$\begin{aligned}&{\mathbb {E}} \parallel v^k_{m+1} \parallel ^2 = {\mathbb {E}} \parallel \nabla f_{i_k}(w^{k}_{m+1}) - \nabla f_{i_k}(\tilde{w}_m) + \tilde{g}_m \parallel ^2\nonumber \\&\quad \le 2{\mathbb {E}} \parallel \nabla f_{i_k}(w^{k}_{m+1}) - \nabla f_{i_k}(w_*) \parallel ^2 \nonumber \\&\qquad + 2{\mathbb {E}} \parallel \nabla f_{i_k}(\tilde{w}_m) - \nabla f_{i_k}(w_*) - \tilde{g}_m \parallel ^2 \nonumber \\&\quad \le 2{\mathbb {E}} \parallel \nabla f_{i_k}(w^{k}_{m+1}) - \nabla f_{i_k}(w_*) \parallel ^2 \nonumber \\&\qquad + 4{\mathbb {E}} \parallel \nabla f_{i_k}(\tilde{w}_m) - \nabla f_{i_k}(w_*) \parallel ^2 \nonumber \\&\qquad + 8 \parallel \frac{1}{2n} \sum _{l=1}^{2n} [\nabla f_{i_l}(\phi _m^l) - \nabla f_{i_l}(w_*)] \parallel ^2 \nonumber \\&\qquad + 8 \parallel \frac{1}{2n} \sum _{l=1}^{2n} \nabla f_{i_l}(w_*) \parallel ^2 \nonumber \\&\quad \le 4L(F(w^{k}_{m+1})-F(w_*)) + 8L(F(\tilde{w}_m)-F(w_*)) \nonumber \\&\qquad + \frac{4}{n} \sum _{l=1}^{2n} \parallel \nabla f_{i_l}(\phi _m^l) - \nabla f_{i_l}(w_*) \parallel ^2\nonumber \\&\quad \le 4L(F(w^{k}_{m+1})-F(w_*)) + \left( 8L + \frac{16L^2}{\mu }\right) (\tilde{G}_m - F(w_*)). \end{aligned}$$

(19)

The first and second inequalities use $\parallel a + b \parallel ^2 \le 2\parallel a \parallel ^2 + 2\parallel b \parallel ^2$. The third inequality uses inequality (17) and $\parallel \sum _{i=1}^n a_n \parallel ^2 \le n\sum _{i=1}^n \parallel a_i \parallel ^2$. The forth inequality uses inequality (18) and $F(\tilde{w}_m)=F(\frac{1}{2n} \sum _{l=1}^{2n}\phi _m^l) \le \frac{1}{2n} \sum _{l=1}^{2n} F(\phi _m^l)$. Note that ${\mathbb {E}} v^k_{m+1} = \nabla F(w^k_{m+1}) - \nabla F(\tilde{w}_m) + \tilde{g}_m$. Conditioned on $w_{m+1}^{k-1}$, we can obtain the bound:

$$\begin{aligned}&{\mathbb {E}} \parallel w_{m+1}^{k} - w_* \parallel ^2 = {\mathbb {E}} \parallel w_{m+1}^{k-1} - \eta v^{k-1}_{m+1} - w_* \parallel ^2\nonumber \\&\quad = \parallel w_{m+1}^{k-1} - w_* \parallel ^2 - 2 \eta (w_{m+1}^{k-1}-w*)^\top {\mathbb {E}} v^{k-1}_{m+1} \nonumber \\&\qquad + \eta ^2 {\mathbb {E}} \parallel v^{k-1}_{m+1} \parallel ^2 \nonumber \\&\quad = \parallel w_{m+1}^{k-1} - w_* \parallel ^2 - 2 \eta (w_{m+1}^{k-1}-w_*)^\top \nabla F(w^{k-1}_{m+1}) \nonumber \\&\qquad - 2 \eta (w_{m+1}^{k-1}-w_*)^\top (\tilde{g}_m - \nabla F(\tilde{w}_m)) + \eta ^2 {\mathbb {E}} \parallel v^{k-1}_{m+1} \parallel ^2\nonumber \\&\quad \le \parallel w_{m+1}^{k-1} - w_* \parallel ^2 - 2 \eta (F(w_{m+1}^{k-1})-F(w_*)) \nonumber \\&\qquad - \mu \eta \parallel w_* - w_{m+1}^{k-1} \parallel ^2 \nonumber \\&\qquad - 2 \eta (w_{m+1}^{k-1}-w_*)^\top (\tilde{g}_m - \nabla F(\tilde{w}_m)) + \eta ^2 {\mathbb {E}} \parallel v^{k-1}_{m+1} \parallel ^2 \nonumber \\&\quad \le \parallel w_{m+1}^{k-1} - w_* \parallel ^2 - 2 \eta (F(w_{m+1}^{k-1})-F(w_*)) \nonumber \\&\qquad - \mu \eta \parallel w_{m+1}^{k-1} - w_* \parallel ^2 \nonumber \\&\qquad - 2 \eta (w_{m+1}^{k-1}-w_*)^\top (\tilde{g}_m - \nabla F(\tilde{w}_m)) \nonumber \\&\qquad + 4 \eta ^2 L (F(w^{k-1}_{m+1})-F(w_*)) \nonumber \\&\qquad + \eta ^2\left( 8L + \frac{16L^2}{\mu }\right) (\tilde{G}_m - F(w_*)) \nonumber \\&\quad = \parallel w_{m+1}^{k-1} - w_* \parallel ^2 - \mu \eta \parallel w_{m+1}^{k-1} - w_* \parallel ^2 \nonumber \\&\qquad - (2 \eta - 4 \eta ^2 L)(F(w^{k-1}_{m+1})-F(w_*))\nonumber \\&\qquad - 2 \eta (w_{m+1}^{k-1}-w_*)^\top (\tilde{g}_m - \nabla F(\tilde{w}_m)) \nonumber \\&\qquad + \eta ^2\left( 8L + \frac{16L^2}{\mu }\right) (\tilde{G}_m - F(w_*)). \end{aligned}$$

(20)

The first inequality uses inequality (11). The second inequality uses inequality (19). The inner product term can be bounded as follows:

$$\begin{aligned}&-2(w_{m+1}^{k-1}-w_*)^\top (\tilde{g}_m - \nabla F(\tilde{w}_m)) \nonumber \\&\quad \le 2\parallel (w_{m+1}^{k-1}-w_*)^\top \parallel \parallel \tilde{g}_m - \nabla F(\tilde{w}_m) \parallel \nonumber \\&\quad \le 2\parallel (w_{m+1}^{k-1}-w_*)^\top \parallel \left[ \frac{1}{2n} \sum _{l=1}^{2n} \parallel \nabla f_{i_l}(\phi _m^l)-\nabla F(\tilde{w}_m) \parallel \right] \nonumber \\&\quad \le 2\parallel (w_{m+1}^{k-1}-w_*)^\top \parallel \left[ \frac{L}{2n} \sum _{l=1}^{2n} \parallel \phi _m^l-\tilde{w}_m \parallel \right] \nonumber \\&\quad = \frac{1}{2n} \sum _{l=1}^{2n} [2\parallel (w_{m+1}^{k-1}-w_*)^\top \parallel L \parallel \phi _m^l-\tilde{w}_m \parallel ]\nonumber \\&\quad \le \parallel w_{m+1}^{k-1}-w_* \parallel ^2 + \frac{L^2}{2n} \sum _{l=1}^{2n} \parallel \phi _m^l-\tilde{w}_m \parallel ^2\nonumber \\&\quad \le \parallel w_{m+1}^{k-1}-w_* \parallel ^2 \nonumber \\&\qquad + \frac{L^2}{2n} \frac{2}{\mu } \sum _{l=1}^{2n}(F(\phi _m^l)-F(\tilde{w}_m)- \nabla F(\tilde{w}_m)(\phi _m^l-\tilde{w}_m))\nonumber \\&\quad = \parallel w_{m+1}^{k-1}-w_* \parallel ^2 \nonumber \\&\qquad + \frac{2L^2}{\mu } [((\tilde{G}_m-F(w_*))- (F(\tilde{w}_m)-F(w_*))]. \end{aligned}$$

(21)

The third inequality uses inequality (10), the fourth inequality uses Cauchy–Schwartz inequality, and the fifth inequality uses inequality (11). We set $F(\tilde{w}_m)-F(w_*)=\alpha (F(w_{m+1}^{k-1})-F(w_*))$, where $\alpha > 0$. Then we can rewrite (20) as:

$$\begin{aligned}&{\mathbb {E}} \parallel w_{m+1}^{k} - w_*\parallel ^2 \nonumber \\&\quad \le \parallel w_{m+1}^{k-1} - w_* \parallel ^2 - (\mu -1 )\eta \parallel w_{m+1}^{k-1} - w_* \parallel ^2 \nonumber \\&\qquad - \left( 2 \eta + \frac{2\eta \alpha L^2}{\mu } - 4 \eta ^2 L)(F(w^{k-1}_{m+1})-F(w_*)\right) \nonumber \\&\qquad + \left( 8 \eta ^2 L + \frac{16 \eta ^2 L^2}{\mu } + \frac{2 \eta L^2}{\mu }\right) (\tilde{G}_m - F(w_*)). \end{aligned}$$

(22)

By summing the previous inequalities over $k=1,\ldots ,2K$, taking expectation with all the history, we obtain

$$\begin{aligned}&{\mathbb {E}} \parallel w_{m+2}^{1} - w_* \parallel ^2 \nonumber \\&\quad \le {\mathbb {E}} \parallel w_{m+1}^{1} - w_* \parallel ^2 - (\mu -1 )\eta \sum _{i=1}^{2K} {\mathbb {E}} \parallel w_{m+1}^{k} - w_* \parallel ^2\nonumber \\&\qquad - \left( 2 \eta + \frac{2\eta \alpha L^2}{\mu } - 4 \eta ^2 L\right) \sum _{i=1}^{2K} {\mathbb {E}}( F(w^{i}_{m+1})-F(w_*))\nonumber \\&\qquad + 2K\left( 8 \eta ^2 L + \frac{16 \eta ^2 L^2}{\mu } + \frac{2 \eta L^2}{\mu }\right) {\mathbb {E}} (\tilde{G}_m - F(w_*))\nonumber \\&\quad \le {\mathbb {E}} \parallel w_{m+1}^{1} - w_* \parallel ^2 - (\mu -1 )\eta \sum _{i=1}^{2K} {\mathbb {E}} \parallel w_{m+1}^{k} - w_* \parallel ^2\nonumber \\&\qquad - 2n\left( 2 \eta + \frac{2\eta \alpha L^2}{\mu } - 4 \eta ^2 L\right) {\mathbb {E}}(\tilde{G}_{m+1}-F(w_*))\nonumber \\&\qquad + (2n-2K) \left( 2 \eta + \frac{2\eta \alpha L^2}{\mu } - 4 \eta ^2 L\right) {\mathbb {E}} (\tilde{G}_{m}-F(w_*))\nonumber \\&\qquad + 2K(8 \eta ^2 L + \frac{16 \eta ^2 L^2}{\mu } + \frac{2 \eta L^2}{\mu }) {\mathbb {E}} (\tilde{G}_m - F(w_*))\nonumber \\&\quad \le (1-\mu \eta +\eta ){\mathbb {E}} \parallel w_{m+1}^{1} - w_* \parallel ^2\nonumber \\&\qquad - 2n\left( 2 \eta + \frac{2\eta \alpha L^2}{\mu } - 4 \eta ^2 L\right) {\mathbb {E}}(\tilde{G}_{m+1}-F(w_*))\nonumber \\&\qquad + (2n-2K) \left( 2 \eta + \frac{2\eta \alpha L^2}{\mu } - 4 \eta ^2 L\right) {\mathbb {E}} (\tilde{G}_{m}-F(w_*))\nonumber \\&\qquad + 2K\left( 8 \eta ^2 L + \frac{16 \eta ^2 L^2}{\mu } + \frac{2 \eta L^2}{\mu }\right) {\mathbb {E}} (\tilde{G}_m - F(w_*)). \end{aligned}$$

(23)

The first inequality is summing the previous inequality over 2K at $m+1$ epoch. Let $\tilde{G}^{-j}_{m+1}$ denote the $\tilde{G}_{m+1}$ removing the selected subset j at $(m+1)$-epoch, which means $\tilde{G}^{-j}_{m+1} = \frac{1}{2n-2K} \sum _{l=1}^{2n-2K} F(\phi _m^l)$. The second inequality uses

$$\begin{aligned}&-\sum _{i=1}^{2K} {\mathbb {E}} (F(w^{i}_{m+1})-F(w_*)) \\&\quad = - {\mathbb {E}}\left[ \sum _{i=1}^{2K} F(w^{i}_{m+1}) + \tilde{G}^{-j} - 2nF(w_*) \right. \\&\qquad \left. - (\tilde{G}^{-j}-(2n-2K)F(w_*))\right] \\&\quad = -2n{\mathbb {E}}(\tilde{G}_{m+1}-F(w_*)) + (2n-2K){\mathbb {E}}(\tilde{G}_{m}-F(w_*)). \end{aligned}$$

The third inequality uses

$$\begin{aligned} -\sum _{i=1}^{2K} {\mathbb {E}} \parallel w_{m+1}^{k} - w_* \parallel ^2 \le -{\mathbb {E}} \parallel w_{m+1}^{1} - w_* \parallel ^2. \end{aligned}$$

Rewriting (23), we can get:

$$\begin{aligned}&{\mathbb {E}} \parallel w_{m+2}^{1} - w_* \parallel ^2 + 2n\left( 2 \eta + \frac{2\eta \alpha L^2}{\mu } - 4 \eta ^2 L\right) {\mathbb {E}}(\tilde{G}_{m+1}-F(w_*))\nonumber \\&\quad \le (1-\mu \eta +\eta ){\mathbb {E}} \parallel w_{m+1}^{1} - w_* \parallel ^2\nonumber \\&\qquad + (2n-2K)\left( 2 \eta + \frac{2\eta \alpha L^2}{\mu } - 4 \eta ^2 L\right) {\mathbb {E}}(\tilde{G}_{m}-F(w_*))\nonumber \\&\qquad + 2K\left( 8 \eta ^2 L + \frac{16 \eta ^2 L^2}{\mu } + \frac{2 \eta L^2}{\mu }\right) {\mathbb {E}} (\tilde{G}_m - F(w_*)). \end{aligned}$$

(24)

So we can get the bound:

$$\begin{aligned}&{\mathbb {E}} \parallel w_{m+2}^{1} - w_* \parallel ^2 + \beta {\mathbb {E}}(\tilde{G}_{m+1}-F(w_*)) \nonumber \\&\quad \le \delta [{\mathbb {E}} \parallel w_{m+1}^{1} - w_* \parallel ^2 + \beta {\mathbb {E}}(\tilde{G}_{m}-F(w_*)], \end{aligned}$$

(25)

When $1-\mu \eta +\eta < 1$ and $\beta > 0$ and $\gamma > 0$ and $\frac{\gamma }{\beta } < 1$, $\delta$ can be limited into (0,1). So Theorem 1 needs to satisfy the following conditions:

$$\begin{aligned}&0<1-\mu \eta +\eta < 1, \\&\quad 2 \eta + \frac{2\eta \alpha L^2}{\mu } - 4 \eta ^2 L> 0,\\&\quad \beta - \gamma = 2\eta K\left( 2 + \frac{2 \alpha L^2}{\mu } - 12 \eta L - \frac{16 \eta L^2}{\mu } - \frac{2 L^2}{\mu }\right) >0, \end{aligned}$$

that is

$$\begin{aligned}&\mu > 1,\\&\quad \eta < \min \left( \frac{1}{\mu -1},\frac{\alpha L^2 + \mu }{2 \mu L},\frac{\mu + \alpha L^2 - L^2}{12\mu L + 16L^2}\right) . \end{aligned}$$

Then the results immediately follows. $\square$