summaries: 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.  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.   The paper applies Anderson mixing to stochastic optimization problems and shows its effectiveness in training neural networks. To make Anderson mixing work in stochastic settings, the paper proposes a series of strategies to handle noise and uncertainty, including damped projection, adaptive regularization, variance reduction, and preconditioned mixing. Proof of convergence is provided, together with numerical evaluations.   The paper under review consider the stochastic Anderson Mixing algorithm and  shows its convergence analysis. An enhanced variance reduction of the algorithms is also derived. The algorithms are tested on extensive data sets with other existing method. This paper seems to be the first one to address the theoretical convergence analysis in this direction. In this paper, the author propose a stochastic version of Anderson acceleration (widely used in fixed point iterations)  to solve non convex optimization problem. Some important modifications of the deterministic  version are introduced in order to stabilize and improve the algorithm. The paper theoretically shows that the new algorithm can achieve O(epsilon^{-2}) sample complexity. Numerical results are provided to verify the performance of SAM.
abstractText: Anderson mixing (AM) is an acceleration method for fixed-point iterations. Despite its success and wide usage in scientific computing, the convergence theory of AM remains unclear, and its applications to machine learning problems are not well explored. In this paper, by introducing damped projection and adaptive regularization to the classical AM, we propose a Stochastic Anderson Mixing (SAM) scheme to solve nonconvex stochastic optimization problems. Under mild assumptions, we establish the convergence theory of SAM, including the almost sure convergence to stationary points and the worst-case iteration complexity. Moreover, the complexity bound can be improved when randomly choosing an iterate as the output. To further accelerate the convergence, we incorporate a variance reduction technique into the proposed SAM. We also propose a preconditioned mixing strategy for SAM which can empirically achieve faster convergence or better generalization ability. Finally, we apply the SAM method to train various neural networks including the vanilla CNN, ResNets, WideResNet, ResNeXt, DenseNet and LSTM. Experimental results on image classification and language model demonstrate the advantages of our method.
authors.name: Fuchao Wei
authors.name: Chenglong Bao
authors.name: Yang Liu
references.authors: Zeyuan Allen-Zhu Elad Hazan
references.title: Variance reduction for faster non-convex optimization
references.venue: In International conference on machine learning,
references.authors: Donald G Anderson
references.title: Iterative procedures for nonlinear integral equations
references.venue: Journal of the ACM (JACM),
references.authors: Wei Bian Xiaojun Chen CT Kelley
references.title: Anderson acceleration for a class of nonsmooth fixed-point problems
references.venue: SIAM Journal on Scientific Computing,
references.authors: Léon Bottou Frank E. Curtis Jorge Nocedal
references.title: Optimization methods for large-scale machine learning
references.venue: SIAM Review,
references.authors: Claude Brezinski Michela Redivo-Zaglia Yousef Saad
references.title: Shanks sequence transformations and Anderson acceleration
references.venue: SIAM Review,
references.authors: Claude Brezinski Stefano Cipolla Michela Redivo-Zaglia Yousef Saad
references.title: Shanks and Anderson-type acceleration techniques for systems of nonlinear equations
references.venue: arXiv preprint arXiv:2007.05716,
references.authors: Richard H Byrd Samantha L Hansen Jorge Nocedal Yoram Singer
references.title: A stochastic quasi- Newton method for large-scale optimization
references.venue: SIAM Journal on Optimization,
references.authors: Eric Cancès Claude Le Bris
references.title: Can we outperform the DIIS approach for electronic structure calculations
references.venue: International Journal of Quantum Chemistry,
references.authors: Xiangyi Chen Sijia Liu Ruoyu Sun Mingyi Hong
references.title: On the convergence of a class of Adam-type algorithms for non-convex optimization
references.venue: In International Conference on Learning Representations,
references.authors: Kai Lai Chung
references.title: On a stochastic approximation method
references.venue: The Annals of Mathematical Statistics,
references.authors: Jing Dong Xin T Tong
references.title: Replica exchange for non-convex optimization
references.venue: Journal of Machine Learning Research,
references.authors: John Duchi Elad Hazan Yoram Singer
references.title: Adaptive subgradient methods for online learning and stochastic optimization
references.venue: Journal of machine learning research,
references.authors: Rick Durrett
references.title: Probability: Theory and examples, volume 49
references.venue: Cambridge university press,
references.authors: Haw-ren Fang Yousef Saad
references.title: Two classes of multisecant methods for nonlinear acceleration
references.venue: Numerical Linear Algebra with Applications,
references.authors: Jerome Friedman Trevor Hastie Robert Tibshirani
references.title: The elements of statistical learning, volume 1. Springer series in statistics
references.venue: New York,
references.authors: Alejandro J Garza Gustavo E Scuseria
references.title: Comparison of self-consistent field convergence acceleration techniques
references.venue: The Journal of chemical physics,
references.authors: Saeed Ghadimi Guanghui Lan
references.title: Stochastic first- and zeroth-order methods for nonconvex stochastic programming
references.venue: SIAM Journal on Optimization,
references.authors: Robert Gower Dmitry Kovalev Felix Lieder Peter Richtárik
references.title: RSN: Randomized subspace Newton
references.venue: In Conference on Neural Information Processing Systems,
references.authors: Robert M Gower Peter Richtárik
references.title: Randomized quasi-Newton updates are linearly convergent matrix inversion algorithms
references.venue: SIAM Journal on Matrix Analysis and Applications,
references.authors: Kaiming He Xiangyu Zhang Shaoqing Ren Jian Sun
references.title: Deep residual learning for image recognition
references.venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,
references.authors: Sepp Hochreiter Jürgen Schmidhuber
references.title: Long short-term memory
references.venue: Neural computation,
references.authors: Gao Huang Zhuang Liu Laurens Van Der Maaten Kilian Q Weinberger
references.title: Densely connected convolutional networks
references.venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,
references.authors: Shaojun Huang Yuanzhang Sun Qiuwei Wu
references.title: Stochastic economic dispatch with wind using versatile probability distribution and L-BFGS-B based dual decomposition
references.venue: IEEE Transactions on Power Systems,
references.authors: Xunpeng Huang Xianfeng Liang Zhengyang Liu Lei Li Yue Yu Yitan Li
references.title: SPAN: A stochastic projected approximate Newton method
references.venue: In Proceedings of the AAAI Conference on Artificial Intelligence,
references.authors: Brian Irwin Eldad Haber
references.title: Secant penalized BFGS: A noise robust quasi-Newton method via penalizing the secant condition
references.venue: arXiv preprint arXiv:2010.01275,
references.authors: Rie Johnson Tong Zhang
references.title: Accelerating stochastic gradient descent using predictive variance reduction
references.venue: Advances in neural information processing systems,
references.authors: Diederik P Kingma Jimmy Ba
references.title: Adam: A method for stochastic optimization
references.venue: arXiv preprint arXiv:1412.6980,
references.authors: M Galina
references.title: Korpelevich. The extragradient method for finding saddle points and other problems
references.venue: Matecon, 12:747–756,
references.authors: Alex Krizhevsky Geoffrey Hinton
references.title: Learning multiple layers of features from tiny images
references.authors: Yann LeCun Léon Bottou Yoshua Bengio Patrick Haffner
references.title: Gradient-based learning applied to document recognition
references.venue: Proceedings of the IEEE,
references.authors: Jian Li Xuanyuan Luo Mingda Qiao
references.title: On generalization error bounds of noisy gradient methods for non-convex learning
references.venue: In International Conference on Learning Representations,
references.authors: Yuanzhi Li Colin Wei Tengyu Ma
references.title: Towards explaining the regularization effect of initial large learning rate in training neural networks
references.venue: Advances in Neural Information Processing Systems,
references.authors: Mitchell P Marcus Mary Ann Marcinkiewicz Beatrice Santorini
references.title: Building a large annotated corpus of English: The penn treebank
references.venue: Computational Linguistics,
references.authors: James Martens
references.title: Deep learning via Hessian-free optimization
references.venue: In Proceedings of the 27th International Conference on Machine Learning,
references.authors: Aryan Mokhtari Alejandro Ribeiro
references.title: Stochastic quasi-Newton methods
references.venue: Proceedings of the IEEE,
references.authors: Arkadi Nemirovski Anatoli Juditsky Guanghui Lan Alexander Shapiro
references.title: Robust stochastic approximation approach to stochastic programming
references.venue: SIAM Journal on optimization,
references.authors: Boris T Polyak
references.title: New stochastic approximation type procedures
references.venue: Automat. i Telemekh,
references.authors: Boris T Polyak Anatoli B Juditsky
references.title: Acceleration of stochastic approximation by averaging
references.venue: SIAM journal on control and optimization,
references.authors: Florian A Potra Hans Engler
references.title: A characterization of the behavior of the Anderson acceleration on linear problems
references.venue: Linear Algebra and its Applications,
references.authors: Phanisri P Pratapa Phanish Suryanarayana John E Pask
references.title: Anderson acceleration of the Jacobi iterative method: An efficient alternative to Krylov methods for large, sparse linear systems
references.venue: Journal of Computational Physics,
references.authors: Ning Qian
references.title: On the momentum term in gradient descent learning algorithms
references.venue: Neural networks,
references.authors: Sashank J Reddi Ahmed Hefny Suvrit Sra Barnabas Poczos Alex Smola
references.title: Stochastic variance reduction for nonconvex optimization
references.venue: In International conference on machine learning,
references.authors: Herbert Robbins Sutton Monro
references.title: A stochastic approximation method
references.venue: The annals of mathematical statistics,
references.authors: Youcef Saad Martin H Schultz
references.title: GMRES: A generalized minimal residual algorithm for solving nonsymmetric linear systems
references.venue: SIAM Journal on scientific and statistical computing,
references.authors: Yousef Saad
references.title: Iterative methods for sparse linear systems
references.authors: Jerome Sacks
references.title: Asymptotic distribution of stochastic approximation procedures
references.venue: The Annals of Mathematical Statistics,
references.authors: Nicol N Schraudolph Jin Yu Simon Günter
references.title: A stochastic quasi-Newton method for online convex optimization
references.venue: In Artificial intelligence and statistics,
references.authors: Damien Scieur Alexandre d’Aspremont Francis Bach
references.title: Regularized nonlinear acceleration
references.venue: In Proceedings of the 30th International Conference on Neural Information Processing Systems,
references.authors: Damien Scieur Francis Bach Alexandre d’Aspremont
references.title: Nonlinear acceleration of stochastic algorithms
references.venue: In Proceedings of the 31st International Conference on Neural Information Processing Systems,
references.authors: Damien Scieur Edouard Oyallon Alexandre d’Aspremont Francis Bach
references.title: Nonlinear acceleration of CNNs
references.venue: In ICLR Workshop track,
references.authors: Damien Scieur Alexandre d’Aspremont Francis Bach
references.title: Regularized nonlinear acceleration
references.venue: Mathematical Programming,
references.authors: Shai Shalev-Shwartz Shai Ben-David
references.title: Understanding machine learning: From theory to algorithms
references.venue: Cambridge university press,
references.authors: Tijmen Tieleman Geoffrey Hinton
references.title: Lecture 6.5-RMSprop: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural networks for machine learning
references.authors: Alex Toth CT Kelley
references.title: Convergence analysis for Anderson acceleration
references.venue: SIAM Journal on Numerical Analysis,
references.authors: Alex Toth J Austin Ellis Tom Evans Steven Hamilton CT Kelley Roger Pawlowski Stuart Slattery
references.title: Local improvement results for Anderson acceleration with inaccurate function evaluations
references.venue: SIAM Journal on Scientific Computing,
references.authors: Henk A Van der Vorst C Vuik
references.title: The superlinear convergence behaviour of GMRES
references.venue: Journal of computational and applied mathematics,
references.authors: Homer F Walker Peng Ni
references.title: Anderson acceleration for fixed-point iterations
references.venue: SIAM Journal on Numerical Analysis,
references.authors: Xiao Wang Shiqian Ma Donald Goldfarb Wei Liu
references.title: Stochastic quasi-Newton methods for nonconvex stochastic optimization
references.venue: SIAM Journal on Optimization,
references.authors: Saining Xie Ross Girshick Piotr Dollár Zhuowen Tu Kaiming He
references.title: Aggregated residual transformations for deep neural networks
references.venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,
references.authors: Pan Xu Jinghui Chen Difan Zou Quanquan Gu
references.title: Global convergence of Langevin dynamics based algorithms for nonconvex optimization
references.venue: In Proceedings of the 32nd International Conference on Neural Information Processing Systems,
references.authors: Zhewei Yao Amir Gholami Sheng Shen Mustafa Mustafa Kurt Keutzer Michael Mahoney
references.title: Adahessian: An adaptive second order optimizer for machine learning
references.venue: In Proceedings of the AAAI Conference on Artificial Intelligence,
references.authors: Sergey Zagoruyko Nikos Komodakis
references.title: Wide residual networks. In British Machine Vision Conference 2016
references.venue: British Machine Vision Association,
references.authors: Michael Zhang James Lucas Jimmy Ba Geoffrey E Hinton
references.title: Lookahead optimizer: k steps forward, 1 step back
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Juntang Zhuang Tommy Tang Yifan Ding Sekhar C Tatikonda Nicha Dvornek Xenophon Papademetris James Duncan
references.title: Adabelief optimizer: Adapting stepsizes by the belief in observed gradients
references.venue: Advances in Neural Information Processing Systems,
sections.heading: 1 Introduction
sections.text: 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.
sections.heading: 2 Methodology
sections.heading: 2.1 Anderson Mixing
sections.text: 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].
sections.heading: 2.2 Stochastic Anderson Mixing
sections.text: 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].
sections.heading: 2.3 Enhancement of Stochastic Anderson Mixing
sections.text: We introduce the variance reduction and preconditioned mixing techniques to further enhance SAM.
  Variance reduction. Variance reduction techniques turn out to be effective if a scan over the full dataset is feasible [1, 44]. Similar to SdLBFGS-VR proposed in [60], we also incorporate SVRG to SAM, which we call SAM-VR (Algorithm 2), for solving (2). To simplify the description, we denote one iteration of SAM in Algorithm 1 as SAM_update(xk, gk), i.e. one update of xk given the gradient estimate gk.
  Preconditioned mixing. Motivated by the great success of preconditioning in solving linear systems and eigenvalue computation [18], we present a preconditioned version of SAM. The key modification is the mixing step (4b). We replace the simple mixing xk+1 = x̄k + βkr̄k with xk+1 = x̄k +M−1k r̄k where Mk approximates the Hessian. Combining it with (8) and (12), we obtain
  xk+1 = xk + ( M−1k − αk ( Xk +M −1 k Rk ) ( RTkRk + δkX T k Xk )† RTk ) rk. (17)
  Setting αk ≡ 1 and δk ≡ 0, (17) reduces to a preconditioned AM update, which can be recast as the solution to the constrained optimization problem: min
  Hk ‖Hk −M−1k ‖F subject to HkRk = −Xk,
  a direct extension of (7). This preconditioned version of AM is related to quasi-Newton updates [20]. We also point out that the action of M−1k can be implicitly done via an update of any optimizer at hand, i.e. xk+1 = optim(x̄k,−r̄k), where optim updates x̄k given the extragradient −r̄k. If Rk = −∇2f(xk)Xk, which is the case in deterministic quadratic optimization, the projection step in preconditioned AM is still a minimal residual procedure. Remark 2. The same as SdLBFGS [60], SAM needs another 2md space to store Xk and Rk. The extra main computational cost for SAM compared with SGD is O(m2d) +O(m3), which accounts for the matrix multiplications (Rm×d × Rd×m) and matrix decomposition of a small Rm×m matrix.
  Since dense matrix multiplication can be ideally parallelized and the cost of gradient evaluations often dominates the computing, the benefit from SAM pays for this extra cost. Also, this cost can be further reduced to O(md) +O(m3) (see Appendix B.2). In our implementation, we incorporate sanity check of the positive definiteness, alternating iteration and moving average, and the details are given in the supplementary materials.
