A Non-asymptotic comparison of SVRG and SGD: tradeoffs between compute and speed

1 INTRODUCTION . Many large-scale machine learning problems , especially in deep learning , are formulated as minimizing the sum of loss functions on millions of training examples ( Krizhevsky et al. , 2012 ; Devlin et al. , 2018 ) . Computing exact gradient over the entire training set is intractable for these problems . Instead of using full batch gradients , the variants of stochastic gradient descent ( SGD ) ( Robbins & Monro , 1951 ; Zhang , 2004 ; Bottou , 2010 ; Sutskever et al. , 2013 ; Duchi et al. , 2011 ; Kingma & Ba , 2014 ) evaluate noisy gradient estimates from small mini-batches of randomly sampled training points at each iteration . The mini-batch size is often independent of the training set size , which allows SGD to immediately adapt the model parameters before going through the entire training set . Despite its simplicity , SGD works very well , even in the non-convex non-smooth deep learning problems ( He et al. , 2016 ; Vaswani et al. , 2017 ) . However , the optimization performance of the stochastic algorithm near local optima is significantly limited by the mini-batch sampling noise , controlled by the learning rate and the mini-batch size . The sampling variance and the slow convergence of SGD have been studied extensively in the past ( Chen et al. , 2016 ; Li et al. , 2017 ; Toulis & Airoldi , 2017 ) . To ensure convergence , machine learning practitioners have to either increase the mini-batch size or decrease the learning rate toward the end of the training ( Smith et al. , 2017 ; Ge et al. , 2019 ) . Recently , several clever variance reduction methods ( Roux et al. , 2012 ; Defazio et al. , 2014 ; Wang et al. , 2013 ; Johnson & Zhang , 2013 ) were proposed to alleviate the noisy gradient problem by using control-variates to achieve unbiased and lower-variance gradient estimators . In particular , the variants of Stochastic Variance Reduced Gradient ( SVRG ) ( Johnson & Zhang , 2013 ) , k-SVRG ( Raj & Stich , 2018 ) , L-SVRG ( Kovalev et al. , 2019 ) and Free-SVRG ( Sebbouh et al. , 2019 ) construct control-variates from previous staled snapshot model parameters . These methods enjoy a superior asymptotic performance in convex optimization compared to the standard SGD . The control-variate techniques are shown to improve the convergence rate of SGD from a sub-linear to a linear convergence rate . These variance reduction methods can also be combined with momentum ( Allen-Zhu , 2017 ) and preconditioning methods ( Moritz et al. , 2016 ) to obtain faster convergence . Despite their strong theoretical guarantees , SVRG-like algorithms have seen limited success in training deep learning models ( Defazio & Bottou , 2018 ) . Traditional results from stochastic optimization focus on the asymptotic analysis , but in practice , most of deep neural networks are only trained for hundreds of epochs due to the high computational cost . To address the gap between the asymptotic benefit of SVRG and the practical computational budget of training deep learning models , we provide a non-asymptotic study on the SVRG algorithms under a noisy least squares regression model . Although optimizing least squares regression is a basic problem , it has been shown to characterize the learning dynamics of many realistic deep learning models ( Zhang et al. , 2019 ; Lee et al. , 2019 ) . Recent works suggest that neural network learning behaves very differently in the underparameterized regime vs the overparameterized regime ( Ma et al. , 2018 ; Vaswani et al. , 2019 ) , characterized by whether the learnt model can achieve zero expected loss . We account for both training regimes in the analysis by assuming a linear target function and noisy labels . In the presence of label noise , the loss is lower bounded by the label variance . In the absence of the noise , the linear predictor can fit each training example perfectly . We summarize the main contributions as follows : • We show the exact expected loss of SVRG and SGD along an optimization trajectory as a function of iterations and computational cost . • Our non-asymptotic analysis provides an insightful comparison of SGD and SVRG by considering their computational cost and learning rate schedule . We discuss the trade-offs between the total computational cost , i.e . the total number of back-propagations performed , and convergence performance . • We consider two different training regimes with and without label noise . Under noisy labels , the analysis suggests SGD only outperforms SVRG under a mild total computational cost . However , SGD always exhibits a faster convergence compared to SVRG when there is no label noise . • Numerical experiments validate our theoretical predictions on both MNIST and CIFAR10 using various neural network architectures . In particular , we found the comparison of the convergence speed of SGD to that of SVRG in underparameterized neural networks closely matches with our noisy least squares model prediction . Whereas , the effect of overparameterization is captured by the regression model without label noise . 1.1 RELATED WORKS . Stochastic variance reduction methods consider minimizing a finite-sum of a collection of functions using SGD . In case we use SGD to minimize these objective functions , the stochasticity comes from the randomness in sampling a function in each optimization step . Due to the induced noise , SGD can only converge using decaying step sizes with sub-linear convergence rate . Methods such as SAG ( Roux et al. , 2012 ) , SVRG ( Johnson & Zhang , 2013 ) , and SAGA ( Defazio et al. , 2014 ) , are able to recover linear convergence rate of full-batch gradient descent with the asymptotic cost comparable to SGD . SAG and SAGA achieve this improvement at the substantial cost of storing the most recent gradient of each individual function . In contrast , SVRG spends extra computation at snapshot intervals by evaluating the full-batch gradient . Theoretical results such as Gazagnadou et al . ( 2019 ) show that under certain smoothness conditions , we can use larger step sizes with stochastic variance reduction methods than is allowed for SGD and hence achieve even faster convergence . In situations where we know the smoothness constant of functions , there are results on the optimal mini-batch size and the optimal step size given the inner loop size ( Sebbouh et al. , 2019 ) . Applying variance reduction methods in deep learning has been studied recently ( Defazio & Bottou , 2018 ) . The authors conjectured the ineffectiveness is caused by various elements commonly used in deep learning such as data augmentation , batch normalization and dropout . Such elements can potentially decrease the smoothness and make the stored gradients become stale quickly . The proposed solution is to either remove these elements or update the gradients more frequently than is practical . Dynamics of SGD and quadratic models Our main analysis tool is very closely related to recent work studying the dynamics of gradient-based stochastic methods . Wu et al . ( 2018 ) derived the dynamics of stochastic gradient descent with momentum on a noisy quadratic model ( Schaul et al. , 2013 ) , showing the problem of short horizon bias . In ( Zhang et al. , 2019 ) , the authors showed the same noisy quadratic model captures many of the essential characteristic of realistic neural networks training . Their noisy quadratic model successfully predicts the effectiveness of momentum , preconditioning and learning rate choices in training ResNets and Transformers . However , these previous quadratic models assume a constant variance in the gradient that is independent of the current parameters and the loss function . It makes them inadequate for analyzing the stochastic variance reduction methods , as SVRG can trivially achieve zero variance under the constant gradient noise . Instead , we adopted a noisy least-squares regression formulation by considering both the mini-batch sampling noise and the label noise . There are also recent works that derived the risk of SGD , for least-squares regression models using the bias-variance decomposition of the risk ( Belkin et al. , 2018 ; Hastie et al. , 2019 ) . We use a similar decomposition in our analysis . In contrast to the asymptotic analysis in these works , we compare SGD to SVRG along the optimization trajectory for any finite-time horizon under limited computation cost , not just the convergence points of those algorithms . Underparameterization vs overparameterization . Many of the state-of-the-art deep learning models are overparameterized deep neural networks with more parameters than the number of training examples . Even though these models are able to overfit to the data , when trained using SGD , they generalize well ( Zhang et al. , 2017 ) . As suggested in recent work , underparameterized and overparameterized regimes have different behaviours ( Ma et al. , 2018 ; Vaswani et al. , 2019 ; Schmidt & Roux , 2013 ) . Given the infinite width and a proper weight initialization , the learning dynamics of a neural network can be well-approximated by a linear model via the neural tangent kernel ( NTK ) ( Jacot et al. , 2018 ; Chizat & Bach , 2018 ) . In NTK regime , neural networks are known to achieve global convergence by memorizing every training example . On the other hand , previous convergence results for SVRG have been obtained in stochastic convex optimization problems that are similar to that of an underparameterized model ( Roux et al. , 2012 ; Johnson & Zhang , 2013 ) . Our proposed noisy least-squares regression analysis captures both the underparameterization and overparameterization behavior by considering the presence or the absence of the label noise . 2 PRELIMINARY . 2.1 NOTATIONS . We will primarily focus on comparing the minibatch version of two methods , SGD and SVRG ( Johnson & Zhang , 2013 ) . Denote Li as the loss on ith data point . The SGD update is written as , θ ( t+1 ) = θ ( t ) − α ( t ) ĝ ( t ) , ( 1 ) where ĝ ( t ) = 1b ∑b i ∇θ ( t ) Li is the minibatch gradient , t is the training iteration , and α ( t ) is the learning rate . The SVRG algorithm is an inner-outer loop algorithm proposed to reduce the variance of the gradient caused by the minibatch sampling . In the outer loop , for every T steps , we evaluate a large batch gradient ḡ = 1N ∑N i ∇θ ( mT ) Li , where N b , and m is the outer loop index , and we store the parameters θ ( mT ) . In the inner loop , the update rule of the parameters is given by , θ ( mT+t+1 ) = θ ( mT+t ) − α ( t ) ( ĝ ( mT+t ) − g̃ ( mT+t ) + ḡ ) ( 2 ) where ĝ ( mT+t ) = 1b ∑b i ∇θ ( mT+t ) Li is the current gradient of the mini-batch and g̃ ( mT+t ) = 1 b ∑b i ∇θ ( mT ) Li is the old gradient . Note that in our analysis , the reference point is chosen to be the last iterate of previous outer loop θ ( mT ) , recommended as a practical implementation of the algorithm by the original SVRG paper Johnson & Zhang ( 2013 ) . 2.2 THE NOISY LEAST SQUARES REGRESSION MODEL . We now define the noisy least squares regression model ( Schaul et al. , 2013 ; Wu et al. , 2018 ) . In this setting , the input data is d-dimensional , and the output label is generated by a linear teacher model with additive noise , ( xi , i ) ∼ Px × P ; yi = x > i θ∗ + i , where E [ xi ] = µ ∈ Rd and Cov ( xi ) = Σ , E [ i ] = 0 , Var ( i ) = σ2y . We assume WLOG θ∗ = 0 . We also assume the data covariance matrix Σ is diagonal . This is an assumption adopted in many previous analysis and it is also a practical assumption as we often apply whitening to pre-process the training data . We would like to train a student model θ that minimizes the squared loss over the data distribution : min θ L ( θ ) : = E [ 1 2 ( x > i θ − yi ) 2 ] . ( 3 ) At each iteration , the optimizer can query an arbitrary number of data points { xi , yi } i sampled from data distribution . The SGD method uses b data points to form a minibatch gradient : ĝ ( t ) = 1 b b∑ i ( xix > i θ ( t ) − xi i ) = XbX > b θ ( t ) − 1√ b Xb b , ( 4 ) where Xb = 1√b [ x1 ; x2 ; · · · ; xb ] ∈ R d×b , and the noise vector b = [ 1 ; 2 ; · · · ; b ] > ∈ Rb . SVRG on the other hand , queries for N data points every T steps to form a large batch gradient ḡ = XNX > Nθ ( mT ) − 1√ N XN N , where XN and N are defined similarly . At each inner loop step , it further queries for another b data points , to form the update in Eq . 2 . Lastly , note that the expected loss can be written as a function of the second moment of the iterate , L ( θ ( t ) ) = 1 2 E [ ( x > i θ ( t ) − i ) 2 ] = 1 2 ( tr ( ΣE [ θ ( t ) θ ( t ) > ] ) + σ2y ) . Hence for the following analysis we mainly focus on deriving the dynamics of the second moment E [ θ ( t ) θ ( t ) > ] , denoted as A ( θ ( t ) ) . When Σ is diagonal , the loss can further be reduced to 1 2 diag ( Σ ) > diag ( E [ θ ( t ) θ ( t ) > ] ) + 12σ 2 y . We denote diag ( E [ θ ( t ) θ ( t ) > ] ) by m ( θ ( t ) ) .

Summary: This paper compares SGD and SVRG (as a representative variance reduced method) to explore tradeoffs. Although the computational complexity vs overall convergence performance tradeoff is well-known at this point, an interesting new perspective is the comparison in regions of interpolation (where SGD gradient variance will diminish on its own) and label noise (which propogates more seriously in SGD vs SVRG). The analysis is done on a simple linear  model with regression, with some experiments on simulations, MNIST, and CIFAR.