Bolstering Stochastic Gradient Descent with Model Building

1 INTRODUCTION . Stochastic gradient descent ( SGD ) is a popular optimization algorithm for machine learning problems which can achieve very fast convergence rates when its stepsize and its scheduling are tuned well for the specific application at hand . This tuning procedure can take up to thousands of CPU/GPU days resulting in big energy costs ( Asi & Duchi , 2019 ) . A number of researchers have studied adaptive strategies for improving the direction and the stepsize choices of the stochastic gradient descent algorithm . Adaptive sample size selection ideas ( Byrd et al. , 2012 ; Balles et al. , 2016 ; Bollapragada et al. , 2018 ) improve the direction by reducing its variance around the negative gradient of the empirical loss function , while stochastic quasi-Newton algorithms ( Byrd et al. , 2016 ; Wang et al. , 2017 ) provide adaptive preconditioning . Recently , several stochastic line search approaches have been proposed . Not surprisingly , some of these work cover sample size selection as a component of the proposed line search algorithms ( Balles et al. , 2016 ; Paquette & Scheinberg , 2020 ) . The Stochastic Model Building ( SMB ) algorithm proposed in this paper is not designed as a stochastic quasi-Newton algorithm in the sense explained by Bottou et al . ( 2018 ) . However , it still produces a scaling matrix in the process of generating trial points , and its overall step at each outer iteration can be written in the form of matrix-vector multiplication . Unlike the algorithms proposed by Mokhtari & Ribeiro ( 2014 ) and Schraudolph et al . ( 2007 ) , we have no accumulation of curvature pairs throughout several iterations . Since there is no memory carried from earlier iterations , the scaling matrices in individual past iterations are based only on the data samples employed in those iterations . In other words , the scaling matrix and the incumbent random gradient vector are dependent . Vaswani et al . ( 2019 ) apply a deterministic globalization procedure on mini-batch loss functions . That is , the same sample is used in all function and gradient evaluations needed to apply the line search procedure at a given iteration . However , unlike our case , they employ a standard line search procedure that does not alter the search direction . They establish convergence guarantees for the empirical loss function under the interpolation assumption , which requires each component loss function to have zero gradient at a minimizer of the empirical loss . Mutschler & Zell ( 2020 ) assume that the optimal learning rate along the negative batch gradient is a good estimator for the optimal learning rate with respect to the empirical loss along the same direction . They test validity of this assumption empirically on deep neural networks ( DNNs ) . Rather than making such strong assumptions , we stick to the general theory for stochastic quasi-Newton methods . Other work follow a different approach to translate deterministic line search procedures into a stochastic setting , and they do not employ fixed samples . In Mahsereci & Hennig ( 2017 ) , a probabilistic model along the search direction is constructed via techniques from Bayesian optimization . Learning rates are chosen to maximize the expected improvement with respect to this model and the probability of satisfying Wolfe conditions . Paquette & Scheinberg ( 2020 ) suggest an algorithm closer to the deterministic counterpart where the convergence is based on the requirement that the stochastic function and gradient evaluations approximate their true values with a high enough probability . With our current work , we make the following contributions . We use a model building strategy for adjusting the stepsize and the direction of a stochastic gradient vector . This approach also permits us to work on subsets of parameters . This feature makes our model steps not only adaptive , but also suitable to incorporate into the existing implementations of deep learning networks . Our method changes the direction of the step as well as its size which separates our approach from the backtracking line search algorithms . It also incorporates the most recent curvature information from the current point . This is in contrast with the stochastic quasi-Newton methods which use the information from the previous steps . Capitalizing our discussion on the independence of the sample batches , we also give a convergence analysis for SMB . Finally , we illustrate the computational performance of our method with a set of numerical experiments and compare the results against those obtained with other well-known methods . 2 STOCHASTIC MODEL BUILDING . We introduce a new stochastic unconstrained optimization algorithm in order to approximately solve problems of the form min x2 < n f ( x ) = E [ F ( x , ⇠ ) ] , ( 1 ) where F : Rn ⇥ Rd ! R is continuously differentiable and possibly nonconvex , ⇠ 2 Rd denotes a random variable , and E [ . ] denotes the expectation taken with respect to ⇠ . We assume the existence of a stochastic first-order oracle which outputs a stochastic gradient g ( x , ⇠ ) of f for a given x . A common approach to tackle ( 1 ) is to solve the empirical risk problem min x2 < n f ( x ) = 1 N NX i=1 fi ( x ) , ( 2 ) where fi : Rn ! R is the loss function corresponding to the ith data sample , and N denotes the data sample size which can be very large in modern applications . As an alternative approach to line search for SGD , we propose a stochastic model building strategy inspired by the work of Öztoprak & Birbil ( 2018 ) . Unlike core SGD methods , our approach aims at including a curvature information that adjusts not only the stepsize but also the search direction . Öztoprak & Birbil ( 2018 ) consider only the deterministic setting and they apply the model building strategy repetitively until a sufficient decent is achieved . In our stochastic setting , however , we have observed experimentally that multiple model steps does not benefit much to the performance , and its cost to the run time can be extremely high in deep learning problems . Therefore , if the sufficient decent is not achieved by the stochastic gradient step , then we construct only one model to adjust the size and the direction of the step . Conventional stochastic quasi-Newton methods adjust the gradient direction by a scaling matrix that is constructed by the information from the previous steps . Our model building approach , however , uses the most recent curvature information around the latest iteration . In the popular deep learning model implementations , model parameters come in groups and updates are applied to each parameter group separately . Therefore , we also propose to build a model for each parameter group separately making the step lengths adaptive . The proposed iterative algorithm SMB works as follows : At step k , given the iterate xk , we calculate the stochastic function value fk = f ( xk , ⇠k ) and the mini-batch stochastic gradient gk = 1 mk Pmk i=1 g ( xk , ⇠k , i ) at xk , where mk is the batch size and ⇠k = ( ⇠k,1 , . . . , ⇠k , mk ) is the realization of the random vector ⇠ . Then , we apply the SGD update to calculate the trial step stk = ↵kgk , where { ↵k } k is a sequence of learning rates . With this trial step , we also calculate the function and gradient values f tk = f ( x t k , ⇠k ) and g t k = g ( x t k , ⇠k ) at x t k = xk + s t k. Then , we check the stochastic Armijo condition f tk fk c ↵kkgkk 2 , ( 3 ) where c > 0 is a hyper-parameter . If the condition is satisfied and we achieve sufficient decrease , then we set xk+1 = xtk as the next step . If the Armijo condition is not satisfied , then we build a quadratic model using the linear models at the points xk , p and xtk , p for each parameter group p and find the step sk , p to reach its minimum point . Here , xk , p and xtk , p denote respectively the coordinates of xk and xtk that corresponds to the parameter group p. We calculate the next iterate xk+1 = xk + sk , where sk = ( sk , p1 , . . . , sk , pn ) and n is the number of parameter groups , and proceed to the next step with xk+1 . This model step , if needed , requires extra mini-batch function and gradient evaluations ( forward and backward pass in deep neural networks ) . For each parameter group p , the quadratic model is built by combining the linear models at xk , p and xtk , p , given by l0k , p ( s ) : = fk , p + g > k , ps and l t k , p ( s s t k , p ) : = f t k , p + ( g t k , p ) > ( s stk , p ) , respectively . Then , the quadratic model becomes mtk , p ( s ) : = ↵ 0 k , p ( s ) l 0 k , p ( s ) + ↵ t k , p ( s ) l t k , p ( s s t k , p ) , where ↵0k , p ( s ) = ( s stk , p ) > ( stk , p ) ( stk , p ) > ( stk , p ) and ↵tk , p ( s ) = s > stk , p ( stk , p ) > stk , p . The constraint ksk2 + ks stk , pk 2 kstk , pk 2 , is also imposed so that the minimum is attained in the region bounded by xk , p and xtk , p . This constraint acts like a trust region . Figure 1 shows the steps of this construction . In this work , we solve a relaxation of this constrained model as explained in ( Öztoprak & Birbil , 2018 , Section 2.2 ) . The minimum value of the relaxed model is attained at the point xk , p + sk , p with sk , p = cg , p ( ) gk , p + cy , p ( ) yk , p + cs , p ( ) s t k , p , ( 4 ) where yk , p : = gtk , p gk , p . Here , the coefficients are given as cg , p ( ) = kstk , pk 2 2 , cy , p ( ) = kstk , pk 2 2 ✓ [ ( y > k , ps t k , p + 2 ) ( s t k , p ) > gk , p + ks t k , pk 2y > k , pgk , p ] , cs , p ( ) = kstk , pk 2 2 ✓ [ ( y > k , ps t k , p + 2 ) y > k , pgk , p + kyk , pk 2 ( stk , p ) > gk , p ] , with ✓ = ⇣ y > k , ps t k , p + 2 ⌘2 kstk , pk2kyk , pk2 and = 1 2 ✓ kstk , pk ✓ kyk , pk+ 1 ⌘ kgk , pk ◆ y > k , pstk , p ◆ , ( 5 ) where 0 < ⌘ < 1 is a constant which controls the size of sk , p by imposing the condition ksk , pk ⌘kstk , pk . Then , the adaptive model step becomes sk = ( sk , p1 , . . . , sk , pn ) . We note that our construction in terms of different parameter groups lends itself to constructing a different model for each parameter subspace . We summarize the steps of SMB in Algorithm 1 . Line 5 shows the trial point , which is obtained with the standard stochastic gradient step . If this step satisfies the stochastic Armijo condition , then we proceed with the next iteration ( line 8 ) . Otherwise , we continue with bulding the models for each parameter group ( lines 10- 12 ) , and move to the next iteration with the model building step in line 13 . Algorithm 1 : SMB : Stochastic Model Building 1 Input : x1 2 Rn , stepsizes { ↵k } Tk=1 , mini-batch sizes { mk } Tk=1 , c > 0 , and ↵max satisfying ( 8 ) 2 for k = 1 , . . . , T do 3 fk = f ( xk , ⇠k ) , gk = 1mk Pmk i=1 g ( xk , ⇠k , i ) ; 4 stk = ↵kgk ; 5 xtk = xk + s t k ; 6 f tk = f ( x t k , ⇠k ) , g t k = 1 mk Pmk i=1 g ( x t k , ⇠k , i ) ; 7 if f tk fk c ↵kkgkk2 then 8 xk+1 = xtk ; 9 else 10 for p = 1 , . . . , n do 11 yk , p = gtk , p gk , p ; 12 sk , p = cg , p ( ) gk , p + cy , p ( ) yk , p + cs , p ( ) stk , p ; 13 xk+1 = xk + sk with sk = ( sk , p1 , . . . , sk , pn ) ;

Summary: The authors proposed a method called stochastic model building (SMB) that uses a combination of existing techniques to get faster convergence in stochastic non-convex optimization. In particular, they use a stochastic adaptation of the model-building globalization strategy from Oztoprak and Birbil (2018), in which the deterministic Armijo condition check is also computed using stochastic gradients. By re-writing their update as a preconditioned SGD step, they are able to bound the spectrum of the preconditioner. This allows them to obtain convergence results for smooth and non-convex objectives by directly invoking the analysis of Wang et al. (2017). Experiments demonstrate that with additional heuristics, their proposed method can outperform state-of-the-art optimizers on common deep learning benchmarks.