A Class of Short-term Recurrence Anderson Mixing Methods and Their Applications
Authors: Fuchao Wei
Chenglong Bao
Yang Liu
Venue: ICLR 2022 Poster (2022)
Keywords: Anderson mixing
sequence acceleration
fixed-point iteration
nonconvex optimization
stochastic optimization

Abstract: Anderson mixing (AM) is a powerful acceleration method for fixed-point iterations, but its computation requires storing many historical iterations. The extra memory footprint can be prohibitive when solving high-dimensional problems in a resource-limited machine. To reduce the memory overhead, we propose a novel class of short-term recurrence AM methods (ST-AM). The ST-AM methods only store two previous iterations with cheap corrections. We prove that the basic version of ST-AM is equivalent to the full-memory AM in strongly convex quadratic optimization, and with minor changes it has local linear convergence for solving general nonlinear fixed-point problems. We further analyze the convergence properties of the regularized ST-AM for nonconvex (stochastic) optimization. Finally, we apply ST-AM to several applications including solving root-finding problems and training neural networks. Experimental results show that ST-AM is competitive with the long-memory AM and outperforms many existing optimizers. 

Published as a conference paper at ICLR 2022
A CLASS OF SHORT-TERM RECURRENCE ANDERSON MIXING METHODS AND THEIR APPLICATIONS
Fuchao Wei Chenglong Bao 3,4 Yang Liu 1.2
Department of Computer Science and Technology, Tsinghua University 2 Innttitute for AI Industry Research, Tsinghua University
3 Yau Mathematical Sciences Center, Tsinghua University
'Yanqi Lake Beijing Institute of Mathematical Sciences and Applications wf wfcl66 60maimmmm....... 1s. ts tsinghua. edu. en, {clbao, htt.://luagnnneestsssooo edu. cn
ABSTRACT
Anderson mixing (AM) is a powerful acceleration method for fixed-point itera- tions, but its computation requires storing many historical iterations. The extra memory footprint can be prohibitive when solving high-dimensional problems in a lesitee machine. To reduce the memory overhead, we propose novel class of short-term recurrence AM methods (ST-AM). The ST-AM methods only store two previous iterations with cheap corrections. We prove that the basic ver- sion of ST-AM is equivalent to the full-memory AM in strongly convex quadratic optimization, and with minor changes it has local linear convergence for solving general nonlinear fixed-point problems. We further analyze the convergence prop- erties of the regularized ST-AM for nonconvex (stochastic) optimization. Finally, we apply ST-AM several applications including solving root-finding problems and training neural networks. Experimental results show that ST-AM is competi- live with the long-memory AM and outperforms many existing optimizers.
INTRODUCTION
1965) is powerful sequence acceleration method (Brezinski al. 2018) for fixed-point iterations and has been widely used in scientific computing LLi et al., Fu et al. 2020; An et al. 2017). e.g., the self-consistent field iterations in electronic structure Specifically, we consider fixed-point 1 Rd is the fixed-point map. By using m historical iterations, AM(m) aims 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 computation (Fang & Saad 2009).
Anderson mixing (AM) (Anderson
2019
computations Garza & Scuseria 2012 Arora et al. iteration Z++1 = g 0, where
2017).
AM can be used a method for solving nonlinear equations (Kelley 2018) the fixed-point problem I = g(z) is equivalent to h(x) := x g(x) = 0. In practice, since computing the 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 the nonlinear conjugate gradient (CG) method (Hager & Zhang 2006), line-search or trust-region technique is used in A.M, 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 GMRES (Saad) & Schultz, 1986). a powerful Krylov subspace method that can exhibit superlincar convergence behaviour 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-ment -memrry
Jacobian of h(()) is commonly difficult
AM (m k) in general (Toth & Kelley 2015 Evans et al. 2020 De Sterck & He 2021).
For the application of AM, one 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.

---

Published as a conference paper at ICLR 2022
information into 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 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 develop the short-term recurrence variant, namely ST-AM. The basic version of ST-AM imposes some orthogonality property on the historical sequence, which inspired by the CG method (Hestenes & Stiefel 1952) that enjoys a three-term recurrence. Fur- thermore, 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; 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 the full- memory AM and the convergence rate is similar to that of the 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 op- timization are consistent with the theoretical results for the basic ST-AM and MST-AM. Furthermore, extensive experiments on training neural networks for image classification and language modeling 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 also known as an extrapolation algorithm in scientific computing (Anderson 2019). A paral- lel method is Shanks transformation (Shanks 1955) which transforms an existing sequence to new sequence for faster convergence. Related classical algorithms include Minimal Polynomial Extrap- olation (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 e-algorithm had been discovered for Shanks transformation for scalar sequence (Wynn] 1956). and was later generalized as the vector e-algorithm (Wynn) 1962) to handle 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 2aa9)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 fixed- point 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 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 gradi- ents can be accessed (Nemirovski & 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 prac- tice. When training neural networks, SGD with momentum (SGDM) Qiaa 1999), and adaptive learning rate methods, e.g. AdaGrad (Duchi et al. 2011), R.Sprrpp (Tillmman & Hinton 2012), Adam (Kingma& & Ba 2014), are very popular optimizers. Our methods have the nature of quasi- Newton methods while the memory footprint is largely reduced to be close to first-order methods. Thus, ST-AM can be competitive optimizer from both theoretical and practical perspectives.
2

---

Published as a conference paper at ICLR 2022
3 METHODOLOGY
In this section, we give the details of the proposed ST-AM. We always assume the objective function as f IR + IR, the fixed-point map g : IR + Id. Moreover, we do not distinguish ri: (ara) and rik = g(zx)) Th in our discussion as V fx) = 0 is equivalent to g(x) = I V fx(x) = x.
3.1 ANDERSON MIXING
The AM finds the fixed point of g via maintaining two sequences of length m (m S k)):
X = [AT my Ark-m+l,
-.x--I))) 1), Rs: = [Ark m. Ark-mttl m.t Ark-1) E Rdxm (1)
where the operator A denotes the forward difference, e.g. Arg 2k+1 Zg. Each update of AM can be decoupled into two steps, namely the projection step and the mixing step:
T& = Jk XR
(Projection step),
Tk+1 = Tx + BaTK, (Mixing step),
(2)
where FF = rk R and B& > is the mixing parameter. The F3 determined by
T3 arg min || TER
R2
(3)
Thus, the full form of AM (Fang & Saad 2009
Walker & Ni 2011 is
Ik-1 = I.k + BETE - (Xx + By R K
(4)
Remark 1. To see the rationality of AM, assume g continuously differentiable, then we have
M(x))-((zj-I) ze 1 (za)((z--jj)) around zs, where h' (xxx is the Jacobian of h(x) := x gx)(x)) So, it is reasonable to assume Rg R -h' (xxx)) and we see ||rk - R&T'I22 R || + + Thus, can recognize 3 as solving h (z
h(*x) in a least- squares sense, where dy XyTg. The mixing step incorporates Tk into the new update Ik+1 if 3x > O. Otherwise, if BB = 0,
then Tk+1 -x is interpolation of the previous iterates, leading a stagnation.
3.2
THE BASIC SHORT-TERM RECURRENCE ANDERSON MIXING
The basic ST-AM is to solve the strongly convex quadratic optimization:
min f(x) =
T 5a² x² Ax 6Tx,
(5)
where > p- Po 90 0 Rdd At the k-th iteration, given the two matrices
9-1
Pk-1 = (Pk-2.Pk-1) E Rdx2 Qi- (gk 2: 11)
E Rdxx and defining = Ig Ik-1 and
q = T3 Tk-1, the basic ST-AM constructs
- = p- Pa- (2-1(( q 1Q-1(0
(6a)
pk = 1111/112 9k q/llål12-
(6b)
Then, we update Ps = (p Pk)., Qx = (qk-1, gk) E Rdx2 Such construction ensures Q Qk = Ia for k 2 and the storage of Ps and Q& is equal to AM(2). With the corrected Ps and Qu. the ST-AM method modifies the projection step and the mixing step accordingly, that is,
F = F4 PETR. (Projection step). Fk+1 = 3x + BEFE: (Mixing step).
where To arg min [|rk- - Qu Q and Fs = r QuTi. Thus, the ST-AM replaces X and R& in (T) by Pa and Qs respectively and imposes the orthogonality condition on Qx. The details of basic ST-AM
given in Algorithm 2 in Appendix C.I Define Pe. (p1, P2.-.. ,P)), Qk
(41, l2..... qk), the Krylov subspace K,,, (A, v) = span e Av, APo.... A ".. the range of X
as range(X), We give the properties of the basic ST-AM in Theorem []]
Theorem 1. Let {xx) be the sequence generated by the basic ST-AM. The following relations hold:
(i) ||:]| 2 > 0, range Pt) ) = rangel Xx) = (.((A,r)), range((.) = range( Re) AKo(A, ro);
(ii) Q+ -APp, Q
Tui
(iii Fv 1 range( (Aag and 3h Do + 2k. where 2 arg min..xx 1|ro A2l12
If /|| = O, then - the exact solution.
3

---

Published as a conference paper at ICLR 2022
The proof is in Appendix C.I
Note that the property (iii) in Theorem exactly describes the relation Ik = zz. where 2x is the output of the k-th iteration of GMRES (Saad & Schultz, 1986). Moreover, let #xMh be the k-th intermediate iterate in the full-memory AM. holds that T% = (See (elendix(C.1),.
Proposition 1 in Appendix [.1 ) which induces that F& = =x = zx.- This equivalence indicates that ST-AM 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)), obtain the convergence rate of the basic ST-AM for solving (5):
Corollary 1. Suppose the eigenvalues of A lie in [u, L] with > O, and lei { be the se- quence generated by the basic ST-AM, then the k-th intermediate residual ix satisfies || S
2 Vifu-1 VLiu+1
||ro|l2. Moreover. the algorithm finds the exact solution in most (d 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] )) the convergence rate of ST-AM has the optimal dependence on the condition number, while ST-AM does not form the Hessian- products explicitly.
3.3 THE MODIFIED SHORT-TERM RECURRENCE ANDERSON MIXING
For general nonlinear fixed-point problems, global convergence may be unavailable for the basic ST- AM, as a coumter-example exists for AM (Mai Johansson 2020). Thus, we propose a modified version of the basic ST-AM (MST-AM) and prove the local linear convergence rate under similar conditions used in (Toth & Kelley, 2015 Evans al. 2020). Concretely, the MST-AM makes three main changes to the basic ST-AM.
Change 1: Instead of applying the normalization 6b the MST-AM constructs pk and qs via
SE = 1Q-
P
9s = 1K-1St
(8)
where " is the Moore- Penrose inverse. Accordingly, we choose T3 = arg min ||, -
12 =
(Q% )x) @TQ... This change relaxes the orthonormality for Qs (k > 2), but keeps the orthogonality
condition: QT 9h = 0. In fact, QT 9h = O in the case of solving 5
Change 2: MST-AM imposes the boundedness constraints on Pa-1 and Sandu--16 P P > cpl|pl12 or QQ=1 > cal1g12, then Ps P Qx Qx-1, where c, > 0, Cg E (0,1) are predefined constants, It is worth mentioning that adding some boundedness condition is common in the analysis of AM (Toth & Kelley. 2015 Evans et al. 2020).
Change 3: MST-AM restarts, i.e. setting Pa = Q+ = 0 E Rd every m iterations. This restart operation is to limit the number of higher-order terms appeared in the residual expansion in our analysis and we can set m to be a large number in practice.
The detailed description of MST-AM is given in Appendix In the next theorem, we establish the convergence rate analysis for the MST-AM.
Theorem 2. Let } be the sequence generated by MST-AM, x²' E IR be fixed point of and m be the restarting period for MST-AM. Suppose that in the ball B(p) == (i IR I| x 1|2 p} for some p > 0. g is Lipschitz continuously differentiable and there are constants K E (0, 1) and R > 0 with (i) (1 - -gxx)l22 S xily- 2llyy= for every x, y E B(p), and (ii) [g'((y) -000ggg(an gx(x)ll2 S kly fely-za xll2 xl12
for every I, y E B(p), where g' is the Jacobian of 9. Assume |1 Sal + KBR I Ko for a constant
Ki E (0, 1). If wo sufficiently close to x", then for rk == g(xx) Ng, the following bound holds:
me
||r S 0x((1 B&l + mBB))lrK
where 02 = [|F/x|12/11 rl < 1 and my = k mod m. Thus, the residuals {rk} converge Q-linearly, and the errors {1|= x" l12 converge R-linearly.
Remark 3. In a local region around x², the convergence rate is determined by the first-order term [] 3x1 + x:xx) +:8x))l&&l2 We can choose 3x 1 such that |1 3xl + KB = K < 1. Since Fk is the orthogonal projection of rak onto the subspace range( )+) Oz has the interpretation of
4

---

