Statistical Adaptive Stochastic Optimization

1 INTRODUCTION . We study adaptive stochastic optimization methods in the context of large-scale machine learning . Specifically , we consider the stochastic optimization problem minimize x∈Rp F ( x ) , Eξ [ fξ ( x ) ] , ( 1 ) where ξ is a random variable representing data sampled from some ( unknown ) probability distribution , x ∈ Rp represents the parameters of the machine learning model ( e.g. , the weight matrices in a neural network ) , and fξ is the loss function associated with data ξ . The objective function F is the expectation of fξ over the distribution of the data . Many stochastic optimization methods for solving problem ( 1 ) can be written in the form of xk+1 = xk − αkdk , ( 2 ) where dk is a ( stochastic ) search direction and αk > 0 is the step size or learning rate . In the classical stochastic gradient descent ( SGD ) method , dk = gk , ∇fξk ( xk ) , ( 3 ) where ξk is a training example ( or a mini-batch of examples ) randomly generated at iteration k. This method traces back to the seminal work of Robbins & Monro ( 1951 ) , and it has become very popular in machine learning ( e.g. , Bottou , 1998 ; Goodfellow et al. , 2016 ) . Many modifications of SGD aim to improve its theoretical and practical performance . For example , Gupal & Bazhenov ( 1972 ) studied a stochastic analog of the heavy-ball method ( Polyak , 1964 ) , where dk = ( 1− βk ) gk + βkdk−1 , ( 4 ) and βk ∈ [ 0 , 1 ) is the momentum coefficient . Sutskever et al . ( 2013 ) proposed to use a stochastic variant of Nesterov ’ s accelerated gradient method ( Nesterov , 2004 ) , where dk = ∇fξk ( xk − αkβkdk−1 ) + βkdk−1 . ( 5 ) Other recent variants include , e.g. , Jain et al . ( 2018 ) and Ma & Yarats ( 2019 ) . Theoretical conditions for the asymptotic convergence of SGD are well studied , and they mostly focus on polynomially decaying learning rates of the form αk = a/ ( b + k ) c for some a , b > 0 and 1/2 < c ≤ 1 ( e.g. , Robbins & Monro , 1951 ; Polyak & Juditsky , 1992 ) . Similar conditions for the stochastic heavy-ball methods are also established ( e.g. , Gupal & Bazhenov , 1972 ; Polyak , 1977 ) . However , these learning rate schedules still require significant hyperparameter tuning efforts in modern machine learning practice . Adaptive rules for adjusting the learning rate and other parameters on the fly have been developed in both the stochastic optimization literature ( e.g. , Kesten , 1958 ; Mirzoakhmedov & Uryasev , 1983 ; Ruszczyński & Syski , 1983 ; 1984 ; 1986a ; b ; Delyon & Juditsky , 1993 ) and by the machine learning community ( e.g. , Jacobs , 1988 ; Sutton , 1992 ; Schraudolph , 1999 ; Mahmood et al. , 2012 ; Baydin et al. , 2018 ) . More recently , adaptive algorithms that use diagonal scaling—replacing αk in ( 2 ) with an adjustable diagonal matrix—have become very popular ( Duchi et al. , 2011 ; Tieleman & Hinton , 2012 ; Kingma & Ba , 2014 , e.g. , ) . Despite these advances , costly hand-tuning efforts are still needed to obtain good ( generalization ) performance in practice ( Wilson et al. , 2017 ) . In this paper , we focus on statistical methods for automatically scheduling the learning rate . This line of work goes back to Kesten ( 1958 ) , who used the change of sign of the inner products of consecutive stochastic gradients as a statistical indicator of slow progress and as a trigger to decrease the learning rate ( see extensions in Delyon & Juditsky , 1993 ) . Pflug ( 1983 ; 1988 ) considered the dynamics of SGD with constant step size for minimizing convex quadratic functions and derived necessary conditions for the resulting process to be stationary . Pflug ( 1983 ; 1990 ) also devised a sequential statistical test to detect stationarity ; when the test fires , the step size is decreased by a constant factor . Ruszczyński & Syski ( 1983 ) used online statistical tests to check if the present step size and momentum constants satisfy certain optimality conditions and adjust them if the conditions are not satisfied . However , Lang et al . ( 2019 ) found that these methods have limited success in machine learning applications due to their reliance on quadratic approximation and/or strong assumptions on the noise models . Most recently , Yaida ( 2018 ) derived fluctuation-dissipation relations ( necessary conditions ) that characterized the stationary behavior of stochastic gradient methods with constant learning rate and momentum . These relations hold exactly for any stationary state ( regardless of the loss function or the noise model ) , so they can be very effective at detecting stationarity and can serve as a reliable indicator of when to decrease the learning rate . Lang et al . ( 2019 ) devised a more rigorous statistical test for Yaida ’ s conditions and obtained robust and competitive performance on common deep learning tasks . Our contributions in this paper are threefold : • We consider a broad family of stochastic optimization methods with constant hyperparameters ( including the learning rate and various forms of momentum ) and derive a more general necessary condition ( than Yaida ’ s ) for the associated learning process to be stationary . Yaida ( 2018 ) established a “ master equation , ” from which different stationarity conditions can be derived for different variants of stochastic gradient methods ( some of them may not admit a usable analytical form ) . We derive a simple “ master condition ” that works for different methods without any change and holds exactly for any stationary process . Because of its conceptual and analytical simplicity , it greatly simplifies implementation and deployment in software packages . • We develop a simple statistical test for checking stationarity based on our “ master condition. ” It is a simple confidence interval test ( checking if it contains zero ) , as opposed to the “ equivalence test ” for Yaida ’ s condition used by Lang et al . ( 2019 ) . This simple test is as robust as the equivalence test and avoids the use of an additional hyperparameter . • A major challenge for almost all adaptive stochastic gradient methods is how to set the initial learning rate appropriately without expensive trial and error . We propose a smoothed stochastic line-search method that can be used to warm up the optimization process . Starting from any safe , small learning rate , this method automatically increases it to a suitably large value for fast initial convergence . Afterwards , the statistical test for stationarity can be applied to gradually reduce the learning rate and obtain finer convergence . Combining the three components above , we obtain a highly autonomous algorithm called SALSA ( Stochastic Approximation with Line-search and Statistical Adaptation ) . We conduct extensive experiments on several deep learning tasks to test its training and testing perfromance as well as its robustness to changes in the hyperparameter settings . SALSA matches the best performance of hand-tuned methods across different deep learning tasks using default hyperparameters . 2 NECESSARY CONDITIONS FOR STATIONARITY . We consider general stochastic optimization methods in the form of ( 2 ) with a constant learning rate : xk+1 = xk − αdk . ( 6 ) We assume that the stochastic search direction dk is generated with time-homogeneous dynamics ; In particular , any additional hyperparameters involved must be constant ( not depending on k ) . As an example , we consider the search directions generated by the family of Quasi-Hyperbolic Momentum ( QHM ) methods ( Ma & Yarats , 2019 ) : hk = ( 1− β ) gk + βhk−1 , dk = ( 1− ν ) gk + νhk , ( 7 ) where 0 ≤ β < 1 and 0 ≤ ν ≤ 1 . With β = 0 or ν = 0 , it recovers the SGD direction ( 3 ) . With ν = 1 , it recovers the stochastic heavy-ball direction ( 4 ) . With 0 < β = ν < 1 , it is equivalent to the direction with Nesterov momentum given in ( 5 ) ( Sutskever et al. , 2013 ; Gitman et al. , 2019 ) . We assume that the dynamics of ( 6 ) , e.g. , driven by the stochastic gradients gk through ( 7 ) , is stable . Stability regions of the hyperparameters in QHM are characterized by Gitman et al . ( 2019 ) . In addition , we assume that the stochastic process { xk } , as k → ∞ , becomes stationary . Recall that { xk } is ( strongly ) stationary if the joint distribution of any subset of the sequence is invariant with respect to simultaneous shifts in the time index ( e.g. , Dembo , 2013 , Section 3.2.3 ) . As a direct consequence , for any deterministic function φ : Rp → R , we have Eπ [ φ ( xk+1 ) ] = Eπ [ φ ( xk ) ] , ∀ k , ( 8 ) where π denotes the stationary distribution of { xk } . If the test function φ is smooth , then we use Taylor expansion to obtain φ ( xk+1 ) = φ ( xk − αdk ) = φ ( xk ) − α 〈 ∇φ ( xk ) , dk 〉 + α 2 2 〈 ∇2φ ( xk ) dk , dk 〉 +O ( α3 ) . Taking expectations on both sides of the above equality and applying ( 8 ) , we have Eπ [ 〈 ∇φ ( xk ) , dk 〉 − α2 〈 ∇2φ ( xk ) dk , dk 〉 ] = O ( α2 ) . ( 9 ) For an arbitrary smooth function φ , it is very hard in practice to characterize or approximate the O ( α2 ) term on the right-hand side . In addition , computing the Hessian-vector product ∇2φ ( xk ) dk can be very costly . Therefore , we only consider simple quadratic functions for which the O ( α2 ) term in ( 9 ) vanishes . In particular , the choice of φ ( x ) = ( 1/2 ) ‖x‖2 results in Eπ [ 〈 xk , dk 〉 − α2 ‖d k‖2 ] = 0 . ( 10 ) This condition holds exactly for any stochastic optimization method of the form ( 6 ) if it reaches stationarity . Indeed , weak stationarity is sufficient since φ is a quadratic function of the state ( e.g. , Dembo , 2013 , Section 3.2.3 ) . Beyond stationarity , it requires no specific assumption on the loss function or noise model for the stochastic gradients . Moreover , it can be applied to different methods without any change . We call this the master condition for stationarity . Yaida ( 2018 ) focused on state perturbation along gk , ∇fξk ( xk ) to obtain his “ master equation ” Eπ [ φ ( xk − αgk ) ] = Eπ [ φ ( xk ) ] , ∀ k. This equation can be combined with the actual state update equation ( 6 ) to obtain more specific stationarity conditions for different algorithms . For example , for the stochastic heavy-ball method with dk = ( 1− β ) gk + βdk−1 , the master equation leads to ( Yaida , 2018 ) Eπ [ 〈 xk , gk 〉 − α2 1+β 1−β ‖d k‖2 ] = 0 . ( 11 ) It is a simple exercise to show that this is equivalent to our master condition ( 10 ) . For the QHM update ( 7 ) , a more complex stationarity condition may also be derived , but it will still be equivalent to ( 10 ) . In practice , the single , simple master condition ( 10 ) is much more preferred . If { xk } starts with a nonstationary distribution and converges to a stationary state , we have lim k→∞ E [ 〈 xk , dk 〉 − α2 ‖d k‖2 ] = 0 . ( 12 ) In the next section , we devise a simple statistical test to determine if the master condition fails to hold . In this case , the learning process has not stalled , and we can continue with the same step size . If we fail to detect non-stationarity ( i.e. , the dynamics may be approximately stationary ) , we would like to reduce the learning rate α to allow for finer convergence . Algorithm 1 : SASA+ : SASA with master condition and simple statistical testing 1 input : x0 ∈ Rp , α0 > 0 , Nmin > 0 , Ktest > 0 , δ ∈ ( 0 , 1 ) , θ ∈ ( 0 , 1 ) , τ ∈ ( 0 , 1 ) 2 α← α0 3 ko ← 0 4 for k = 0 , ... , T − 1 do 5 Randomly sample ξk and compute stochastic search direction dk ( e.g. , using QHM ) 6 xk+1 ← xk − αdk 7 ∆k ← 〈 xk , dk 〉 − α2 ‖d k‖2 8 N ← dθ ( k − ko ) e 9 if N > Nmin and k mod Ktest == 0 then 10 ( µ̂N , σ̂N ) ← sample mean and BM/OLBM variance of { ∆k−N+1 , . . . , ∆k } 11 if 0 ∈ µ̂N ± t∗1−δ/2 σ̂N√ N then 12 α← τα 13 ko ← k 14 end 15 end 16 end 17 output : xT ( or the average of last epoch )

Summary: The authors explore how stationarity tests can be leveraged to automatically tune the learning rate during training. Their algorithm also add a robust line search algorithm, to reduce the need to tune the initial learning rate. The paper is clear and the literature review is honest and thorough. However, it is unclear to me if the contribution of the authors is enough, as the method used and its presentation are very close to Lang and al. In particular: