Stochastic Anderson Mixing for Nonconvex Stochastic Optimization

1 Introduction . Stochastic optimization is important in various areas such as statistics [ 15 ] , machine learning [ 4 , 54 ] and power systems [ 25 ] . In this paper , we consider the following stochastic optimization problem : min x∈Rd f ( x ) = Eξ [ F ( x ; ξ ) ] , ( 1 ) where F : Rd × Ξ → R is continuously differentiable and possibly nonconvex and the random variable ξ ∈ Ξ may follow an unknown probability distribution . It is assumed that only noisy information about the gradient of f is available through calls to some stochastic first-order oracle ( SFO ) . One special case of ( 1 ) is the empirical risk minimization problem : min x∈Rd f ( x ) def = 1 T T∑ i=1 fξi ( x ) , ( 2 ) where fξi : Rd → R is the loss function corresponding to the i-th data sample and T denotes the number of data samples . T can be extremely large such that it prohibits the computation of the full gradient∇f . Thus designing efficient and effective numerical algorithms for solving problem ( 1 ) or ( 2 ) with rigorous theoretical analysis is a challenging task . ∗Corresponding author . clbao @ mail.tsinghua.edu.cn . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . One classical approach for solving ( 1 ) is the stochastic gradient descent ( SGD ) method [ 45 ] . It mimics the gradient descent ( GD ) method with noisy gradients and exhibits optimal convergence rate for some strongly convex stochastic problems [ 10 , 48 ] . Some early related works of SGD in convex optimization can be found in [ 38 , 39 , 40 ] . For nonconvex cases , Ghadimi and Lan [ 17 ] proposed a randomized stochastic gradient ( RSG ) method that randomly selects a solution x̄ from previous iterates . To ensure x̄ satisfying E [ ‖∇f ( x̄ ) ‖22 ] ≤ , the total number of SFO-calls2 needed by RSG is O ( −2 ) . For training neural networks , adaptive learning rate methods were proposed to accelerate SGD , e.g . Adagrad [ 12 ] , RMSprop [ 55 ] and Adam [ 29 ] , though their theoretical properties in nonconvex optimization were unknown until recently [ 9 ] . Regarding to second-order optimization methods [ 49 , 36 , 37 ] , one notable work is the framework of stochastic quasi-Newton ( SQN ) methods proposed by Wang et al . [ 60 ] , which covers a class of SQN methods and has theoretical guarantees in nonconvex stochastic optimization . However , these second-order methods usually demand more gradient evaluations in every iteration and less noisy gradient information to achieve actual acceleration [ 4 ] . There are also some recent works inspired from the dynamics perspective , such as the Langevin dynamics based algorithms [ 62 , 33 , 11 ] , but many of them focus on the theory and their applications in deep learning are still under-explored . In practice , the choices of optimizers can vary for different applications [ 63 ] and their practical performances are still unsatisfactory in terms of convergence rate or generalization ability [ 66 ] . For these reasons , we develop a novel second-order method based on Anderson mixing ( AM ) [ 2 ] . Our analysis and experiments show the proposed method is competent from both theoretical and practical perspectives and performs well in training various neural networks in different tasks . AM is a sequence acceleration method in scientific computing [ 5 ] . It is widely used to accelerate the slow convergence of nonlinear fixed-point iterations arising in computational physics and quantum chemistry , e.g . the self-consistent field iteration in electronic structure calculations [ 16 , 8 ] , where the function evaluation is costly . It turns out that AM is closely related to multisecant quasi-Newton methods in nonlinear problems [ 14 , 5 ] or the GMRES method [ 46 ] in linear problems [ 59 , 41 ] . Inspired by the great success of AM in accelerating fixed-point iterations , it is natural to ask whether AM can be applied to accelerate nonlinear optimization since the gradient descent with constant stepsize is a fixed-point iteration . This idea has been explored in [ 50 , 51 ] , but the proposed Regularized Nonlinear Acceleration ( RNA ) method is built on the minimal polynomial extrapolation ( MPE ) approach [ 5 ] , a sequence transformation method that has subtle difference from AM . Also , their methods rely heavily on the contraction assumption of the fixed-point map and the strong convexity . For AM , although current research has proved local linear convergence of AM for fixed-point iterations under some conditions [ 56 , 57 , 3 ] , there exists no version of AM that guarantees convergence for nonconvex optimization , let alone stochastic optimization . In this paper , we develop a stochastic version of AM . Due to the nonconvexity and noise inside the problems and lack of line search or trust-region , a straightforward migration of AM to nonconvex stochastic optimization is infeasible . As a result , we make several fundamental modifications to AM . We highlight the main contributions of our works as follows : 1 . We develop a stochastic version of AM , namely Stochastic Anderson Mixing ( SAM ) , by introducing damped projection and adaptive regularization . We prove its almost sure convergence to a stationary point and analyze its work complexity . When a randomly chosen iterate xR is returned as the output of SAM , we prove that the worst-case SFO-calls complexity to guarantee E [ ‖∇f ( xR ) ‖22 ] ≤ is O ( −2 ) . ( See Theorems 1-4 . ) 2 . We give a variance reduced extension of SAM by borrowing the stochastic variance reduced gradient ( SVRG ) [ 28 ] technique and analyze its SFO-calls complexity . ( See Theorem 5 . ) We also propose a preconditioned mixing strategy for AM and obtain the preconditioned SAM method which can empirically converge faster or generalize better . ( See Section 2.3 . ) 3 . Extensive experiments on training Convolutional Neural Network ( CNN ) , ResNet , WideResNet , ResNeXt , DenseNet , and LSTM on different tasks and datasets show the faster convergence or better generalization ability of our method compared with the state-of-the-art methods . ( See Section 4 . ) 2One call to the SFO in ( 1 ) or one evaluation of ∇fξi ( x ) in ( 2 ) is counted as one SFO-call . 2 Methodology . 2.1 Anderson Mixing . AM is proposed for acceleration of fixed-point iterations . We assume the fixed-point iteration is xk+1 = g ( xk ) def = xk + rk , where rk def = −∇f ( xk ) . Then g ( xk ) − xk = rk . Here , we adopt the description of AM in [ 14 , 59 ] . Let ∆ denote the forward difference operator , say , ∆xk = xk+1 − xk . Let Xk and Rk record the most recent m ( m ≤ k ) iterations : Xk = [ ∆xk−m , ∆xk−m+1 , · · · , ∆xk−1 ] , Rk = [ ∆rk−m , ∆rk−m+1 , · · · , ∆rk−1 ] . ( 3 ) AM can be decoupled into two steps . We call them the projection step and the mixing step : x̄k = xk −XkΓk , ( Projection step ) ( 4a ) xk+1 = x̄k + βkr̄k , ( Mixing step ) ( 4b ) where βk is the mixing parameter , and r̄k def = rk −RkΓk is reminiscent of extragradient [ 30 ] . Γk is determined by solving Γk = arg min Γ∈Rm ‖rk −RkΓ‖2 . ( 5 ) Combining ( 4a ) and ( 4b ) , we obtain the full form of AM [ 14 , 59 , 5 , 42 ] : xk+1 = xk + βkrk − ( Xk + βkRk ) Γk . ( 6 ) Remark 1 . To see the rationality of AM , we assume f is twice continuously differentiable . Then a quadratic approximation of f implies∇2f ( xk ) ( xj − xj−1 ) ≈ ∇f ( xj ) −∇f ( xj−1 ) in a local small region around xk , so it is reasonable to assume Rk ≈ −∇2f ( xk ) Xk . Thus we see ‖rk −RkΓ‖2 ≈ ‖rk +∇2f ( xk ) XkΓ‖2 . Hence , we can recognize ( 5 ) as solving ∇2f ( xk ) pk = ∇f ( xk ) in a leastsquares sense , where pk = XkΓk . When the quadratic approximation is exact , solving ( 5 ) is a minimal residual procedure , thus being verified as a residual projection method [ 47 ] . Moreover , let Hk be the solution to the constrained optimization problem [ 14 ] : min Hk ‖Hk − βkI‖F subject to HkRk = −Xk , ( 7 ) then iterate ( 6 ) is xk+1 = xk+Hkrk , which is indeed a multisecant quasi-Newton method . Note that a key simplification in AM is using differences of historical gradients Rk to approximate−∇2f ( xk ) Xk , which reduces the heavy cost to compute Hessian-vector products [ 7 , 19 , 26 ] . 2.2 Stochastic Anderson Mixing . We describe our method Stochastic Anderson Mixing ( SAM ) in this section . At the k-th iteration , let Sk ⊆ [ T ] def = { 1 , 2 , . . . , T } be the sampled mini-batch and the corresponding objective function value is fSk ( xk ) = 1 |Sk| ∑ i∈Sk fξi ( xk ) . Then rk def = −∇fSk ( xk ) and the noisy Rk is defined correspondingly ( cf . ( 3 ) ) . Due to the instability and inaccurate estimation of Rk , we stabilize the projection step by proposing damped projection and adaptive regularization techniques . Algorithm 1 is a sketch of our method . Now , we elaborate the mechanism of this algorithm . Damped projection . From Remark 1 , we see the determination of Γk in ( 5 ) relies on the local quadratic approximation of ( 2 ) , which can be rather inexact in general nonlinear optimization . To improve the stability , we propose a damped projection method for ( 4a ) . Let αk be the damping parameter , we obtain x̄k via x̄k = ( 1− αk ) xk + αk ( xk −XkΓk ) = xk − αkXkΓk . ( 8 ) Combining ( 8 ) and ( 4b ) and noting that r̄k = rk − αkRkΓk , the new iterate xk+1 is given by xk+1 = xk + βkrk − αk ( Xk + βkRk ) Γk . ( 9 ) It is worth noting that βk and αk in ( 9 ) behave like stepsize or learning rate in SGD , and the extra term ( αkXk + αkβkRk ) Γk can be viewed as a generalized momentum term . Algorithm 1 Stochastic Anderson Mixing ( SAM ) Input : x0 ∈ Rd , m = 10 , αk = 1 , βk = 1 , µ ∈ ( 0 , 1 ) , max_iter > 0 . Output : x ∈ Rd 1 : for k = 0 , 1 , . . . , max_iter do 2 : rk = −∇fSk ( xk ) 3 : if k = 0 then 4 : xk+1 = xk + βkrk 5 : else 6 : mk = min { m , k } 7 : Xk = [ ∆xk−mk , ∆xk−mk+1 , · · · , ∆xk−1 ] 8 : Rk = [ ∆rk−mk , ∆rk−mk+1 , · · · , ∆rk−1 ] 9 : Check Condition ( 14 ) and use smaller αk if ( 14 ) is violated . 10 : xk+1 = xk + βkrk − ( αkXk + αkβkRk ) ( RTkRk + δkX T k Xk ) † RTk rk 11 : end if 12 : Apply learning rate schedule of αk , βk 13 : end for 14 : return xk Adaptive regularization . Since Rk may be rank deficient and no safeguard method is used in AM , the least squares problem ( 5 ) can be unstable . A remedy is to add regularization [ 6 ] to ( 5 ) . One well-known choice is the constant regularization introduced in [ 50 , 53 ] , which can be viewed as forcing ‖Γk‖2 to be small [ 53 ] , leading to a penalty term to ( 5 ) : Γk = arg min Γ∈Rm ‖rk −RkΓ‖22 + δ‖Γ‖22 , ( 10 ) where δ ≥ 0 is the penalty constant . The solution of ( 10 ) is Γk = ( RTkRk + δI ) † RTk rk , where “ † '' denotes the Penrose-Moore inverse . This regularized variant of AM is named as RAM and serves as a baseline in the experiments . Here , we propose a new regularization , namely adaptive regularization , to better suit the stochastic optimization . Since −XkΓk = x̄k − xk denotes the update from xk to x̄k , a large magnitude of ‖XkΓk‖2 tends to make the intermediate step x̄k overshoot the trust region around xk . Thus it is more reasonable to force ‖XkΓk‖2 rather than ‖Γk‖2 to be small . We formulate this idea as min Γ ‖rk −RkΓ‖22 + δk‖XkΓ‖22 , ( 11 ) where δk ≥ 0 is a variable determined in each iteration . Explicitly solving ( 11 ) leads to Γk = ( RTkRk + δkX T k Xk ) † RTk rk . ( 12 ) We call AM with this regularization and damped projection as SAM , i.e . the prototype algorithm given in Algorithm 1 . Positive definiteness . From ( 9 ) and ( 12 ) , the SAM update is xk+1 = xk + Hkrk , where Hk = βkI − αkYkZ†kRTk , Yk = Xk + βkRk , Zk = RTkRk + δkXTk Xk . Hk is generally not symmetric . A critical condition for the convergence analysis of SAM is the positive definiteness of Hk , i.e . pTkHkpk ≥ βkµ‖pk‖22 , ∀pk ∈ Rd , ( 13 ) where µ ∈ ( 0 , 1 ) is a constant . Next , we give an approach to guarantee it . Let λmin ( · ) denote the smallest eigenvalue , λmax ( · ) denote the largest eigenvalue . Since pTkHkpk = 1 2p T k ( Hk +H T k ) pk , Condition ( 13 ) is equivalent to λmin ( 1 2 ( Hk +H T k ) ) ≥ βkµ . With some simple algebraic operations , we obtain λmin ( 1 2 ( Hk +H T k ) ) = βk − 12αkλmax ( YkZ † kR T k + RkZ † kY T k ) . Let λk def = λmax ( YkZ † kR T k +RkZ † kY T k ) , then Condition ( 13 ) is equivalent to αkλk ≤ 2βk ( 1− µ ) . ( 14 ) To check Condition ( 14 ) , note that λk = λmax ( ( Yk Rk ) ( 0 Z†k Z†k 0 ) ( Y Tk RTk ) ) = λmax ( ( Y Tk RTk ) ( Yk Rk ) ( 0 Z†k Z†k 0 ) ) . ( 15 ) Algorithm 2 Stochastic Anderson Mixing with variance reduction ( SAM-VR ) Input : x̃0 ∈ Rd ; βkt , αkt , δkt for SAM_update ( xkt , gkt ) ; Batch size n ≥ 1 . Output : x ∈ Rd 1 : for k = 0 , . . . , N − 1 do 2 : xk0 = x̃k 3 : ∇f ( x̃k ) = 1T ∑T i=1∇fξi ( x̃k ) 4 : for t = 0 , . . . , q − 1 do 5 : Sample a subset K ⊆ [ T ] with |K| = n 6 : gkt = ∇fK ( xkt ) −∇fK ( x̃k ) +∇f ( x̃k ) where∇fK ( xkt ) = 1|K| ∑ i∈K∇fξi ( xkt ) 7 : xkt+1 = SAM_update ( x k t , g k t ) 8 : end for 9 : Set x̃k+1 = xkq 10 : end for 11 : return Iterate x chosen uniformly random from { { xkt } q−1 t=0 } N−1 k=0 Since ( Y Tk RTk ) ( Yk Rk ) , ( 0 Z†k Z†k 0 ) ∈ R2m×2m , and m d , λk can be computed efficiently , say , using an eigenvalue decomposition algorithm with the time complexity of O ( m3 ) . This cost is negligible compared with those to form XTk Xk , R T kRk . After that , to guarantee the positive definiteness , we check if αk satisfies ( 14 ) and use a smaller αk if necessary . AdaSAM method . We choose the δk in ( 11 ) as δk = max { c1‖rk‖22 ‖xk − xk−1‖22 + , c2β −2 k } , ( 16 ) where c1 , c2 , are constants . Such choice reflects the curvature change in the vicinity of xk . A large ‖rk‖2 indicates a potential dramatic change in landscape , suggesting using a precautious tiny stepsize . The denominator in ( 16 ) behaves like annealing , which can measure the noise in gradients like that in secant penalized BFGS [ 27 ] .

Summary: This paper proposes a stochastic Anderson mixing method to solve non-convex stochastic optimization problems. The main contribution is a  convergence theory and the applications of this method for deep learning problems. Adaptive, variance reduced, and preconditioned versions of this methods are also studied.