Published as a conference paper at ICLR 2022
Algorithm RST-AM for stochastic programming
Input: xp RR, Bx > 0, on E [0, 1,, /
0, %
>
Output: x E Rdd
1: Po. Qo = 0 E R Pn. 40 = 0 E R
2: for k 0. ,.,.,. until convergence, do
3: T = - -v,s, (#x)
4:
i > 0 then
5:
p = Ik 2b-1, == -k-1
1 1 -6 (1 Pk_ 11
6:
7:
ON - 1h-1 = P
8:
Pk Pk), Qx = (9
9:
end if
Check Condition 13 and use smaller og 13 is violated
11:
P& = (Q% Qk +
6 P Pt)
12:
Fk 3K Og PaT
k,
Th
Q&Q&T
13:
FA+1 = Fy + F
Apply learning rate schedule of or 3k
15: end for
return 2%
the direction-sine between Tkk and the subspace range(Qk). When 02 is small, e.g., TІ nearly lies in range( (Qx). the acceleration by MST-AM significant. Compared to AM(m). MST-AM incorporates historical information with orthogonalization. In the SPD linear case and without restart, the global orthogonality property holds, iee. Ti: range(Qu). which means there is no loss of historical information, while AM(m) (Evans al. 2020) does not have such property in this ideal case.
3.4 THE REGULARIZED SHORT-TERM RECURRENCE ANDERSON MIXING
Inspired by the recent work on stochastic Anderson mixing (SAM) method (Wei et al. 202 we
develop a regularized ST-AM (RST-AM) for solving nonconvex stochastic optimization problems.
Consider the nonconvex optimization problem min f(x) := + fç. (z). where fe. : Rd
IR is the loss function corresponding to i-th data sample and T is the number of data samples. In mini-batch training, the gradient evaluated for fs, (Ixx = mu Eies, (rrx), where S₁ S [7] = (1.2...... TT is the sampled mini-batch, and ny := SS is the batch size. In this case, we set Fh Vs, (Irx) (Line 3 in Algorithm D which is unbiased estimate of the negative gradient.
arg min || Qu-1Sll2 q Ark-1 by definition. Since ST-AM is based on local quadratic approximation (5) in a small region around Ik. large magnitude of [|P tends make the change from Ark-1 pk Ara-1 Pk-1Ca too aggressively, which may lead to instability. Consequently, we add penalty term in the above least squares problem, i.e.
Recalling from 8). CL = ( (T
CA = arg min |Ark-1 - Qa- + (2) 1
(10)
where ((4) > The same as SAM Wei et al. 2021 ) we also add regularization term for comput- ing I' via
D = arg min |r - Q & |
(11)
where 5(2) > O. and a damping term is used as shown in Line 12 in Algorithm In practice, we choose the two regularization parameters as
cilTra|
max 02 1ella 11
(1))
Ar 11 + ((2)
CB=
(12)
[pel
where C1, 02 C away from zero
constants, and eo 0 is small constant to bound the denominators Assuming || ~ | = O1]A2z=111)) the choices of (12 make
(1o Pr-lle O(1irx11)) and (b((0) p Palla O(Irx)11)) aware of the change of the local curvature: large (small) |rxl|22 tends to lead large (small) regularization.
5

---

Published as a conference paper at ICLR 2022
Remark 4. One update of Iı given by Line 11-13 in Algorithm Can be formulated Ik+1
Ik + HYTK, where H2 = -3 as o 2 Q Yv = Ps + BrQu, Zz = Q% Qu + 0% Pk Pk- To guarantee the positive definiteness of H2, we follow the same procedure in SAM. Let As be the largest eigenvalue of 2 + Qu Z Ifa; satisfies
OR S (1-)).
(13)
A then 86 HASK 2 Beu][ssIle, Vag E Rd, where fs E (0, 1) is constant. Note that X5. can be cheaply obtuined by computing the largest eigenvalue of a matrix of RAX (see Appendix C.3.1
We summarize the RST-AM in Algorithm D and establish its convergence properties here. First, we impose the same assumptions on the objective function f as those in (Wei al., 202I).
Assumption 1. f IR + IR is continuously differentiable. f(x) 2 flow > -o0 for any x E IR... f is globally L-Lipschitz continuous; namely (x)
L|lx-yl22 for any x, y E Rdd
Assumption 2. For any iteration k. the stochastic gradient ffç, (xx) satisfies Es, fV. (Ix)] =
V f(xx), Ee, IV (ax) - v ( S o², where a > 0, and Ex: 0, l,... .. are independent samples that are independent of {
The diminishing condition about 3x is
+
+ 32 < +00.
Σ
3x = +00,
Σ k==0
(14)
k=
We give the convergence properties of RST-AM in nonconvex (stochastic) optimization and proofs are deferred to Appendix C.3.2
Theorem 3. Suppose Assumption hold and {5x] is the sequence generated by full-batch RST-AM,
i.e. ni = T. Let 3x = 3 E (0, 20(1C 15) be constant, ag E [0, 1] and satisfies (13) then
N-1
2(f(zo)- - flom)
(15)
N
Nu3
in the N iterations. To ensure +
E IV (Ix)) < e, the number of iterations is O(L/e).
Theorem 4. Suppose Assumptions
7 and
2 hold and {xx} is the sequence generated by RST-AM
with batch size ni = S T. If Be E (0, T 1y) and satisfies 14), ag E min 8 H and
satisfies (13, then
lim inf [V = 0 with probability and 3M, > 0 + [[f((]) & Mj (16)
R-+00
[, V Se, (xx)
[2] < Mg Vk, where Mg > is constant, have [V(()))) = O with probability 1.
lim k
(17)
Theorem 5. Suppose Assumptions
7 and 2 hold and {
N-1 is the first iterations generated by
D RST-AM with fixed batch size ng. Let Bx minf 1 o/N ) where D is a problem-
independent constant; are [0, min(1, 18 H and satisfies (13) Let R be random variable follow- ing PR(K) == P = k} = 1/N, then
[[|VV[(X))1 16D, L(1 + C-1)
(4D, D
4(L + 1 +1+++=))+++++ )
(18)
2
AVN
n
where Dj == f( and the expectation is taken with respect to R and (5 N-1 To ensure
B[IV [(x))]]] S the number of iterations is (1)/e)).
Remark 5. The proofs of Theorem and 5 are based on the analysis of SAM (Wei et al. 202 The theorems show that the convergence of RST-AM is no worse than SGD (Robbins & Monro 1951}. There are two key differences between RST-AM and SAM: RST-AM is based on short-term recurrences while SAM usually maintains longer historical sequences to ensure effectiveness; RST- AM uses additional correction and regularization terms (Line 5-8 in Algorithm 7 to incorporate historical information while SAM simply discards the oldest iteration to make space for A and Ark-1- The reduced memory requirement in RST-AM makes it applicable for solving more challenging problems in machine learning.

---

Published as a conference paper at ICLR 2022
4 EXPERIMENTS
We validated the effectiveness of our proposed ST-AM methods in various applications in fixed-point iterations and nonconvex optimization, including linear and nonlinear problems, deterministic and stochastic optimization. Specifically, we first tested ST-AM in linear problems, cubic-regularized quadratic minimization (Carmon Duchi 2020) and a multiscale deep equilibrium (MDEQ) model (Bai et al., 2020). Then we applied RST-AM to train neural networks and compared them with several first-order and second-order optimizers. Experimental details are in Appendix
4.1 EXPERIMENTS ABOUT THE BASIC ST-AM AND MST-AM
We verified the properties of ST-AM declared in Theorem D and by solving four problems (details are in Appendix D.I): (I) strongly convex quadratic optimization (corresponding to Theorem (II) solving a nonsymmetric linear system Ar b (corresponding to A = 0 in Theorem 2)
cubic-regularized quadratic minimization mingeR f(x) ) = || Ar b12 + 3
M ||=|| (corresponding to
A 0 in Theorem 2): (IV) root-finding problems in MDEQ on CIFAR-10 (Krizhevsky al. 2009)
The compared methods were gradient descent (GD), fixed-point iteration (FP), conjugate residual method (CR) (Saad) 2003), BFGS (Nocedal & Wright 2006), Broyden's 's method (Broyden 1965), and the full-memory AM (AM). We used the basic ST-AM to solve Problem I and II, and MST-AM (cp Cg 1) solve Problem III and IV.
10-2
10-1
10-3
10
10-7
10
10-3
10-41
10
10:11
Q2
44
10-11
JPGS
10
20
30
40
50
10
50
10
20
30
40
50
ieration
(a) Problem
(b) Problem II
(c) Problem III (M 0.1)
Broydoo XM
Broydon AM
Black Blackwaards Iattil
Forewrdt ([Ac)))f11
2.6
MSTAAM
MST-MM
R. S
Goodin
step
Step
eooch
(d) Forward process
(c) Backward process
(1) Test accuracy and loss
Figure 1: (a) ||rx]|// /roll/ of GD and CR, and || /roll///rolle 12 of AM and ST-AM for solving Prob- lem I; (b) [|x]|/// rooll for solving Problem II; (c) ||r 2 ||ro ll for solving Problem III (M 0.1); (d)(e) relative residuals of the forward and backward root-finding processes in MDEQ, and shaded areas correspond the standard deviations; (f) test accuracy and loss MDEQ/CIFAR-10.
The numerical results shown in Figure I demonstrate the power of ST-AM as a variant of Krylov subspace methods. It significantly accelcrates the slow convergence of the GD or FP method, and can outperform AM. Figure a clearly verifies the correctness of Theorem I within the machine precision, the intermediate residual Fs of ST-AM coincides with the residual of CR. Note that AM fails to coincide with CR and ST-AM due to the intrinsic numerical weakness to solve 3 as also pointed out in (Walker & Nij 2011). Figure (b) shows that ST-AM can outperform CR. though both methods enjoy short-term recurrences and are equivalent for solving SPD linear systems. Fig- ure TTc) also shows MST-AM surpasses BFGS in solving cubic- cubic-regularized problems. The tests in MDEQ/CIFAR- 10 indicate that MST-AM is comparable to the full- full-memory methods the forward root-finding process and converges faster in the backward process. The accuracy is also comparable.
7

---

Published as a conference paper at ICLR 2022
10"
nom gradien:
10t
10
10-2
Adagred
10
5GD
Mam
Adam
20"
Adagred
10-4
5.81221
10-4
RSTAN
AST-NM
100
(a) Train loss (w/o preconditioning)
(b) SNG (w/o preconditioning)
Train loss (w/ preconditioning)
Figure 2: Experiments MNIST. (a)(b) Training loss and the squared norm of gradient (SNG) (w/o
preconditioning for SAM, RST-AM); (c) Training loss (w/ preconditioning for SAM, RST-AM).
4.2 EXPERIMENTS ABOUT RST-AM
We applied RST-AM to train neural networks, with full-batch training on MNIST (LeCun et al. 1998), and mini-batch training on CIFAR RAAARIO/CIFAR-OO and Penn Treebank (Mareus et 1993).
Experiments on MNIST. We trained a convolutional neural network (CNN) on MNIST to see the convergence behaviour of RST-AM in nonconvex optimization (ef. Theorem 3. for which we were only concerned about the training loss. Figureßaa(h) show that the short-memory SAMs (m 2) hardly show any improvement over the first-order optimizers SGD and Adam, while RST-AM can close the gap of the long-memory (m 10) and the short-memory methods. We also considered the effect of preconditioning on RST-AM (see Appendix A.3). The notation "A_B" means B method preconditioned by A method. Figure 2e indicates that preconditioning also works much better for RST-AM than SAM(2), and RMSprop RST-AM can outperform the non- non-preconditioned RST-AM.
Table 1:
Experiments on CIFAR1O/CIFARION "L" means failing to complete the test in our device
due memory limit. "** indicates numbers published in (Wei et al. 2021
(a) Final TOP1 test accuracy (mean + standard deviation) (%) for training 160 epochs.
Test accuracy on CIFARIO
Test accuracy on CIFARI00
Method
ResNet18
ResNet20
ResNet32
ResNos44
ResNet56
WRNI6-4
ResNet18
ResNeXt
DenseNet
SGDM
94.821.15
92.03...16
92.86.1.15
93.10.1.33
93.47
94.901.09
77.27.109
78.41.154
78.49.1.12
Adam"
93.032.07
91.171.13
92.03-128
92.28:.62
22.39-223
92.452.11
72411..7
73.574.11
70.80:1.23
AdaBound
94.25+.31
90.77+.08
91.73+.06
92.00=.18 18
92.44-04
93,50=.12
75.07+
75.74+:20
76.06.t.1
AdaBelief
94.65+.13
91.15+-21
92.15++11
92.79--24
93.30-07
94.46+.13
76.254.06 06
78.27+t116
78.83-+.15
Lookaheud"
92.071.04
92.86.1.15
93.36-..13
AdaHessian"
94.3611.09
91 92:1.32
92. 18.18:1.18
92.742.11
92.40-066
94,041.12
76.591.42
SAM12)
95.07:.04
92.142.33
23
93.462.09
93.66-006
95.07=.16
77.51:224
79.024421
80.00:1:29
SAM(IO)*
15
92.43+.19
03.22+3-22
93.57=214
93.77-.12
98.23-.07
78.13+.14
79.31
80.09:t.5
RST-AM
95.27+.04
92.39+-11
93.241.36
93.52 02
93.69
95.21
77.91+.22
79.53-134
80.36
(b) The memory and computation cost compared with SGDM. The notations "m","tle" and "t" are abbrevia- tions of memory, per-epoch time and total running time, respectively.
Cost
CIFARI0/RsNee118
CIFARIO/WRNI6-4
CIFARI00ResNee1500
CIPAR100/DensNee111
(X SGDM)
t/c
U/c
SGDM"
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
SAM(IO)*
1.73
1.78
1.00
1.26
1.28
0.80
1.30
1.16
0.58
1.16
1.19
0.60
RST.AM
Experiments on CIFAR. We trained ResNe118//0//2/444/1 (He et al. 2016). WideResNet16-4 Zagoruyko & Komodakis] 2016) (abbr. WRN16-4) on CIFAR-TO, and ResNetT8, ResNeXi50 (Xie)
ct al. 2017 7). DenseNet121 (Huang et al. 2017) on CIFAR- 100. The baseline optimizers were
SGDM, Adam, AdaBound (Luo et al. 2018), AdaBelief (Zhuang et al. 2020), Lookahead (Zhang et al. 2019). AdaHessian (Yao et al. 2021) and SAM. Here, some results of the baselines in Wei
used for reference since the experimental settings were the same. Table [(a] shows RST-AM improves SAM(2) and has comparable test accuracy to SAM(10). RST-AM also outper-
al. 2021
forms other baseline optimizers. Table ((b) reports the memory and computation cost, where
8

---

Published as a conference paper at ICLR 2022
used SGDM the baseline and other optimizers were terminated when achieving comparable or better test accuracy than SGDM. indicates that RST-AM introduces Σ 5% extra memory over- head compared with SGDM, and significantly reduces the memory footprint of AM. Since RST-AM needs fewer training epochs, the total running time is less than SGDM.
220
100
100
SCDN
Adam
Azabeld
SAM121
SAMI21
Validatio
www.mm
Validaaion
Valldstion
05
320
2033
110
116
100
1.00
200
020
opoch
epach
opoch
(a) 1- L-Layer LSTM
(b) 2-Layer LSTM
(c) 3-Layer LSTM
Figure 3: Validation perplexity of training 1,2,3-layer LSTM on Penn Treebank.
Experiments Penn Treebank. We trained LSTMs with 1-3 layer(s) on Penn Treebank and
Table 2: Test perplexity of training 1,2,3-layer
LSTM on Penn Treebank. Lower is better.
report the validation perplexity in Figure 3 and test perplexity in Table (lower is better). The results suggest that RST-AM comparable to or even better than SAM(10). The improvement of RST-AM over other optimizers is also sig- nificant. We report the computation and mem- ory cost in Appendix D.2.4] RST-AM can still surpass Adam while using much fewer epochs, thus reducing the total running time.
Method
I-Layer
2-Layer
3-Layer
SGDM
83.481.03
65.89.1.18
61.88.123
Adam
80.334.15
64.32:1.6
59.72:2.13
AdaBelief
11.2
64.68.1.10
60.46.1.07
SAM(2)
80.791.19
65.52.1.29
61.1311.2
SAM(10)
78.782.11
62.46::.11
58.932.09
RST-AM
78.41
62.46.1.08
58.31 1.23
Table 3: Test accuracy (%) for adversarial training.
Table 4:
FID score for SN-GAN.
Mathod
Adam
AdaBelief
RST-AM
Optinizer
Cless
PGSM
PGD-20
C&W
Clean
FISM
PCID-20 CRW
Best FID
13.074.88
12,80-109
12.05
SGD
12
63.23
$1.91
30.22
39 25
39.76
30.92
29:00
Final FID
13.341.14
13.591.21
12.501.29
RM
N2.53
63.75
$2.43
50.52
69.45
40.41
31.20
29.52
2018 as the outer- optimizer and compared it with SGD by the clean test accuracy and robust test accuracy. The results on CIFAR 10/ARRet/ and CIFAR 100/DenseNer121 are reported in Table It can be seen that RST-AM can achieve both higher clean test accuracy and higher robust test accuracy. More results can be found in Appendix D.2.5
Adversarial training. We applied RST-AM to adversarial training (Madry et al.
Generative adversarial network (GAN). We tested RST-AM by training a GAN equipped with spectral normalization (SN-GAN) (Miyato al. 2018), where the generator and discriminator net- works were ResNets and the dataset was CIFAR- 10. Table shows that RST-AM can achieve lower FID score (better accuracy) than Adam and AdaBelief.
CONCLUSION
In this paper, address the memory issue of Anderson mixing (AM), we develop a novel class of short-term recurrence AM methods (ST-AM) and test it in various applications. including solv- ing linear and nonlinear problems, deterministic and stochastic optimization. We give a complete theoretical analysis of the proposed methods. We prove that the basic ST-AM is equivalent to the full-memory AM in strongly convex quadratic optimization. With some minor changes, it has local linear convergence for solving general fixed-point problems under some common assumptions. We also introduce the regularized form of ST-AM and analyze its convergence properties. The numerical results show that the ST-AM methods comparable while consuming less memory. The regularized ST-AM also outperforms many existing optimizers in training neural networks in various tasks.
even better than the mam-mmoory AM

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
This work was supported by the National Key R&D Program of China (No. 2021YFAJ001300 National Natural Science Foundation of China (No.61925001,, Tsinghua University Initiative Scien- tific Research Program, National Natural Science Foundation of China (No. 1991119,333), and Huawei Noah's ss Ark Lab. We thank all anonymous reviewers for their valuable comments and suggestions on this work.
REFERENCES
Hengbin An, Xiaowei Jia, and Homer Walker. Anderson acceleration and application to the three- temperature energy equations. Journal of Computational Physics, 347:1-19, 2017.
Donald G Anderson. Iterative procedures for nonlinear integral equations. Journal of the ACM (JACM), 12(4)5547-560, 1965.
Donald G Anderson.
Comments on " "Anderson acceleration, mixing and extrapolation". Numerical Algorithms, 80(1):135-234 2019.
Marcin Andrychowicz, Misha Denil, Sergio Gómez Colmenarejo, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando de Freitas Learning to learn by gradient descent by gradient descent. In Proceedings of the 30th International Conference Neural Information Processing Systems, pp. 3988-3996, 2016.
Akash Arora, David C Morse, Frank S Bates, and Kevin D Dorfman. Accelerating self-consistent field theory of block polymers in variable unit cell. The Journal of Chemical Physics, 146(24): 244902, 2017.
Shaojie Bai, J Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances in Neural Information Processing Systems, 32:690-701, 2019.
Shaojie Bai, Vladlen Koltun, and J. Zico Kolter. Multiscale deep equilibrium models. In Advances in Neural Information Processing Systems (NeurlPS), 2020.
Albert S Berahas, Frank E Curtis, and Baoyu Zhou. Limited-mennn memory BFGS with displacement aggregation. Mathematical Programming, pp. 1-37, 2021.
Lon Bottou, Frank E. Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM Review, 60(2):223 311, 2018.
Claude Brezinski and Michela Redivo ZZglia. Shanks function transformations in vector space. Applied Numerical Mathematics, 1 16:57-63, 2017.
Claude Brezinski, Michela Redivo-Zaglia, and Yousef Saad. Shanks sequence transformations and Anderson acceleration. SIAM Review, 60(3):646-669, 2018.
Charles G Broyden. A class methods for solving nonlinear simultaneous equations. Mathematics of Commuutaton, 19(92):577-593, 1965.
Richard H Byrd, Samantha L Hansen, Jorge Nocedal, and Yoram Singer. A stochastic quasi-Newton method for large-scale optimization. SIAM Journal on Optimization, 26(2):1008-10311 2016.
Stan Cabay and LW Jackson. polynomial extrapolation method for finding limits and antilimits of vector sequences. SIAM Journal Numerical Analysis, 13(5):734-752, 1976.
Nicholas Carlini and David Wagner. Towards evaluating the robustness of neural networks. In 2017 IEEE Symposium on Security and Privacy (SP), pp. 39-57. IEEE, 2017.
Yair Carmon and John Duchi. First-order methods for nonconvex quadratic minimization. SIAM Review, 22(2):395-436 2020.
Mauro Cettolo, Jan Niehues, Sebastian Stüker, Luisa Bentivogli, and Marcello Federico. Report on the 11th IWSLT evaluation campaign, IWSLT 2014. Proceedings of the International Workhhop Spoken Language Translation, Hanoi, Vietnam, volume 57, 2014.
10

---

Published as a conference paper at ICLR 2022
Hans De Sterck and Yunhui He. On the asymptotic linear convergence speed of Anderson accelera- tion, Nesterov acceleration, and nonlinear GMRES. SIAM Journal on Scientific Computing, (0): S21-S46, 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Fei-Fei. ImageNet: A large-scale hier- archical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp. 248-255. IEEE, 2009.
Yunbin Deng. Deep learning on mobile devices: A review. In Mobile Multimedia/lnage Processing, Security, and Applications 2019, volume 10993, pp. 109930A. International Society for Optics and Photonics, 2019.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(7), 2011.
lain Duff, Albert Maurice Erisman, and John Ker Reid. Direct methods for sparse matrices.
Oxford University Press, 2017.
Rick Durrett. Probability: Theory and examples, volume 49. Cambridge University Press, 2019.
P Eddy, Extrapolating to the limit of vector sequence. In Peter C.C. Wang, Arthur L. Schoenstadt, Bert I. Russak, and Craig Comstock (eds.), Information Linkage Between Applied Mathematics and Industry, pp. 387-396. Academic Press, 1979.
Claire Evans, Sara Pollock. Leo G Rebholz, and Mengying Xiao. A proof that Anderson acceler- ation improves the convergence rate linearly converging fixed-point methods (but not those converging quadratically). SIAM Journal on Numerical Analysis, 58(1):788 8100 2020.
Haw-ren Fang and Yousef Saad. Two classes of multisecant methods for nonlinear acceleration. Numerical Linear Algebra with Applications, 16(3):197-221, 2009.
Anqi Fu, Junzi Zhang, and Stephen Boyd. Anderson accelerated Douglas-Rachford splitting. SIAM Journal on Scientific Computing, 23360-3350-A3355 2020.
Alejandro Garza and Gustavo E Scuseria. Comparison of self-consistent field convergence accel-
eration techniques. The Journal of Chemical Physics, 137(5):054110, 2012.
Saced Ghadimi and Guanghui Lan. Stochastic first-and zeroth-order methods for nonconvex stochas- tic programming. SIAM Journal on Optimization, 23(4):2341-2368 2013.
Gene H Golub and Charles F Van Loan. Matrix computations, 4th. Johns Hopkins, 2013.
Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. arXiv preprint 112.6::41 2014.
William W Hager and Hongchao Zhang. A survey of nonlinear conjugate gradient methods. Pacific
Journal of Optimization, 2(1):35-58, 2006.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pp- 770-778, 2016.
Magnus R. Hestenes and Eduard Stiefel. Methods of conjugate gradients for solving linear systems. Journal of Research of the National Bureau of Standards, 49:409 435,, 1952.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by two time-scale update rule converge local nash equilibrium. Advances in Neural Information Processing Systems, 30, 2017.
Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinherger. Densely connected convolutional networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 4700-4708, 2017.
Carl T Kelley. Numerical methods for nonlinear equations, Acta Numerica, 27:207-287, 2018.
11

---

Published as a conference paper at ICLR 2022
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Tamara G Kolda, Dianne P O'leary, and Larry Nazareth. BFGS with update skipping and varying memory. SIAM Journal on Optimization, 8(4):1060-1083, 1998.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, (12(111::272235553 1998.
Lin Lin, Jianfeng Lu, and Lexing Ying. Numerical methods for Kohn-Sham density functional theory, Acta Numerica, 28:405-539,, 2019.
Dong Liu and Jorge Nocedal. On the limited memory BFGS method for large scale optimization. Mathematical Programming. 45(1):503-52, 1989.
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. In International Conference Learning Representations, 2019.
Liangchen Luo, Yuanhao Xiong, Yan Liu, and Xu Sun. Adaptive gradient methods with dynamic bound of learning rate. International Conference on Learning Representations, 2018.
Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. In International Conference on Learning Representations, 2018.
Vien Mai and Mikael Johansson. Anderson acceleration of proximal gradient methods. Interna- tional Conference on Machine Learning, pp. 6620-6629, PMLR, 2020.
Mitchell Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz. Building large annotated corpus of English: The Penn Treebank. 1993.
Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. In International Conference on Learning Representations, 2018.
Aryan Mokhtari, Mark Eisen, and Alejandro Ribeiro. IQN:
An incremental quasi-Newton method with local superlinear convergence rate. SIAM Journal on Optimization, 2872)11670-168, 2018.
Ion Necoara, Yu Nesterov, and Francois Glineur. Linear convergence of first order methods for non-strongly convex optimization. Mathematical Programming, 175(1):69-107, 2019.
Arkadij Semenovič Nemirovski and David Borisovich Yudin. Problem complexity and method efficiency in optimization. 1983.
Jorge Nocedal and Stephen Wright. Numerical optimization. Springer Science & Business Media, 2006.
Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. BLEU: A method for automatic evaluation of machine translation. In Proceedings of the 40th Annual Meeting of the Association for Computational Linguistics, pp. 311-318, 2002.
Florian A Potra and Hans Engler. A characterization of the behavior of the Anderson acceleration on linear problems. Linear Algebra and Its Applications, 102-3)1002-1011, 2013.
Ning Qian. On the momentum term in gradient descent learning algorithms. Neural Networks, 12 (1):145-151, 1999.
Leslie Rice, Eric Wong, and Zico Kolter. Overfitting in adversarially robust deep learning. In International Conference Machine Learning, pp. 8093-8104, PMLR, 2020.
12

---

Published as a conference paper at ICLR 2022
Herbert Robbins and Sutton Monro. A stochastic approximation method. The Annals of Mathemat- ical Statistics, pp. 400-407, 1951.
Youcef Saad and Martin H Schultz. GMRES: A generalized minimal residual algorithm for solving nonsymmetric linear systems. SIAM Journal on Scientific and Statistical Computing, 7(3):856- 869, 1986.
Yousef Saad. Iterative methods for sparse linear systems. SIAM, 2003.
Damien Scieur, Alexandre dAspremont, and Francis Bach. Regularized nonlinear acceleration. Mathematical Programming, 179(1):47-83, 2020.
Daniel Shanks. Non-linear transformations of divergent and slowly convergent sequences. Journal of Mathematics and Physics, 34(1-4):1-42 1955.
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initial- ization and momentum in deep learning. In International Conference on Machine Learning, pp- 1139-1147, 2013.
Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-RMSprop: Divide the gradient by running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2):66 31, 2012.
Alex Toth and CT Kelley. Convergence analysis for Anderson acceleration. SIAM Journal on Numerical Analywis, 53(2):805-819) 2015.
Henk Van der Vorst and Vuik. The superlinear convergence behaviour of GMRES. Journal of Computational and Applied Mathematics, 48(3):327-341, 1993.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In Advances Neural Infor- mation Processing Systems, pp. 5998-6008, 2017.
Homer Walker and Peng Ni. Anderson acceleration for fixed-point iterations. SIAM Journal on Numerical Analysis, 49(4):1715-11) 1735, 2011.
Xiao Wang, Shiqian Ma, Donald Goldfarb, and Wei Liu. Stochastic quasi-Newtor methods for nonconvex stochastic optimization. SIAM Journal on Optimization, 27(2):927-956, 2017.
Fuchao Wei, Chenglong Bao, and Yang Liu. Stochastic Anderson mixing for nonconvex stochastic optimization. Advances in Neural Information Processing Systems, 34, 2021.
Peter Wynn. On a device for computing the em (S,) transformation. Mathematical Tables and Other Aids Computation, pp. 91-96, 1956.
Peter Wynn. Acceleration techniques for iterated vector and matrix problems. Mathematics of Computation, 16(79):301-322, 1962.
Saining Xie, Ross Girshick, Piotr Dollár, Zhuowen Tu, and Kaiming He. Aggregated residual trans- formations for deep neural networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 1492-1500, 2017.
Yunan Yang. Anderson acceleration for seismic inversion. Geophysics, 86(1):R99-RI0) 2021.
Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael Mahoney. AdaHessian: An adaptive second order optimizer for machine learning. Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 10665-10673, 2021.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In British Muchine Vision Conference 2016. British Machine Vision Association, 2016.
Michael Zhang, James Lucas, Jimmy Ba, and Geoffrey E Hinton. Lookahead optimizer: k steps forward, step back. In Advances in Neural Information Processing Systems, pp. 9597-9608, 2019.
Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar C Tatikonda, Nicha Dvomek, Xenophon Pa- pademetris, and James Duncan. AdaBelicf optimizer: Adapting stepsizes by the belief in observed gradients. Advances in Neural Information Processing Systems, 33, 2020.
13

---

Published as a conference paper at ICLR 2022
A ADDITIONAL PRELIMINARIES
We provide some additional preliminaries in this section for readers that are not familiar with An- derson mixing, fixed-point iterations and some techniques mentioned in the main paper.
A.1 FIXED-POINT ITERATION
The fixed-point problem and the optimization problem are the main application scenarios of our methods. It is worth pointing out that there are some minor differences between these two prob- lems that make algorithm designs different. The key difference is that for a fixed-point problem, the Jacobian (if exists) is generally not symmetric while for optimization, the Hessian is naturally symmetric. In principle, a fixed-point solver can also be applicable for optimization problem since the frrs- necessary condition of min f(r), where f R - R, Vf =
Consider a contraction mapping g R Rdd i.e. for some K < 1, l[g(z) g(y)ll S =l|l ull2, Var, E Rd. According to the contraction mapping theorem, unique fixed point x' exists for g and the iterates generated by the iteration 12+1 = g(zx) converge to x², starting from any To E Rd, In practice, the damped fixed-point iteration is also commonly used: given the B& (0, 2). the update is 26+1 = (1 - Ba) + 3kg(xx) = 24 + B&TK, where тІ == g(rx) Ik is called the residual.
The fixed-point iteration, also known as Picard iteration in some areas, can converge very slowly in practice. Anderson mixing is method to improve the convergence.
A.2 ANOTHER FORM OF ANDERSON MIXING
The derivation of Anderson mixing (AM) in Section 3.1 explicitly interprets AM as two-step procedure. the literature, there is another equivalent form of AM.
Let the projection coefficients FÅ (1) 2 (m)yT E RT. Define the auxiliary coefficients (1)
(06) m 0 0
(2,00)
- ( : 1). O(!) = 7k
(11) then Σド= 0
76
,000
1 and Fs =
g Tk Tk-m+j. Hence the least squares problem (3 can be reformulated as Eo 0f5) Tk-m s.t. Ej 0
constrained problem min (o( ymm
1, which indicates that
as linear combination of the historical residuals (ry k Fa is minimal terms of the Ly-norm.
Also, the projection step and the mixing step 2 can be reformulated as 3 Ey= 86 Ik-m+j
and 1 = (1 - 3x) Σø= of I'a m 3x > 1=n 0(1) g respectively. Such formulation is also adopted in the literature (Toth & Kelley/ 2015 Mai & Johansson 2020 Scicur et al. 2020).
Let H3 be the solution to the constrained optimization problem (Fang & Saad 2009)
min || H - 3r I || subject to Hy Rg = -Xx,
(19)
then the update (4) is Tk+1 Ix HETK- k- It suggests that AM is multisecant quasi-Newton method FFng & Saad 2009).
A.3 PRECONDITIONED ANDERSON MIXING
Like preconditioning for Krylov subspace methods (Saad 2003). preconditioning can also be incor- porated into Anderson mixing to mitigate the ill-conditioning of the original problem (Wei et al.) 2021). The idea is to replace the mixing step in (2) via a preconditioned mixing.
Suppose that there is basic solver preconditioneer(na. sk), which works as black-box procedure that updates 2& given the residual sk. i.c. xa+1 = preconditioner (2 sk). then the preconditioned mixing of RST-AM is
Tt preconditioner((22 (@g, Fk))
which substitutes for the Line 13 in Algorithm I The simple mixing (Line 13 in Algorithm can be seen as a special case by defining preconditioner((gg (xg. = 2g + Sg i.e. precondi- tioned by damped fixed-point iteration. Moreover, if write the preconditioning operation as
14

---

Published as a conference paper at ICLR 2022
preconditioner sR) == 23 SkSk, where G₁ is the matrix to approximate the inverse Jaco- bian, then the preconditioned AM is
Ek+1 = T4 + GETK - (X2 + GRRA)(R T
(20)
(cf. the definitions of XR, R& in Section 3.1
The matrix Hx G2 (Xx + (R R
forms low-rank updated upproximation to the inverse Jacobian: it solves
min H4 - G || subject to HyR, -X H2
which is direct extension of 19
For the stochastic Anderson mixing, using damped projection can be helpful, i.e. 7y (1 - + ag - = Ia ax'' and Fu Tk ar RxRR''s correspondingly. Then the preconditioned AM with damped projection is
Th+1 Tk GaTE ax Ga
(21)
B ADDITIONAL DISCUSSION
B. THE MEMORY AND COMPUTATIONAL EFFICIENCY
Since the ST-AM methods only need to store two previous iterations, the memory and computational cost can be reduced to be close first-order methods.
Memory cost. Assume that the iteration number is k and the model parameter size is d. The full- memory AM stores all previous iterations, thus the additional memory 2kd. To reduce the memory overhead, the limited- AM(m) only maintains the most recent m iterations while discarding the older historical information (cf. (IP). Hence the additional memory of AM(m) 2md. Choosing (Walker & Ni, 2011). and it is often suggested that > (Ma & Johansson 2020 Fu et al. 2020 to avoid too much historical information being discarded. There is equivalence between AM(m) and the full-memory AM. On the other side, ST-AM methods have the same memory footprint as that AM(2), i.c. only introducing 4d additional memory. For a large model that contains millions or even billions of parameters, such reduction in memory requirement can be significant.
suitable can be problem-depender
Computational cost. For ST-AM, besides the gradient evaluations, the main additional computa- tional cost comes from vector corrections, the projection and mixing step. These operations are very cheap: matrix multiplications of 2x x Rdx2, R2xd x Rdx1, and IR Rdx22 R2211; the pseudo- inverse can be exactly solved as the size of the matrix R2x2, Also, the 1 Qx-1 and Pk_1Pk Pk-1 in Line in Algorithm [ can reuse the results in the previous iteration (cf. Line 11). So the total additional computational cost of RST-AM is (d).
B.2 APPLICABILITY
In the main paper, we develop a class of short-term recurrence Anderson mixing. Among them, the basic ST-AM can serve new linear solver for linear systems; MST-AM can be used fixed-point solver for nonlinear equations; RST-AM is a new method for optimization. These meth- ods are based Anderson mixing and have close relationship between each other. The theoretical analysis and numerical results show the great potential of ST-AM methods for various applications. Here, we give some comparisons between the ST-AM methods and some other classical methods.
ST-AM versus current solvers for linear systems. ST-AM is an iterative solver. It is equivalent to the full- full-memory AM and GMRES in strongly convex quadratic optimization (or SPD linear sys- tems), and can also have linear convergence rate for general nonsymmetric linear systems. Since it is based on short-term recurrences, like the CG method, the memory and per-iteration cost ST-AM is economical. On the other side, the LU factorization based methods (Golub & Van Loan 2013) are direct solvers that can incur overwhelming memory and computation cost for large sparse linear systems. Although sparse direct solvers (Duff et al. 2017) alleviate overhead, they often more complicated to implement and difficult for parallel computing. Moreover, iterative solvers can benefit from the preconditioning technique that improves convergence (Sadd 2003). Also, ST-AM
15

---

Published as a conference paper at ICLR 2022
has advantages over other iterative methods since it does not need to directly access the matrix and only the residual is required. In the case that explicitly accessing the matrix is difficult, this property is appealing. Hence, unlike nonlinear CG that relies line search, it is direct to extend ST-AM to unconstrained optimization where the gradient is commonly available while the Hessian can be too costly obtain. Moreover, ST-AM is very flexible and any iterative solver can be viewed black-box iterative process to be accelerated by ST-AM.
MST-AM versus other quasi-Newton methods for nonlinear equations. MST-AM has the na- ture of quasi-Newton methods since is built upon AM that is recognized as a multisecant quasi- Newton method (Fang & Saad 2009). In scientific computing, AM has been successfully applied to
solve many difficult nonlinear problems arising from many areas (An al. 2017 Lin et al. 2019 Fu et al. 2020 Yang 2021). Since AM only manipulates residuals and does not use line-search or trust-region technique, it is efficient to apply AM to accelerate a slowly convergent black-box iterative process. A comprehensive discussion about the applicability of AM for nonlinear problems and the relation between AM and Broyden's methods can be found in (Fang & Saad 2009).
One of the biggest issues of AM and other quasi-Newton methods is the additional memory over- head, because they need to store historical iterations to form the secant equations. To make a com- promise, limited-memory quasi-Newton methods such as L-BFGS (Liu & Nocedal, 1989) are pro- posed in which only limited number of historical iterations are stored. In each update, the oldest iteration is discarded make space for the up-to-date secant pair. As result, the limited- memory quasi-Newton methods can lose the local superlinear convergence properties achieved by the full- memory schemes (Berahas al. 202T). Also, as stated by Berahaa al. (2021), the choice of the number of historical iterations (i.e. historical length) is problem-dependent, and one does not know a priori the best choice when solving a particular problem.
Unlike the limited-memory quasi-Newton methods whose performance can be sensitive to the his- torical length, our MST-AM method only needs store two corrected historical iterations. MST- AM carefully incorporates historical information through orthogonalization. In the ideal case, i.e. strongly convex quadratic optimization (or SPD linear systems), it is equivalent to the full-memory AM which means there is no loss of historical information. Our experiments verify the theoretical properties of MST-AM and indicate MST-AM can be a competitive method for nonlinear equations.
RST-AM versus first-order methods for stochastic optimization. For stochastic optimization such as training deep neural networks in machine learning, the first-order methods have come to dominate the field due to the low memory and per-iteration cost. The RST-AM is an extension of ST-AM and MST-AM to tackle this challenging problem. RST-AM has theoretical guarantees in deterministic/stochsstic optimization and also inherits several advantages of AM and ST-AM:
Fast convergence in quadratic optimization. It is important to possess such property for an optimizer since smooth function can be approximated by quadratic function in the local region around the optima and many techniques such as trust region (Nocedal & Wrightt 2006 rely on this local approximation. Adaptive learning rate methods such AdaGrad (Duchi et al. 201T). Adam (Kingma & Ba 2014) use a diagonal approximation of the Fisher information matrix that is approximation of the Hessian. However, in the sim- ple quadratic case, these methods can only roughly match the performance of the Jacobi method for solving linear systems (Saad 2003). which better than gradient descent but far inferior to the powerful Krylov subspace methods. The momentum method mimics CG method by incorporating a historical iteration into the search direction. However, the choice of momentum and stepsize can be an art (Sutskever et al.) 2013). For RST-AM, there no need to determine the stepsize and the fast convergence rate of ST-AM can be recovered by simply setting a₃ = and 8( = 5(2) 0. For more difficult functions, the damping and regularization in RST-AM can be enabled to improve stability.
Theoretical guarantee in stochastic optimization. If only first-order information can be ac- cessed, the SGD GGaadimi & Lan 2013) achieves optimal convergence rate (l1ee)) to obtain an e-accurate solution (Nemrovski & Yudin| 1983). In such case, it seems to be a big mismatch for current second-order methods, because there is no theoretical improvement albeit with more memory and computation resource. Such mismatch may also account for the popularity of first-order methods. Since RST-AM has very limited additional memory
16

---

Published as a conference paper at ICLR 2022
overhead. and also achieves the (1)/²²) complexity, it can be applied to many applications that are dominated by first-order methods.
Flexibility in use. The application of RST-AM can be very flexible. In principle, RST-AM can be applied to improve any slowly convergent black-box iterative process by viewing the latter as a fixed-point iteration. For example, consider accelerating solver of a com- mercial software, where we have access to the underlying codes to provide our custom implementation, or some case rewriting the codes is too cumbersome. Moreover, RST- AM can efficiently incorporate the preconditioning technique. Hence any optimizer, even an optimizer built upon neural networks (Andrychowicz et al.) 2016), can be used as a preconditioner for RST-AM. Preconditioning largely enhances the applicability of RST- AM for various applications. For example, RST-AM can be preconditioned by Adam for the language task and SGDM for image classification. So any fine-tuned first-order opti- mizer can be combined with RST-AM to achieve an overall improvement. As RST-AM is light, this additional cost is marginal and can be largely countcracted by the actual improvement. So in some sense, the purpose of RST-AM is not to totally replace current off-the-shelf optimizers but to achieve collaborative effectiveness: RST-AM is aware of the second-order information while the first-order method can mitigate the ill-conditioning of the problem.
Overall, the ST-AM methods have wide applicability and can competent methods from both theoretical and practical perspectives.
C
PROOFS
We give more details about ST-AM and the proofs of the theorems in the main paper.
C.I THE BASIC ST-AM FOR STRONGLY CONVEX QUADRATIC OPTIMIZATION
Recall that the strongly convex quadratic optimization is formulated
min f(x) = 2 x''Az Az -bxx,
(22)
zgRd
where A Rdxd is SPD, Rd. Solving
22 is equivalent to solving the SPD linear system
Ax b.
(33)
The detail of the basic ST-AM is given in Algorithm
We first state the relationship of AM with GMRES in the following proposition. Similar results can also be found in (Walker & Ni 2011 Wei et al. 2021).
= Let of r := 6 - A denote the k-th GMRES iterate and residual, respectively, and Kg v) span (v, Av,.. . A denotes the k-th Krylov subspace generated by A and v. Define e := 1)T E R for j > 1. Let range(X) denote the linear space spanned by the columns of X. The main results The Proposition| is the same the Proposition 2 in (Wei et al. 202I), and we restate it here for completeness. The Proposition 2 is new as far as we know.
the full-memory AM are stated Proposition | and Proposition 2
Proposition 1 (General lincar system). For solving a general linear system Ar = b with the full- memory AM (m k), suppose that 3x O and the fixed-point map is g(x) (1 A)x b. If the initial point of AM is Fo = I and rank R&) = m, then the intermediate iterate Ty satisfies Is G
Proof. The definition of the fixed- fixen-point map suggests that the residual F = g(xx) xs = Ar
Since R -AX and A is nonsingular, we have rank(X;) m. We first show
range(Xx, (Xx) Kr ())
(24)
by induction. We abbreviate (A,A,rg) Kr in this proof.
17

---

Published as a conference paper at ICLR 2022
Algorithm 2 ST-AM for strongly convex quadratic optimization
Input: E _R_ 5x > 0,0 < maa iter S d.
Output: x E R
1: Po. Qo = 0 E Rdx2
Po. 40 = 0 E R
2: for k 0, o......, maz iter do
3: T Vf(*k)
4:
if k 0 then
5:
p Ix -k-1- Th kk-
(Compute A Ark-1)
6:
q = q - Qk- Q1 k 19) p Ph-1(QE-19))
( I Qx-1) 1)
7:
PA ///li/22 qs 9/11@/2
(Igs 12 1)
8:
Pk 1k-1, Qk -, 9kb
(Q Q+ = 12)
9:
end if
10:
I's Q Tk
3k = 2g P = nh Q&T
(Projection step: Fk 1 Q2) (Mixing step)
12:
Ek+1 = 3g + FB
if 1
Othen
14: break
end if
16:
end for
17:
return 2
First, Azo = Soro- = forg since = zo + Soro. If = 1, then the proof is complete. Then, suppose that k > 1 and, as an inductive hypothesis, that range(Xx--)) Kg-1- With we have
- Zg
Ik
B (X R
Bk-1(b - Arx 1) - (X+-1 - - 10-
3x-1b Bx-1A(I0 AZo + .... +................... Arx (Xk-1 - BE- AXx-1){
B5-1TO - 3x- -X k- (Xx-1 A
(25)
Since E KG-1. and by the inductive hypothesis range(X:- C KA-1 which also implies range(AXg-a (1X C KR. we know ArÆ-1 E KR, which implies range( S Kr+ Since we assume rank(X;) = = & which implies dim Xx)) dim(K))) we have range(Xa) = 1c2, thus completing the induction. As result, we also have
range(Ra) range(AX,, AK₁ (A,
(26)
Recalling that to determine Fs we solve the least squares problem and Rs = - AX, We have
I' arg min ||rk + AX,1'122-
(27)
TERm
Since rann(AAg)) = rank(Xx) = m, 27 has unique solution. Also, since = b Az =
b - A(zo ) ro A we have ra + AX ro AXte* + A = AX,T,
where r = et -T. I.Soo So I' solves (7 if and only if I e I& solves
min ||ro AX
(28)
FERm
According to 24 (28 equal to min,ek., (4.rg) l|ro - Azl12 which the GMRES minimization
problem. Since the solution of (28) also unique, we have
Fs = Fs X
26 Xx(e* f&)) = 2o + x2fx=
In Proposition I the assumption that R& has full column rank critical to ensure no stagnation occurs in AM for solving general linear system. In fact, for SPD linear systems convex quadratic optimization (22), when AM breaks down, ie. R2 is rank deficient, AM obtains the exact solution, as shown in the next proposition.
(3) strongly
18

---

Published as a conference paper at ICLR 2022
Proposition 2 (SPD). For applying the full-memory AM to minimize strongly convex quadratic
problem (22) or equivalently, solve a SPD linear system (23) suppose that Bg > 0 and the fixed- k holds for 1 S k < while failing to hold for
point map (zx (I A)x + b. If ffran((R, k= == 8, where & 2 1, then the residual of AM satisfies r, = f=-1=
Proof. The definition of ! suggests that the residual g(xx)) Th = b- Arx. The relation Rs: -Ax holds during the iterations and the nonsingularity of A implies rank(X;) rank(Rz).
2o + Boro. the assumption rank = 0 implies that For = 1, since the first step of AM is rank(ro) rank(X1) 0, i.e. r1- Fo 0.
For 8 > 1, because Ar,-1 Is - Ia-I -X,- + 85-17 the rank deficiency of
X₂ implies Ar,-1 range -1) which further implies e--1 E vange(X3-1) So there exists
Ç E R=- such
such that F- = X.-1Ç. Note that according to 3 F- 1 Ra-1 -A, so we
have
0 0= - (X,-1() - -T
-x,-1
(29)
Because rank(X, )- - and A is SPD, we know A is also SPD. So Ç = 0, which implies F--1 = 0. Hence Vg 1 and r. F--1
Now we give the proof of Theorem
Proof of Theorem 7 Besides relations (i)-(ii),, we add an auxiliary relation here:
(iv) Ta = + Qul'k E (A, where I'g E 12
We prove the relations (i)-(iv) by induction.
For k = 1. since io + 0. according to Proposition rank Aro) rank(X,) = 1, rank( (Ar))
rank R1) = 1, so q * 0, which implies Line 7 in A Algorithm@ia iss well-defined. The relation (i) holds. Arg, and Aro -AAro, the equality Q1 -AP, also holds. Due
Since q= 9 = Aro, to the normalization in Line 7, Q1Q1= = Since r1 = -n Ar and range (1)) = range(Arr), is clear that 1ı ro- Q 1 (2( namely relation (iv). Due the projection step Line 11, Ti I Lrange(Q (Q1) ) = range(Q1 Also, F1 = r Q 1- = To- Ar - T = ro - Qin. where the last equality due span(Arr)
range(Q1) (Q1). For G ro Azı, where 21 arg min,ek, (A.ro) ||ro- Az]12. it holds rf 1 AK, = range (Q1). As result, both F1 and r are the orthogonal projections of ro the subspace range(Q:)+ which implies F1 = r So 5i = σ xo + because their residuals are equal and A is nonsingular. Hence relation (iii) holds.
range(Q1)
Suppose that A > 1, and as an inductive hypothesis, the relations (i)-ivv hold for j = I,... k - 1. Consider the k-th iteration. From Line in Algorithm 7 E range(Ark-1) Qx-1), and p We first prove that q A 0 by contradiction.
(A
q = then from Line in Algorithm 2. Ark-1 E range(Qu-i) S range(Q(=1). which implies A E range( 1 C range( Pk-1) as Ark-1 = - -A and Qk-1 -APk-1 and A is
nonsingular. From Line and Line 12, we have
A I& Fk-1 - + B
(30)
So fk-1 E range( P1- C range( Pa-1) since Ark- E Pange(Pa-1) Hence there exists
< R such that Fk-1 Pa-IC. From the inductive hypothesis, we know Fk-1 1 QK-1-
- -APp-1, SO we have
T RK AP 1aa0--1-P------ (Pr-1
Since Qf Qk 1--1, 1, we know rank(Qu-) 1)= k- 1,, which implies rank(Pa-1) = k: due
to QK-1 = - AP Hence PT AP:-1 is also SPD. Then Ç = 0 which implies Ft-1 1 = 0. is
impossible otherwise Algorithm well-defined.
2 has terminated in the (I- 1) iteration. So q * 0 and Line is
Since FK-1 = T4-1 Qk PT-1, and T6-1 E (A, rangeeanuaaaa S range(Qu-1) I ro) as the inductive hypothesis, we have FR-1 KC (A,ro). which together with (30
19

---

Published as a conference paper at ICLR 2022
and rango(PP 1 C range(Pk-1) = KA-1(A, ro) infers A (A, Hence Ark-1 - AA 1 E
So range Q+)) To), As a result, Ok = 9/ //ll2 E Ark--1, S (A, ro). range(Qu-14)
_ AKx(A, ro). Moreover, 9k & range(Qk-1l otherwise q E rangel Qk-1) that implies Ark-1 =q E (ange(Qk-1) which is impossible following the former proof of q # 0. So we have range(Q1)) A (A)
Because Ark-1 = -AA and QK-1 = -APk-1 due to Q - - Line 6 in Algo- rithm inffrss q = Aÿ, which implies qh - Apk. So Qk APk. Since A is nonsingular and range( (Qx) - (A, we have range(Ps) (x)) - Kr (A,
As range( Xx) (A, ro) and range(Re) relation (i) holds for the k-th iteration.
AK& (A, (A) has been proved in Proposition | the
To prove Q
Lx. it suffices to show qi Q2-1, as the equalities
Qk- = 4-1 and
||g 122 = has already held. It is equivalent to prove 1 QK-1- From the construction of q in Line 6 in Algorithm 2
we know 0 0, so 7 1 span(x 2, (k-1) 1) (for h = 2, q 1 4a = clearly holds). To further prove L range (k 2 4), note that
Ark-1 = -A = AP
- 3x- A = - QK-IT 3k-1
where the second equality is direct substitution with 30 Therefore,
- 3
BA-1 Arx k- = 0 3x 1 (AQk-s) (
7k-1 =0 (31)
where the second equality is due to range QQ--)) spanssa 91. 11- 1 range (ooge(Qi-uu (Qa-3) and A is
SPD, the third equality is due to FR-1 range (Q+-1) = A A9Kg-a(A,ro) S. 1 As a result, noting that 0- Ar we obtain
and rango(A AQu=3) =
Qh = Q4-31 - 1Q 0,
which is due to (31b and range(Qk=1) (k=1) span (panf4k-20 96-1)) 96-13
1 range Qk Therefore, we show that Q = Ix. which along with Qu = -APh proves relation (ii) in the k-th iteration.
Next, we prove the relation (iv). We have
rk rk-1 1 - AFk1
rk-1 1 - Qu-IIk-1 3h A(rx - (R-I
rk- - 1 -k= Ark-1 + Bk 1
rr + +Q--ITa-1- - Qk- h----- Bk-1(Aro + AQu-ITa-)) + AR--AQ
where the last equality due to rk- = ro + px by the inductive hypothesis. Since (ange(Qk-))
(A, rango AK,-1(A,ro) S. (K, ro), range(Qk-i) S range( spann() S
S A² S AK2 (o. and (Qa) (A, it is clear that F& ro + Qul's (x+1(A,ro) for some I's E R#. The relation (iv) is proved.
S rango(AQ
Finally, we prove the relation (iii). For proving T. . range(Qr) (Q+). note that Fx 1 span 9k range(Qs) already holds due to the projection step (Line 10 and Line 11 in Algorithm 2). It suffices to prove Fk 1 range(QR=2) In fact, since we have Fg Fk Q&Ir. k. can prove that r 1 range(( (ange(Qx-2) Qx-2) and & 1. rango(Qk-2)
Since range (Qx) spanfgk-1, 9 1 (Qnge(Qe-2) as induced from Q.Qr = Ix. it is clear that QuT's & 1 range(Q:-2). For rk. according to Line 12 in Algorithm 2 we have Tи- PA-11 - We have Fk-1 1 range(Qu- (Qxx1) 2 range(Qu- by the inductive hypothesis. Also, A (AQu-2)T
0 due to range(AQu-aa
A2,,==(AArr)) S 1
range(Qu (Qu-1).
Therefore, we obtain 7x 1 range( 2), which along with Ps 1 span(gu-1, | implies FA 1 range(Qu))
To prove Ft = 26 = xo + 2k. where 21 = arg min (A.ro) Ilro - Azl12. first we have Tk
ro + Q.Tx. where Fx E R². Hence F = r% Q&Ta = ro + Q&Ta Q&Ta ro Quek. where
Th E RR Since F 1 Qu, Fx is the orthogonal projection of rn onto the subspace range(Qx)
On the other side, for GMRES, rg = ro Azk 1 rK)) = range(Qu), so r is also the
20

---

Published as a conference paper at ICLR 2022
orthogonal projection of onto the subspace range(Q:)). So Fkk rk. which further indicates F& 2 Hence, the relation (iii) holds.
With relations (i)-(iv) being proved in the k-th iteration, we complete the induction.
C.2 MODIFIED ST-AM FOR GENERAL HIXED-POINT ITERATIONS
Algorithm is suitable for analysis and implementation in the linear case. For general nonlinear fixed-poine iterations, we adopt an alternative form described in Algorithm B which discards the normalization of ğ in each iteration (Line 7 in Algorithm 2). In Line 7 in Algorithm 3 the orthogonal projection of Ar is checked to ensure - "less linearly dependent" on (Qn).(Qu)). which ensures l|ge is bounded away from zero; the check of PP-- l2 ensures that Il Ar S cpll p01Aag which also important since a large deviation from Azg-1 can make llp& > P (p is the radius introduced in Theorem 2) When this condition cannot be satisfied, the algorithm simply reuses the old 1 i -1 The main procedure of MST-AM restarts every m iterations, i.e. Pk. Qx 0. Such restart mechanism is to restrict the higher-order terms in the residual expansion, as shown in ( Also, restart can flush out the outdated historical information that may weaken the quality of Ps and Qx that are used to pursue a local first-order approximation of g in MST-AM.
Algorithm MST-AM for nonlinear fixed-point problems
Input: zp E [Rd 3g E (0, [], Go > 0, c E (0, 1), m 0.
Output: x E Rd
1: Po. Qo = 0 E Rdx2 Po: 40 = 0 E Rd
2: for k 0, 1..... until convergence do
3:
g( -
4: if mod m + 0 then
5:
p Ih- 2k-1. rg Tk-1
6:
CA = Q
7:
if P S lllppll and 1|Qt-1 S cylllll then
8:
Pk = p- P5-1 qE = q- Qx-1
9:
Ps (pt-1.Pe. Qu 1M-1,9h 4k)
(gk Qk-1)
10:
else
11:
Ps Pk-1, Qx = -Q--1
12:
end if
13:
else
14:
Pk. Qk = E Rilx2 Pk: 9k = = Rdd
15:
end if
16:
TE
(Q
17:
Th = Pal',, = Qullk
(Projection step: FR 1 Q②) (Mixing step)
3k++1 Eh BRFK
19:
end for
return 2%
(The notation "1" is the Moore-Penrose pseudoinverse.)
In the linear case, Algorithm 2 and Algorithm 3 (with m = -o) are equivalent. Similar to Algo-
rithm 2, have the following properties held for MST-AM:
Claim 1. In the k-th iteration (k > 0) of Algorithm I applied to minimize a strongly corver
quadratic problem (5). assuming Cp 00, Cq 1,m
the following relations hold:
(i) ligullz > 0. rango PR) range X₂) = K4 (A), range (xx) range(R, = AK & ro):
(ii) Q+ APr,, qi (1 k);
(iii) Fk range(Qu and 3 = Da + 2b. where 2 arg minmiinn,,,.... (A.ro) I|ro A2
If [|[x = 0, then F the exact solution.
The proof of Claim i essentially the same as the proof of Theorem I with a special care that Q Qk 1t is replaced by the relation that columns of Qu are orthogonal to each other, other words, QT diag [| 1 l12:
21

---

Published as a conference paper at ICLR 2022
For minimizing nonlinear functions accelerating nonlinear fixed-point iterations, the long-temn
relation that Q7 Qx = diagf//qu112: |/.............................. |/ (2] generally cannot hold, while the orthogonalization procedure in Line 6 still leads to short-term orthogonality relation: qi 1 93 for i - 2.
Hence, Q diag ( * 0)1/1l/--10 1 & 0)1/ (2), where Z(-) is the indicator function that
diag || p 10x (1) for h 1. Thus the pseudoinverse Qz)
I(x) -
{
1 !
x is true, is false.
Now, we give the proof of Theorem
2
Proof of Theorem?
2 For convenience, we restate the main assumptions of here:
(i) (g(()) - g(>)] - xlly 2|12, x E (0, 1), for Vx, y E B(p), (ii (v) g (2) - &lyy 21|2, A > 0, for Vx, y B(p).
(32a)
(32b)
know 3, 0s bounded, i.e. Bx < 3 where 3 that ||rx||2 S ||ro|22 by induction.
Also, since |1- + 1 x
is constant.
The proof is based on the two lemmas given in Lemma 1 and Lemma 2 Besides
we also prove
For h = 0, Øo = zo E B(p). and due to (48b). ||=]] x" 222 S 1 xx x2112 + Bollrolla S |
xx 122 + B(1 + =r] x= 2 S p provided 20
xx 1|2 S p/ 3(1 + k)). Since
r1 g(x1) xı g((1) (ar (oro) = g(m)) (w0 + ro) + (1 3rrro
g(()) - g(x0) + (1 Bo)ro:
it follows that
||r/|/2 s (1(()) g20r0))2 + |1 Blllrolla
5 1x00l1rol12 + 3o foollo (noo + |1 Bol)llooll.
Also note that 0p ||ooll//|1oolz 1. Thus p
holds. Because x30 + |1 3ol S K0 < 1, we have
||rilla < ||r
Now, suppose that (9) and ||rxlla S ||rull2 hold for k 2 0. We establish the results for k + 1.
Let I'x = (7 7k follows that
1) IR2. Since I's = =QT Qx) Q. Q. Qx diag (iag((19x-1112.. | [1))
r T
(2) I(gu + 0) -h
(1) I(qu-1 1 + 0)
9k-1
7
(33)
q a 9
q qk
Therefore,
||ral12
1|rk 96-111 112
[11 < I(gk-1 + 0)
11 - 7/8 S max I(9k-1 = 0), * #0)
-1
(1)
1gg--112
196-1l12
||rk 1|2
-1-)) I max I(gx = 0),I(q #0) ||, 9xl12
I/ < I(gx # 0)
|1 11-
(34)
|194
||9x 1|2
Define c (1-x)(1-c,) We have
1tep
||PTT = 1| (2)
+ PkP 1 7 1| < || //
(1)
(2)
|| ||pp- 17k-
(1)
l|2
|| l|
< I((g + 0) pk
+ I(gk-1 + 0) |k 1|
S 2c|[ral12,
(35)
101112
[k
11
where the second inequality is due to 334 and the third inequality is due
449 Then
|i+ x= 1|2 || P&T x 1|2 S || x || +
P.
<
|x ll |all
+ 2c||ru|l2 = (T Kr
+
22

---

Published as a conference paper at ICLR 2022
and due to F ||r 12, it holds that
| x" 1l2 = || + BEFK x=112 - || - x=112 + 8 Fxl1/
- || - x x=112 + 31|r=l12 = II- +2c + 8)11ral12.
By the inductive hypothesis that ||r 12 ||rull2, and (48b), it has
1 1 + 2c (1 + k)]l20 x'112.
1 1- 22)]|roll2 | x 2=112- S
(36)
and
| - x
(T-m= + 2c + 3 |roll S
$T-==
+ 20 + 3) (1 x) 11:0 x2112. (37)
As a result, we can choose ||zا - x1122 sufficiently small to ensure 3y E B(p) and Ix+1 E B(p), which ensure the (āxx) and 1(*x41) well defined.
At the end of the k-thh iteration of Algorithm
B
we have
Tk+1 = g(xx+1) - Ik+1
g( g(āx) + ((xx) (Ix + BxFU)
= (0(0x+1) - -((xx)) + (g(āx) - 3 fk) + (1 BBk
(38)
Let La = g((x+1) g(xx) + (1 BB))FK, H4 == g(ox) 3x -kk then
| 2|2 S x|| Ix l1 + FAll1FFlle
KBx 1 - 111-BxllFFaa
Og (k3x + |1-3x)]]rxlle
(39)
which bounds the linear part of the residual 1h+1- For the higher-order terms He. we have
Hr g(tx) - (Ix - P&Tx Tk - -
g((x) g(xx) (P& + Qu)l'R.
g((x)) - g(xx) + +px + mk (2)
+ (px-1 + qu qn-1) (1)
(40)
According to the formula fo g((2 + fly - =x)) - x)dt = g(u) - g(r)), we have
g(āx)= - g(xx)) = g(āx) - -[(x -47k (2) + g P (2)) - g(xx)
g 1x pa (2) 1p - 17k PPk - 17k dl
(1)
(1)
f-' -g'(s (Ix 1px7) (2)
(2) dt.
PP
(41)
Also, by Lemma
2 we have
Pk + Us = (' - to--tpppple + Ar
Σ
O(IAr,lla),
j=k-ma mp
( -'pox---taa-- tpo-1)pe-1dt Ar
Σ O(1Ar))la);
Pk-1 + 9k-1
where =kk) denotes that the latest update of gu by Line 8 occurred in the π (k)-th iteration and e(k) 7 (k<kk()))) 1) marks that qk-1 records (o(k) that is the penultimate update by Line occurring in the -kh)-th iteration.
23

---

Published as a conference paper at ICLR 2022
By substituting these relations to 40 it follows that
1/'0 g (Ik tpx tp + A))Ar=(0)-
Hk =
Σ
o 0(1rr,lla) 7k
j-k-m
S g tpx 1px- (#x-1-tpa-
k-2
Σ 0(1A))122) rie
+
+ A (o(k)
j=k-ma
g' g'(zk -Pa7k (2) - 17 (1) 17k 20 dt - S/¹ g((2x -tpwne
12) )arc dt
S. (gf (0x tpx) g (xxx= tp (2) )) pk
(2) dt
+ S( (g'(0x11 tpx-1) g' (xx
p47k tpk-17k (1)))))))pppp (1) dt
(2)
k-1
+ (1)rra())-lll
Σ
j=k-ma
k-2
+ A
Σ
0(1Ar,112)% (1)
j=
= A4 + B& + C DK-
(42)
Then can bound each terms of He follows (Here we assume T) F 0 and Φx-1 F 0 as qk leads to 7 = and -k-1 = 0 leads to 7 = where the result is trivial.):
| = 1/,0 (g (xx - tpx) g( -tpx7k (2) ))kk) % dt
11g' (x - -tpa) - -g' (xxx tp (2) Illll (2) 12dt
$
I R l|tpk tpk (2) 11 11 || (2) dt
R
-- (1(1)
2
2 1 ((1rxllell - gxll2)
R Σ
(11r)11))
(43)
1
j=k-mæ
where the third equality is due to (34). and the last equality is due to 449 and (50b
|| [Blll2 = 1.1. (g((zx-1 11) g' (xx -P (2) - ttp-117h })pp 7/ dt
(1)
$ / 1 - tpk-1) g' (xx - (2) tpk-17 (1) 12211 (1) 112dt
||AT--1 - Pk7ik 1pk-17 (1) + tpk 1pl-1ll2lllp-111 112 ( dt
S A
S
|| AA 11/2 + 11px|1217 (2) + ト
11p--1112 - 2 Ilpk 112 (2)
||all22
| 1l12 l|r - gh-1112 1p--1l12-
||
< A
|| 11l2 + 1lalll
12
11
1
11
R||ræ|12 (((m))))) O(IIr= - q l|2)
0(1rml12)
R Σ
(11r)11)
(44)
j=k-m mm
24

---

Published as a conference paper at ICLR 2022
where the last inequality is due (3). and the second equality is due to 48a
49
50b
k-1
|C& = A (1)
ol1Arylean
(2)
||all2
I k Ar (r)))
O(1Ar,112)
110x
j=k-m
||r/12
A Ai|lrxx()) Ar))r(())
Σ
o(
[
j=k-me
0(1r,)1)),
(45)
j=k k 17
where the first inequality is from (34 and the second equality due to (50b):
h-2
|D = A Ar (r
Σ
O(1Arylla)7h (1)
j=k-ma
k-2
S R|| Aro))r 11/2
Σ
0(1Ar,112)
1
km
||mall2
k-2 Σ
A|] 11/2
o(1Ar,
1 122
j-k-ma ma
O(1r))1))
(46)
where the first inequality is from (34). and the second equality due to (50).. Then with the bounds
(45) and (46), we obtain
Hall2 = (11rs11))
(47)
j=k-ma
Combining 339-m and (47 to (38)
we obtain
|r |l |[ax|l + |176x 1 I (x(B +11 11 BB)llrll R
Σ
O(Ir,)1B)
j=k-m
as desired. Since my m, the higher-order terms limited. Note that x, + |1 3xl S Kg <1
by assumption. Then, for |=:0-201 - 2l2 sufficiently small, the residuals {ra} are Q-linearly convergent, which infers ||ræll < ||rol22. Therefore, we complete the induction.
Remark 6. There may be some concern about whether R can be counteracted by the constant hidden in the Big-O notation. In fact, since R is the Lipschitz constant of g' and the constants in (() are composed of K. Cp. Cqp it follows that R is unrelated to O(-). Hence small i. can lead small uniform boundedness of the higher-order terms. In the exireme case where R = 0, e.g. is a linear map, the residual only consists of the first-order term Lx-
Lemma 1. Under the same assumptions of Theorem 2 for k 2 1, we have the following bounds:
(1 A 1|2 < || < (1 + x)]l 1|2,
(48a)
(1 x) |
x 221 < ||ral12 < (1 + x) || x ||2
(48b)
qk + 0, then
l 1|
1 Ep
<
(49)
1 l|
(1 - (1 Cg)
25

---

Published as a conference paper at ICLR 2022
the condition in Line 7 is tme, then
|| (1 + (p)ll
(50a)
[k # 0, (1 - (-) Ark 11 S 11rr-11l2
(50b)
Proof. From the assumption (32a) of g, we have
(1 - 1 -
| - Tk 1 [(((+))
[1g(2x) g( (z xx-1))l2
1|rk - rk-1ll2 2 = |
[Ig()))) g( )|2 | xk--1l
$ (1 + k))l -xx-1
(1 x)]laa 0"112 I | x² 112 (1g(2k)) g(xx))l2
< llg(a)) - g(x") (zx - x ) 12- = ||ra
$ l1g(xx) - g(x*)) lll + +1158 l|2
$ (1 k))l x x2112-
If the condition in Line is true, i.e., S ll S 1111A 1122:112 1Q--1kkl1z S collArkall2. have
1l - || - Pa PR-ICA 1 1 + Pa 1 - (1+ (
119 = ||Ark-1 Q 11 1Qu~1 こ (1 Cr)llAr-1
The inequality 19 12 I Ar is due to the fact that Q is the orthogonal projection of Ark-1 onto (ange(Qk-1) Also, qs * O must hold otherwise q QK- Cx which violates the condition 11Qx-1Cel12 S Cyl Ark- ll2 as Gg E 1).
If qx: O, then 9s must be updated by Line in previous iteration. We assume the latest update by Line 8 occurred in the j-th iteration, i.e., qj Pk = Pj. hence
llpllle lpyl22 <
(1 +1+ (1
1 cp
llq
14 l2 (1 Ar) (1 - c1-cc)11-1
where the first inequality due to (50a) and (50b) and the second inequality due to 48a
Lemma 2. Under the same assumptions of Theorem 2 the k-th iteration (k 2 0) of the restarted MST-AM (Algorithm 3. we have
Pk + 4h S/' g'(xx t-pp)ppddt -())
Σ
O(1Ar,112).
(51)
j=k-ma m.
where my k mod m, Ark-ma-1 == Th-my- (k) denotes that the latest update of qu by Line occurred in the T (k)) iteration and (kk) = k-m if Line 8 is never executed up to the k-iteration.
Proof. We prove 51 by induction. Denote ÇA (ch) (())).
For A 0, the relation trivially holds.
For 1, pi Aro, Q = Aro. It follows that
Pı ạ1 g((1) g(xo) f g( tp tpi)padt.
So 51 holds for l: = 1.
Suppose that (5I) holds for 0 i 1 where 2 2. For k: 2 2f if := mod m = 0 or 1,
51 holds because is the same the of k 0 or k Now consider the nontrivial cases.
26

---

Published as a conference paper at ICLR 2022
For my # 0 and my * 1, if the condition in Line 7 in Algorithm B true, then TTk) k. With the convention that Ark-ma-1 = m we have
Pk + 9k = A 15 + Ark-1 Qk-15k
= gg(xx) - g((x-1) (px-1 + 9
(2)
+ (1)
= -f0 g -A 1
S.' 8' -ca-1-1pp-ill d1 + Ar
k-2
Σ
jmn- m
- S.' g' 0' 100-2-1pp-20000 1pk-2)Px-25h (2) di + Rll Rll
k-3
Σ
OllAr, , 1
J 2
(52)
where ukk - 1) = ( - 1) - 1) denotes that qk-2 records (o(k- 1) that is the penultimate update by Line 8 occurring in the v(k: 1 iteration. Considering the terms in pk + Ik. we know
g((x - 1A A g' (rk-1 tpk pp
(2) g (xx-2 tpk 2pk-2)pk-25$ (1)
14
=g'(z = 14-1 (2) - (1)
+ (#x t - g (Ix-1 tp: 1 1) 1
+ (f - 1Axk- 1) - 9 (x 1pk 2 pk 2
(53)
For Cx
(0 1x- + Ark-1 because Q Qx- diag(l1 21:::211::: 11--1113 (12), it follows
that
c = I(qk-2 # 0)
-- 9K- (2) I(qx-1 + 0)
Art
4- 1-19--1
16-29k=2
9k-
19-->
Therefore,
c S I(qk-2 + 0)
|Ark-
(22) S I(q + #0)
1 /rrr-1ll2
(54)
196-2112
119:-1112
Now, to bound the terms in (53). we have
( (Ark-1) g' (IR-1 (pp----pp--)-ppp 1/22
S Allæz Ik- - t ( P 1|2
S A(((1 - 1
Ark
s A((1 +) + -
* 0)
S ATR-
Σ
O (l2) = Ar/Arkill 12
O(1
(55)
j=k-
5=
where the third inequality is due to 54 and the last inequality is due to 48a). the constants are absorbed into the big-O notation. Similarly, it follows that
50a and
(49
where
(Z tATx-1) - g (IX-2 -pp---pp-))ppp5556 (1) ll2
S R1| + - +
S A((1 1) ||2 + |
Ar ||2
s AA(11 -0) 1)
1 * 0)
11
S Al|Arrzlll2
Σ
= = Σ
01 112))
(56)
j=k-2 m
j=k-m
27

---

Published as a conference paper at ICLR 2022
For the remaining terms in (52). have
2-2
A (1)x l2
Σ
0(1Ar,112)5 + AillAr 1
(2)
Σ
olIAr, 12)5k
5 - m-ma
Am 2-
Ark- 112
R-2
$ k| [a]]Arn(k-a) (11)) +0) 0)
Σ
O(1Ar,112)
11
j=R-I-ma-1
AArk-1ll2
k-3
+ R| Ar 1o #0)
Σ
O(1Ary112)
1
1=
2
S AA|AAr-1ll k=
O(1Ar)112),
(57)
where the last inequality is due
506
and that 9k-1 - 97(k- 9k-2 = 9v(k-1)-
With relations (52) 53).
55
556
and 57 and noting that pk: AT P4- (2)
Pk- (1)
and
11,0 (g'(zz tAr AT g (I - tpk 1P c/ g' xxxxx22 - tpk 13pk-2)pp 2k 25k )dt
)tt
L. lox 1A g (2x11 -px 1) 1)pk- Sk
(2)
g (xxx_ tpx-2) 2)pk 2 lladt,
can estimate Pk + Ik as
Pk + Th = S. g( 14xx-1)Pudll + R AllArr-ille
Σ
O(1Ar,1122)
(58)
j=k
To further obtain 51 notice that the difference
1./ o' (Ix tpk)pıdt { I (Ix +tAxx
1'' (g' - -tpa) - -xx - (t
(. I1o' (o* tpx) g'(xxx - 1Am=-))l22lpa 1latt
S R / 01lpa-14 ((2) + 1 [px 22
A cp 1all--lllllllla///
R (1
where the last inequality is due to the condition I.PA-ICAll S cpllllla and inequalities (48a) and 50a). Hence the difference can be absorbed in the Big-() notation in (58). Thus (51) holds when the condition in Line 7 is true for the k-th iteration.
We consider the case where the condition in Line 7 is false for the k-th iteration. Then Pe Pk-1, Qs = Qk-1- In other words, the memory recording Pk, keeps unchanged from the (T(k) + 1)-th to k-th iteration. Since (kk) < k and π (T(k)) π (k), with the inductive hypothesis, we have
Pk + Qh- = Pm(k) + (m(k)
= f.' g' (F - (P(()) Rl| 111(1)1112
x(k) 1
Σ
((1Ar,112)
j=m(k)- mm
= S.' g' (II(K) - tpx)padt + Ar
k-1
Σ o(1Ar,112).
(59)
j=k-m mm
28

---

Published as a conference paper at ICLR 2022
To further obtain
51 note that the difference
1.' (ax tpx)pedt S.¹ g (FTT()) tpp (ppdtt
11 (ark - tpx) I (x) - tpx))12 11ll2/pxl122t
I Rl 2 *=(x)ll2llpall
k-1
k-1
A
|Aayllalleell = A||Arx
Σ j=m(k)
O(1Ar,112)
(-)
can be absorbed into the Big-O notation in 59 Thus 51 also holds when the condition in Line false for the k-th iteration.
As result, complete the induction in the k-th iteration.
C.3 REGULARIZED SHORT-TERM RECURRENCE ANDERSON MIXING
C.3.1 CHECK OF POSITIVE DEFNITENESS
We describe the check of positive definiteness in RST-AM, which follows the same procedure as that of SAM (Wei et al., (021]. From Line 11-13 in Algorithm [ one update of Fy RST-AM is Tk+1 I + HaTK, where H4 = 3 - au Y o Y2 Pa + QrQ 2x Of Qx + 6 Pa. Hs generally not symmetric. For the convergence analysis of RST-AM, critical condition the positive definiteness of Hg, i.e.
$ H 2 Basl|ssl12 Vsa E Rd
(60)
where μ E (0, 1) is a constant. Next, we show how to guarantee
Let Amin(.) (.) denote the smallest eigenvalue, and Xmar(.) denote the largest eigenvalue. Since
of H4SS = 1 (H H ss., Condition (60) is equivalent to Amin (1 (H + H.)) 2 3x By
some simple algebraic operations. we obtain Amain & (H₂ + +H+) Bx 2ax QT +
Q Let d == Amaz (Y + Z ), then Condition (60) is equivalent to
ax I 28 - 4):
(61)
namely, 133 in Remark 4
To check Condition 66]. note that
A Amar (or
zt
z 7 0 )
Qk)
Aman
(Y
Qk)
(62)
z
0
Z/
Since
(Yn Qx),
E R2x1,
>x can be computed cheaply. This cost is negligible
compared with those to form P Pa, Q7 Qu. which need (d) flops. Then, to guarantee the positive definiteness of Hy, we check whether og satisfies (61 and use smaller og if necessary, e.g. ou 21, /)
C.3.2 PROOFS OF THE THEOREMS
We first give the proof of the boundedness of || || and 11Qx-19112
Lemma 3. For P. Q E dd m (d 2 m), 8 0. and Z QTQ 8P P, have
|PZ
(63a)
(63b)
Proof. We first consider the that Z is nonsingular, then PZ PT QT QZ It can be seen that PTP and QTQ are symmetric positive semidefinite, and Z is SPD as it is assumed to be nonsingular. Also, have 8p' P Z and QTQ x Z, where the notation "-"
29

---

Published as a conference paper at ICLR 2022
denotes the Loewner partial order, i.e. A 3 B with A,B E R=x" means that B - A is positive semidefinite. Hence, we have
z --26p= Pz 11, 10
which implies
|z (2 56
: (Q) z - 11 1
Let Amar (-) denote the largest eigenvalue, we have
IQz='Q"1lz = Amaz (QZ- = Amaz (Q" = Amar (z- 1
2
Amar (PZ-
Amax
p=PZ--Q"QZ-")
Amax
< 1 (
1 (PTP) z z- 1121 (Q' Z 11
Therefore, || For the case that Z is singular, it can be proved that
11 s 5 and IQZ-'Q" Q') 2| 1 1.
ker = {ar RM Zx = 0} = {x E Pm and Qx = 0}.
(64)
In fact, if Pr = and Qx = 0. it is obvious that Z5 0. On the other side, if Zx 0, then x. Zx = i.e. x² xTQ Qz + 8x pT Px = O. Since 0 < PTP.0 = Q Q 5 > 0, it follows that oFCQ. FQT Qx = and 8x pT Px = 0, which further implies Qx = 0 and Px = 0. Hence (64) holds.
Let U1 satisfy UFU, = I and rUn) ker(Z), i.e. the orthonormal basis of ker(Z), and Un satisfy U2 know PU1 0, QUi O. Define U (U1, U2) E Rmxm then U Im and by direct computation, we have
and U3
0, i.e. the orthonormal basis of ( With the equality (64).
UTZU = (o
0
U ZU2)
where UUZU2 (QU + 8(PU,) PU2 is nonsingular according to the definition of U2. So
zt= 1 ( (UTZZU)) UT.
As a result, we can further compute PT and QZ!Q" as
P = (PU2) (PU))0 (QU,))
QuiQ_ = [QU5)0U
Then let P- PU2, = QUZ, and 2 = QTQ + 5PT P. and noting that nonsingular, we can
obtain |PZ 11 = 1 11 < 8- , Qt = 11Q2-1Q" < 1.
As result of Lemma
3
| S4 1Q in Algorithm I are bounded by O()))x-ill))
shown in the following corollary.
Corollary 2.
Qu - 1 <l12 Algorithm 7 Mare bounded, i.e.
P 2|2 < (6()) =11Ar= 11Qu-15Hll2 S |1ra-1l12.
(65)
where Ce 10 (1) P Pa Pk +F F Ark-n
30

---

Published as a conference paper at ICLR 2022
Now, we turn to the proofs of the theorems about RST-AM in Section 3.4 For brevity, use 6ş to denote 8(2), i.e. óy = 6622 The proofs follow those of SAM (Wei et al. (021). Nonetheless, since RST-AM uses different historical sequences compared with SAM, we give the detailed proof' for RST-AM for completeness.
From Assumption for the mini-batch gradient (2x) the following properties hold:
L TE
Eies, (1a)), where ne = |
E (z) V((aa).
(66a)
o²
E[IV (T&x) (+(+)))(( S
(66b)
Consider the update of RST-AM. From Line 11-13 in Algorithm I it can be written as Ik+1 Ik + Hark, where rk- -V fs, (zx), and for l 2 0,
H = B - an (P& + BxQu) (Q 6x oF.
(67)
To prove the theorems, the critical points are (i) the positive definiteness of the approximate Hessian Hy and (ii) an adequate suppression of the noise from the gradient estimates in the stochastic case.
We first give a lemma related to the projection step.
Lemma 4. Suppose that { is generated by RST-AM. If ax 2 0, 3x > then for any Uy Rd, we have
| 2 (5) (1 2a -20x) + t0,0x
||0x 11
(68)
Proof. The result holds when k: 0 as Ho = Bol. For 2 1,
Hg = 3R - (a + xak32Qk))
(69)
where F& (QT + 6x P P. ) Q. 2y solves
min 12 Q 5 P.
(70)
By direct computation, ||uk - + P = || - 0 thus
Q Uk.
IV QxTx11 P 12.11
(71)
Therefore,
H Uk 12
11. uh (a + +
1| (vx an (xTT) - a4 P
= || (1 - ax)²k + 3e (Ux Qxl's) 0 28 PAT'el3
S (5) - au)² + BL + + (1|02 + 110x - +
< (L) (1 + 20 - 200) + + (1us + 1vx111
2 (8) (1 + 20 20x) + ag ||ux112.
(72)
In the above, the first inequality uses the inequality
0x
Σ
|aل|||i|la
<
x 12
(73)
where a; E RR x4 E Rd. The second inequality is based on inequality
711
With Lemma 4 can prove the deterministic case of RST-AM.
31

---

Published as a conference paper at ICLR 2022
Proof of Theorem 3
Since 1 1+2af. -2ax < for 0 - ay I and 5x 2 C=!3², with Lemma 4
we have
| 223 cr1)
Since ag satisfies
61.. we have
rh Hark rul1
Then, under Assumption I have
[(++11) S ((xR) + VS(TA) ( "(x+(11=11111 - Ikk + L
| -xrkl
f(xx) HATE + L
2
S f(x)) 9p11rx11. - C1=
f(x)) 3 (n -3L(1 () ||xl12
S f(x)) 2 Bu. [V
(74)
where the last inequality is due 0 < B v 2u1+c-1y Thus, [(x++1) (1xx) s
[.Bul[V[(az)) (axx)).
Summing both sides of this inequality for & {O....,N N - 1} and recalling f(x) flow in Assumption gives
N flow f(xo) S f(xN)= - f(zo) S 3p
[f
2
k=0
Rearranging and dividing further by N yields 15
The next lemmas and proofs about the stochastic case.
Lemma 5. Suppose that Assumption
2 holds for { generated by RST-AM. In addition, if Bv > 0,
and 0g 2 0 and satisfies (73. then
+ of
6n ) (IVS(EN)B++
-
Es. HATR < 2
8 (1 2a - 20x)
(75a)
nk
1 o 1/2 + B&) o²
1 B f +
(f(x))
[Hrrk] [Hr
(75b)
2
n2k
where μ > O is the constant introduced in Remark 4
Proof. (i) From Lemma
we have
a
Es, Her 22 8
(1((1 (1+ 20 -2ax)
Es, [1)x
(76)
8k
From Assumption
have
Es, rx 1
Es. Illr Es. [reel121 [Es [rw]ll [V o² ///
(77)
With 76)..
obtain (55)
(ii) The result holds for & since Ho = Bot. Consider 2 1. Define 6 = Vis. (xx)
V((ak)
-Tk V (zx). Then HATE- = H₁ (-EH vf(zx)) Since ag satisfies 13 it has
Amin (2 (H& HT)) Sep. Thus
V((zx) 14 Vf(x))" (Hh + HH) V((xx) 2 (.
2
which implies
Es, (TTR)) H H&V((=))) 2 V
(78)
32

---

Published as a conference paper at ICLR 2022
Let M₁ a4 (P& + (xQu) 6 QF then H2 B&I M2. With the assumption
66a). i.e. Es, ((x) = 0, we have
Es, V ((rx)""Hee)) HRER] Es. (xx)) T (BRES - MaEx]]
((x)) " Esyles] Es, (zx) [M Es (V)))))
M
Using the Cauchy- Schhwar inequality with expectations, we obtain
Es, (T)( Hec = [S,
- Es. ( Es, M
(79)
We now bound ||M For brevity, let Z QE + P Pk. and N1 Pr
2. then
||ael2 I|au (N1 + N2) ll S on (1 l2 (N2ll2 (x + Bk),
(80)
where the last inequality is from Lemma 3
With (80, have || 2l2 S (n + 3%) 1l)/222 which implies
(5p + Ba) Es, (l (ī2] S a² 3 =Bx)2
Es. M
(81)
n
where the last inequality due to 66bb) Now we bound [Rs, (xx)) (HA as follows (cf. (7h):
Esp V fx HREall
[V Es. M&
(n $ + 3w] V fzxx) IEs, 11
<a + 3k) v
[x(x))]l
as (5
3kk σ TTk
Biau
1
10 Bz)) o² Bal
2 3ap (1
(82)
7kk
With the inequality
78 and 82
obtain
v (Ixx) Es, [ArK]]
(((x)) [He (Ex [[(xx))
Es. / HH Rs,, [V (1xx H&CK]
S - Es.. ( H (
+ Es.. V (Tx)) T HEE4]]
10 + +zx)² 0²
< -B
2 3ull [/(xx))))) [xxx)11
3ul
T1p
1 a 1 + Bx) o²
(83)
B&P
By imposing one more restriction on au, we can obtain a convenient corollary: Corollary 3. Suppose that Assumption
2 holds for { generated by RST-AM. C > 0 is the
constant in 12 If 3x 0, 0 < as < min B } and satisfies 61 then
Es, ||| Hark < 2,32 (1 + C-1) (IV/(Cn))1
o²
(84a)
Thk
1 Bapl (xxx)) + 1 (1 c-1) o²
(f(x)) [Hrr <
(84b)
n1k
33

---

Published as a conference paper at ICLR 2022
Proof. The first result (84a) is clear by considering (55g) and noticing that 1 2a - 2aa; 2 1 when
Ok [0, 1] and 0%
2 C Sg. Since as 5 5x 2 C and (1+C-3p C- 2(1 C-
have
1 a 1 (6x
+ Bx)²
o²
1x 1 + 3x)²
o²
2 3xll
nkk
7
0²
(C- à +1)2
2
71k
0²
- B3 (1
71k
Substituting it into (7b)). we obtain (84b).
With Corollary we establish the descent property of RST-AM:
Lemma 6. Suppose that Assumptions
7 and 2 hold for {a*} generated by RST-AM. C
the
constant 120. 150 Bx S
alifo-- S an min(1, { and satisfies 61), then
o²
Es. [f(++11) S (Zx)
Bal [f(xx)] & ((L # (1+
(85)
71
Proof. According to Assumption
we have
f(xx11) & f(x)) + Vf(xx) - 2k) + I.
1x - 2x
2
L _ -1
fxxx))+ + ff(x)) HETK +
(86)
Taking expectation with respect to the mini-batch S₃ on both sides of (86) and using Corollary obtain
3 we
Bs, ((
L
$ {(xx) + Vf(xx) FT,, [HRTR] + E.s, |
2
o2 T. + c-1)
o²
< {(xx) Bau V/(x))) (1 + c-
7
n
f(x) Bx () ( - [V)(2))
0²
(87)
n
Then 87 combined with the assumption Bg S wateo-y implies 85
Lemma 6 Suggests that the term related to the noise from gradient estimates is bounded as a second-
order term (i.e. () Thus, with the diminishing stepsize, the effect of noise also diminishes. To establish the global convergence, we introduce the definition of supermartingale following the proofs in (Wang et al. 2017] Wei et al. (021).
Definition 1. Let {Fi) be an increasing sequence of a-algebras. If (Xx) is a stochastic process satisfying (i) [[[Xx]] oo, (ii) Xş E for all k, and (iii) [[[ Fu] S Xx for all k, then {Xx}
called a supermartingale.
Proposition 3 (Supermartingale convergence theorem, see, e.g., Theorem 4.2.12 in (Durrett 2019 If (X&) is nonnegative supermartingale, then limg-soo Xx - X almost surely and IR X [[Xo].
S
Now, we prove the convergence theory of RST-AM in the nonconvex stochastic case.
Proof of Theorem 4 Define di == Den [>x)]] and == (I+ +
7 fxx) + o. E B. Let Fx be the o-algebra measuring da. 7k, and zx.. From (85) we know that
34

---

Published as a conference paper at ICLR 2022
for any k.
o²
20
= R [F) [B[[(+++1)]
-.-
13
n
1=++1
(((x) - 1 - B 112 +
20
8 i=k
7/k 6k.
(88)
n
which implies that E[78+11 plow FF.] 7k - flow th_ Since Φx 2 O, we have 0 < E flow] I 70 -Foo < +00. As the diminishing condition (4) holds, obtain (xf() - My f for
some constant My > 0. According to Definition [ f7H [low supermartingale. Therefore, Proposition 3 indicates that there exists a y such that limg- 7h with probability I, and Ehn] S [o Note that from 88
we have [[Øu] S [B[m]] 117++1] Thus,
E
ΦK
hhx E(7x+1)) +00,
which further yields that
Σ
0
3& V too with probability 1.
(89)
Since Ek=o 3h +00, it follows that 16) holds.
(ii) the noisy gradient is bounded, i.e.,
V S Mg: where M, > 0 is a constant, then a stronger result can be obtained.
(90)
For any give O, according 16. there exist infinitely many iterates oy such that ||V (xx) s
Then if (17) does not hold, there must exist two infinite sequences of indices {st), (4) with 8/, such that for 0, 0, k 8 I....,....
V 2 2e, |I f(r,))) < [f((m)ll 2c.
(91)
Then from it follows that
+00
Be V f(xx)
Be with probability 1,
which implies that
3R + 0 with probability a, +00.
(92)
k=s.
According to and (2)
we have
[[[| [
E Hu]|Hurk|llzz
$ 2 (38 (1+ +af - -20x) + azoz [[[[[/llla
S 3x /2 |12
S BA (B
< 8x 2
(93)
where the last inequalitics are due to Cauchy-Schwac inequality and 90). Then it follows from that
93
Zt, Fe, 11211] C2
Σ SE.
Amma
which together with 92) implies that || - Fea 1|2 + O with probability 1, as + +00. Hence, from the Lipschitz continuity of f.. it follows that ff(zt,) fx 0 with probability as i - +00. However, this contradicts (9I). Therefore, the assumption that (17) does not hold is not true.
35

---

Published as a conference paper at ICLR 2022
Proof of Theorem
5
According to 87 in Lemma 6 we have
N
Σ
3k
2
3x 3xL(1+C-1)) E8|
[
N-1
0²
- f(0o)
flow
+
$ + + (1
(94)
k=0
where the expectation is taken with respect to { N-N-1 Define
Bx (5 Bn L(1 + (-+))
Pn(k) def Prob ( kt)
Σラ Bj (3/1 33 + ( k= ON 11,
(95)
then
-ー11 BI - (1+ [[V
E [
Σラ= 83 B3 +
< DJ + o² +=
(96)
Σ/= B, -,- (1,+-,,
Let D be cidbemm-n constant. If we choose 3x 3 minf =11==--) a. and
D
n. = n, then the definition of Pn simplifies to Pn(k) = 1/N. From 96 we have
(L+ C No
ΣØ 8
+ oZ + / )1 + 0 Nga
N3 7
4D₃
d²(L (1+ 1--
N8p
Tmp
4Ds
4L(1+ + +C-1)
OVN
40 + 1 +1++) D
max
Nu
D
n/44
oVN
4D₃
4L((
C-1)
OVN
40 (1)(+ (1
Nu
D
npv/N
16D L(1+C-)
4D
(1+++)))++++10
2
PVN
D
n
Therefore, to ensure [ & e, the number of iterations is ((1/e).
D EXPERIMENTAL DETAILS
Our main codes were written based on the PyTorch framework and one GeForce RTX 2080 Ti GPU was used for the tests in training neural networks. Our methods are ST-AM (MST-AM) and the regularized version RST-AM.
D.1 EXPERIMENTAL DETAILS ABOUT ST-AM/MST-AM
The experiments about ST-AM conducted verify the main theorems, i.e. Theorem 1 Corol- lary and Theorem Four types of problems were used for the experiments. The conjugate residual (CR) method is short-term recurrence version the full-memory GMRES and needs matrix-vector products to fulfill the algorithm. We give the pseudocode of the CR method (Algorithm 6.20 in (Saad 2003 here for readers who are not familiar with this numerical algorithm.
Information about this framework found in https://pytorch.org
36

---

Published as a conference paper at ICLR 2022
Algorithm 4 CR Algorithm for solving a SPD linear system Av: b.
Input: E I
Output: 2 E R
1: To = b- -Axo Po = To
2: for k = =0.1,... until convergence, do
3: aę = rh Ar
(Apx)* Aps
4:
#k+1 = Ik + ORPR
5:
rk+1 Tk Apk
Ark11
6:
B& Ar
7: Pk+1 Tk+1 + 3EPE
8: Ap&+1 = Ark+1 + 3r Aps
9: end for
10: return 2²
According to Theorem [] the CR method is essentially equivalent ST-AM for solving strongly convex quadratic optimization. However, it should be pointed out that the CR method needs to directly assess the matrix A, and the residual update (Line 5) based on the linear assumption, which makes it inapplicable for nonlinear problems the case that A is unavailable. Also, even though finite difference technique can be used to construct the matrix-vector products, the number of gradient evaluations is twice of that of ST-AM.
D.1.1 STRONGLY CONVEX QUADRATIC OPTIMIZATION
To construct a case of we considered the least squares problem:
zE min f(z) == - Az - blβ2
1
(97)
where A E Redd b E g,, which can be reformulated as form of 5
We generated a dense random matrix E 10000x000 and dense random vector b E R500 for test. The gradient descent used fixed stepsize S TAI that can guarantee convergence, and the same 7) was used as the 3& for the full-memory AM and ST-AM.
We also conducted additional tests about solving the problem 997 with different condition numbers, where the eigenvalues of A were set be uniformly distributed. The results with different condition numbers (cond( A A) = 102, 10a 10%) are shown in Figure It can be observed that the curves of CR. AM and ST-AM nearly coincide, which verify the correctness Theorem T Also, since the cigenvalues uniformly distributed, the superlinear convergence behaviour may not happen. Nonetheless, CR, AM and ST-AM are still much faster then the GD method.
IKILAÍII DIGIAIAAI [Palbeieee
10
10-1
DSGEI/MA AM: [P>I)//)
10-2
ST AM. MANAAMI
ST AM: [PA]LА/
10-2
10-3
10
10-3
CR [[lbbiial
AM (SA15/19
AM ([bl-)])al
10-4
10
10
20
30
40
50
0
10
20
30
40
50
10
20
30
40
Beration
iteration
iteration
(a) cond(A"A))11 1 102
(b) cond(ATA) 1 x 10
(c) cond(A A) 10
Figure 4: Solving 97 with different condition numbers: cond(AFA) Amax (AT Amim (AT A).
A)/Aminn
37

---

Published as a conference paper at ICLR 2022
D.1.2 SOLVING A NONSYMMETRIC LINEAR SYSTEM
For the solution of a nonsymmetric linear system
Az
(98)
where E RRdd . E IRd, the fixed point iteration (FP) is Ik+1 = g(ax) == I 17(6 Axx), Theorem 2 requires g to be contractive map, i.e. ///- < 1. We used test matrix "fidap029" from the Matrix Market P]] This matrix is a banded matrix and not symmetric. We used Jacobi preconditioner for all the tested iterative methods. Since solving linear systems is not the main focus of this work, we did not do a thorough test of applying ST-AM to solve various nonsymmetric linear systems.
D.1.3 CUBIC-REGULARIE QUADRATIC MINIMIZATION
The concerned cubic-regularized quadratic minimization
min f(z - 11 b12 x =E2"
M |=112
(99)
where A Bexd b R., and M is the regularization parameter. It can be computed that
vf(x)= = (Az 6) |||2<|22 ( 1 M|| A.
Hence, for the gradient descent Fx+1 = g( := FA - ( the Jacobian of g n (|[r]l) ax. + || which has i 0 in the local region B(p)-
We generated a dense random matrix A E 500 100 and dense vector b E 12500 for the test. The official implementation of L-BFGS PyTorch was used for comparison. The historical length of L-BFGS was 50, i.e., BFGS was actually used in the test.
MICOO
137
AM MeE
10-1
10
04210 2000
MST-AM
(h) ST-AM
MST-AM
(d) ST-AM
Figure 5: Solving 99 with different M. (a) 12//1////ll// of MST-AM: (b) | of ST-AM;
(c) P / Ark lla and 1|Qu-1 S2 Ark-1l of MST-AM: (d) (PA-1Gall//lAaag
and Q: 1 Ara-1l22 of ST-AM.
We also conducted the tests related to different regularization parameters M. and the cases that M = 0.01, 1, 100 are shown in Figure Figure that with the large M = 100, both AM and MST-AM converge faster. An ablation study was also conducted about the boundedness restriction
of Pa-1Cel12,, 1 in Theorem 2 In Figure ((b)) we show the convergence behaviour of 100, ST- AM does not show a monotone decrease of the residual. To further investigate the cause, we plot the Figura and Figure d). It can magnitude of P P A 1/12: 11te:l12.111xx1444.. 11 | QQ 1 are quite oscilla- be observed that the evolutions of P 11/4/1/A//AAallll/ 11laaarallll11QAA-G......
ST-AM without the boundedness check. In the case of the rather large regularization
tory in ST-AM, while being roughly bounded below in MST-AM. This phenomenon phen-mmmmnnn may accounts for the more stable convergence behaviour of MST-AM and verifies the necessity of the changes of MST-AM compared with ST--MM. In our experiments, we also found the restarting period can be set quite large and has little effect on ST-AM.
38

---

Published as a conference paper at ICLR 2022
1.1.4 ROOT-FINDING PROBLEMS IN THE MULTISCALE DEEP EQUILIBRIUM MODEL
The multiscale deep equilibrium (MDEQ) model is recent extension of the deep equilibrium (DEQ) model (Bai et al. (2119) for computer vision. One of the central engines for these DEQ models is the coot-finding problem:
f0(2;;)) (=;xr) - > find =" s.t. fo(=";z) ;x)) 0,
(100)
where 8 and x are parameters and the input representation, respectively. In (Bai et al.) 2020), the Broyden's 's method is employed to solve this problem. Since MST-AM is suitable for solving non- linear systems of equations, we can apply MST-AM to solve (100).
We implemented the MST-AM method and integrated it into the MDEQ framework The task was image classification on CIFAR-10 10 and we used the small model for test. We followed the suggested experimental setting of the framework. The optimizer was Adam with learning rate of 0.001 and the weight-decay was 2.5 x 10 The batch size was 128 and the number of epochs was 50. The tested fixed-point solver was used for the forward root-finding process and for the backward root- moot-finding process. The threshold of the steps for the forward process was 18 and the threshold of the steps for the backward process was 20.
We used the built-in AM method and Broyden's method as the baseline methods. The m for AM was 20, i.e. using the full-memory AM.
D.2 EXPERIMENTAL DETAILS ABOUT RST-AM
Our experiments RST-AM focused training neural networks. Since ST-AM can be regarded
as a special case of RST-AM with 6()) = x(2) = 0, the basic ST-AM also covered. In the
Line 10 in Algorithm 1 ag should be adjusted meet the positive definiteness check (13) which
can be simplified to (xxx) ( Tk 2 Sap |rx]3 in practice. The adjustment of ag can be (1) 0 =
minfar, 23 (1 /a J. or (ii) an O. We used the option (ii) since the violation of 13) seldom
happened in our tests and option (ii) is more simple to apply.
In the experiments MNIST, Penn Treebank, incorporated preconditioning (described in Ap- pendix A.3) into the baseline method SAM and our method. Preconditioning is found to be effective for difficult problems, e.g. mini-batch training with very small batch sizes and the scaling of the model's 's parameters being important.
D.2.1 HYPERPARAMETER SETTING OF RST-AM
The hyperparameters of RST-AM are easy to tune. The only hyperparameters that need to be care- fully tuned are the regularization parameters C₁, c2 in of and τ(2) We found RST-AM is more sensitive to C1, possibly due to the fact that it influences the construction of secant equations. 02 can be quite small in our tests. The hyperparameter C' in a(2) was set to be very small (C = 10-16)
so as to ensure g(2)
<glirvll almost all the time. is introduced to prevent the denominators 1|
from being zero. We found 11 > 0 and || > O always held the tests, so eg can be omitted. In the experiments except for deterministic optimization on MNIST, we kept the setting C1 = 12 = x 10 unchanged and found such setting is quite robust.
The other hyperparameters are ake, Su. which can be always initially set as 1. So RST-AM has the same number of hyperparameters to tune as SGDM, since SGDM needs to tune learning rate and the momentum.
D.2.2 EXPERIMENTS ON MNIST
The experiments MNIST aimed to validate the effectiveness of RST-AM in deterministic opti- mization, so we were only concerned about the training loss by regarding it as a nonlinear function to be optimized. To facilitate the full-batch training, used subset of the training dataset by randomly selecting 10k images from the total 60k images.
hhtpp://github.coom htgggghhuh..oo//oouss////eeeee
Based on the official PoTorch hwppm/awwh.iiegggiioooooooooooom
39

---

Published as a conference paper at ICLR 2022
The baselines were SGDM, Adam, Adagrad, RMSprop and SAM. We tried our best to ensure that the baselines had the best performance in the tests. We tuned the learning rates by log-scale grid- searches from 10 33 to 100. For SGDM, Adam, Adagrad, and RMSprop, the learning rates were 0.2, 0,001, 0.01, 0.001, respectively. For SAM, we used the hyperparameter setting recommended in (Wei et al. 2021).
For RST-AM, we set ç₁ = 0.05, 02 = x 10-7, 000 = 8k = 1. When (Amx) rk Σ 0 occurs, 1 F2 + 0.2rk was used the new update.
For the preconditioned RST-AM, we set c =
x 10 = x 10-7 for the Adagrad- preconditioned RST-AM and e₁ 1 10 c2 1 x 10 for the RMSSpro--rreconditioned
RST-AM.
For all these tests, we trained the model for 200 epochs.
10
10'
100-2
10-3
Adam
Adamt
1001
1M55rtp
RMSpppp
Adom, SAMU?
Adam, SAMES
1M5prrp 54MI2
R35pppp: SAM02
10-4
10
spach
(a) Train loss
(b) Squared norm of gradient
Figure 6: Training MNIST with the Adam/RMSprop preconditioner.
In Figure 2 in the main paper, we report the RMSpropp/Adagrad preconditioned SAM/RST-AM. We give the result about using Adam the preconditioner in Figure 6 It suggests that Adam does not perform as well as the RMSprop method to serve as a preconditioner for SAM/RST-AM in this task. Nevertheless, Adam_ RSS-AMSTTTMA is still better than Adam and Adam_SAMC which demonstrates the effect of our proposed short-term recurrence scheme.
D.2.3 EXPERIMENTS ON CIFAR
This group of experiments were the same as those in (Wei al. 2021) so that we can have direct comparison between RST-AM and SAM. We still give the details here for completeness.
We followed the same way of training ResNet (He et al. 2016). The batchsize was be 128. For N iterations of training, the learning rate of each optimizer was decayed at the (L / I)-th and the N iterations. Here, for SAM and RST-AM, the ag and 3 serve the learning rates, so the learning rate decay denotes decaying 0k Bg simultancously. The experiments were run with random seeds.
The baseline optimizers were SGDM, Adam, AdaBound (Luo et al. 2018), AdaBelicf Zhhang et al. 2020), Lookahead (Zhang al. 2019) and AdaHessian (Yao et al. 2021). Adam and the recently proposed optimizers AdaBound and AdaBelief are adaptive learning rate methods which different learning rates for different model parameters. Lookahead k-step method and has an inner optimizer. In each cycle of Lookahead, the inner-optimizer oplim is used to iterate for steps and then the first iterate and the last iterate are interpolated to obtain the starting point of the next cycle. Compared to Adam, AdaHessian uses Hessian-vector products to construct diagonal approximation of the Hessian.
For fair comparison, the hyperparameters of all the optimizers (including RST-AM) were tuned For each optimizer, the hyperparameter setting that attained the highest final test accuracy on the other networks on CIFAR. We found the results hyperparameter tuning were consistent with those reported in (Wei et al.] 2021). For completeness, we list the settings of hyperparameters here (learning rate is abbreviated as Ir, and =*g indicates the same setting as that in (Wei et al. 202Ih).
through experiments on CIFAR-
was kept unchanged for training
SGDM': Ir 0.1, momentum 0.9, weight-decay x 10-4, Ir-decay 0.1.
40

---

Published as a conference paper at ICLR 2022
3
Test Accuracy %
Test Accuracy %
70-
220
140
100
Adam
SANILOI
SAMOO11
RST-AM
KET-MM
140
180
20
1.25
340
185
opachs
epochs
(a) CIFAR-IO/eeSee118
(b) CIFAR-IOWRNI644
03-
62-
60-
Test Accuracy
Test Accuracy
62-
120
160
120
S0DN
23-
Acaboand
SAM
SAM
R51-AM
102
320
362
302
320
opechs
opechs
(c) CIFAR-100/esNeeXSO0
(d) 200AA-000/DensNee1221
Figure Test accuracy training ResNet18 and WideResNet16-4 CIFAR-10 and training ResNeXt50 and DenseNet121 on CIFAR-100.
Adam": Ir 0.001, (81,92) = (0.9, 0.999), weight-decay = x 10-4, Ir-decay 0.1.
AdaBound: = 0.001, ($1,33) = 0.999, final_Ir = 0.1, gamma = 0.001, weight-
decay = 5 10 Ir-decay = 0.1.
AdaBelief": Ir 0.001, (81, (2) (0.9, 0.999), cps 10 10-8, weight-decay = 10-4,
Ir-decay = 0.1.
Lookahead optim: SGDM (Ir 0.1, momentum 0.9, weight-decay 10-3), a 0.8, steps 10, Ir-decay 0.1.
AdaHessian" :: 1.15, (81, (3) = 0.9.9..99)... eps=1 10 hessian-power: weight-
decay = 5 10 1011/0.15, Ir-decay 0.1.
SAM(2): oplim: SGDM (Ir 0.1, momentum = 0, weight-decay = 1.5 10 10-3),
1.0, 84 = 1.0, c1 = 0.01, = 1, m 2, weight =.5 10 Ir-decay = 0.06.
SAM(10)": optim: SGDM (Ir = 0.1, momentum 0, weight-decay = 1.5 x 3), ag
1.0, 8r = 1.0, c = 0.01, 1, m 10, weight-decay x 10 Ir-decay 0.06.
RST-AM: e 1, x 10-7
.00 10 Bo weight-decay 1 10-3 Ir-decay 0.1.
For the tests of SGDM, Adam, AdaBelief, Lookahead, AdaHessian and SAM(10), also had consistent numerical results with those reported in (Wei et al. 2021), so we reported their results of these methods in the main paper for reference.
shows the curves test accuracy for training ResNet18 and WRN16-4 on CIFAR-10 -00 and training Xt50 and DenseNet121 on CIFAR-100. The numerical results of final test accuracy can found in Table [(a] in the main paper. can be found Figure 7 that the convergence behaviour of RST-AM is less erratic than SAM(10) and SAM(2). In the first 80 epochs, RST-AM
Figure 7
41

---

Published as a conference paper at ICLR 2022
converges faster than SAM, partly due using smaller weight decay. The learning rate schedule has a large impact on the convergence of each optimizer. Similar to SAM(10), RST-AM can always climb up and stabilize to higher test accuracy in the final 40 epochs. It is found that RST-AM is comparable to SAM(10), while improving the short-memory SAM(2).
Trein Accuracy
Test Accuracy
SGDM
SGDM
RST AST c2-0.1
RST AK: C2-0.1
RST-AM (2=0,001
c2-10-2
RST AN: c2-10-
RST-AM: c2+10
RST-AM: c2=10-f
epochs
(a) Train accuracy
(b) Test accuracy
Figure 8: Train accuracy and test accuracy RST-AM with different c2.
100
90
Teat Accuracy
70
SGDM
SGDM
Train
RST AM: CL-10
c1=0.1
ci-0.1
RSTAM: c1-0.01
40
AST c2-0.01
140
epechs
(a) Train accuracy
(b) Test accuracy
Figure 9: Train accuracy and test accuracy RST-AM with different C1.
Test Accuracy %
SCOM
0c000.005
RST-AM: wd-0.0000
30
RST-MM: wc0000008
wdl-0.0008
RST AM: wd=0.001
RST-AM: www.
40-
RST-AM: wd-0.0015
20
60
D0
100
120
140
260
epochs
epochs
(a) Train accuracy
(b) Test accuracy
Figure 10: Train accuracy and test accuracy of RST-AM with different weight-decays.
Since our hyperparameter tuning was conducted on CIFARIO/R<sNe100, we give some results about the hyperparameters on this model.
Effeet of 02 in RST-AM. Figure g shows the effect of c2, where we kept e₁ = 1, weight-decay=5 x
10 fixed, and 0 10 10-1, 10 10 1007.10 10-0. It indicates that RST-AM not sensitive to C2.
Effect of Ci in RST-AM. Figure shows the effect of c1, where kept e2 10-7, weight- decay=5 x 10 fixed, and c = 0.0I, 0.1, 1, 10. It suggests that with smaller cj RST-AM converges faster in terms of train accuracy, but the test accuracy is worse.
42

---

Published as a conference paper at ICLR 2022
Effect of generalization of each optimizer. In Figure I0 we show the convergence behaviour of RST-AM with different weight-decays. It suggests that with too small weight-decay, RST-AM tends to be overfitting in the test dataset.
weight decay in RST-AM. Weight-decay is common hyperparameter that affect the
$9.92
64.02
24:2
Test Accuracy $5
Thitt Accuracy W
SCON: spoch=200
520M: 12501=202
SCOM spoch=100
SCOM: sco01=100
SAM 192e7
SAM 13:m
RST.A.M apect-000
RSTAN, opoch -00
EST-AM opect-100
ESTAM: epeth-102
151.4M) 120
IStaM cocch-195
20
60
100
122
140
161
20
60
100
122
3/0
162
epochs
epochs
(a) Test accuracy CIFARIO/ResNet18
(b) Test accuracy on CIFARIOWiiddRsNeello.
02
60
44
Test Accuracy
53-
62-
50
120
120
240
120
220
320
145
120
5CONE spach
SGOM
SCONe spoch
SCDM epoch 503
2000 spod
spocl (m
R:T JM cocch-13t
RSSTAN c3651-306
8:5T-AM: opoch-122
R.ST.AM e30ch-160
100
320
162
302
320
epochs
epechs
(c) Test accuracy on CIFARIO01ReNeXI5O
(d) Test accuracy on CIFAR
Figure 11: Training deep neural networks for 80,100,160 epochs. The results of final test accuracy of RST-AM for training 80, 160 epochs and the final test accuracy of SGDM for training 160 epochs are shown in the nested figures for comparison.
The cost and final test accuracy compared with SGDM. The notations "m", "te", "e", "t" and "a" are abbreviations of memory, per-epoch time, training epochs, total running time, and accuracy, respectively. "g* indicates numbers published in (Wei al. 2021).
Table 5:
Cost (x SGDM) accuracy
CIFAR10/ResNe118
CIFARIO/WRNI64
m
t/c
e
a( (%)
m
t/c
a(%)
SGDM"
1.00
1.00
1.00
1.00
94.82
1.00
1.00
1.00
1.00
94.90
SAM(10)*
1.73
1.78
0.56
1.00
94.81
1.26
1.28
0.63
0.80
94.94
RST-AM
1.05
1.46
0.56
0.82
94.84
1.03
1.14
0.63
0.71
94.95
Cost (x SGDM)
CIFAAR000/eeNeX1550
(IFAAR00/Denseeee112
& accuracy
t/e
e
a (%)
m
t/e
a(%)
SGDM
1.00
1.00
1.00
1.00
78.41
1.00
1.00
1.00
1.00
78.49
SAM(10)
1.30
1.16
0.50
0.58
78.37
1.16
1.19
0.50
0.60
78.84
RST-AM
1.04
1.07
0.50
0.54
78.39
1.01
1.11
0.50
0.55
78.90
Table I(a) in the main paper reports the test accuracy each optimizer when training for the same epochs. In fact, shown in Figure [T] within 100 epochs, RST-AM can achieve better test
43

---

Published as a conference paper at ICLR 2022
accuracy than SGD. So if the running time of the training process matters, it is expected that RST- AM can use less total running time due to the large number of reduction in training epochs. In Table [(b] we set SGD as the baseline, and compare the computation and memory cost with SGD. we give more details about the saving in training epochs and the final test accuracy. achieve comparable better test accuracy than SGDM with less computation time, while reducing the memory footprint of SAM.
In Table 5
be seen that RST-AM
We also tried using Adam preconditioner for RST-AM but found the final test accuracy was often worse. For example, the test accuracy of 90.79%. We suppose is the worse generalization ability of Adam (Luo al.) 2018) that makes Adam not suitable as a preconditioner for RST-AM in the image classification task.
for CIFAR- 1eR is only
D.2.4 EXPERIMENTS ON PENN TREEBANK
This group of experiments were the same as those in Wei et in Table 2 were measured with 3 random seeds. The parameter settings of the LSTM models were the same as those in (Zhuang et al. 2020 Wei et al. (2021). The baseline optimizers were SGDM, Adam, AdaBelief, and SAM. The validation dataset used for tuning hyperparameters.
2021) for direct comparison. Results
For SGDM, the learning rate (abbr. Ir) was tuned via grid-search in {1, 10, 30, 100}. We set Ir= 10, momentum 0.9 for the 2-layer/3-layer LSTM, and Ir 30, momentum for the 1-layer LSTM.
For Adam, the learning rate was tuned via grid-search in {1 x 10-3, x 10x 10-3, 1,,881003111 10 2 2,2 10 ") We found the setting that Ir 10 performs best.
5
For AdaBelief, we set Ir = 5 10-3 which is better than the recommended settings of the official implementation"
For SAM, we used the recommended setting in (Wei et al. 2021 and used the baseline Adam as the preconditioner. The cases 2 and 20 are denoted as SAM(2) and SAM(10), respectively.
For our method RST-AM, we kept c₁ = 1, c2 = 1 x 10-7 unchanged and used the same precondi- tioner (Adam) as SAM.
The batch size was 20. We trained the model for 500 epochs and the learning rate was decayed by 0.1 the 250th epoch and the 375th epoch.
The cost to achieve comparable results of Adam. The notations "m","te" and "t" abbreviations of memory, per-epoch time and total running time, respectively.
Table 6:
Cost
L-Layer t/e
2-Layer t/e
3-Layer t/e
(x Adam)
Adam
1.00
1,00
1.00
1.00
1.00
1.00
1.00
1,00
1.00
SAM(10)
1.20
1.84
0.74
1.36
1.88
0.75
1.90
1.67
0.67
RST-AM
1.06
1.73
0.69
1.15
1.78
0.71
1.11
1.53
0.61
Table 7: Test perplexity of training 1,2.3-layer LSTM on Penn Treebank for 200 epochs. Lower is
better. **** indicates numbers published in (Wei al. 202T).
Method
I-Layer
2-Layer
3-Layer
Adam
80.8
64.54
60.344.22
SAM(2)
81.82+.09
66.624.26
61.55+.11
SAM(10)
79.30
63.211.02
59.47+.07
RST-AM
79.49+.11
63.614.26
59.344.23
In Table 6
we report the memory footprint and per-epoch running time SAM(10) and RST-AM compared with Adam. It indicates that RST-AM also largely reduces the memory overhead of the
hhtpg//ywtttccmooooooomommmaamai 0.2.0/PyTorch
EXperiments/I.STM O2.MPeooc
44

---

Published as a conference paper at ICLR 2022
long-memone SAM(10) in the language modeling task. Since the batch size is very small, the cost of stochastic gradient evaluation is quite cheap, which makes the additional cost of matrix computation in RST-AM and SAM(10) become considerable. However, if we consider achieving comparable validation/test perplexity to that of Adam, RST-AM does not need to train for the same number of epochs Adam. Table 7 the test perplexity of Adam, SAM(2), SAM(10), and RST-AM for training 200 epochs. By comparing the results of Table Table 7 we see RST-AM is better than Adam with much fewer training epochs, thus RST-AM can save large amount of training time, as shown in Table 6
D.2.5 EXPERIMENTS ON ADVERSARIAL. TRAINING
The adversarial training considers the empirical adversarial risk minimization (EARM) problem:
min EER
max (x), 2e
(101)
where & is the i-th adversarial data within the e-ball centered at &, We followed the standard PGD adversarial training in (Madry et al. 2018), using projection gradient descent (PGD) to solve the inner maximization problem and the tested optimizers (SGD) and RST-AM) solve the outer mini- CIFAR10/WRN16
mization problem. The experiments were conducted on CIFAR 100/ResNet18, and C 1e1... We trained the neural networks for 200 epochs and decayed the learning rate at the 100th and 150th epoch.
Since adversarial training is much more time consuming than the ordinary training in Section D.2.3 we tuned the hyperparameters in CIFARI0/ResNe120,, and applied the same hyperparameters to other models. We found the setting that c₁ 1, 02 1 10 and weight-decay 0.001 still suitable for this task.
The CIFAR- (CIFAR-100) dataset contains 50K images for training and 10K images for testing. Since it is found that the phenomenon of overfitting is severer (Rice et al.] (020) in adversarial training, we randomly selected 5K images from the total 50K training dataset as the validation dataset (the other 45K images remained as the training dataset), and chose the best checkpoint model the validation set to evaluate on the test dataset. We consider two types of test accuracy:
(i the clean test accuracy, where clean data was used for model evaluation;
(ii) the robust test accuracy, where corrupted data was used for model evaluation. The attacking methods are FGSM (Goodfellow et al., 2014), PGD-20 (Madry al. 2018), and C&W oo attack
(Carlini & Wagner 2017).
Table 8:
Test accuracy (%) for adversarial training.
CIFAR 1
CIFAR 100/eenscNNt111
Optimizer
Clean
FGSM
PGD-20
C&W 0c
Clean
FGSM
PGD-20
C&W
SGD
82.16
63.23
51.91
50.22
59.45
39.76
30.92
29,00
RST-AM
82.53
63.78
52.43
50.52
60.48
40.41
31.20
29.52
CIFARIO/WRNI64
CIFAR100/ReNNe118
Optimizer
Clean
FGSM
PGD-20
C&W mc
Clean
FGSM
PGD-20
C&W x
SGD
80.84
60.97
49.29
47.62
55.42
36.17
28.18
26.31
RST-AM
81.36
61.38
49.93
47.95
56.49
37.00
28.50
26.66
In Tabel g we report the average results of tests with three different random seeds. shows that RST-AM can achieve both higher clean test accuracy and higher robust test accuracy across various models on CIFAR1O/CIFAL 100. To see the convergence behaviour of SGD and RST-AM, we plot the curves of the clean validation accuracy and the PGD-10 attacked validation accuracy in Fig- ure 12, We can observe the phenomenon robust overfitting (Rice et al. (2020) from these curves, which justifies our experimental setting with validation set for the checkpoint selection. Nonetheless, the numerical results suggest that RST-AM can still be better than SGDM in adversarial training.
It is also found that due to the heavy cost of gradient evaluations in PGD adversarial training, the additional computational cost incurred by RST-AM is negligible and the per-epoch running time of SGD and RST-AM is roughly the same. So we do not report it here.
45

---

Published as a conference paper at ICLR 2022
55
SGDM RST-AM
SGDM
50
50
RST-AM
Validation Accuracy
70
45
60
AdvValidation
40-
50-
40
25.
30
25
50
75
100
125
150
175
200
25
75
100
125
150
175
200
epochs
epochs
(a) CIFARIO/ReSNe118
(b) CIFARIO/ReSNe118
SGDM
Valiaaiii Accuracy %
SGDM
RSTAM
RST AM
Validation Accuracy
45
40
35
50
30
40
30
0
25
50
75
100
125
150
175
200
0
25
50
75
100
125
150
175
200
cpochs
epochs
(c) CIFARIOWWRNI644
(d) CIFARIOWRNI6-4
SGDM
SGDM RST-AM
50
RSTAAM
Adv-Validation Accuracy
25
Validation Accuracy
40
20
20
10
0
25
50
75
100
125
150
175
200
0
25
50
75
100 epochs
125
150
175
200
epochs
(c) JIFAR100/ResNett
(f) CIFAR N00/8
SGDM
wwwo
Adv-Validation Accuracy %
30-
SGDM
Validation Accuracy %
50
20
30-
10
epochs
epochs
(g) 21
(h) CIPAR1000DenneNe1112
Figure 12: Clean accuracy and PGD-10 attacked accuracy on the validation in training different neural networks.
46

---

Published as a conference paper at ICLR 2022
D.2.6 EXPERIMENTS ON TRAINING GENERATIVE ADVERSARIAL NETWORK
We describe our setting of training generative adversarial network (GAN) here. Like the adver- sarial training, the GAN training process is also min-max problem. The stability of an optimizer critical for the training process. To demonstrate the applicability of RST-AM, we conducted experiments GAN which was equipped with spectral normalization (Miyato et al. 2018) (SN- GAN). The experimental setting was the same as that of AdaBelief (Zhuang et al. 2020): the dataset was CIFAR- 10; ResNets were used the generator and the discriminator networks; the steps for optimization in the discriminator and the generator per iteration were and I, respectively; the min- batch size 64 and the total iteration number was 100000. The Frechet Inception Distance (FID) Heusel et al. 2017) was used as the evaluation metric: lower FID score means better accuracy.
The baseline optimizer were Adam and AdaBelief as they perform well in this task (Zhuang al. 2020 We also used the recommended hyperparameter settings for the two optimizers. For our RST-AM method, due to the ill-conditioning of the min-max problem, we used the AdaBelief the preconditioner and set the damping parameter ogk = 0.6. The regularization parameters Cj = 1 and were still kept unchanged just as the previous experiments.
22
10
45
AdaBellet
RST-AM
score
35
FID
15
20008
60000
80000
100000
iteration
Figure 13: FID score for training SN-GAN on CIFAR-10.
In Figure 13 we show the curve of FID score for each optimizer, which is the average of three independent runs. It indicates that the RST-AM method stable for this min-max optimization problem.
Table 9: The effect of og for RST-AM.
Method
Adam
AdaBelief
a = 0.8
á = 0.6
a: = 0.5
à =
0.4
à = 0.2
a = 0.1
FID score
13.07
12.80
12.48
12.05
12.75
13.13
13.07
12.59
Since we only tuned the damping parameter ag for RST-AM, we report the FID scores of other choices of oge in Table 9 during our experiment. It shows that even with the suboptimal choices of ak, e.g. Og = 0.1, 0.5, 0.8, RST-AM can still outperform the baselines.
D.3 ADDITIONAL EXPERIMENTS
To further test the performance of our method in training neural networks on larger datasets or different models, we conducted additional experiments of the image classification task ImageNet DDng al. 2009) and the Transformer (Vaswani et al. (007)) based neural machine translation task
in the IWSTL14 DE-EN (Cettolo et al. 2014) dataset.
D.3.1 EXPERIMENTS ON IMAGENET
We trained ResNet50 on ImageNet with SGDM and RST-AM. We used the built-in ResNet50 model in PyTorch. We ran the tests of each optimizer with three random seeds and four GeForce RTX 2080
47

---

Published as a conference paper at ICLR 2022
Ti GPUs were used for each test. The hyperparameters of SGDM setting in PyTorch. For RST-AM, the hyperparameters were kept the same as those in the CIFAR experiments. The weight-decay was 1 x 10 The number of the total training epochs is 90. The learning rate decay of SGD was at the 30th and 60th epochs. For RST-AM, since the experiments on CIFAR show it can often converge faster early learning rate decay strategy recommended in (Zhang et al. 2019): decay the and 3y for RST-AM the 30th, 50th and 70th epochs.
the recommended
acceptable solution than SGDM, adopted the
Table 10: TOP 1 test accuracy (%) w.r.t. epoch, the best TOPI test accuracy (%), and the cost. The memory, per-epoch time and total time of SGDM are set the units. The total time is the time to first achieve the accuracy 2 75.90%.
Method
epoch 72
epoch 88
epoch 90
best
memorry
per-epuch time
total time
SGDM
75.75
75.90
75.81 75.98
75.93 76.04=.06
1.00
1.00
1.00
RST-AM
75.90
75.95
1.06
1.01
0.83
In Table 10 we report the TOPI accuracy in the validation dataset during training. Note that we also report the epoch number for each optimizer to achieve the accuracy equal or exceeding 75.90%. shows RST-AM needs 72 epochs while SGDM needs 88 epochs. The curves of the training accuracy and test accuracy are shown Figure [4 The results suggest that RST-AM is still competitive optimizer in training larger model in a larger dataset.
SGOM
SGOM
Accuracy
20-
Accuracy
22
d0
epochs
eparhs
(a) Train accuracy
(b) Test accuracy
Figure 14: Train and test accuracy for training ImageNet/ResNet50..
D.3.2 TRAINING TRANSFORMER
We conducted the neural machine translation task with Transformer. We implemented our RST-AM The basic experimental setting set as
method and integrated it into the fairseq framework [ the recommended setting in (Yo et al. 202T). We trained the model for 50 epochs and the BLEU (Papineni al. 2002) score was calculated using the average model of the last five checkpoints. We added a baseline optimizer RAdam (Liu et al. 2019) which was inspired by the warmup procedure training Transformer.
Table 11: The BLEU score of training Transformer on IWSLT14.
Method
SGD
Adam
AdaBelief
RAdam
RST-AM
BLEU score 28. 18.14=.08 00
35.71#.03
35.15 14
35.60+.06
35.89+.02
http:://hhhu..com/yyrrrfffaaeg
48

---

Published as a conference paper at ICLR 2022
The numerical results reported in Table 11 show that Adam is well-suited for this neural machine translation task, though it does not perform well in the image classification task. Also, the results demonstrate that RST-AM still outperform Adam in this task.
Table 12: BLEU score evaluated the 40/45/50-th epoch, and the cost. The memory, per-epoch time and total time of Adam are set as the units. The total time is the time of RST-AM to achieve a BLEU score matching the final BLEU of Adam: BLEU BLEU(Adam) I 0.03, where 0.03 is the standard deviation of the results of Adam.
Method
epoch
epoch = 45
epoch
memory
per-epoch time
total time
Adam
35.42=.10
35.54+.08
35.71=.03
1.00
1.00
1.00
RST-AM
35.59=..4
35.691.06
35.89:.02
1.16
1.00
0.90
In Table 12 we report the BLEU scores evaluated in the test dataset at the 40th, 45th and 50th epochs. The results show the better performance of RST-AM over Adam and the per-epoch computational cost is nearly the same. To achieve comparable solution to Adam, RST-AM can save 10% training time.
49


Reviews:
[{"review_id": "gcE3fPl4ahs", "reviewer": "Reviewer_xxcF", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper tackles an ambitious goal, which is to design better optimizers for machine learning workloads than the first order methods which have come to dominate the field. It makes an interesting contribution with several well-motivated new algorithms, establishing convergence proofs and running detailed experiments to check whether the theory matches the practice.\nThe paper is generally well-written and nice to read, with special care taken to cite appropriate earlier works.\n\nHere are a few outstanding issues, however.\n1. The paper is perhaps not as self-contained as one could hope for, and assumes a strong familiarity with many topics: e.g. Anderson Mixing, the properties of first order methods on fixed point problems, preconditioning (which is only made explicit far in the Appendix). It would make the paper more readable and easier to understand to add more context to it.\n\n2. A claim is made that the memory footprint of the algorithms is close to that of GD or SGD, but a more precise theoretical comparison would be helpful. For instance, we see that the memory footprint seems to only increase by 5% over SGD in some experiments, but are bigger than Adam in others which leaves the reader wondering how these seemingly inconsistent results came about.\n\n3. Finally, the biggest issue with the paper is its experimental section. While thorougly reported (which is appreciated), it relies on toyish datasets which make the claims of the paper more difficult to support. Another issue is that the authors chose RNNs for language tasks while transformers have obtained SOTA results for years, and Adam seems to be more critical to good transformer performance than good RNN performance. All told, it seems like the experimental section, while well-executed, misses key points to make it really relevant to the community.\n\nAll told, this is a nice paper which clears the bar for publication at ICLR this year, but which could be much improved with a revamped, more ambitious experimental section.\n\nDetails\n- section 3.4, first paragraph: optimizaiton --> optimization\n- fig 1 caption : Probrem --> Problem\n- page 34 precondtioner --> preconditioner", "summary_of_the_paper": "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.\nThe 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.", "summary_of_the_review": "Interesting new algorithms, well-motivated, delivered with good theoretical analysis. The experimental section is a bit lacking, but that's not sufficiently concerning not to recommend the paper for publication.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "drhPKwMwJC2", "reviewer": "Reviewer_i7Ws", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I begin by noting that I am not an expert on the finer details of the theory of (non)linear solvers, so my review is primarily focused on the application side of things.\n\nThis appears to be an excellent paper. The quality of the writing is uniformly high throughout. I appreciated the various remarks throughout, used to elucidate the occasional mathematical detail.\n\nMy first comment is that I think it would be really beneficial to more explicitly state the delta between this work and previous work. When is it worth the effort to implement the introduced methods as compared to some off-the-shelf competing method? (LU solvers for linear problems; quasi-Newton for nonlinear problems; Adam for supervised learning etc.) For the audience interested in applications, rather than theory, I believe a section describing this would be especially valuable.\n\nI would describe the experiments as still having room to improve. (And I apologise for being \"that reviewer\" who asks for more experiments, in a paper that already features several.) Numerical methods for the solution of (non)linear systems are applied widely across disparate fields, so in this case I think it is particularly important to check how broadly applicable any new technique is. I would regard the following as being particularly important extra comparisons:\n\n- Section 4.1: how does the proposed method compare to Newton and standard quasi-Newton methods (BFGS, L-BFGS, chord, Levenberg–Marquardt, etc.)? For those problems for which the Jacobian is available -- an increasingly common case following the popularisation of autodifferentiation -- these are still the go-to technique in many/most regimes.\n- Section 4.1: the test problems considered are all very simple. What about more complicated test problems? In particular I would be interested in finding the fixed point of \"complicated\" nonlinear functions (e.g. a neural network, or the system arising from an implicit ODE solver).\n  - Focusing on the implicit ODE case: at present this is ubiquitously done via Newton's method or the (quasi-Newton) chord method. (And in particular not via Picard iterations, which can introduce an undesirable loss of stability to the overall numerical method.)\n- Training GANs. The choice of optimiser can have a *very* dramatic impact when training GANs: easily making the difference between convergent and divergent training, despite the same change in optimiser having minimal-to-no effect in the supervised learning case. A few examples:\n  - [1] consider a case in which GAN training with SGD/Adam fails completely, but training with Adadelta produces stable/convergent results.\n  - [2] derive optimisers by solving the gradient flow ODE. [Usually not a great idea in the supervised learning case, but the unreliability of other optimisers in the adversarial case motivates a return to this idea.]\n  - [3, 4] consider using negative or complex momentum within the optimiser.\n  - [5] consider taking Cesàro means over the parameters being optimised, during training. (This is commonly also done in theoretical studies of optimisers.)\n\nAn in addition the following would be nice as extra comparisons, but I think they are less important than those suggested above:\n\n- Section 4.2: MNIST: SAM and RST-AM preconditioned with Adam is noticeably absent.\n- Section 4.2: CIFAR: SAM and RST-AM with Adadelta/Adagrad/Adam/RMSprop preconditioners. These seemed to help a lot on MNIST. (Figure 2(c))\n\n## Minor points / typos\n\nEquation (1): I think specifying $X_k \\in \\mathbb{R}^{m \\times d}$, $R_k \\in \\mathbb{R}^{m \\times d}$ would improve readability slightly.\n\nRemark 1: I'd suggest including here that the mixing step is essentially a gradient descent step ($\\bar{r}_k$ is some linear combination of $\\nabla f(x_j)$; c.f. also the standard condition given by equation (14)).\n\nTop of page 4: both $\\bar{x}_k^G$ and $x_k^G$ are used. Neither notation is really defined precisely, but I am guessing that $\\bar{x}_k^G$ is meant to be $x_k^G$.\n\nTop of page 4: it is stated that $\\bar{x}_k = \\bar{x}_k^F = x_k^G$ . This would seem to indicate that the $\\bar{x}_k$ iterates of ST-AM compute exactly the same values as AM (with the corollary that the actual output $x_k$ is almost identical), and as such ST-AM is a strict improvement  over AM due to its lower memory requirements. This seems at odds with the rest of the text, which seems to suggest that the use of ST-AM represents a trade-off.\n\nTheorem 2: \"constant $\\kappa$...\" -> \"constants $\\kappa$...\"\n\nSection 3.4, end of second paragraph: $\\nabla_{S_k} f(x_k)$ -> $\\nabla f_{S_k}(x_k)$.\n\nEquation (14): some space (\\qquad) between the equations would be nice.\n\nEnd of page 8: \"less training epochs\" -> \"fewer training epochs\"\n\nPage 9: \"less epochs\" -> \"fewer epochs\"\n\n## References\n\n[1] Kidger et al. \"Neural SDEs as Infinite-Dimensional GANs\" NeurIPS 2021 https://arxiv.org/abs/2105.13493\n\n[2] Qin et al. \"Training Generative Adversarial Networks by Solving Ordinary Differential Equations\" NeurIPS 2020 https://arxiv.org/abs/2010.15040\n\n[3] Gidel et al. \"Negative Momentum for Improved Game Dynamics\" AISTATS 2019 https://arxiv.org/abs/1807.04740\n\n[4] Lorraine et al. \"Complex Momentum for Optimization in Games\" 2021 https://arxiv.org/abs/2102.08431\n\n[5] Yazici et al. \"The Unusual Effectiveness of Averaging in GAN Training\" ICLR 2019 https://arxiv.org/abs/1806.04498", "summary_of_the_paper": "The paper introduces three variants of Anderson acceleration: short-term Anderson mixing (ST-AM), modified short-term Anderson mixing (MST-AM), and regularised short-term Anderson mixing (RST-AM). These may be used to obtain the solution to linear systems, nonlinear systems, and optimisation problems.", "summary_of_the_review": "I recommend acceptance due to the high quality of the paper. My main two concerns are (a) readability wrt applications and (b) experiments across a broader range of regimes.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "PI-GsMduop7", "reviewer": "Reviewer_Axgn", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I think overall this is a paper with solid theory and experiments. The theoretical and experiments can be separated into two parts, one part is for the basic version of ST-AM designed mostly for nice problems like strongly convex quadratic ones, the other part is for RST-AM designed for neural network training. In my perspective, the first part may not be of too much interest to the audience of ICLR and is more suitable for optimization journals. Yet, the part of the MST-AM algorithm could be of great interest since it is applicable to training deep models. For the weakness, I feel the discussion on improvement of MST-AM over SAM (Wei et al., 2021) could be improved. The paper mentioned MST-AM is motivated by SAM  and talked about which steps are different. It could be better if the authors could provide a discussion on intuitions why MST-AM outperforms SAM in the paper. Also, the content of the paper is quite dense, it might be better to move some less important results into appendix and provide more discussion on intuitions in the main paper.", "summary_of_the_paper": "The paper proposes a new class of memory-efficient Anderson mixing (AM) methods. Compared with classical Anderson mixing which requires saving m historical iterates, the new variants require only storing two historical iterates while keeping good performance. Convergence analyses are given to the proposed variants showing a variant for strongly convex quadratic problem enjoys the same guarantee as full-memory AM, and another variant converges with a similar convergence rate as SGD (O(1/\\sqrt{number of iterations})) on nonconvex problems. Experiments on MNIST, CIFAR-10, and PENN TREEBANK with some popular deep models validates the superior performance of one of the proposed method.", "summary_of_the_review": "In summary, I think this is a relatively strong paper due to the technical depth and the empirical studies, though some parts could be improved.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers found this work well-motivated and the additional experiments conducted during the response phase were greatly appreciated. Anderson's acceleration appears to be a simple device that may be of great value to this field, and therefore this work is a timely contribution. The presented theoretical results justify the authors' modifications, although at times it felt more comparisons would be welcome: (a) Section 3.2 could have compared to a lot of three-term recurrences that lead to the optimal dependence on the condition number, including Chebyshev's polynomial and conjugate gradient, as well as the results in Brezinski et al. (2018); (b) Section 3.3 would benefit from some comparison with "Evans, Claire, Sara Pollock, Leo G. Rebholz, and Mengying Xiao (2020). “A Proof That Anderson Acceleration Improves the Convergence Rate in Linearly Converging Fixed-Point Methods (But Not in Those Converging Quadratically)”. SIAM Journal on Numerical Analysis, vol. 58, no. 1, pp. 788–810." (c) the results in Section 3.4 seem to be a bit preliminary and it would be great if the authors could compare to standard rates of SGD. 

Overall we believe this work will generate more interest on memory-based optimization techniques in deep learning, and we encourage the authors to thoroughly polish their draft by incorporating the reviewers' comments and the responses during the discussion phase.
Decision: Accept (Poster)