sections.heading: 3 Theory
sections.text: In this section, we analyze the convergence and complexity of SAM. All the proofs are deferred to the Appendix A.1. We first give two assumptions about the objective function f that are widely used in stochastic programming [17, 60, 44]. Assumption 1. f : Rd → R is continuously differentiable. f(x) ≥ f low > −∞ for any x ∈ Rd. ∇f is globally L-Lipschitz continuous; namely ‖∇f(x)−∇f(y)‖2 ≤ L‖x−y‖2 for any x, y ∈ Rd. Assumption 2. For any iteration k, the stochastic gradient ∇fξk(xk) satisfies Eξk [∇fξk(xk)] = ∇f(xk), Eξk [‖∇fξk(xk)−∇f(xk)‖22] ≤ σ2, where σ > 0, and ξk, k = 0, 1, . . . , are independent samples that are independent of {xj}kj=0.
  Throughout our analysis, we always assume both Assumptions 1 and 2 hold. Moreover, we state the following conditions about the three parameters αk, βk, δk of SAM.
  +∞∑ k=0 βk = +∞, +∞∑ k=0 β2k < +∞, (18a)
  δk ≥ Cβ−2k , αk ∈ [0,min{1, β 1 2
  k }] and statisfies (14), (18b) where C > 0 is a constant.
  Convergence and complexity. We give the convergence of SAM by analyzing the iterations {xk} generated by Algorithm 1. Theorem 1. Let {xk} be the sequence generated by Algorithm 1 with fixed batchsize n. Given a positive constant C, if βk ∈ (0, µ4L(1+C−1) ], then under the conditions (18a) and (18b), it has
  lim inf k→∞
  ‖∇f(xk)‖2 = 0 with probability 1. (19)
  Moreover, there exists a positive constant Mf such that
  E[f(xk)] ≤Mf ∀k. (20) Theorem 2. Under the same conditions as Theorem 1, if we require that the noisy gradient is bounded, i.e., Eξk [‖∇fξk(xk)‖22] ≤Mg, (21) where Mg > 0 is a constant, we have
  lim k→∞
  ‖∇f(xk)‖2 = 0 with probability 1. (22)
  Now, we give the iteration complexity of SAM. Theorem 3. Let {xk} be the sequence generated by Algorithm 1 with batchsize n. Given constants C > 0 and r ∈ (0.5, 1), we choose βk = µ4L(1+C−1) (k + 1) −r. Under the condition (18b), it has
  1
  N N−1∑ k=0 E‖∇f(xk)‖22 ≤ 16L(1 + C−1)(Mf − f low) µ2 Nr−1 + (1 + L−1µ−1)σ2 (1− r)n (N−r − rN−1),
  (23)
  where n is the batchsize, andN is the iteration number. Moreover, for a given ∈ (0, 1), to guarantee that 1N ∑N−1 k=0 E‖∇f(xk)‖22 < , the number of iterations N needed is at most O( − 11−r ).
  If the output xR is randomly chosen from previous iterates according to the certain probability distribution, SAM has the same complexity O( −2) as RSG [17] and SQN [60].
  Theorem 4. Let {xk} be the sequence generated by Algorithm 1 with batchsize n and xR be the random output of the first N iterations where the random variable R follows the uniform distribution PR(k) def = Prob{R = k} = 1/N . Given constant C > 0, we choose βk = µ4L(1+C−1) and assume condition (18b) holds. For any > 0, there exist some n and N = O( −2)/n such that we have E [ ‖∇f(xR)‖22 ] ≤ , where the expectation is taken with respect to R and {Sj}N−1j=0 .
  In fact, from the proof of Theorem 4, we can find the explicit choice of the batchsize. Define Df def = f(x0) − f low and D̃ is a problem-independent positive constant. The batchsize n can be
  chosen as n = ⌈ min { N̄ ,max { 1, σL √ N̄ D̃ }}⌉ ,where N̄ is the total number of SFO-calls satisfying
  N̄ ≥ max { C21 2 + 4C2 , σ2
  L2D̃
  } , and C1 = 32Df (1+C −1)σ
  µ2 √ D̃
  + (L+ µ−1)σ √ D̃, C2 = 32DfL(1+C −1)
  µ2 . Thus, to ensure E [ ‖∇f(xR)‖22 ] ≤ , the total number of SFO-calls is O( −2). Moreover, we give the SFO-calls complexity of the variance reduced version of SAM. Theorem 5. Let x be the output of Algorithm 2 with batchsize n. Given a positive constant C, suppose ν, µ0 ∈ (0, 1) to be two constants satisfying
  µ0µ 2(1 + C−1)1/2 − µ
  2 0(2L+ µ −1)(e− 1) L − 2µ20n− 2µ30(2L+ µ −1)(e− 1)n L ≥ ν,
  where e is the Euler’s number. Set βkt = β := µ0n L(1+C−1)1/2T 2/3 , δkt ≥ Cβ−2, αkt ∈ [0,min{1, β 1 2 }] and satisfies (14), and the number of iteration q = ⌊
  T µ0nd0 ⌋ , where d0 = 6 + 2L(1+C−1)1/2 . Then,
  E [ ‖∇f(x)‖22 ] ≤ T 2/3L
  ( f(x0)− f low ) qNnν . (24)
  To ensure E [ ‖∇f(x)‖22 ] ≤ , the total number of SFO-calls is O(T + (T 2/3/ )). Remark 3. It is worth noting that the approximated HessianHk in SAM depends on the data samples {ξi}i∈Sk of the current mini-batch, which violates the assumption AS.4 in [60].
  Analysis of AdaSAM. All the results in Theorems 1-5 are applicable for AdaSAM as (16) satisfies the condition (18b). We further discuss the rationality of the first term in (16). Since rTkHkrk ≥ βkµ‖rk‖22, and ‖Hkrk‖22 ≤ 2 ( β2k + α 2 kδ −1 k ) ‖rk‖22 ≤ 2β2k ( 1 + C−1 ) ‖rk‖22 ifαk ≤ 1, δk ≥ Cβ−2k (see Corollary 1 in Appendix A.1), it is sensible to suppose ‖xk − xk−1‖2 ≈ ‖xk+1 − xk‖2 = ‖Hkrk‖2 = βkhk‖rk‖2, where hk ∈ (0, h) (h < +∞) is related to Hk. Therefore,
  c1‖rk‖22 ‖xk − xk−1‖22 ≈ c1‖rk‖ 2 2 β2kh 2 k‖rk‖22 ≥ c1h−2β−2k , (25)
  which coincides with the requirement that δk ≥ Cβ−2k . A further discussion is in Appendix D.3.
