A Class of Short-term Recurrence Anderson Mixing Methods and Their Applications

1 INTRODUCTION . Anderson mixing ( AM ) ( Anderson , 1965 ) is a sequence acceleration method ( Brezinski et al. , 2018 ) in scientific computing . It is widely used to accelerate the slow convergence of fixed-point iterations ( Lin et al. , 2019 ; Fu et al. , 2020 ; An et al. , 2017 ) , e.g. , the self-consistent field iterations in electronic structure computations ( Garza & Scuseria , 2012 ; Arora et al. , 2017 ) . Specifically , we consider a fixed-point iteration xk+1 = g ( xk ) , k = 0 , 1 , . . . , where g : Rd 7→ Rd is the fixed-point map . By usingm historical iterations , AM ( m ) aims to extrapolate a new iterate that satisfies certain optimality property . When the function evaluation is costly , the reduction of the number of iterations brought by AM can save a large amount of computation ( Fang & Saad , 2009 ) . AM can be used as a method for solving nonlinear equations ( Kelley , 2018 ) as the fixed-point problem x = g ( x ) is equivalent to h ( x ) : = x − g ( x ) = 0 . In practice , since computing the Jacobian of h ( x ) is commonly difficult or even unavailable ( Nocedal & Wright , 2006 ) , AM can be seen as a practical alternate for Newton ’ s method ( An et al. , 2017 ) . Also , compared with the classical iterative methods such as the nonlinear conjugate gradient ( CG ) method ( Hager & Zhang , 2006 ) , no line-search or trust-region technique is used in AM , which is preferable for large-scale unconstrained optimization . Empirically , it is observable that AM can largely accelerate convergence , though its theoretical analysis is still under-explored . It turns out that in the linear case ( Walker & Ni , 2011 ; Potra & Engler , 2013 ) , the full-memory AM ( m = k ) is essentially equivalent to GMRES ( Saad & Schultz , 1986 ) , a powerful Krylov subspace method that can exhibit superlinear convergence behaviour in solving linear systems ( Van der Vorst & Vuik , 1993 ) . For general nonlinear problems , AM is recognized as a multisecant quasi-Newton method ( Fang & Saad , 2009 ; Brezinski et al. , 2018 ) . As far as we know , only local linear convergence has been obtained for the limited-memory AM ( m < k ) in general ( Toth & Kelley , 2015 ; Evans et al. , 2020 ; De Sterck & He , 2021 ) . For the application of AM , one of the major concerns is the historical length m , a critical factor related to the efficiency of AM ( Walker & Ni , 2011 ) . A larger m can incorporate more historical ∗Corresponding author . information into one extrapolation , but it incurs heavier memory overhead since 2m vectors of dimension d need to be stored in AM ( m ) . The additional memory footprint can be prohibitive for solving high-dimensional problems in a resource-limited machine ( Deng , 2019 ) . Using small m can alleviate the memory overhead but may deteriorate the efficacy of AM since much historical information is omitted in the extrapolation ( Walker & Ni , 2011 ; Evans et al. , 2020 ) . To address the memory issue of AM , we deeply investigate the properties of the historical iterations produced by AM and leverage them to develop the short-term recurrence variant , namely ST-AM . The basic version of ST-AM imposes some orthogonalization property on the historical sequence , which is inspired by the CG method ( Hestenes & Stiefel , 1952 ) that enjoys a three-term recurrence . Furthermore , to better suit the more difficult nonconvex optimization , a regularized short-term form is introduced . We highlight the main contributions of our work as follows . 1 . We develop a novel class of short-term recurrence AM methods ( ST-AM ) , including the basic ST-AM , the modified ST-AM ( MST-AM ) , and the regularized ST-AM ( RST-AM ) . The basic ST-AM is applicable for linear systems , MST-AM can solve general fixed-point problems , and RST-AM aims for solving stochastic optimization . An important feature of ST-AM is that all methods only need to store two previous iterations with cheap corrections , which significantly reduces the memory requirement compared with the classical AM . 2 . A complete theoretical analysis of the ST-AM methods is given . When solving strongly convex quadratic optimization , we prove that the basic ST-AM is equivalent to full-memory AM and the convergence rate is similar to CG method . We also prove that MST-AM has improved local linear convergence for solving fixed-point problems . Besides , we establish the global convergence property and complexity analysis for RST-AM when solving stochastic optimization problems . 3 . The numerical results on solving ( non ) linear equations and cubic-regularized quadratic optimization are consistent with the theoretical results for the basic ST-AM and MST-AM . Furthermore , extensive experiments on training deep neural networks for image classification and language models show that RST-AM is competitive with the long-memory AM and outperforms many existing optimizers such as SGD and ADAM . 2 RELATED WORK . AM is also known as an extrapolation algorithm in scientific computing ( Anderson , 2019 ) . A parallel method is Shanks transformation ( Shanks , 1955 ) which transforms an existing sequence to a new sequence for faster convergence . Related classical algorithms include Minimal Polynomial Extrapolation ( Cabay & Jackson , 1976 ) and Reduced Rank Extrapolation ( Eddy , 1979 ) , and a framework of these extrapolation algorithms including AM is given in ( Brezinski et al. , 2018 ) . Note that an elegant recursive algorithm named -algorithm has been discovered for Shanks transformation for scalar sequence ( Wynn , 1956 ) , and was later generalized as the vector -algorithm ( Wynn , 1962 ) for handling vector sequences , but this short-term recurrence form is not equivalent to the original Shanks transformation in general ( Brezinski & Redivo-Zaglia , 2017 ) . Since AM is closely related to quasi-Newton methods ( Fang & Saad , 2009 ) , there are also some works trying to derive equivalent forms of the full-memory quasi-Newton methods using limited memory ( Kolda et al. , 1998 ; Berahas et al. , 2021 ) , while no short-term recurrence is available . To the best of our knowledge , ST-AM is the first attempt to short-term recurrence quasi-Newton methods . Recently , there have been growing demands for solving large-scale and high-dimensional fixedpoint problems in scientific computing ( Lin et al. , 2019 ) and machine learning ( Bottou et al. , 2018 ) . For these applications , Newton-like methods ( Byrd et al. , 2016 ; Wang et al. , 2017 ; Mokhtari et al. , 2018 ) are less appealing due to the heavy memory and computational cost , especially in nonconvex stochastic optimization , where only sublinear convergence can be expected if only stochastic gradients can be accessed ( Nemirovskij & Yudin , 1983 ) . On the other side , first-order methods ( Necoara et al. , 2019 ) stand out for their low per-iteration cost , though the convergence can be slow in practice . When training neural networks , SGD with momentum ( SGDM ) ( Qian , 1999 ) , and adaptive learning rate methods , e.g . AdaGrad ( Duchi et al. , 2011 ) , RMSprop ( Tieleman & Hinton , 2012 ) , Adam ( Kingma & Ba , 2014 ) , are very popular optimizers . Our methods have the nature of quasiNewton methods while the memory footprint is largely reduced to be close to first-order methods . Thus , ST-AM can be a competitive optimizer from both theoretical and practical perspectives . 3 METHODOLOGY . In this section , we give the details of the proposed ST-AM . We always assume the objective function as f : Rd → R , the fixed-point map g : Rd 7→ Rd . Moreover , we do not distinguish rk = −∇f ( xk ) and rk = g ( xk ) − xk in our discussion as∇f ( x ) = 0 is equivalent to g ( x ) = −∇f ( x ) + x = x . 3.1 ANDERSON MIXING . The AM finds the fixed point of g via maintaining two sequences of length m ( m ≤ k ) : Xk = [ ∆xk−m , ∆xk−m+1 , · · · , ∆xk−1 ] , Rk = [ ∆rk−m , ∆rk−m+1 , · · · , ∆rk−1 ] ∈ Rd×m , ( 1 ) where the operator ∆ denotes the forward difference , e.g . ∆xk = xk+1 − xk . Each update of AM can be decoupled into two steps , namely the projection step and the mixing step : x̄k = xk −XkΓk , ( Projection step ) , xk+1 = x̄k + βkr̄k , ( Mixing step ) , ( 2 ) where r̄k : = rk −RkΓk and βk > 0 is the mixing parameter . The Γk is determined by Γk = arg min Γ∈Rm ‖rk −RkΓ‖2 . ( 3 ) Thus , the full form of AM ( Fang & Saad , 2009 ; Walker & Ni , 2011 ) is xk+1 = xk + βkrk − ( Xk + βkRk ) Γk . ( 4 ) Remark 1 . To see the rationality of AM , assume g is continuously differentiable , then we have h ( xj ) −h ( xj−1 ) ≈ h′ ( xk ) ( xj−xj−1 ) around xk , where h′ ( xk ) is the Jacobian of h ( x ) : = x−g ( x ) . So , it is reasonable to assume Rk ≈ −h′ ( xk ) Xk , and we see ‖rk −RkΓ‖2 ≈ ‖rk + h′ ( xk ) XkΓ‖2 . Thus , we can recognize ( 3 ) as solving h′ ( xk ) dk = h ( xk ) in a least-squares sense , where dk = XkΓk . The mixing step incorporates rk into the new update xk+1 if βk > 0 . Otherwise , xk+1 = x̄k is an interpolation of the previous iterates , leading to a stagnating sequence . 3.2 THE BASIC SHORT-TERM RECURRENCE ANDERSON MIXING . The basic ST-AM is to solve the strongly convex quadratic optimization : min x∈Rd f ( x ) : = 1 2 xTAx− bTx , ( 5 ) where A 0 . Initialized with p−1 = q−1 = p0 = q0 = 0 ∈ Rd , at the k-th iteration , given the two matrices Pk−1 = ( pk−2 , pk−1 ) ∈ Rd×2 , Qk−1 = ( qk−2 , qk−1 ) ∈ Rd×2 and defining p = xk−xk−1 and q = rk − rk−1 , the basic ST-AM constructs p̃ = p− Pk−1 ( QTk−1q ) , q̃ = q −Qk−1 ( QTk−1q ) , ( 6a ) pk = p̃/‖q̃‖2 , qk = q̃/‖q̃‖2 . ( 6b ) Then , we update Pk = ( pk−1 , pk ) , Qk = ( qk−1 , qk ) ∈ Rd×2 . Such construction ensures QTkQk = I2 for k ≥ 2 and the storage of Pk and Qk is equal to AM ( 2 ) . With the corrected Pk and Qk , the ST-AM method modifies the projection and mixing step accordingly , that is , x̄k = xk − PkΓk , ( Projection step ) , xk+1 = x̄k + βkr̄k , ( Mixing step ) , ( 7 ) where Γk = arg min ‖rk−QkΓ‖2 = QTk rk and r̄k = rk−QkΓk . Thus , the ST-AM replacesXk and Rk in ( 1 ) by Pk and Qk respectively and imposes the orthogonalization property on Qk . The details of basic ST-AM are given in Algorithm 2 in Appendix C.1 . Define P̄k = ( p1 , p2 , . . . , pk ) , Q̄k = ( q1 , q2 , . . . , qk ) , the Krylov subspace Km ( A , v ) ≡ span { v , Av , A2v , . . . , Am−1v } , the range of X as range ( X ) . We give the properties of the basic ST-AM in Theorem 1 . Theorem 1 . Let { xk } be the sequence generated by the basic ST-AM . The following relations hold : ( i ) ‖q̃‖2 > 0 , range ( P̄k ) = range ( Xk ) = Kk ( A , r0 ) , range ( Q̄k ) = range ( Rk ) = AKk ( A , r0 ) ; ( ii ) Q̄k = −AP̄k , Q̄Tk Q̄k = Ik ; ( iii ) r̄k ⊥ range ( Q̄k ) and x̄k = x0 + zk , where zk = arg minz∈Kk ( A , r0 ) ‖r0 −Az‖2 . If ‖r̄k‖2 = 0 , then xk+1 is the exact solution . The proof is in Appendix C.1 . Note that the property ( iii ) in Theorem 1 exactly describes the relation x̄k = x G k , where x G k is the output of the k-th iteration of GMRES ( Saad & Schultz , 1986 ) . Moreover , let x̄AMk be the k-th intermediate iterate in full-memory AM ( m = k ) . We prove that x̄ AM k = x G k ( See Proposition 1 in Appendix C.1 . ) , which induces that x̄k = x̄AMk = x G k . This equivalence indicates that ST-AM is more efficient than AM and GMRES since only two historical iterations need to be stored . Moreover , by directly applying the convergence analysis of GMRES ( Corollary 6.33 in ( Saad , 2003 ) ) , we obtain the convergence rate of the basic ST-AM for solving ( 5 ) : Corollary 1 . Suppose the eigenvalues of A lie in [ µ , L ] with µ > 0 , and let { xk } be the sequence generated by the basic ST-AM , then the k-th intermediate residual r̄k satisfies ‖r̄k‖2 ≤ 2 ( √ L/µ−1√ L/µ+1 ) k ‖r0‖2 . Moreover , the algorithm finds the exact solution in at most ( d+ 1 ) iterations . Remark 2 . The GMRES can be simplified to an elegant three-term recurrence algorithm called the conjugate residual ( CR ) method ( Algorithm 6.20 in ( Saad , 2003 ) ) when solving ( 5 ) . Thus , a similar simplification for AM is expected to exist . Like CG and Chebyshev acceleration ( Algorithm 12.1 in ( Saad , 2003 ) ) , the convergence rate of ST-AM has the optimal dependence on the condition number , while ST-AM does not form the Hessian-vector products explicitly .

Summary: The paper introduces 3 new variants of Anderson Mixing methods relying on very limited short term memory. This makes these methods more attractive for usual machine learning workloads. The paper also provides detailed analysis of the performance of each of these AM algorithms, as well as thorough experiments, showcasing improved performance of select neural network training.