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: In this manuscript, the authors proposes a Stochastic Anderson Mixing (SAM) scheme to solve nonconvex optimization problems. Under the smoothness, unbiased gradient estimate with bounded variance assumptions, the authors show different variants of the scheme (SAM, SAM with variance reduction, Ada-SAM) have asymptotic and non-asymptotic convergence. Experiments on MNIST, CIFAR and Penn Treebank seem to show the superior performance over other baselines.