Stochastic Anderson Mixing for Nonconvex Stochastic Optimization
Year: 2021

Abstract:
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.

heading: 1 Introduction
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.

heading: 2 Methodology

heading: 2.1 Anderson Mixing
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].

heading: 2.2 Stochastic Anderson Mixing
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].

heading: 2.3 Enhancement of Stochastic Anderson Mixing
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.

heading: 3 Theory
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.

heading: 4 Experiments
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.

heading: 5 Discussion
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.

heading: 6 Conclusion
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.

Review 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.