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:  This paper proposes an alternative to stochastic line search which is based on forwarding step model building which corrects the direction of move and its magnitude at the same time. In its proposed algorithm it first checks if the given step size satisfies the stochastic line search. If yes then just use the step size and do the SGD update. Otherwise, it builds linear models around two points and combines these two models and minimizes this new model and this becomes the new iterate value.