sections.heading: 4 Experiments
sections.text: We used AdaSAM (AdaSAM-VR) for the experiments, which is a special case of SAM (SAM-VR) with δk specified as (16). The pseudocodes are given in Appendix B.1. To evaluate the effectiveness of our methods, we compared them with several first-order and second-order optimizers for mini-batch training of neural networks on different machine learning tasks. The datasets were MNIST [32], CIFAR-10/CIFAR-100 [31] for image classification and Penn Treebank [35] for language model. The experimental details and hyper-parameter tuning are referred to Appendix C.
  Experiments on MNIST. We trained a simple convolutional neural network (CNN) 3 on MNIST, for which we were only concerned about the minimization of the empirical risk (2), i.e. the training loss, with large batch sizes. The training dataset was preprocessed by randomly selecting 12k images from the total 60k images to facilitate large mini-batch training. Neither weight-decay nor dropout was used. We compared AdaSAM with SGDM [43] (SGD with momentum), Adam [29], SdLBFGS [60], and RAM (cf. (10)). The learning rate was tuned and fixed for each optimizer. The historical lengths for SdLBFGS, RAM and AdaSAM were set as 20. δ = 10−6 for RAM and c1 = 10−4 for AdaSAM.
  3Based on the official PyTorch implementation https://github.com/pytorch/examples/blob/master/mnist.
  Figure 1 (a) and (b) show the curves of training loss when training 100 epochs with batch sizes of 6K and 3K, which indicate that AdaSAM can significantly minimize the empirical risk in large mini-batch training. The comparison with RAM verifies the benefit of adaptive regularization. We notice that there hardly exists any oscillation in AdaSAM during training except for the first few epochs, which suggests that AdaSAM is tolerant to noise. We also tested the effectiveness of variance reduction and preconditioning introduced in Section 2.3. The variance reduced extension AdaSAMVR was compared with AdaSAM for different batch sizes. The variants of AdaSAM preconditioned by Adagrad [12] and RMSprop [55] are denoted as Adagrad-AdaSAM and RMSprop-AdaSAM, respectively. Though AdaSAM-VR is more costly per iteration and the preconditioned variants seem to deteriorate the final training loss, we point out that AdaSAM-VR can achieve lower training loss (10−9) and the preconditioned variants converge faster to an acceptable training loss (e.g. 10−3).
  Experiments on CIFAR. For CIFAR-10 and CIFAR100, both datasets have 50K images for training and 10K images for test. The test accuracy at the final epoch was reported as the evaluation metric. We trained ResNet18/20/32/44/56 [21] and WideResNet16-4 [64] on CIFAR-10, and ResNet18, ResNeXt50 [61] and DenseNet121 [24] on CIFAR-100. The baseline optimizers were SGDM, Adam, AdaBelief [66], Lookahead [65], AdaHessian [63] and RNA [52]. The hyperparameters were kept unchanged across different tests. We trained 160 epochs with batch size of 128 and decayed the learning rate at the 80th and 120th epoch. For AdaSAM/RAM, αk and βk were decayed at the 80th and 120th epoch.
  Table 1 demonstrates the generalization ability of AdaSAM (AdaHessian ran out of memory for training ResNeXt50 and DenseNet121). It shows AdaSAM consistently outperforms other optimizers for various neural
  networks on CIFAR10/CIFAR100. Compared with SGDM/Lookahead, AdaSAM is built on a noisy quadratic model to extrapolate historical iterates more elaborately, which may account for its effectiveness. We also conducted tests on training for 120 epochs and 80 epochs. Figure 2 shows AdaSAM
  can achieve comparable or even lower test loss than SGDM/Lookahead when training with fewer epochs, thus saving a large number of iterations. More results can be found in Appendix C.3.
  To compare the total training time, we set SGDM as the baseline and terminated the training of AdaSAM when the accuracy difference AdaSAM − SGDM > −0.05%. The memory, per-epoch computational cost and total training time of training ResNet20/WideResNet16-4 (abbr. WResNet) on CIFAR10 and ResNeXt50/DenseNet121 on CIFAR100 are reported in Table 2. It shows AdaSAM is more efficient than AdaHessian and saves more than 15% training time than SGDM in these models.
  Experiments on Penn Treebank. We trained LSTM [22] on Penn Treebank and reported the perplexity on the validation set in Figure 3 and perplexity on the test set in Table 3, where pAdaSAM denotes the variant of AdaSAM preconditioned by Adam. The experimental setting was the same as that in AdaBelief [66]. In our practice, we found that the vanilla AdaSAM with default hyperparameter setting is not suitable for this task. Nevertheless, a suitable preconditioner (e.g. Adam) can largely
  improve the behaviour of AdaSAM. Conversely, AdaSAM can also enhance a optimizer when the latter is used as a preconditioner.
sections.heading: 5 Discussion
sections.text: In this section, we provide discussions that compare SAM with SQN [60] and SGD.
  SAM versus SQN. Our theoretical results are similar to those of SQN [60], but there are some fundamental differences:
  (i) The basic assumptions are different. (a) Recall that the update is xk+1 = xk − Hkgk, where gk = ∇fSk(xk). Let S[k−1] := (S0, . . . , Sk−1) denote the random samplings in the first k iterations. A key assumption in [60] is that Hk depends only on S[k−1], which guarantees E[Hkgk|S[k−1]] = Hk∇f(xk) that can simplify the proof. In contrast, in our methods SAM and SAM-VR, Hk depends on Sk since rk = −gk participates in constructing Hk, thus violating this assumption and complicating our analysis. (b) The proposed algorithm SdLBFGS in [60] needs twice continuously differentiable assumption of f and the boundedness of∇2f , which are not required in our algorithm.
  (ii) The ways to construct the secant equations are different. In the deterministic case, a secant equation is Hksk = yk, where sk = xk − xk−1 and yk = gk − gk−1. In the stochastic case, various choices exist in forming yk. For SAM, the choice of yk is the same as that of the deterministic case, thus making Hk depend on Sk. For the SQN methods in [60], yk = ∇fSk−1(xk)−∇fSk−1(xk−1), which ensures Hk is independent of Sk. Such choice can reduce variance of noise in secant equations but demands evaluating gradients on the same mini-batch Sk−1 twice. Some other SQN methods [7] use subsampled Hessian-vector products to form yk = ∇2fS′k(xk − xk−1), where S ′ k denotes the data samples independent of Sk, which also ensures Hk is independent of Sk.
  SAM versus SGD. We compare SAM with SGD from theoretical and practical perspectives.
  (i) Theoretical perspective. (a) For general nonconvex stochastic optimization, SAM achieves the same worst-case complexity O(1/ 2) as SGD. (b) There are some specific cases that SAM can be theoretically better than SGD. For strongly convex quadratic minimization, we prove that SAM is essentially equivalent to GMRES (see Appendix A.2) that can exhibit superlinear convergence [58]. For nonlinear optimization, SAM can also have superlinear convergence rate in some cases as it is closely related to the multisecant quasi-Newton method [14], as also pointed out in Remark 1.
  (ii) Practical perspective. (a) With the same number of epochs as SGD, it is shown in Table 1 and Table 3 that our method has better generalization performance than SGD in both image classification and language tasks. (b) Our method can largely reduce the number of epochs to achieve a comparable solution as SGD, as shown in Figure 2 in Section 4 and Figure 8 in Appendix C.3. The total running time can reduce correspondingly, as shown in Table 2. (c) SAM needs additional memory and computational cost, but it is still much more efficient and economical than other second-order optimizers, as shown in Table 2 in Section 4 and Figure 22 in Appendix D.8.
sections.heading: 6 Conclusion
sections.text: In this paper, we develop an extension of Anderson mixing, namely Stochastic Anderson Mixing, for nonconvex stochastic optimization. By introducing damped projection and adaptive regularization, we establish the convergence theory of our new method. We also analyze its work complexity in terms of SFO-calls and show it can achieve the O(1/ 2) complexity for an -accurate solution. We also give a specific form of adaptive regularization. Then we propose two techniques to further enhance our method. One is the variance reduction technique, which can further improve the work complexity of our method theoretically and help achieve lower empirical risk in our experiments. The other one is the preconditioned mixing strategy that directly extends Anderson mixing. Experiments show encouraging results of our method and its enhanced versions in terms of convergence rate or generalization ability in training different neural networks in different machine learning tasks. These results confirm the suitability of Anderson mixing for nonconvex stochastic optimization.
  Acknowledgments and Disclosure of Funding
  This work was supported by the National Key R&D Program of China (No. 2018YFB1005103), Technology and Innovation Major Project (No. 2020AAA0108403), National Natural Science Foundation of China (No.61925601, No. 61772302, No. 11901338) and Huawei Noah’s Ark Lab. We thank all anonymous reviewers for their valuable comments and suggestions on this work.
title: Stochastic Anderson Mixing for Nonconvex Stochastic Optimization