A C LASS OF S HORT - TERM R ECURRENCE A NDERSON M IXING M ETHODS AND T HEIR A PPLICATIONS
A BSTRACT
Anderson mixing (AM) is a powerful acceleration method for ﬁxed-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 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 ﬁxed-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-ﬁnding problems and training neural networks.
Experimental results show that ST-AM is competitive with the long-memory AM and outperforms many existing optimizers
I NTRODUCTION
Anderson mixing (AM) (Anderson, 1965) is a sequence acceleration method (Brezinski et al., 2018) in scientiﬁc computing.
It is widely used to accelerate the slow convergence of ﬁxed-point iterations (Lin et al., 2019; Fu et al., 2020; An et al., 2017), e.g., the self-consistent ﬁeld iterations in electronic structure computations (Garza & Scuseria, 2012; Arora et al., 2017).
Speciﬁcally, we consider a ﬁxed-point iteration x k +
= g ( x k ) , k = 0 , 1 , . . .
, where g : R d (cid:55)→ R d is the ﬁxed-point map.
By using m historical iterations, AM( m ) aims to extrapolate a new iterate that satisﬁes certain optimality property.
When the function evaluation is costly, the reduction of the number of iterations brought by AM can save a large amount of computation (Fang & Saad, 2009).
AM can be used as a method for solving nonlinear equations (Kelley, 2018) as the ﬁxed-point problem x
= g ( x ) is equivalent to h ( x ) :
= x − g ( x ) = 0 .
In practice, since computing the
Jacobian of h ( x ) is commonly difﬁcult or even unavailable (Nocedal & Wright, 2006), AM can be seen as a practical alternate for Newton’s method (An et al., 2017).
Also, compared with the classical iterative methods such as the nonlinear conjugate gradient (CG) method (Hager & Zhang, 2006), no line-search or trust-region technique is used in AM, which is preferable for large-scale unconstrained optimization.
Empirically, it is observable that AM can largely accelerate convergence, though its theoretical analysis is still under-explored.
It turns out that in the linear case (Walker & Ni, 2011;
Potra & Engler, 2013), the full-memory AM ( m = k ) is essentially equivalent to GMRES (Saad & Schultz, 1986), a powerful Krylov subspace method that can exhibit superlinear convergence behaviour in solving linear systems (Van der Vorst & Vuik, 1993).
For general nonlinear problems,
AM is recognized as a multisecant quasi-Newton method (Fang & Saad, 2009; Brezinski et al., 2018).
As far as we know, only local linear convergence has been obtained for the limited-memory
AM ( m < k ) in general (Toth & Kelley, 2015; Evans et al., 2020; De Sterck & He, 2021).
For the application of AM, one of the major concerns is the historical length m , a critical factor related to the efﬁciency of AM (Walker & Ni, 2011).
A larger m can incorporate more historical information into one extrapolation, but it incurs heavier memory overhead since 2 m vectors of dimension d need to be stored in AM( m ).
The additional memory footprint can be prohibitive for solving high-dimensional problems in a resource-limited machine (Deng, 2019).
Using small m can alleviate the memory overhead but may deteriorate the efﬁcacy of AM since much historical information is omitted in the extrapolation (Walker & Ni, 2011; Evans et al., 2020).
To address the memory issue of AM, we deeply investigate the properties of the historical iterations produced by AM and leverage them to develop the short-term recurrence variant, namely ST-AM.
The basic version of ST-AM imposes some orthogonalization property on the historical sequence, which is inspired by the CG method (Hestenes & Stiefel, 1952) that enjoys a three-term recurrence.
Furthermore, to better suit the more difﬁcult nonconvex optimization, a regularized short-term form is introduced.
We highlight the main contributions of our work as follows.
We develop a novel class of short-term recurrence AM methods (ST-AM), including the basic ST-AM, the modiﬁed 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 ﬁxed-point problems; RST-AM aims for solving stochastic optimization.
An important feature of STAM is that all methods only need to store two previous iterations with cheap corrections, which signiﬁcantly reduces the memory requirement compared with the classical AM.
A complete theoretical analysis of the ST-AM methods is given.
When solving strongly convex quadratic optimization, we prove that the basic ST-AM is equivalent to full-memory AM and the convergence rate is similar to that of the CG method.
We also prove that MST-AM has improved local linear convergence for solving ﬁxed-point problems.
Besides, we establish the global convergence property and complexity analysis for RST-AM when solving stochastic optimization problems.
The numerical results on solving (non)linear equations and cubic-
regularized quadratic optimization are consistent with the theoretical results for the basic ST-AM and MST-AM.
Furthermore, extensive experiments on training neural networks for image classiﬁcation and language modelling show that RST-AM is competitive with the long-memory AM and outperforms many existing optimizers such as SGD and ADAM.
R ELATED WORK
AM is also known as an extrapolation algorithm in scientiﬁc computing (Anderson, 2019).
A parallel method is Shanks transformation (Shanks, 1955) which transforms an existing sequence to a new sequence for faster convergence.
Related classical algorithms include Minimal Polynomial Extrapolation (Cabay & Jackson, 1976) and Reduced Rank Extrapolation (Eddy, 1979), and a framework of these extrapolation algorithms including AM is given in (Brezinski et al., 2018).
Note that an elegant recursive algorithm named (cid:15) -algorithm had been discovered for Shanks transformation for scalar sequence (Wynn, 1956), and was later generalized as the vector (cid:15) -algorithm
(Wynn, 1962) for handling vector sequences, but this short-term recurrence form is not equivalent to the original
Shanks transformation in general (Brezinski & Redivo-Zaglia, 2017).
Since AM is closely related to quasi-Newton methods (Fang & Saad, 2009), there are also some works trying to derive equivalent forms of the full-memory quasi-Newton methods using limited memory (Kolda et al., 1998; Berahas et al., 2021), while no short-term recurrence is available.
To the best of our knowledge, ST-AM is the ﬁrst attempt to short-term recurrence quasi-Newton methods.
Recently, there have been growing demands for solving large-scale and high-dimensional ﬁxedpoint problems in scientiﬁc computing (Lin et al., 2019) and machine learning (Bottou et al., 2018).
For these applications, Newton-like methods (Byrd et al., 2016; Wang et al., 2017; Mokhtari et al., 2018) are less appealing due to the heavy memory and computational cost, especially in nonconvex stochastic optimization, where only sublinear convergence can be expected if only stochastic gradients can be accessed (Nemirovskij & Yudin, 1983).
On the other side, ﬁrst-order methods (Necoara et al., 2019) stand out for their low per-iteration cost, though the convergence can be slow in practice.
When training neural networks, SGD with momentum (SGDM) (Qian, 1999), and adaptive learning rate methods, e.g. AdaGrad (Duchi et al., 2011), RMSprop (Tieleman & Hinton, 2012),
Adam (Kingma & Ba, 2014), are very popular optimizers.
Our methods have the nature of quasiNewton methods while the memory footprint is largely reduced to be close to ﬁrst-order methods.
Thus, ST-AM can be a competitive optimizer from both theoretical and practical perspectives.
M ETHODOLOGY
In this section, we give the details of the proposed ST-AM.
We always assume the objective function as f : R d → R , the ﬁxed-point map g : R d (cid:55)→ R d .
Moreover, we do not distinguish r k = −∇ f ( x k ) and r k = g ( x k ) − x k in our discussion as ∇ f ( x ) = 0 is equivalent to g ( x )
=
−∇ f ( x ) + x = x .
3. A NDERSON MIXING
The AM ﬁnds the ﬁxed point of g via maintaining two sequences of length m ( m ≤ k ) :
[Equation]
where the operator ∆ denotes the forward difference, e.g. ∆ x k = x k +1 − x k .
Each update of AM can be decoupled into two steps, namely the projection step and the mixing step :
[Equation]
where ¯ r k :=
r k − R k Γ k and β k > 0 is the mixing parameter.
The Γ k is determined by
[Equation]
Thus, the full form of AM (Fang & Saad, 2009; Walker & Ni, 2011) is
[Equation]
Remark 1.
To see the rationality of AM, assume g is continuously differentiable, then we have h ( x j ) − h
( x j − 1 ) ≈ h (cid:48)
( x k )( x j − x j − 1 ) around x k , where h (cid:48) ( x k ) is the Jacobian of h ( x ) :
= x − g ( x ) .
So, it is reasonable to assume R k ≈ − h (cid:48)
( x k ) X k , and we see (cid:107) r k − R k Γ (cid:107)
≈ (cid:107) r k
+
h
(cid:48)
( x k ) X k Γ (cid:107) 2 .
Thus, we can recognize (3) as solving h (cid:48)
( x k ) d k = h ( x k ) in a least-squares sense, where d k =
X k Γ k .
The mixing step incorporates r k into the new update x k +1 if β k > 0 .
Otherwise, x k +1 =
¯ x k is an interpolation of the previous iterates, leading to a stagnation.
3. T HE BASIC SHORT - TERM RECURRENCE A NDERSON MIXING
The basic ST-AM is to solve the strongly convex quadratic optimization:
[Equation]
where A (cid:31) 0 . Let p − 1 = q − 1 = p 0 = q 0 = 0 ∈ R d . At the k -th iteration, given the two matrices P k − 1 =
( p k − 2 , p k − 1 ) ∈ R
d × 2 , Q k − 1 =
( q k − 2 , q k − 1 ) ∈ R d × 2 and deﬁning p = x
k − x k − 1 and q = r k − r k − 1 , the basic ST-AM constructs
[Equation]
(6b
[Equation]
Then, we update P k
=
( p k − 1 , p k ) , Q k = ( q k − 1 , q k ) ∈
R
d × 2 .
Such construction ensures Q T k Q k =
I 2 for k ≥ 2 and the storage of P k and Q k is equal to AM(2).
With the corrected P k and Q k , the
ST-AM method modiﬁes the projection step and the mixing step accordingly, that is,
[Equation]
where Γ k = arg min (cid:107) r k − Q k Γ (cid:107) 2 = Q T k r k and ¯ r k = r k − Q k Γ k .
Thus, the ST-AM replaces X k and
R k in (1) by P k and Q k respectively and imposes the orthogonalization property on Q k .
The details of basic ST-AM are given in Algorithm 2 in Appendix C.1.
Deﬁne ¯ P k
=
( p 1 , p 2 , . . .
, p k ) , ¯ Q k = ( q 1 , q 2 , . . . , q k ) , the Krylov subspace K m ( A, v ) ≡ span { v, Av, A 2 v, . . .
A m − 1 v } , the range of X as range(X) .
We give the properties of the basic ST-AM in Theorem 1.
Theorem 1.
Let { x k } be the sequence generated by the basic ST-AM.
The following relations hold:
(i) (cid:107) ˜ q (cid:107) 2 > 0 , range( ¯ P k )
=
range( X k ) =
K k ( A, r 0 ) , range( ¯ Q k )
=
range( R k )
=
A K k ( A, r 0 ) ; (ii) ¯ Q k = − A ¯ P k , ¯ Q T k ¯ Q k
= I k ; (iii) ¯ r k ⊥ range( ¯ Q k ) and
¯ x k = x 0 + z k , where z k = arg min
z ∈K k
( A,r 0 ) (cid:107) r 0 − Az (cid:107) 2 .
If (cid:107) ¯ r k (cid:107) 2 = 0 , then x k +1 is the exact solution.
The proof is in Appendix C.1.
Note that the property (iii) in Theorem 1 exactly describes the relation ¯ x k = x
G k , where x G k is the output of the k -th iteration of GMRES (Saad & Schultz, 1986).
Moreover, let ¯ x
AM k be the k -th intermediate iterate in full-memory AM ( m = k ).
We prove that ¯ x AM k
= x
G k
(See
Proposition 1 in Appendix C.1.), which induces that ¯ x k = ¯ x AM k
= x G k .
This equivalence indicates that ST-AM is more efﬁcient than AM and GMRES since only two historical iterations need to be stored.
Moreover, by directly applying the convergence analysis of GMRES (Corollary 6.33 in (Saad, 2003)), we obtain the convergence rate of the basic ST-AM for solving (5):
Corollary 1.
Suppose the eigenvalues of A lie in [ µ, L ] with µ > 0 , and let { x k } be the sequence generated by the basic ST-AM, then the k -th intermediate residual ¯ r k satisﬁes (cid:107) ¯ r k (cid:107) 2 ≤ (cid:18) √
L/µ − 1 √
L/µ + (cid:19) k (cid:107) r 0 (cid:107) 2 .
Moreover, the algorithm ﬁnds the exact solution in at most ( d +1) iterations.
Remark 2.
The GMRES can be simpliﬁed 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 simpliﬁcation for AM is expected to exist.
Like CG and Chebyshev acceleration (Algorithm 12.1 in (Saad, 2003)), the convergence rate of ST-AM has the optimal dependence on the condition number, while ST-AM does not form the Hessian-vector products explicitly.
3. T HE MODIFIED SHORT - TERM RECURRENCE ANDERSON MIXING
For general nonlinear ﬁxed-point problems, global convergence may be unavailable for the basic
ST-AM, as a counter-example exists for full-memory AM (Mai & Johansson, 2020).
Thus, we propose a modiﬁed version of the basic ST-AM (MST-AM) and prove the local linear convergence rate under similar conditions used in (Toth & Kelley, 2015; Evans et 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 p k and q k via
[Equation]
where “ † ” is the Moore-Penrose inverse.
Accordingly, we choose Γ k = arg min (cid:107) r k − Q k Γ (cid:107)
= ( Q T k Q k ) † Q T k r k .
This change relaxes the orthonormality for Q k , but keeps the orthogonal condition:
Q T k − 1 q k = 0 .
In fact, ¯ Q T k − 1 q k = 0 in the case of solving (5).
Change 2 : MST-AM imposes the boundedness constraints on P k − 1 ζ k and Q k − 1 ζ k : If (cid:107) P k − 1 ζ k (cid:107) 2 >
c p (cid:107) p (cid:107) 2 or (cid:107) Q k − 1 ζ k (cid:107) 2 > c q (cid:107) q (cid:107) 2 , then P k
=
P
k − 1 ,
Q k = Q k − 1 , where c p > 0 , c q ∈ (0 , 1) are predeﬁned constants.
It is worth mentioning that the boundedness condition is common in the analysis of AM (Toth & Kelley, 2015; Evans et al., 2020).
Change 3 : MST-AM restarts, i.e. setting P k = Q k = 0 ∈
R d × 2 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 C.2.
In the next theorem, we establish the convergence rate analysis for the MST-AM.
Theorem 2.
Let { x k } be the sequence generated by MST-AM, x
∗ ∈ R
d be a ﬁxed point of g and m be the restarting period for MST-AM.
Suppose that in the ball B ( ρ ) :
= { x ∈ R d |(cid:107) x − x ∗ (cid:107) 2 < ρ } for some ρ > 0 , g is Lipschitz continuously differentiable and there are constants κ ∈ (0 , 1) and ˆ κ >with (i) (cid:107) g
( y ) − g ( x ) (cid:107) 2 ≤ κ (cid:107) y − x (cid:107) 2 for every x, y ∈ B ( ρ ) , and (ii) (cid:107) g (cid:48)
( y ) − g (cid:48)
( x ) (cid:107) 2 ≤ ˆ κ (cid:107)
y − x
(cid:107)for every x, y ∈ B ( ρ ) , where g (cid:48)
( x ) is the Jacobian of g .
If (cid:107) x 0 − x ∗ (cid:107) 2 is sufﬁciently small, then for the k -th residual r k = g
( x k ) − x k , the following bound holds:
[Equation]
where θ k = (cid:107) ¯ r k (cid:107) 2 / (cid:107) r k (cid:107) 2 ≤ 1 and m k = k mod m .
Thus, the residuals { r k } converge Q -linearly, and the errors {(cid:107) x k − x ∗ (cid:107) 2 } converge R -linearly.
Remark 3.
In a local region around x ∗ , the convergence rate is determined by the ﬁrst-order term θ k ( | 1 − β k | + κβ k ) (cid:107) r k (cid:107) 2 .
Since ¯ r k is the orthogonal projection of r k onto the subspace range( Q k ) ⊥ , θ k has the interpretation of the direction-sine between r k and the subspace range( Q k ) .
When θ k is
Input : x 0 ∈ R d , β k > 0 , α k ∈ [0 , 1] , δ (1) k > 0 , δ (2) k > 0 .
Output : x ∈ R d 1: P 0 , Q 0 =
∈ R
d × 2 , p 0 , q 0 =
∈ R d 2: for k = 0 , 1 , . . .
, until convergence, do 3: r k = −∇ f S k
( x k ) 4:
if k > 0 then 5: p = x
k − x k − 1 , q = r k − r k −6: ζ k =
( Q T k − 1 Q k − 1 + δ (1) k
P T k − 1 P k − 1 ) † Q T
k − 1 q 7: q k = q − Q k − 1 ζ k , p k
= p − P k − 1 ζ k 8:
P k
=
[ p k − 1 , p k ] , Q k =
[ q k − 1 , q k ] 9: end if 10:
Check Condition (13) and use smaller α k if (13) is violated 11:
Γ k = ( Q T
k Q k
+ δ (2) k P T k P k ) †
Q T k r k 12: ¯ x k = x
k − α k P k Γ k , ¯ r k = r k − α k Q k Γ k 13: x k +
=
¯ x k + β k ¯ r k 14:
Apply learning rate schedule of α k , β k 15: end for 16: return x k small, e.g., r k nearly lies in range( Q k ) , the acceleration by MST-AM is signiﬁcant.
Unlike AM( m ) that simply discards the oldest iteration to make space for ∆ x k − 1 and ∆ r k − 1 , MST-AM incorporates historical information with orthogonalization.
In the SPD linear case and without restart, the global orthogonalization property holds, i.e. ¯ r k ⊥ range ( ¯ Q k ) , which means there is no loss of historical information, while AM( m ) (Evans et al., 2020) does not have such property in this ideal case.
3. T HE REGULARIZED SHORT - TERM RECURRENCE A NDERSON MIXING
Inspired by the recent work on stochastic Anderson mixing (SAM) method (Wei et al., 2021), we develop a regularized ST-AM (RST-AM) for solving nonconvex stochastic optimization problems.
Consider the nonconvex optimization problem min x ∈ R d f
[Equation]
x
[Equation]
:=
[Equation]
x ) , where f ξ i : R d → R is the loss function corresponding to i -th data sample and T is the number of data samples.
In mini-batch training, the gradient is evaluated for f S k ( x k ) :
= 1 n k (cid:80) i ∈ S k f ξ i ( x k ) , where S k ⊆
[ T ] := { 1 , 2 , . .
, T } is the sampled mini-batch, and n k :=
| S k | is the batch size.
In this case, we set r k = −∇ f S k ( x k )
(Line 3 in Algorithm 1), which is an unbiased estimate of the negative gradient.
Recalling from (8), ζ k =
( Q T k − 1 Q k − 1 ) † Q T k − 1 ∆ r k − 1 = arg min (cid:107) ∆
r k − 1 − Q k − 1 ζ (cid:107) 2 as q = ∆ r k −by deﬁnition.
Since ST-AM is based on a local quadratic approximation (5) in a small region around x k , a large magnitude of (cid:107) P k − 1 ζ k (cid:107) 2 tends to make p k = ∆ x k − 1 − P k − 1 ζ k overshoot the trust region around x k , thus leading to instability.
Consequently, we add a penalty term in the above least squares problem, i.e.
[Equation]
(10) where δ (1) k > 0 .
The same as SAM (Wei et al., 2021), we also add a regularization term for computing Γ k via
[Equation]
where δ (2) k > 0 , and a damping term is used as shown in Line 12 in Algorithm 1.
In practice, we choose the two parameters as
[Equation]
where c 1 , c 2 , C > 0 are constants, and (cid:15) 0 > 0 is a small constant to bound the denominators away from zero.
Assuming (cid:107) p k − 1 (cid:107)
≈ (cid:107) p k (cid:107)
= O ( (cid:107) ∆ x k − 1 (cid:107) 2 ) , the choices of (12) make (cid:107) δ (1)
k P T k − 1 P k − 1 (cid:107)
≈ O ( (cid:107) r k
(cid:107) 22 ) and (cid:107) δ (2) k P T k P k (cid:107)
≈ O ( (cid:107) r k
(cid:107) 22 ) aware of the change of the local curvature: large (small) (cid:107) r k (cid:107) 2 tends to lead to a large (small) regularization.
Remark 4.
One update of x k given by Line 11-13 in Algorithm 1 can be formulated as x k +
= x k + H k r k , where H k = β
k
I − α k Y
k Z † k Q T k ,
Y k
= P k + β k Q k , Z k = Q T
k Q k
+ δ (2) k P T k P k .
To guarantee the positive deﬁniteness of H k , we follow the same procedure in SAM.
Let λ k be the largest eigenvalue of Y k Z † k Q T k
+ Q k Z † k Y T k .
If α k satisﬁes
[Equation]
13) then s T
k H k s k ≥ β k µ (cid:107) s k (cid:107) 22 , ∀ s k ∈ R d , where µ ∈ (0 , 1) is a constant.
Note that λ k can be cheaply obtained by computing the largest eigenvalue of a matrix of R 4 × 4 (see Appendix C.3.1).
We summarize the RST-AM in Algorithm 1 and establish its convergence properties in the following context.
First, we impose the following assumptions on the objective function f .
Assumption 1.
f : R d → R is continuously differentiable .
f ( x ) ≥ f low >
−∞ for any x ∈ R d .
∇ f is globally L -Lipschitz continuous; namely (cid:107)∇ f ( x ) −∇ f ( y ) (cid:107) 2 ≤ L (cid:107) x − y (cid:107) 2 for any x, y ∈ R d .
Assumption 2.
For any iteration k , the stochastic gradient ∇ f ξ k ( x k ) satisﬁes E ξ k [ ∇ f ξ k ( x k )]
= ∇ f ( x k ) , E ξ k [ (cid:107)∇ f ξ k ( x k ) − ∇ f ( x k ) (cid:107) 22 ] ≤ σ 2 , where σ > 0 , and ξ k , k = 0 , 1 , . .
. are independent samples that are independent of { x j } kj =0 .
The diminishing condition about β
[Equation]
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 1 hold and { x k } is the sequence generated by full-batch RST-AM, i.e. n k = T .
Let β k = β ∈
(0 , µ 2 L (1+ C − 1 ) ] be a constant, α k ∈ [0 , 1] and satisﬁes (13), then
[Equation]
in the N iterations.
To ensure 1 N (cid:80) N − 1 k =0 (cid:107)∇ f ( x k ) (cid:107) 22 < (cid:15) , the number of iterations is O (1 /(cid:15) ) .
Theorem 4.
Suppose Assumptions 1 and 2 hold and { x k } is the sequence generated by RST-AM with batch size n k = n ≤ T .
If β k ∈ (0 , µ 4 L (1+ C − 1 ) ] and satisﬁes (14), α k ∈ [0 , min { 1 , β 12 k } ] and satisﬁes (13), then
[Equation]
probability
[Equation]
If E ξ
[Equation] [Equation]
22 ] ≤ M g , ∀ k, where M g > 0 is a constant, we have
[Equation]
probability
[Equation]
Theorem 5.
Suppose Assumptions 1 and 2 hold and { x k } N − 1 k =0 is the ﬁrst N iterations generated by
RST-AM with ﬁxed batch size n k = n .
Let β k = min { µ 4 L (1+ C − 1 ) ,
˜ Dσ √ N } , where ˜ D is a problemindependent constant; α k ∈ [0 , min { 1 , β 12 k } ] and satisﬁes (13).
Let R be a random variable following P R ( k ) :
= Prob { R = k } = 1 /N , then
[Equation]
where D f := f
( x 0 ) − f low and the expectation is taken with respect to R and { S j } N − 1 j =0 .
To ensure
E [ (cid:107)∇ f ( x R ) (cid:107) 22 ] ≤ (cid:15) , the number of iterations is O (1 /(cid:15) 2 ) .
Remark 5.
The proofs of Theorem 4 and 5 are based on the analysis of SAM (Wei et al., 2021).
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; RSTAM uses additional correction and regularization terms (Line 5-8 in Algorithm 1) to incorporate historical information while SAM simply discards the oldest iteration to make space for ∆ x k −and ∆ r k − 1 .
The reduced memory requirement in RST-AM makes it applicable for solving more challenging problems in machine learning.
E XPERIMENTS
We validated the effectiveness of our proposed ST-AM methods in various applications in ﬁxed-point iterations and nonconvex optimization, including linear and nonlinear problems, deterministic and stochastic optimization.
Speciﬁcally, we ﬁrst 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 ﬁrst-order and second-order optimizers.
Experimental details are in Appendix D.
4. E XPERIMENTS ABOUT THE BASIC ST-AM AND MST-AM
We veriﬁed the properties of ST-AM declared in Theorem 1 and 2 by solving four problems (details are in Appendix D.1): (I) strongly convex quadratic optimization (corresponding to Theorem 1); (II) solving a nonsymmetric linear system Ax = b (corresponding to ˆ κ = 0 in Theorem 2); (III) cubic-regularized quadratic minimization
min x ∈ R d f ( x ) := (cid:107)
Ax − b (cid:107) 22 + M 3 (cid:107) x (cid:107) 32 (corresponding to ˆ κ > 0 in Theorem 2); (IV) root-ﬁnding problems in MDEQ on CIFAR-10 (Krizhevsky et al., 2009).
The compared methods were gradient descent (GD), ﬁxed-point iteration (FP), conjugate residual method (CR) (Saad, 2003), BFGS (Nocedal & Wright, 2006), Broyden’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 ( c p =
c q = 1 ) to solve Problem III and IV.
[Figure]
: || f
[Figure]
) ||/|| z
||
[Figure]
The numerical results shown in Figure 1 demonstrate the power of ST-AM as a variant of Krylov subspace methods.
It signiﬁcantly accelerates the slow convergence of the GD or FP method, and can outperform AM.
Figure 1(a) clearly veriﬁes the correctness of Theorem 1: within the machine precision, the intermediate residual ¯ r k of ST-AM coincides with the residual r k 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 & Ni, 2011).
Figure 1(b) shows that ST-AM can outperform CR, though both methods enjoy short-term recurrences and are equivalent for solving SPD linear systems.
Figure 1(c) also shows MST-AM surpasses BFGS in solving cubic-regularized problems.
The tests in
MDEQ/CIFAR-10 indicate that MST-AM is comparable to the full-memory methods in the forward root-ﬁnding process and converges faster in the backward process.
The accuracy is also comparable.
[Figure]
4. E XPERIMENTS 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-10/CIFAR-100 and Penn Treebank (Marcus et al., 1993).
Experiments on MNIST .
We trained a convolutional neural network (CNN) on MNIST to see the convergence behaviour of RST-AM in nonconvex optimization (cf.
Theorem 3), for which we were only concerned about the training loss.
Figure 2(a)(b) show that the short-memory SAMs ( m = 2 , 5 ) hardly show any improvement over the ﬁrst-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 2(c) indicates that preconditioning also works much better for
RST-AM than SAM(2), and RMSprop RST-AM can outperform the non-preconditioned RST-AM.
[Table]
[Table]
Experiments on CIFAR .
We trained ResNet18/20/32/44/56 (He et al., 2016), WideResNet16- (Zagoruyko & Komodakis, 2016)
(abbr.
WRN16-4) on CIFAR-10, and ResNet18, ResNeXt50 (Xie et al., 2017), 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.
Table 1(a) reports the ﬁnal TOP1 test accuracy for training 160 epochs (AdaHessian ran out of memory for training ResNeXt50, DenseNet121).
It shows RST-AM improves SAM(2) and has comparable test accuracies to SAM(10).
RST-AM also outperforms other baseline optimizers.
Table 1(b) reports the memory and computation cost, where we used SGDM as the baseline and other optimizers were terminated when achieving a comparable or better test accuracy than SGDM.
It indicates that RST-AM introduces ≤ 5% extra memory overhead compared with SGDM, and signiﬁcantly reduces the memory footprint of AM.
Since RST-AM needs fewer training epochs, the total running time is less than SGDM
[Figure] [Figure]
[Table]
Experiments on Penn Treebank .
We trained
LSTMs with 1-3 layer(s) on Penn Treebank and report the validation perplexity in Figure 3 and test perplexity in Table 2 (lower is better).
The results suggest that RST-AM is comparable to or even better than SAM(10).
The improvement of RST-AM over other optimizers is also signiﬁcant.
We report the computation and memory cost in Appendix D.2.4.
RST-AM can still surpass Adam while using much fewer epochs, thus reducing the total running time.
[Table] [Table]
Adversarial training .
We applied RST-AM to adversarial training (Madry et al., 2018) as the outeroptimizer and compared it with SGD by the clean test accuracy and robust test accuracy.
The results on CIFAR10/ResNet18 and CIFAR100/DenseNet121 are reported in Table 3.
It can be seen that
RST-AM can achieve both higher clean test accuracy and robust test accuracy.
More results can be found in Appendix D.2.5.
Generative adversarial network (GAN) .
We tested RST-AM by training a GAN equipped with spectral normalization (SN-GAN) (Miyato et al., 2018), where the generator and discriminator networks were ResNets and the dataset was CIFAR-10.
Table 4 shows that RST-AM can achieve lower
FID score (better accuracy) than Adam and AdaBelief.
C ONCLUSION
In this paper, to 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 solving 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 ﬁxed-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 are comparable to or even better than the long-memory AM while consuming less memory.
The regularized ST-AM also outperforms many existing optimizers in training neural networks in various tasks.
A CKNOWLEDGMENTS
This work was supported by the National Key R&D Program of China (No. 2021YFA1001300),
National Natural Science Foundation of China (No.61925601), Tsinghua University Initiative Scientiﬁc Research Program, National Natural Science Foundation of China (No.11901338) and Huawei
Noah’s Ark Lab.
We thank all anonymous reviewers for their valuable comments and suggestions on this work.
A A DDITIONAL PRELIMINARIES
We provide some additional preliminaries in this section for readers that are not familiar with Anderson mixing, ﬁxed-point iterations and some techniques mentioned in the main paper.
A. F IXED - POINT ITERATION
The ﬁxed-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 problems that make algorithm designs different.
The key difference is that for a ﬁxed-point problem, the Jacobian (if exists) is generally not symmetric while for optimization, the Hessian is naturally symmetric.
In principle, a ﬁxed-point solver can also be applicable for an optimization problem since the ﬁrst-order necessary condition of min x ∈ R d f ( x ) , where f : R d → R , is ∇ f ( x ) = 0 .
Consider a contraction mapping g : R d (cid:55)→ R d , i.e. for some κ < 1 , (cid:107) g
( x ) − g ( y ) (cid:107) 2 ≤ κ (cid:107) x − y (cid:107) 2 , ∀ x, y ∈ R d .
According to the contraction mapping theorem, a unique ﬁxed point x ∗ exists for g and the iterates generated by the iteration x k
+
=
g ( x k ) converge to x ∗ , starting from any x 0 ∈ R d .
In practice, the damped ﬁxed-point iteration is also commonly used: given the β k ∈ (0 , 2) , the update is x k +
=
(1 − β k ) x k + β k g ( x k ) =
x k + β k r k , where r k :=
g ( x k ) − x k is called the residual .
The ﬁxed-point iteration, also known as Picard iteration in some areas, can converge very slowly in practice.
Anderson mixing is a method to improve the convergence.
A. A NOTHER FORM OF A NDERSON MIXING
The derivation of Anderson mixing (AM) in Section 3.1 explicitly interprets AM as a two-step procedure.
In the literature, there is another equivalent form of AM.
Let the projection coefﬁcients Γ k
= ( γ (1) k , . .
, γ ( m ) k ) T ∈ R m .
Deﬁne the auxiliary coefﬁcients { θ ( j ) k } mj =0 as θ (0) k
= γ (1) k , θ ( j ) k = ∆ γ ( j ) k
( j = 1 , . .
, m − 1) , θ ( m ) k
= 1 − γ ( m ) k , then (cid:80) mj
=
θ ( j ) k
= 1 and ¯ r k = (cid:80) mj
=
θ ( j ) k
r k − m + j .
Hence the least squares problem (3) can be reformulated as a constrained problem min { θ ( j ) k } mj =0 (cid:107) (cid:80) mj
=
θ ( j ) k
r k − m + j (cid:107) 2 s.t.
(cid:80) mj
=
θ ( j ) k
= 1 , which indicates that as a linear combination of the historical residuals { r j } mj = k − m , ¯ r k is minimal in terms of the L 2 -norm.
Also, the projection step and the mixing step (2) can be reformulated as ¯ x k = (cid:80) mj
=
θ ( j ) k x k − m + j and x k +
=
(1 − β k ) (cid:80) mj
=
θ ( j ) k x k − m + j + β k (cid:80) mj
=
θ ( j ) k g
( x k − m + j ) , respectively.
Such formulation is also adopted in the literature (Toth & Kelley, 2015; Mai & Johansson, 2020; Scieur et al., 2020).
Let H k be the solution to the constrained optimization problem (Fang & Saad, 2009)
min
[Equation]
I
[Equation]
F subject to
[Equation]
= − X k , (19) then iterate (4) is x k +
= x k + H k r k .
Hence AM is a multisecant quasi-Newton method (Fang &
Saad, 2009).
A. P RECONDITIONED A NDERSON MIXING
Like preconditioning for Krylov subspace methods (Saad, 2003), preconditioning can also be incorporated 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 a basic solver preconditioner ( x k , s k ) , which works as a black-box procedure that updates x k given the residual s k , i.e. x k +
=
preconditioner ( x k , s k ) , then the preconditioned mixing of RST-AM is
[Equation]
which substitutes for the Line 13 in Algorithm 1.
The simple mixing (Line 13 in Algorithmcan be seen as a special case by deﬁning preconditioner ( x k , s k ) =
x k + β k s k , i.e. preconditioned by a damped ﬁxed-point iteration.
Moreover, if we write the preconditioning operation as preconditioner ( x k , s k ) :
= x k + G k s k , where G k is the matrix to approximate the inverse Jacobian, then the preconditioned AM is
[Equation]
20) (cf.
the deﬁnitions of X k , R k in Section 3.1).
The matrix H k
= G k − ( X k + G k R k )
( R T k R k ) † R T k forms a low-rank updated approximation to the inverse Jacobian: it solves
[Equation]
subject to H k R k = − X k , which is a direct extension of (19).
For the stochastic Anderson mixing, using damped projection can be helpful, i.e. ¯ x k =
(1 − α k ) x k + α k
( x k − X k Γ k ) =
x k − α k X k Γ k and ¯ r k = r k − α k R k Γ k correspondingly.
Then the preconditioned
AM with damped projection is
[Equation]
B A DDITIONAL DISCUSSION B. T HE 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 to ﬁrst-order methods.
Memory cost.
Assume that the iteration number is k and the model parameter size is d .
The fullmemory AM stores all previous iterations, thus the additional memory is 2 kd .
To reduce the memory overhead, the limited-memory AM( m ) only maintains the most recent m iterations while discarding the older historical information (cf.
Hence the additional memory of AM( m ) is 2 md .
Choosing a suitable m can be problem-dependent (Walker & Ni, 2011), and it is often suggested that m ≥(Mai & Johansson, 2020; Fu et al., 2020) to avoid too much historical information being discarded.
There is no equivalence between AM( m ) and the full-memory AM.
On the other side, our ST-AM methods have the same memory footprint as AM(2), i.e. only introducing 4 d additional memory.
For a large model that contains millions or even billions of parameters, such reduction in memory requirement can be signiﬁcant.
Computational cost.
For ST-AM, besides the gradient evaluations, the main additional computational cost comes from vector corrections, the projection and mixing step.
These operations are very cheap: matrix multiplications of R 2 ×
d × R d × 2 , R 2 ×
d × R d × 1 , and R d × 2 × R 2 × 1 ; the pseudoinverse can be exactly solved as the size of the matrix is R 2 × 2 .
Also, the Q T k − 1 Q k − 1 and P T k − 1 P k −in Line 6 in Algorithm 1 can reuse the results in the previous iteration (cf.
Line 11).
So the total additional computational cost of RST-AM is O ( d ) .
B. A PPLICABILITY
In the main paper, we develop a class of short-term recurrence Anderson mixing.
Among them, the basic ST-AM can serve as a new linear solver for linear systems; MST-AM can be used as a new ﬁxed-point solver for nonlinear equations; RST-AM is a new method for optimization.
These methods are based on 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-memory AM and GMRES in strongly convex quadratic optimization (or SPD linear systems), 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 of 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) can alleviate overhead, they are often more complicated to implement and difﬁcult for parallel computing.
Moreover, iterative solvers can beneﬁt from the preconditioning technique that improves convergence (Saad, 2003).
Also, ST-AM has advantages over other iterative methods since it does not need to directly access the matrix and only the residual is required.
For applications that explicitly accessing the matrix is difﬁcult, this property is appealing.
Hence, unlike nonlinear CG that relies on 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 to obtain.
Moreover, ST-AM is very ﬂexible and any iterative solver can be viewed as a black-box iterative process to be accelerated by ST-AM.
MST-AM versus other quasi-Newton methods for nonlinear equations.
MST-AM has the nature of quasi-Newton methods since it is built upon AM that is recognized as a multisecant quasiNewton method (Fang & Saad, 2009).
In scientiﬁc computing, AM has been successfully applied to solve many difﬁcult nonlinear problems arising from many areas (An et 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 efﬁcient 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 overhead, because they need to store historical iterations to form the secant equations.
To make a compromise, limited-memory quasi-Newton methods such as L-BFGS (Liu & Nocedal, 1989) are proposed in which only limited number of historical iterations are stored.
In each update, the oldest iteration is discarded to make space for the up-to-date secant pair.
As a result, the limited-memory quasi-Newton methods can lose the local superlinear convergence properties achieved by the fullmemory schemes (Berahas et al., 2021).
Also, as stated by Berahas et al.
(2021), one does not know a priori what number of historical iterations (i.e. historical length) should be maintained to attain a good performance when solving a particular problem.
Unlike the limited-memory quasi-Newton methods whose performance can be sensitive to the historical length, our MST-AM method only needs to store two corrected historical iterations.
MSTAM 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 ﬁrst-order methods for stochastic optimization.
For stochastic optimization such as training deep neural networks in machine learning, the ﬁrst-order methods have come to dominate the ﬁeld 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/stochastic 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 a smooth function can be approximated by a quadratic function in the local region around the optima and many techniques such as trust region (Nocedal & Wright, 2006) rely on this local approximation.
Adaptive learning rate methods such as AdaGrad (Duchi et al., 2011),
Adam (Kingma & Ba, 2014) use a diagonal approximation of the Fisher information matrix that is an approximation of the Hessian.
However, in the simple quadratic case, these methods can only roughly match the performance of the Jacobi method for solving linear systems (Saad, 2003), which is 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 is no need to determine the stepsize and the fast convergence rate of ST-AM can be recovered by simply setting α k = 1 and δ (1) k
= δ (2) k
= 0 .
For more difﬁcult functions, the damping and regularization in RST-AM can be enabled to improve stability.
• Theoretical guarantee in stochastic optimization.
If only ﬁrst-order information can be accessed, the SGD (Ghadimi & Lan, 2013) achieves an optimal convergence rate O (1 /(cid:15) 2 ) to obtain an (cid:15) -accurate solution (Nemirovskij & 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 ﬁrst-order methods.
Since RST-AM has very limited addi-
tional memory overhead, and also achieves the O (1 /(cid:15) 2 ) complexity, it can be applied to many applications that dominated by ﬁrst-order methods.
• Flexibility in use.
The application of RST-AM can be very ﬂexible.
In principle, RST-AM can be applied to improve any slowly convergent black-box iterative process by viewing the latter as a ﬁxed-point iteration.
For example, consider accelerating a solver of a commercial software, where we have no access to the underlying codes to provide our custom implementation, or some case rewriting the codes is too cumbersome.
Moreover, RSTAM can efﬁciently incorporates 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 RSTAM for various applications.
For example, RST-AM can be preconditioned by Adam for the language task and SGDM for image classiﬁcation.
So any ﬁne-tuned ﬁrst-order optimizer can be combined with RST-AM to achieve an overall improvement.
As RST-AM is light-weight, this additional cost is marginal and can be largely counteracted 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 ﬁrst-order method can mitigate the ill-conditioning of the problem.
Overall, the ST-AM methods have wide applicability and can be competent methods from both theoretical and practical perspectives
C P ROOFS
We give more details about ST-AM and the proofs of the theorems in the main paper.
C. T HE BASIC ST-AM FOR STRONGLY CONVEX QUADRATIC OPTIMIZATION
Recall that the strongly convex quadratic optimization is formulated as
[Equation]
(22) where A ∈ R d × d is SPD, b ∈ R d .
Solving (22) is equivalent to solving the SPD linear system
[Equation]
The detail of the basic ST-AM is given in Algorithm 2.
We ﬁrst state the relationship of AM with GMRES in the following proposition.
Similar results can also be found in (Walker & Ni, 2011).
Let x G k , r G k :=
b − Ax G k denote the k -th GMRES iterate and residual, respectively, and K k ( A, v ) :
= span { v, Av, . . .
, A k − 1 v } denotes the k -th Krylov subspace generated by A and v .
Deﬁne e j :=
, 1) T ∈ R j for j ≥ 1 .
Let range( X ) denote the linear space spanned by the columns of X .
The main results of the full-memory AM are stated in Proposition 1 and Proposition 2.
Proposition 1 (General linear system) .
For solving a general linear system Ax = b with the fullmemory AM ( m = k ), suppose that β k > 0 and the ﬁxed-point map is g ( x ) =
( I − A ) x + b .
If the initial point of AM is x 0 = x G 0 and rank( R k ) = m , then the intermediate iterate ¯ x k satisﬁes ¯ x k = x
G k .
Proof.
The deﬁnition of the ﬁxed-point map suggests that the residual r k = g ( x k ) − x k =
b − Ax k .
Since R k = − AX k and A is nonsingular, we have rank( X k ) =
m .
We ﬁrst show
[Equation]
24) by induction.
We abbreviate K k ( A, r G 0 ) as K k in this proof.
Input : x 0 ∈ R d , β k > 0 , 0 < max iter ≤ d .
Output : x ∈ R d 1: P 0 , Q 0 =
∈ R
d × 2 , p 0 , q 0 =
∈ R d 2: for k = 0 , 1 , . .
, max iter do 3: r k = −∇ f ( x k ) 4:
if k > 0 then 5: p = x k − x k − 1 , q = r k − r k −(Compute ∆ x k − 1 , ∆ r k − 1 ) 6: ˜ q = q − Q
k − 1 ( Q T k − 1 q ) , ˜ p
= p − P k − 1 ( Q T k − 1 q ) ( ˜ q ⊥ Q k − 1 ) 7: p k = ˜
p/ (cid:107) ˜ q (cid:107) 2 , q k = ˜
q/ (cid:107) ˜
q
(cid:107)( (cid:107) q k (cid:107) 2 = 1 ) 8:
P k
=
[ p k − 1 , p k ] , Q k =
[ q k − 1 , q k ] ( Q T k Q k
= I 2 ) 9: end if 10: Γ k = Q T
k r k 11: ¯ x k = x
k − P k Γ k , ¯ r k = r k − Q k Γ k
(Projection step: ¯ r k ⊥ Q k ) 12: x k
+1 =
¯ x k + β k ¯ r k
(Mixing step) 13: if (cid:107) ¯ r k (cid:107) 2 = 0 then 14: break 15: end if 16: end for 17: return x k
First, ∆ x 0 = β 0 r 0 = β 0 r G 0 since x 1 = x 0 + β 0 r 0 .
If k = 1 , then the proof is complete.
Then, suppose that k > 1 and, as an inductive hypothesis, that range( X k − 1 ) =
K k − 1 .
With (4) we have ∆ x
[Equation]
Since r 0 ∈ K k − 1 , and by the inductive hypothesis range( X k − 1 ) ⊆
K k − 1 which also implies range( AX k − 1 ) ⊆
K k , we know ∆ x k − 1 ∈ K k , which implies range( X k ) ⊆
K k .
Since we assume rank( X k )
= m
= k which implies dim(range( X k ))
=
dim( K k ) , we have range( X k ) =
K k , thus completing the induction.
As a result, we also have
[Equation]
Recalling that to determine Γ k , we solve the least squares problem (3) and R k = − AX k .
We have
[Equation]
Since rank( AX k ) = rank( X k )
= m , (27) has a unique solution.
Also, since r k = b − Ax k = b − A ( x 0 + X k e k ) =
r 0 − AX k e k , we have r k + AX k Γ = r 0 − AX k e k
+ AX k Γ = r 0 − AX k ˜Γ ,
where ˜Γ =
e k − Γ .
So Γ k solves (27) if and only if ˜Γ k = e k − Γ k solves
[Equation]
According to (24), (28) is equal to min z ∈K m ( A,r G 0 ) (cid:107) r 0 − Az (cid:107) 2 which is the GMRES minimization problem.
Since the solution of (28) is also unique, we have
[Equation]
In Proposition 1, the assumption that R k has full column rank is critical to ensure no stagnation occurs in AM for solving a general linear system.
In fact, for SPD linear systems (23) or strongly convex quadratic optimization (22), when AM breaks down, i.e. R k is rank deﬁcient, AM obtains the exact solution, as shown in the next proposition.
Proposition 2 (SPD) .
For applying the full-memory AM to minimize a strongly convex quadratic problem (22), or equivalently, solve a SPD linear system (23), suppose that β k > 0 and the ﬁxedpoint map is g ( x )
=
( I − A ) x + b .
If rank( R k ) =
k holds for 1 ≤ k < s while failing to hold for k = s , where s ≥ 1 , then the residual of AM satisﬁes r s =
¯ r s − 1 = 0 .
Proof.
The deﬁnition of g suggests that the residual r k = g
( x k )
− x k
= b − Ax k .
The relation R k = − AX k holds during the iterations and the nonsingularity of A implies rank( X k )
=
rank( R k ) .
For s = 1 , since the ﬁrst step of AM is x 1 = x 0 + β 0 r 0 , the assumption rank( R 1 ) =
implies that rank( r 0 )
=
rank( X 1 ) = 0 , i.e. r 1 = ¯ r 0 := 0 .
For s > 1 , because ∆ x s − 1 = x s − x s − 1 = − X s − 1 Γ s − 1 + β s − 1 ¯ r s − 1 , the rank deﬁciency of X s implies ∆ x s − 1 ∈ range( X s − 1 ) , which further implies ¯ r s − 1 ∈ range( X s − 1 ) .
So there exists ζ ∈ R s − 1 , such that ¯ r s − 1 =
X s − 1 ζ .
Note that according to (3), ¯ r s − 1 ⊥ R s − 1 = − AX s − 1 , so we have
[Equation]
Because rank( X s − 1 ) = s − 1 and A is SPD, we know X T s − 1 AX s − 1 is also SPD.
So ζ = 0 , which implies ¯ r s − 1 = 0 .
Hence x s = ¯
x s − 1 and r s = ¯ r s − 1 = 0 .
Now we give the proof of Theorem 1.
Proof of Theorem 1 .
Besides relations (i)-(iii), we add an auxiliary relation here:
(iv) r k
=
r
+ ¯ Q k ¯Γ k ∈ K k +
( A, r 0 ) ,
where
¯Γ k ∈ R k
We prove the relations (i)-(iv) by induction.
For k = 1 , since ¯ r 0 (cid:54) = 0 , according to Proposition 2, rank(∆ x 0 ) =
rank( X 1 ) = 1 , rank(∆ r 0 )
=
rank( R 1 )
= 1 ,
so ˜ q (cid:54) = 0 , which implies Line 7 in Algorithm 2 is well-deﬁned.
The relation (i) holds.
Since ˜ q = q = ∆ x 0 , ˜ p = p = ∆ r 0 , and ∆ r 0 = − A ∆ x 0 , the equality ¯ Q 1 = − A ¯ P 1 also holds.
Due to the normalization in Line 7, ¯ Q T1 ¯ Q 1 = 1 .
Since r 1 = r 0 − β
Ar 0 and range( ¯ Q 1 )
=
range( Ar 0 ) , it is clear that r 1 = r 0 − ¯ Q 1 ¯Γ 1 ∈ K 2 ( A, r 0 ) , namely relation (iv).
Due to the projection step Line 11, ¯ r 1 ⊥ range( Q 1 ) =
range( ¯ Q 1 ) .
Also, ¯ r 1 = r 1 − Q 1 Γ 1 = r 0 − β
Ar 0 − Q 1 Γ 1 = r 0 − ¯ Q 1 η 1 , where the last equality is due to span { Ar 0 } =
range( Q 1 ) =
range( ¯ Q 1 ) .
For r G 1 = r 0 − Az 1 , where z 1 = arg min z ∈K 1 ( A,r 0 ) (cid:107) r 0 − Az (cid:107) 2 , it holds r G 1 ⊥
A K 1 ( A, r 0 ) =
range( ¯ Q 1 ) .
As a result, both ¯ rand r G 1 are the orthogonal projections of r 0 onto the subspace range( ¯ Q 1 ) ⊥ , which implies ¯ r 1 = r G 1 .
So ¯ x 1 = x G 1 = x 0 + z 1 because their residuals are equal and A is nonsingular.
Hence relation (iii) holds.
Suppose that k > 1 , and as an inductive hypothesis, the relations (i)-(iv) hold for j = 1 , . .
, k − 1 .
Consider the k -th iteration.
From Line 6 in Algorithm 2, ˜ q ∈ range(∆ r k − 1 , Q k − 1 ) , and ˜ p
∈
range(∆ x k − 1 , P k − 1 ) .
We ﬁrst prove that ˜ q (cid:54) = 0 by contradiction.
If ˜ q = 0 , then from Line 6 in Algorithm 2, ∆ r k − 1 ∈ range( Q k − 1 ) ⊆
range( ¯ Q k − 1 ) , which implies ∆ x k − 1 ∈ range( P k − 1 ) ⊆
range( ¯ P k − 1 ) as ∆ r k − 1 = −
A ∆ x k − 1 and ¯ Q k − 1 = − A ¯ P k − 1 and A is nonsingular.
From Line 11 and Line 12, we have
[Equation]
So ¯ r k − 1 ∈ range( P k − 1 ) ⊆
range( ¯ P k − 1 ) since ∆ x k − 1 ∈ range( P k − 1 ) .
Hence there exists ζ ∈ R k − 1 , such that ¯ r k − 1 =
¯ P k − 1 ζ .
From the inductive hypothesis, we know ¯ r k − 1 ⊥ ¯ Q k − 1 = −
A ¯ P k − 1 , so we have
[Equation]
Since ¯ Q T k − 1 ¯ Q k − 1 =
I k − 1 , we know rank( ¯ Q k − 1 ) =
k − 1 , which implies rank( ¯ P k − 1 ) =
k − 1 due to ¯ Q k − 1 = −
A ¯ P k − 1 .
Hence ¯ P T k − 1 A ¯ P k − 1 is also SPD.
Then ζ = 0 which implies ¯ r k − 1 = 0 .
It is impossible otherwise Algorithm 2 has terminated in the ( k − 1) -th iteration.
So ˜ q (cid:54) = 0 and Line 7 is well-deﬁned.
Since ¯ r k − 1 = r k − 1 − Q k − 1 Γ k − 1 , and r k − 1 ∈ K k
( A, r 0 ) , range( Q k − 1 ) ⊆
range( ¯ Q k − 1 ) =
A K k − 1 ( A, r 0 ) as the inductive hypothesis, we have ¯ r k − 1 ∈ K k
( A, r 0 ) , which together with (30) and range( P k − 1 ) ⊆
range( ¯ P k − 1 )
=
K k − 1 ( A, r 0 ) infers ∆ x k − 1 ∈
K k ( A, r 0 ) .
Hence ∆ r k − 1 = − A ∆ x k − 1 ∈
A K k ( A, r 0 ) .
As a result, q k = ˜
q/ (cid:107) ˜ q (cid:107) 2 ∈ range(∆ r k − 1 , Q k − 1 ) ⊆
A K k
( A, r 0 ) .
So range( ¯ Q k )
=
range( ¯ Q k − 1 , q k ) ⊆
A K k
( A, r 0 ) .
Moreover, q k / ∈ range( ¯ Q k − 1 ) , otherwise ˜ q ∈ range( ¯ Q k − 1 ) that implies ∆ r k − 1 = q ∈ range( ¯ Q k − 1 ) , which is impossible following the former proof of ˜ q (cid:54) = 0 .
So we have range( ¯ Q k ) =
A K k ( A, r 0 ) .
Because ∆ r k − 1 = − A ∆ x k − 1 and Q k − 1 = − AP k − 1 due to ¯ Q k − 1 = − A ¯ P k − 1 , Line 6 in Algorithm 2 infers ˜ q = − A ˜ p , which implies q k = − Ap k .
So ¯ Q k = − A ¯ P k .
Since A is nonsingular and range( ¯ Q k ) =
A K k ( A, r 0 ) , we have range( ¯ P k ) =
K k ( A, r 0 ) .
As range( X k ) =
K k ( A, r 0 ) and range( R k ) =
A K k ( A, r 0 ) has been proved in Proposition 1, the relation (i) holds for the k -th iteration.
To prove ¯ Q T k ¯ Q k
= I
k , it sufﬁces to show q k ⊥ ¯ Q k − 1 , as the equalities ¯ Q T k − 1 ¯ Q k − 1 =
I k − 1 and (cid:107) q k (cid:107) 2 = 1 has already held.
It is equivalent to prove ˜ q ⊥ ¯ Q k − 1 .
From the construction of ˜ q in
Line 6 in Algorithm 2, we know Q T k − 1 ˜ q = 0 , so ˜ q ⊥ span( q k − 2 , q k − 1 ) (for k = 2 , ˜ q ⊥ q 0 =clearly holds).
To further prove ˜ q ⊥ range( ¯ Q k − 3 )( k ≥ 4) , note that
[Equation]
where the second equality is a direct substitution with (30).
Therefore,
[Equation] [Equation]
where the second equality is due to range( Q k − 1 )
= span( q k − 2 , q k − 1 ) ⊥ range( ¯ Q k − 3 ) and A is
SPD, the third equality is due to ¯ r k − 1 ⊥ range( ¯ Q k − 1 ) =
A K k − 1 ( A, r 0 ) and range( A ¯ Q k − 3 ) =
A 2 K k − 3 ( A, r 0 ) ⊆
A K k − 1 ( A, r 0 ) .
As a result, noting that q = ∆ r k − 1 , we obtain
[Equation]
which is due to (31) and range( Q k − 1 )
=
span { q k − 2 , q k − 1 } ⊥ range( ¯ Q k − 3 ) .
Therefore, we show that ¯ Q T k ¯ Q =
I k , which along with ¯ Q k = −
A ¯ P k proves relation (ii) in the k -th iteration.
Next, we prove the relation (iv).
We have
[Equation]
where the last equality is due to r k −
= r 0 + ¯ Q k − 1 ¯Γ k − 1 by the inductive hypothesis.
Since range( ¯ Q k − 1 )
=
A K k − 1 ( A, r 0 ) ⊆
A K k
( A, r 0 ) , range( Q k − 1 ) ⊆
range( ¯ Q k ) , span { Ar 0 } ⊆
A K k ( A, r 0 ) , range( AQ k − 1 ) ⊆
range( A ¯ Q k − 1 ) ⊆
A 2 K k − 1 ( A, r 0 ) ⊆
A K k
( A, r 0 ) , and range( ¯ Q k ) =
A K k ( A, r 0 ) , it is clear that r k = r
+ ¯ Q k ¯Γ k ∈ K k +
( A, r 0 ) for some ¯Γ k ∈ R k .
The relation (iv) is proved.
Finally, we prove the relation (iii).
For proving ¯ r k ⊥ range( ¯ Q k ) , note that ¯ r k ⊥ span { q k − 1 , q k } = range( Q k ) already holds due to the projection step (Line 10 and Line 11 in Algorithm 2).
It sufﬁces to prove ¯ r k ⊥ range( ¯ Q k − 2 ) .
In fact, since we have ¯ r k = r k − Q k Γ k , we can prove that r k ⊥ range( ¯ Q k − 2 ) and Q k Γ k ⊥ range( ¯ Q k − 2 ) :
Since range( Q k ) = span { q k − 1 , q k } ⊥ range( ¯ Q k − 2 ) as induced from ¯ Q T k ¯ Q k
= I
k , it is clear that Q k Γ k ⊥ range( ¯ Q k − 2 ) .
For r k , according to Line 12 in Algorithm 2, we have r k = ¯
r k − 1 − β k − 1 A
¯ r k − 1 .
We have ¯ r k − 1 ⊥ range( ¯ Q k − 1 ) ⊇ range( ¯ Q k − 2 ) by the inductive hypothesis.
Also, ¯
Q T k −
A ¯ r k − 1 =
( A ¯ Q k − 2 )
T ¯ r k − 1 =
due to range( A ¯ Q k − 2 )
=
A 2 K k − 2 ( A, r 0 ) ⊆
A K k − 1 ( A, r 0 ) =
range( ¯ Q k − 1 ) .
Therefore, we obtain ¯ r k ⊥ range( ¯ Q k − 2 ) , which along with ¯ r k ⊥ span { q k − 1 , q k } implies ¯ r k ⊥ range( ¯ Q k ) .
To prove ¯ x k = x
G k := x 0 + z k , where z k = arg min
z ∈K k
( A,r 0 ) (cid:107) r 0 − Az (cid:107) 2 , ﬁrst
we have r k = r
+ ¯ Q k ¯Γ k , where ¯Γ k ∈ R k .
Hence ¯ r k = r k − Q k Γ k
= r 0 + ¯ Q k ¯Γ k − Q k Γ k
= r 0 − ¯ Q k η k , where η k ∈ R k .
Since ¯ r k ⊥ ¯ Q k , ¯ r k is the orthogonal projection of r 0 onto the subspace range( ¯ Q k ) ⊥ .
On the other side, for GMRES, r G k = r 0 − Az k ⊥
A K k ( A, r 0 )
=
range( ¯ Q k ) , so r G k is also the orthogonal projection of r 0 onto the subspace range( ¯ Q k ) ⊥ .
So ¯ r k = r G k , which further indicates ¯ x k = x
G k .
Hence, the relation (iii) holds.
With relations (i)-(iv) being proved in the k -th iteration, we complete the induction.
C. M ODIFIED ST-AM FOR GENERAL FIXED - POINT ITERATIONS
Algorithm 2 is suitable for analysis and implementation in the linear case.
For general nonlinear ﬁxed-point iterations, we adopt an alternative form as described in Algorithm 3 which discards the normalization of ˜ q in each iteration (Line 7 in Algorithm 2).
In Line 7 in Algorithm 3, the orthogonal projection of ∆ r k − 1 is checked to ensure ∆ r k − 1 is “less linearly dependent” on range( Q k ) , which ensures (cid:107) q k (cid:107) 2 is bounded away from zero; the check of (cid:107) P k − 1 ζ k (cid:107) 2 ensures that (cid:107) p k − ∆ x
k − 1 (cid:107) 2 ≤ c p (cid:107) ∆ x k − 1 (cid:107) 2 , which is also important since a large deviation from ∆ x k − 1 can make (cid:107) p k (cid:107) 2 > ρ ( ρ is the radius introduced in Theorem 2).
When this condition cannot be satisﬁed, the algorithm simply reuses the old P k − 1 , Q k − 1 .
The main procedure of MST-AM restarts every m iterations, i.e.
P k , Q k = 0 .
Such restart mechanism is to restrict the higher-order terms in the residual expansion, as shown in (9).
Also, restart can ﬂush out the outdated historical information that may weaken the quality of P k and Q k that are used to pursue a local ﬁrst-order approximation of g in MST-AM.
Input : x 0 ∈ R d , β k ∈ (0 , 1] , c p > 0 , c q ∈ (0 , 1) , m > 0 .
Output : x ∈ R d 1: P 0 , Q 0 =
∈ R
d × 2 , p 0 , q 0 =
∈ R d 2: for k = 0 , 1 , . . .
, until convergence do 3: r k = g ( x k ) − x k 4: if k mod m (cid:54) =
then 5: p = x
k − x k − 1 , q = r k − r k −6: ζ k =
( Q T k − 1 Q k − 1 ) † Q T k − 1 q 7: if (cid:107) P k − 1 ζ k (cid:107) 2 ≤ c p (cid:107) p (cid:107) 2 and (cid:107) Q k − 1 ζ k (cid:107) 2 ≤ c q (cid:107) q (cid:107) 2 then 8: p k
= p − P k − 1 ζ k , q k = q − Q k − 1 ζ k 9:
P k
=
[ p k − 1 , p k ] , Q k =
[ q k − 1 , q k ] ( q k ⊥ Q k − 1 ) 10: else 11:
P k
=
P
k − 1 , Q k = Q k −12: end if 13: else 14:
P k , Q k =
∈ R d × 2 , p k , q k = 0 ∈ R d 15: end if 16: Γ k
= ( Q T k Q k ) † Q T k r k 17: ¯ x k = x
k − P k Γ k , ¯ r k = r k − Q k Γ k
(Projection step: ¯ r k ⊥ Q k ) 18: x k
+1 =
¯ x k + β k ¯ r k
(Mixing step) 19: end for 20: return x k
(The notation “ † ” is the Moore-Penrose pseudoinverse.)
In the linear case, Algorithm 2 and Algorithm 3 (with m = ∞ ) are equivalent.
Similar to Algorithm 2, we have the following properties held for MST-AM:
Claim 1.
In the k -th iteration ( k > 0) of Algorithm 1 applied to minimize a strongly convex quadratic problem (5), assuming c p = ∞ , c q = 1 , m = ∞ , the following relations hold:
(i) (cid:107) q k (cid:107) 2 > 0 , range( ¯ P k )
=
range( X k ) =
K k ( A, r 0 ) , range( ¯ Q k )
=
range( R k )
=
A K k ( A, r 0 ) ; (ii) ¯ Q k = − A ¯ P k , q i ⊥ q j (1 ≤ i (cid:54) =
j ≤ k ) ; (iii) ¯ r k ⊥ range( ¯ Q k ) and ¯ x k = x 0 + z k , where z k = arg min
z ∈K k
( A,r 0 ) (cid:107) r 0 − Az (cid:107) 2 .
If (cid:107) ¯ r k (cid:107) 2 = 0 , then x k +1 is the exact solution.
The proof of Claim 1 is essentially the same as the proof of Theorem 1, with a special care that ¯
Q T k ¯ Q k
= I k is replaced by the relation that columns of ¯ Q k are orthogonal to each other, in other words, ¯ Q T k ¯ Q k = diag {(cid:107) q 1 (cid:107) 22 , . . .
, (cid:107) q k (cid:107) 22 } .
For minimizing nonlinear functions or accelerating nonlinear ﬁxed-point iterations, the long-term relation that ¯ Q T k ¯ Q k = diag {(cid:107) q 1 (cid:107) 22 , . . .
, (cid:107) q k (cid:107) 22 } generally cannot hold, while the orthogonalization procedure in Line 6 still leads to a short-term orthogonalization relation:
q i ⊥ q j for | i − j | ≤ 2 .
Hence, Q T k Q k = diag {(cid:107) q k − 1 (cid:107) 22 , (cid:107) q k (cid:107) 22 } for k ≥ 1 .
Thus the pseudoinverse ( Q T k Q k ) † = diag {I ( (cid:107) q k (cid:107) 2 (cid:54) =
0)1 / (cid:107) q k (cid:107) 22 , I ( (cid:107) q k (cid:107) 2 (cid:54) =
0)1 / (cid:107) q k (cid:107) 22 } , where I ( · ) is the indicator function that
I ( x ) =
(cid:26)x is true , x is false .
Now, we give the proof of Theorem 2.
Proof of Theorem 2 .
For convenience, we restate the main assumptions of g here:
[Equation]
The proof is based on the two lemmas given in Lemma 1 and Lemma 2.
With (45b), we know if
[Equation]
2 is small, then (cid:107) x k − x ∗ (cid:107) 2 is small as well.
Hence we can always ensure x k ∈ B ( ρ ) if we choose (cid:107) x 0 − x ∗ (cid:107) 2 to be sufﬁciently small and (cid:107) r k (cid:107) 2 is non-increasing during the iterations.
At the end of the k -th iteration of Algorithm 3, we have
[Equation]
Let L k :=
g ( x k +1 )
− g (¯ x k ) + (1 − β k )¯ r k , H k :=
g (¯ x k ) − ¯ x k − ¯ r k , then
[Equation]
34) which bounds the linear part of the residual r k +1 .
For the higher-order terms H k , we have
[Equation]
k , (35) where Γ k = ( γ (1) k , γ (2) k ) T ∈ R 2 .
Since Γ k
=
( Q T k Q k ) † Q T k r k , Q T k Q k = diag {(cid:107) q k − 1 (cid:107) 22 , (cid:107) q k (cid:107) 22 } , it follows that
[Equation]
Therefore,
[Equation]
According to the formula (cid:82) 1 0 g (cid:48) ( x + t ( y − x ))
( y − x ) dt = g ( y )
− g ( x )
[Equation]
we have
[Equation]
Also, by Lemma 2, we have
[Equation]
where π ( k ) denotes that the latest update of q k by Line 8 occurred in the π ( k ) -th iteration and υ ( k )
=
π ( π ( k ) − 1) marks that q k − 1 records q υ ( k ) that is the penultimate update by Line 8 occurring in the υ ( k ) -th
iteration
[Equation]
By substituting these relations to (35), it follows that
[Equation]
Then we can bound each terms of H k as follows (Here we assume q k (cid:54)
=
and q k − 1 (cid:54) =
as q k =leads to γ (2) k
= 0 and q k − 1 = 0 leads to γ (1) k
= 0 where the result is trivial.):
[Equation]
[Table]
[Equation]
where the last inequality is due to (37)
[Equation]
and the second equality is due to (45a), (46), (47b);
[Equation] [Equation]
where the ﬁrst inequality is from (37)
[Equation]
and the second equality is due to (47b);
[Equation]
where the ﬁrst inequality is from (37)
[Equation]
and the second equality is due to (47b).
Then with the bounds (40), (41), (42) and (43), we obtain
[Equation]
Combining (34) and (44) to (33), we obtain
[Equation]
2 + ˆ κ
[Equation]
as desired.
Since m k ≤ m , the higher-order terms are limited.
Then, for (cid:107) x 0 − x ∗ (cid:107) 2 sufﬁciently small, residual r k is non-increasing during the iterations.
Remark 6.
There may be some concern about if ˆ κ can be counteracted by the constant hidden in the Big- O notation.
In fact, since ˆ κ is the Lipschitz constant of g (cid:48) and the constants in O ( · ) are composed of κ, c p , c q , it follows that ˆ κ is unrelated to O ( · ) .
Hence a small ˆ κ can lead to a small uniform boundedness of the higher-order terms.
In the extreme case where ˆ κ = 0 , e.g. g is a linear map, the residual only consists of the ﬁrst-order term L k .
Lemma 1.
Under the same assumptions of Theorem 2, for k ≥ 1 , we have the following bounds:
[Equation]
If q k (cid:54) = 0 , then
[Equation]
If the condition in Line 7 is true, then
[Equation]
Proof.
From the assumption (32a) of g
[Equation]
we have
[Equation]
If the condition in Line 7 is true,
i
[Equation]
e.,
[Equation]
P k −
[Equation]
∆ r k − 1 (cid:107) 2 , we have
[Equation]
∆
[Equation]
k
[Equation]
The inequality (cid:107) q k (cid:107) 2 ≤ (cid:107) ∆ r k − 1 (cid:107) is due to the fact that q k is the orthogonal projection of ∆ r k −onto range( Q k − 1 ) ⊥ .
Also, q k (cid:54) = 0 must hold otherwise q = Q
k − 1 ζ k which violates the condition (cid:107) Q k − 1 ζ k (cid:107) 2 ≤ c q (cid:107) ∆ r k − 1 (cid:107) 2 as c q ∈ (0 , 1) .
If q k (cid:54) = 0 , then q k must be updated by Line 8 in some previous iteration.
We assume the latest update by Line 8 occurred in the j -th iteration, i.e., q k = q j , p k =
p j , hence
[Equation]
p k
[Equation]
where the ﬁrst inequality is due to (47a) and (47b) and the second inequality is due to (45a).
Lemma 2.
Under the same assumptions of Theorem 2, in the k -th iteration ( k ≥ 0 ) of the restarted
MST-AM (Algorithm 3), we have
[Equation]
where m k = k mod m , ∆ r k − m k − 1 :
=
r k − m k , π ( k ) denotes that the latest update of q k by Lineoccurred in the π ( k ) -th iteration and π ( k ) =
k − m k
if Line 8 is never executed up to the k -iteration.
Proof.
We prove (48) by induction.
Denote ζ k
=
( ζ (1) k , ζ (2) k ) T .
For k = 0 , the relation trivially holds.
For k = 1 , p 1 = ∆ x 0 , q 1 = ∆ r 0 .
It follows that
[Equation]
So (48) holds for k = 1 .
Suppose that (48) holds for 0 ≤ j ≤ k − 1 , where k ≥ 2 .
For k ≥ 2 , if m k :
=
k mod m = 0 or 1, (48) holds because it is the same as the case of k = 0 or k = 1 .
Now consider the nontrivial cases.
For m k (cid:54) =
and m k (cid:54) =
1 , if the condition in Line 7 in Algorithm 3 is true, then π ( k ) =
k .
With the convention that ∆ r k − m k − 1 =
r k − m k , we have
[Equation] [Equation]
where υ ( k − 1) =
π ( π ( k − 1) − 1) denotes that q k − 2 records q υ ( k − 1) that is the penultimate update by Line 8 occurring in the υ ( k − 1) -th iteration.
Considering the terms in p k + q k , we know g
[Equation]
For ζ k = ( Q T k − 1 Q k − 1 ) † Q T k − 1 ∆ r k − 1 , because Q T k − 1 Q k − 1 = diag {(cid:107) q k − 2 (cid:107) 22 , (cid:107) q k − 1 (cid:107) 22 } , it follows that
[Equation]
Therefore,
[Equation]
Now, to bound the terms in (50)
[Equation]
we have
[Equation]
where the third inequality is due to (51)
[Equation]
and the last inequality is due to (45a)
[Equation]
(47a) and (46)
[Equation]
where the constants are absorbed into the big- O notation.
Similarly, it follows that
[Equation]
For the remaining terms in (49), we have
[Equation]
where the last inequality is due to (47b) and that q k − 1 = q π ( k − 1) , q k − 2 = q υ ( k − 1) .
With relations (49), (50), (52), (53) and (54), and noting that p k = ∆ x k − 1 − p k − 1 ζ (2) k − p k − 2 ζ (1) k , and (cid:13)(cid:13)(cid:13)(cid:13)(cid:90)
[Equation]
we can estimate p k + q k as
[Equation]
To further obtain (48)
[Equation]
notice that the difference
[Equation]
where the last inequality is due to the condition (cid:107) P k − 1 ζ k (cid:107) 2 ≤ c p (cid:107) p (cid:107) 2 and inequalities (45a) and (47a).
Hence the difference can be absorbed in the Big- O notation in (55).
Thus (48) 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 P k =
P k − 1 , Q k
=
Q
k − 1 .
In other words, the memory recording p k , q k keeps unchanged from the ( π ( k )+ 1) -th to k -th iteration.
Since π ( k ) < k and π ( π ( k ))
=
π ( k ) , with the inductive hypothesis, we have
[Equation]
To further obtain (48), note that the difference
[Equation]
can be absorbed into the Big- O notation in (56).
Thus (48) also holds when the condition in Lineis false for the k -th iteration.
As a result, we complete the induction in the k -th iteration.
C. R EGULARIZED SHORT - TERM RECURRENCE A NDERSON MIXING C.3. C HECK OF POSITIVE DEFINITENESS
We describe the check of positive deﬁniteness in RST-AM, which follows the same procedure of that of SAM (Wei et al., 2021).
From Line 11-13 in Algorithm 1, one update of x k RST-AM is x k +
= x k + H k r k , where H k = β
k
I − α k Y
k Z † k Q T k ,
Y k
= P k + β k Q k , Z k = Q T
k Q k + δ k P T k P k .
H k is generally not symmetric.
A critical condition for the convergence analysis of RST-AM is the positive deﬁniteness of H k , i.e.
[Equation] [Equation]
where µ ∈ (0 , 1) is a constant.
Next, we give an approach to guarantee it.
Let λ min ( · ) denote the smallest eigenvalue, λ max ( · ) denote the largest eigenvalue.
Since s T
k H k s k
= 12 s T k
( H k + H T k ) s k , Condition (57) is equivalent to λ min (cid:0) 12 (cid:0)
H k
+ H T k (cid:1)(cid:1) ≥ β k µ .
With some simple algebraic operations, we obtain λ min (cid:0) 12 (cid:0)
[Equation]
k
[Equation]
Let λ k def = λ max
( Y k Z † k Q T k
+ Q k Z † k Y T k ) , then Condition (57) is equivalent to
[Equation]
≤
[Equation] [Equation]
) , (58) namely, (13) in Remark 4.
To check Condition (58), note that λ k = λ max
[Equation]
Since (cid:18)
[Equation]
Z
[Equation]
(cid:19) ∈ R 4 × 4 , λ k can be computed cheaply.
This cost is negligible compared with those to form P T k P k , Q T k Q k , which need O ( d ) ﬂops.
After that, to guarantee the positive deﬁniteness, we check if α k satisﬁes
C.3. P ROOFS OF THE THEOREMS
We ﬁrst give the proof of the boundedness of (cid:107) P k − 1 ζ k (cid:107) 2 and (cid:107) Q k − 1 ζ k (cid:107) 2 .
Lemma 3.
For P, Q ∈ R d × m ( d ≥ m ) , δ > 0 , and Z = Q T Q + δP T P , we have
[Equation]
Proof.
We ﬁrst consider the case that Z is nonsingular, then PZ † Q T =
PZ − 1 Q T , QZ † Q T =
QZ − 1 Q T .
It can be seen that P T P and Q T Q are symmetric positive semideﬁnite, and Z is SPD as it is assumed to be nonsingular.
Also, we have δP T P (cid:22) Z and Q T Q (cid:22) Z , where the notation “ (cid:22) ” denotes the Loewner partial order , i.e., A (cid:22) B with A, B ∈ R m × m means that B − A is positive semideﬁnite.
Hence, we have
[Equation]
, which implies
[Equation]
Let λ max ( · ) denote the largest eigenvalue, we have
[Equation]
≤ δ −
[Equation]
Therefore, (cid:107) PZ − 1 Q T (cid:107) ≤ δ − 12 and (cid:107)
QZ − 1 Q T (cid:107) 2 ≤ 1 .
For the case that Z is singular, it can be proved that
[Equation]
In fact, if Px = 0 and Qx = 0 , it is obvious that
Zx
= 0 .
On the other side, if Zx = 0 , then x T Zx = 0 , i.e. x T Q T Qx + δx T P T Px = 0 .
Since 0 (cid:22) P T P, 0 (cid:22) Q T Q, δ > 0 , it follows that x T Q T
Qx
= 0 and δx T P T Px = 0 , which further implies Qx = 0 and Px = 0 .
Hence (61) holds.
Let U 1 satisfy U T 1 U
=
I and range( U 1 ) =
ker( Z ) , i.e. the orthonormal basis of ker( Z ) , and Usatisfy U T
U
=
I and U T2 U 1 = 0 , i.e. the orthonormal basis of ker( Z ) ⊥ .
With the equality (61), we know PU 1 = 0 , QU 1 = 0 .
Deﬁne U
=
( U 1 , U 2 ) ∈
R m × m , then U T U =
I m and by direct computation, we have
[Equation]
where U T2 ZU 2 =
[Equation]
PU
[Equation]
PU 2 is nonsingular according to the deﬁnition of U 2 .
So
[Equation]
As a result, we can further compute PZ † Q T and QZ † Q T as
[Equation]
Then let ˆ P = PU 2 , ˆ Q = QU 2 , and ˆ Z
= ˆ Q T ˆ Q
+ δ ˆ P T ˆ P , and noting that ˆ Z is nonsingular, we can obtain (cid:107) PZ † Q T (cid:107)
[Equation]
(cid:107) ˆ P
[Equation]
(cid:107) QZ
[Equation]
Q
[Equation]
(cid:107)
[Equation]
≤ 1 .
As a result of Lemma 3, (cid:107) P k − 1 ζ k (cid:107) 2 , (cid:107) Q k − 1 ζ k (cid:107) 2 in Algorithm 1 are bounded by O ( (cid:107) ∆ r k − 1 (cid:107) 2 ) , as shown in the following corollary.
Corollary 2.
(cid:107) P k − 1 ζ k (cid:107) 2 , (cid:107) Q k − 1 ζ (cid:107) 2 in Algorithm 1 are bounded, i.e.
[Equation]
where ζ k =
( Q T k − 1 Q k − 1 + δ (1) k
P k − 1 P k − 1 ) † Q T k − 1 ∆ r k − 1 .
Now, we turn to the proofs of the theorems about RST-AM in Section 3.4.
For brevity, we use δ k to denote δ (2) k , i.e. δ k ≡ δ (2) k .
From Assumption 2, for the mini-batch gradient f S k ( x k ) = 1 n k (cid:80) i ∈ S k f ξ i ( x k ) , where n k = | S k | , we have
[Equation]
Note that the update of RST-AM (Line 11-13 in Algorithm 1) can be written as x k +
= x k + H k r k , where r k = −∇ f S k ( x k ) , and for k ≥ 0 ,
[Equation]
To prove the theorems, the critical points are (i) the positive deﬁniteness of the approximate Hessian
H k and (ii) an adequate suppression of the noise in the gradient estimation for the stochastic case.
We ﬁrst give a lemma related to the projection step.
Lemma 4.
Suppose that { x k } is generated by RST-AM.
If α k ≥ 0 , β k > 0 , δ k > 0 , then for any v k ∈ R d , we have
[Equation]
Proof.
The result clearly holds when k = 0 as H
= β 0 I .
For k ≥ 1 ,
[Equation]
where Γ k = ( Q T k Q k +
δ k P T k P k ) † Q T k v k solves
[Equation]
[Equation]
By direct computation, (cid:107) v k − Q k
[Equation]
k
[Equation]
22 + δ
[Equation]
P k Γ
[Equation]
22 =
[Equation] [Equation]
Q
[Equation]
Q
[Equation]
+ δ k P
[Equation]
k P k ) † · Q T k v k , thus
[Equation]
≤
[Equation]
Therefore,
[Equation]
In the above
[Equation]
the ﬁrst inequality uses the inequality (cid:107)
[Equation]
(70) where a i ∈ R , x i ∈ R d .
The second inequality is based on inequality (68).
With Lemma 4, we can prove the deterministic case of RST-AM.
Proof of Theorem 3 .
Since 1+2 α 2 k − 2 α k ≤ 1 for 0 ≤ α k ≤ 1 , and δ − 1 k ≤ C − 1 β 2 , with Lemma 4, we have
[Equation]
Since α k satisﬁes (58), we have
[Equation]
Then, under Assumption 1, we have
[Equation]
71) where the last inequality is due to 0 < β ≤
[Equation]
2 L (
[Equation]
C −
[Equation]
Thus, f ( x k +
[Equation]
f
[Equation]
x
[Equation]
≤
[Equation]
Summing both sides of this inequality for k ∈ { 0 , . .
, N − 1 } and recalling f ( x ) >
f low in
Assumption 1 gives
[Equation]
Rearranging and dividing further by N yields (15).
The next lemmas and proofs are about the stochastic case.
Lemma 5.
Suppose that Assumption 2 holds for { x k } generated by RST-AM.
In addition, if β k > 0 , δ k > 0 , and α k ≥ 0 and satisﬁes (13), then
[Equation]
where µ > 0 is the constant introduced in Remark
[Equation]
Proof.
(i) From Lemma
[Equation]
we have
[Equation] [Equation]
From Assumption
[Equation]
we have
[Equation]
With (73), (74), we obtain (72a).
(ii) Recalling that H
= β
I , the result holds for k = 0 .
Deﬁne (cid:15) k
= ∇ f S k ( x k ) − ∇ f ( x k ) = − r k − ∇ f ( x k ) , then H k r k = H k ( − (cid:15) k − ∇ f ( x k )) .
Since α k satisﬁes (13), we can ensure λ min (cid:0) 1 2 (cid:0) H k
+ H T k (cid:1)(cid:1) ≥ β k µ.
Thus
[Equation]
which implies
[Equation]
Let M k
= α k ( P k + β k Q k ) (cid:0)
Q T k Q k
+ δ k P T k P k (cid:1) † Q T k , then
H k
=
β
k
I − M k .
With the assumption (63a), i.e. E S k [ (cid:15) k ] = 0 , we have
[Equation]
Using the Cauchy-Schwarz inequality with expectations , we obtain
[Equation]
We now bound (cid:107) M k
(cid:15) k (cid:107) 22 .
For brevity, let
Z k = Q T
k Q k + δ k P T k P k , and N 1 =
P k
Z † k Q T k , N
=
β
k Q k Z † k Q T k , then
[Equation]
where the last inequality is from Lemma 3.
With (77)
[Equation]
we have (cid:107) M
[Equation]
k (cid:107) 2 , which implies
[Equation]
where the last inequality is due to (63b).
Now we can obtain the bound of | E S k [ ∇ f ( x k ) T H k (cid:15) k ] | as follows (cf.
[Equation]
[Equation]
With the inequality (75) and (79)
[Equation]
we obtain
[Equation]
By imposing more restrictions to α k
[Equation]
β k
[Equation]
δ k
[Equation]
we can obtain a convenient corollary:
Corollary 3.
Suppose that Assumption 2 holds for { x k } generated by RST-AM.
C > 0 is the constant in (12).
If β k > 0 , 0 ≤ α k ≤ min { 1 , β 12 k } and satisﬁes (58) , then
[Equation]
Proof.
The ﬁrst result (81a) is easy to obtain by considering (72a) and noticing that 1+
α 2 k − 2 α k ≤ 1 when α k ∈ [0 , 1] and δ − 1 k ≤ C − 1 β 2 k .
Since α k ≤ β 12 k , δ k ≥ Cβ − 2 k
and (1+ C − 12 ) 2 ≤ 2(1+ C − 1 ) we have
[Equation]
Substituting it into (72b), we obtain (81b).
Using Corollary 3 we obtain the descent property of RST-AM:
Lemma 6.
Suppose that Assumptions 1 and 2 hold for { x k } generated by RST-AM.
C > 0 is the constant in (12).
If 0 < β k ≤
[Equation]
(1+ C
[Equation] [Equation]
α k ≤ min { 1 , β 12 k } and satisﬁes (58), then
[Equation]
Proof.
According to Assumption
[Equation]
we have
[Equation]
Taking expectation with respect to the mini-batch S k on both sides of (83) and using Corollary 3 we obtain
[Equation]
[Equation]
Then (84) combined with the assumption β k ≤
[Equation]
implies (82).
Following the proofs in Wang et al.
(2017), we introduce the deﬁnition of a supermartingale .
Deﬁnition 1.
Let {F k } be an increasing sequence of σ -algebras.
If { X k } is a stochastic process satisfying (
i) E
[ | X k | ] < ∞ , (ii) X k ∈ F k for all k , and (iii) E [ X k +1 |F k ] ≤ X k for all k , then { X k } is called a supermartingale.
Proposition 3 (Supermartingale convergence theorem, see, e.g., Theorem 4.2.12 in (Durrett, 2019)) .
If { X k } is a nonnegative supermartingale, then lim k →∞ X k → X almost surely and E [ X ] ≤ E
[ X 0 ] .
Now, we prove the convergence theory of RST-AM in the nonconvex stochastic case.
Proof of Theorem 4 .
(i) Deﬁne φ k :=
β k µ 4 (cid:107)∇ f ( x k ) (cid:107) 22 and ˜ L :=
( L + µ − 1 )(1 + C − 1 ) , γ k :=
f
( x k )+ ˜ L σ 2 n (cid:80) ∞ i = k β 2 i .
Let F k be the σ -algebra measuring φ k , γ k , and x k .
From (82) we know that for any k ,
[Equation]
which implies that E
[ γ k +1 − f low |F k ] ≤ γ k − f low − φ k .
Since φ k ≥ 0 , we have 0 ≤ E
[ γ k − f low ] ≤ γ 0 − f low < + ∞ .
As the diminishing condition (14) holds, we obtain E [ f ( x k )] ≤ M f for some constant M f > 0 .
According to Deﬁnition 1, { γ k − f low } is a supermartingale.
Therefore,
Proposition 3 indicates that there exists a γ such that lim k →∞ γ k
= γ with probability 1, and E [ γ ] ≤ E [ γ 0 ] .
Note that from (85) we have E [ φ k ] ≤ E [ γ k ] − E
[ γ k +1 ] .
Thus,
[Equation]
which further yields that
[Equation] [Equation]
Since (cid:80) ∞ k =0 β k = + ∞ , it follows that (16) holds.
(ii)
If the noisy gradient is bounded, i.e.,
[Equation]
where M g > 0 is a constant, then a stronger result can be obtained.
For any give (cid:15) > 0 , according to (16), there exist inﬁnitely many iterates x
k such that (cid:107)∇ f ( x k ) (cid:107) 2 ≤ (cid:15) .
Then if (17) does not hold, there must exist two inﬁnite sequences of indices { s i } , { t i } with t i >
s i , such that for i = 0 , 1 , . . .
, k = s i + 1 , . . .
t i − 1 ,
[Equation]
Then from (86) it follows that
[Equation]
β k with probability 1 , which implies that
[Equation]
probability
[Equation]
, as i → + ∞ .
[Equation] [Equation]
According to (74) and (69)
[Equation]
we have
[Equation]
where the last inequalities are due to Cauchy-Schwarz inequality and (87).
Then it follows from (90) that
[Equation]
which together with (89) implies that (cid:107) x
t i − x s i (cid:107) 2 → 0 with probability 1, as i → + ∞ .
Hence, from the Lipschitz continuity of ∇ f , it follows that (cid:107)∇ f ( x t i ) − ∇ f
( x s i ) (cid:107) 2 → 0 with probability 1 as i → + ∞ .
However, this contradicts (88).
Therefore, the assumption that (17) does not hold is not true.
Proof of Theorem 5 .
According to (84) in Lemma 6, we have
[Equation]
where the expectation is taken with respect to { S j } N − 1 j =0 .
Deﬁne
[Equation]
then
[Equation]
Let ˜ D be a problem-independent constant.
If we choose β k = β := min { µ 4 L (1+ C − 1 ) , ˜ Dσ √ N } , and n k = n , then the deﬁnition of P R simpliﬁes to P R ( k ) = 1 /N .
From (93) we have
E [ (cid:107)∇ f ( x R
[Equation]
[Equation] [Equation]
Nβ
[Equation]
Therefore, to ensure E [ (cid:107)∇ f ( x R ) (cid:107) 2 2 ] ≤ (cid:15) , the number of iterations is O (1 /(cid:15) 2 ) .
D E XPERIMENTAL DETAILS
Our main codes were written based on the PyTorch framework 1 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. E XPERIMENTAL DETAILS ABOUT ST-AM/MST-AM
The experiments about ST-AM were conducted to verify the main theorems, i.e. Theorem 1, Corollary 1 and Theorem 2.
Four types of problems were used for the experiments.
The conjugate residual (CR) method is a short-term recurrence version of the full-memory GMRES and needs matrix-vector products to fulﬁll the algorithm.
We give the pseudocode of CR method here for readers who are not familiar with this numerical algorithm.
Input : x 0 ∈
Output : x ∈ R d 1: r 0 =
b − Ax 0 , p 0 = r2: for k = 0 , 1 , . . .
, until convergence, do 3: α
k = r T k
Ar k
( Ap k ) T
Ap k 4: x k +
=
x k + α k p k 5: r k +
=
r k − α k
Ap k 6: β k = r T k
+
Ar k +1 r T k
Ar k 7: p k +
=
r k +
+ β k p k 8:
Ap k +
=
Ar k
+
+ β k
Ap k 9: end for 10: return x k
It should be pointed out that CR method needs to directly assess the matrix A , and the residual update (Line 5) is based on the linear assumption, which makes it inapplicable for nonlinear problems or the case that A is unavailable.
Also, even though ﬁnite 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. S TRONGLY CONVEX QUADRATIC OPTIMIZATION
To construct a case of (5), we considered the least squares problem: min
[Equation]
where A ∈ R (cid:96) ×
d , b ∈ R (cid:96) , which can be reformulated as a form of (5).
We generated a dense random matrix A ∈ R 500 × 100 and a dense random vector b ∈ R 500 for the test.
The gradient descent used a ﬁxed stepsize η ≤
[Equation]
A
[Equation]
22 that can guarantee convergence, and the same η was used as the β k for the full-memory AM and ST-AM.
We also conducted additional tests about solving the problem (94) with different condition numbers, where the eigenvalues of A were set to be uniformly distributed.
The results with different condition numbers ( cond( A T A ) =
10 2 , 10 4 , 10 6 ) are shown in Figure 4.
It can be observed that the curves of CR, AM and ST-AM nearly coincide, which verify the correctness of Theorem 1.
Also, since the eigenvalues are uniformly distributed, the superlinear convergence behaviour may not happen.
Nonetheless, CR, AM and ST-AM are still much faster then the GD method.
D.1. S OLVING A NONSYMMETRIC LINEAR SYSTEM For the solution of a nonsymmetric
linear system
[Equation]
[Figure]
where A ∈ R d × d , b ∈ R d , the ﬁxed-point iteration (FP) is x k +
= g ( x k ) :
= x k + η ( b −
Ax k ) .
Theorem 2 requires g to be a contractive map, i.e. (cid:107) I − ηA (cid:107) 2 < 1 .
We used a test matrix “ﬁdap029” from the Matrix Market 2 .
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. C UBIC -R EGULARIZED QUADRATIC MINIMIZATION
The concerned cubic-regularized quadratic minimization is
[Equation]
where A ∈ R (cid:96) × d , b ∈ R (cid:96) , and M ≥ 0 is the regularization parameter.
It can be computed that
[Equation]
Hence, for the gradient descent x k +
=
g ( x k ) :
= x k − η ∇ f ( x k ) , the Jacobian of g is I −
ηA T A − ηM ( (cid:107) x
(cid:107) − 1 2 xx T + (cid:107) x (cid:107) 2 I ) , which has ˆ κ > 0 in the local region B ( ρ ) .
We generated a dense random matrix A ∈ R 500 × 100 and a dense vector b ∈ R 500 for the test.
[Figure]
We also conducted the tests related to different regularization parameters M and the cases that
M = 0 .
01 , 1 , 100 are shown in Figure 5.
Figure 5 shows that with the large M = 100 , both AM and
MST-AM converge faster.
An ablation study was also conducted about the boundedness restriction of (cid:107) P k − 1 ζ k (cid:107) 2 , (cid:107) Q k − 1 ζ k (cid:107) 2 in Theorem 2.
In Figure 5(b), we show the convergence behaviour of
ST-AM without the boundedness check.
In the case of the rather large regularization M = 100 , STAM does not show a monotone decrease of the residual.
To further investigate the cause, we plot the magnitude of (cid:107) P k − 1 ζ k (cid:107) 2 / (cid:107) ∆ x k − 1 (cid:107) 2 , (cid:107) Q k − 1 ζ k (cid:107) 2 / (cid:107) ∆ r k − 1 (cid:107) 2 in Figure5(c) and Figure 5(d).
It can be observed that the evolutions of (cid:107) P k − 1 ζ k (cid:107) 2 / (cid:107) ∆ x k − 1 (cid:107) 2 , (cid:107) Q k − 1 ζ k (cid:107) 2 / (cid:107) ∆ r k − 1 (cid:107) 2 are quite oscillatory in ST-AM, while being roughly bounded below 1 in MST-AM.
This phenomenon may accounts for the more stable convergence behaviour of MST-AM and veriﬁes the necessity of the changes of
MST-AM compared with ST-AM.
In our experiments, we also found the restarting period can be set quite large and has little effect on ST-AM.
D.1. R OOT - FINDING PROBLEMS IN THE MULTISCALE DEEP EQUILIBRIUM MODEL
The multiscale deep equilibrium (MDEQ) model is a recent extension of the deep equilibrium (DEQ) model (Bai et al., 2019) for computer vision.
One of the central engines for these DEQ models is the root-ﬁnding problem:
[Equation]
where θ and x are parameters and the input representation, respectively.
In (Bai et al., 2020), the
Broyden’s method is employed to solve this problem.
Since MST-AM is suitable for solving nonlinear systems of equations, we can apply MST-AM to solve (97).
We implemented the MST-AM method and integrated it into the MDEQ framework 3 .
The task was image classiﬁcation on CIFAR-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 × 10 − 6 .
The batch size was 128 and the number of epochs was 50.
The tested ﬁxed-point solver was used for the forward root-ﬁnding process and for the backward root-ﬁnding 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. E XPERIMENTAL DETAILS ABOUT RST-AM
Our experiments on RST-AM focused on training neural networks.
Since ST-AM can be regarded as a special case of RST-AM with δ (1) k
=
[Equation]
= 0 , the basic ST-AM is also covered.
In the
Line 10 in Algorithm 1, α k should be adjusted to meet the positive deﬁniteness check (13), which can be simpliﬁed to (∆ x k ) T r k ≥ β k µ (cid:107) r k (cid:107) 22 in practice.
The adjustment of α k can be (i) α k = min { α k , 2 β k (1 − µ ) /λ k } , or (ii) α k = 0 .
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 on MNIST, Penn Treebank, we incorporated preconditioning (described in Appendix A.3) into the baseline method SAM and our method.
Preconditioning is found to be effective for difﬁcult problems, e.g. mini-batch training with very small batch sizes and the scaling of the model’s parameters being important.
D.2. H YPERPARAMETER SETTING OF RST-AM
The hyperparameters of RST-AM are easy to tune.
The only hyperparameters that need to be carefully tuned are the regularization parameters c 1 , c 2 in δ (1) k
and δ (2) k .
We found RST-AM is more sensitive to c 1 , possibly due to the fact that it inﬂuences the construction of secant equations.
c 2 can be quite small in our tests.
The hyperparameter C in δ (2) k was set to be very small ( C = 1 × 10 − 16 )
so as to ensure δ (2) k
[Equation]
(cid:107) p k (cid:107)
+
cid:15) 0 almost all the time.
(cid:15) 0 only serves as the safe-guard to prevent the denominator from being zero.
We found (cid:107) p k (cid:107) 2 > 0 and (cid:107) ∆ x k − 1 (cid:107) 2 > 0 always held in the tests, so (cid:15) 0 can be omitted.
In the experiments except for deterministic optimization on MNIST, we kept the setting c 1 = 1 , c 2 = 1 × 10 − 7 unchanged and found such setting is quite robust.
The other hyperparameters are α k , β k , 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. E XPERIMENTS ON MNIST
The experiments on MNIST 4 aimed to validate the effectiveness of RST-AM in deterministic optimization, so we were only concerned about the training loss by regarding it as a nonlinear function to be optimized.
The training dataset was preprocessed by randomly selecting 10k images from the total 60k images to facilitate the full-batch training.
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 gridsearches from 10 − 3 to 100.
The learning rates of SGDM, Adam, Adagrad, RMSprop 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 c 1 = 0 .
05 , c 2 = 1 × 10 − 7 , α k
= β k
= 1 .
When (∆ x k ) T r k ≤ 0 occurs, x k
+1 =
x k + 0 .
2 r k was used as the new update.
For the preconditioned RST-AM, we set c 1 = 1 × 10 − 2 , c 2 = 1 × 10 − 7 for the Adagradpreconditioned RST-AM and c 1 = 1 × 10 − 2 , c 2 = 1 × 10 − 8 for the RMSprop-preconditioned
RST-AM.
For all these tests, the model was trained for 200 epochs.
[Figure]
In Figure 2 in the main paper, we report the RMSprop/Adagrad preconditioned SAM/RST-AM.
We give the result about using Adam as 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 RST-AM is still better than Adam and Adam SAM(2), which demonstrates the effect of our proposed short-term recurrence scheme.
D.2. E XPERIMENTS ON CIFAR
For this group of experiments, we followed the same setting of training ResNet in (He et al., 2016).
The batchsize was set to be 128 as commonly suggested.
When training for N iterations, the learning rate was decayed at the ( (cid:98) N 2 (cid:99) ) -th and the ( (cid:98) 34 N (cid:99) ) -th iterations.
For SAM and RST-AM, the learning rate decay means decaying α k , β k simultaneously.
The experiments were run with 3 random seeds.
The baseline optimizers were SGDM, Adam, AdaBound (Luo et al., 2018), AdaBelief (Zhuang et al., 2020), Lookahead (Zhang et al., 2019) and AdaHessian (Yao et al., 2021).
AdaBound and
AdaBelief are recently proposed adaptive learning rate methods to improve Adam.
Lookahead is a k -step method, which can be seen as a simple sequence interpolation method.
In each cycle,
Lookahead iterates with an inner-optimizer optim for k steps and then interpolates the ﬁrst and the last iterates to give the starting point of the next cycle.
Compared to Adam, AdaHessian uses
Hessian-vector products to construct a diagonal approximation of the Hessian.
Since tuning hyperparameter in large models is very time-consuming, the hyperparameters of each optimizers were tuned through experiments on CIFAR-10/ResNet20.
For each optimizer, the hyperparameter setting that has the best ﬁnal test accuracy on CIFAR-10/ResNet20 was kept unchanged and used for the other tests.
We list the hyperparameters of all the tested optimizers here.
(Learning rate is abbreviated as lr.)
[Figure]
• SGDM : lr = 0.1, momentum = 0.9, weight-decay = 5 × 10 − 4 , lr-decay = 0.1.
• Adam : lr = 0.001, ( β 1 , β 2 ) =
(0 . 9 , 0 . 999) , weight-decay = 5 × 10 − 4 , lr-decay = 0.1.
AdaBound : lr = 0.001, ( β 1 , β 2 ) =
(0 . 9 , 0 . 999) , ﬁnal lr = 0.1, gamma = 0.001, weightdecay = 5 × 10 − 4 , lr-decay = 0.1.
• AdaBelief : lr = 0.001, ( β 1 , β 2 ) =
(0 . 9 , 0 . 999) , eps = 1 × 10 − 8 , weight-decay = 5 × 10 − 4 , lr-decay = 0.1.
Lookahead : optim : SGDM (lr = 0.1, momentum = 0.9, weight-decay = 1 × 10 − 3 ), α = 0 .
8 , steps = 10, lr-decay = 0.1.
AdaHessian : lr = 0.15, ( β 1 , β 2 ) =
(0 . 9 , 0 . 999) , eps= 1 × 10 − 4 , hessian-power: 1, weightdecay = 5 × 10 − 4 / 0 .
15 , lr-decay = 0.1.
• SAM(2) : optim : SGDM (lr = 0.1, momentum = 0, weight-decay = 1 .
5 × 10 − 3 ), α k = 1 . 0 , β k = 1 . 0 , c 1 = 0 .
01 , p = 1 , m = 2 , weight-decay = 1 . 5 × 10 − 3 , lr-decay = 0.06.
• SAM(10) : optim : SGDM (lr = 0.1, momentum = 0, weight-decay = 1 .
5 × 10 − 3 ), α k = 1 . 0 , β k = 1 . 0 , c 1 = 0 .
01 , p = 1 , m = 10 , weight-decay = 1 . 5 × 10 − 3 , lr-decay = 0.06.
• RST-AM : c 1 = 1 , c 2 = 1 × 10 − 7 , α
=
β 0 = 1 ,
weight-decay = 1 × 10 − 3 , lr-decay = 0.1.
For the tests of SGDM, Adam, AdaBelief, Lookahead and SAM(10), we found consistent numerical results as those in (Wei et al., 2021), so we report their results of these methods in the main paper.
Figure 7 shows the test accuracy of different optimizers for training ResNet18/WRN16-4 on CIFAR10 and ResNeXt50/DenseNet121 on CIFAR-100.
The full results of ﬁnal test accuracy are listed in
Table 1(a) in the main paper.
From Figure 7, we ﬁnd that the convergence behaviour of RST-AM is less erratic than SAM(10) and SAM(2).
In the ﬁrst 80 epochs, RST-AM converges faster than SAM, partly due to using a 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 a higher test accuracy in the ﬁnal 40 epochs.
It is found that RST-AM is comparable to SAM(10)
[Figure]
while improving the short-memory SAM(2)
[Figure] [Figure]
[Figure]
Since our hyperparameter tuning was conducted on CIFAR10/ResNet20, we give some results about the hyperparameters on this model.
Effect of c 2 in RST-AM .
Figure 8 shows the effect of c 2 , where we kept c 1 = 1 , weight-decay= 5 × 10 − 4 ﬁxed, and c 2 = 10 − 1 , 10 − 3 , 10 − 5 , 10 − 7 , 10 − 9 .
It indicates that RST-AM is not sensitive to c 2 .
Effect of c 1 in RST-AM .
Figure 9 shows the effect of c 1 , where we kept c 2 = 10 − 7 , weightdecay= 5 × 10 − 4 ﬁxed, and c 1 = 0 .
It suggests that with smaller c 1 , RST-AM converges faster in terms of train accuracy, but the test accuracy is worse.
Effect of weight decay in RST-AM .
Weight-decay is a common hyperparameter that can affect the generalization of each optimizer.
In Figure 10, we show the convergence behaviour of RST-AM with different weight-decays.
It suggests that with a too small weight-decay, RST-AM tends to be overﬁtting in the test dataset
[Figure]
[Figure] [Figure]
[Table]
Table 1(a) in the main paper reports the test accuracy of each optimizer when training for the same epochs.
In fact, as shown in Figure 11, within 100 epochs, RST-AM can achieve a better test accuracy than SGD.
So if the running time of the training process matters, it is expected that RSTAM can use less total running time due to the large number of reduction in training epochs.
In
Table 1(b)
, we set SGD as the baseline, and compare the computation and memory cost with SGD.
In Table 5, we give more details about the saving in training epochs and the ﬁnal test accuracy.
It can be seen that RST-AM can achieve a comparable or better test accuracy than SGDM with less computation time, while reducing the memory footprint of SAM.
We also tried using Adam as a preconditioner for RST-AM but found the ﬁnal test accuracy was often worse.
For example, the test accuracy of Adam RST-AM for CIFAR-10/ResNet20 is only 90.79%.
We suppose it is the worse generalization ability of Adam (Luo et al., 2018) that makes
Adam not suitable as a preconditioner for RST-AM in the image classiﬁcation task.
D.2. E XPERIMENTS ON P ENN T REEBANK
Our experimental setting on training LSTM models on Penn Treebank dataset was based on the ofﬁcial implementation of AdaBelief (Zhuang et al., 2020).
Results in Table 2 were measured with 3 random seeds.
The parameter setting of the LSTM models were the same as that of AdaBelief.
The baseline optimizers were SGDM, Adam, AdaBelief, and SAM.
We tuned hyperparameters on the validation dataset for each optimizer.
For SGDM, we tuned the learning rate (abbr.
lr) via grid-search in { 1 , 10 , 30 , 100 } and found that lr=10 performs best on 2,3-layer LSTM.
For 1-layer LSTM, we set lr=30 and momentum=0 as that in AdaBelief because we found such setting is better.
For Adam, we tuned the learning rate via grid-search in { 1 × 10 − 3 , 2 × 10 − 3 , 5 × 10 − 3 , 8 × 10 − 3 , 1 × 10 − 2 , 2 × 10 − 2 } and found 5 × 10 − 3 performs best.
For AdaBelief, we tuned the learning rate and found 5 × 10 − 3 is better than 1 × 10 − 2 used in (Zhuang et al., 2020).
For SAM, we used the recommended setting in (Wei et al., 2021) and used the baseline Adam as the preconditioner for SAM.
Historical lengths of m = 2 and m = 10 were tested for SAM, i
[Table]
e. denoted as SAM(2), SAM(10), respectively.
For our method RST-AM, we kept c 1 = 1 , c 2 = 1 × 10 − 7 unchanged and used the same preconditioner (Adam) as SAM.
The batch size was 20.
We trained for 500 epochs and decayed the learning rate by 0.1 at the 250th and 375th epoch.
[Table] [Table]
In Table
[Table]
we report the memory footprint and per-epoch running time of SAM(10) and RST-AM compared with Adam.
It indicates that RST-AM also largely reduces the memory overhead of the long-memory SAM(10) in the language modelling 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 a comparable validation/test perplexity to that of Adam, RST-AM does not need to train for the same number of epochs as Adam.
Table 7 shows the test perplexity of Adam, SAM(2), SAM(10), and RST-AM for training 200 epochs.
By comparing the results of Table 2 and Table 7, we see RST-AM is better than
Adam with much fewer training epochs, thus RST-AM can save a large amount of training time, as shown in Table 6.
D.2. E XPERIMENTS ON ADVERSARIAL TRAINING
The adversarial training considers the empirical adversarial risk minimization (EARM) problem:
[Equation]
where ¯ ξ i is the i -th adversarial data within the (cid:15) -ball centered at ξ i .
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) to solve the outer minimization problem.
The experiments were conducted on CIFAR10/ResNet18, CIFAR10/WRN16-4,
CIFAR100/ResNet18, and CIFAR100/DenseNet121.
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 CIFAR10/ResNet20, and applied the same hyperparameters to other models.
We found the setting that c 1 = 1 , c 2 = 1 × 10 − 7 and weight-decay = 0 .
001 is still suitable for this task.
The CIFAR-10 (CIFAR-100) dataset contains 50K images for training and 10K images for testing.
Since it is found that the phenomenon of overﬁtting is severer (Rice et al., 2020) in adversarial training
[Table]
we randomly selected 5K images from the total 50K training dataset as the validation dataset (the other 45K images remained as the training dataset)
[Table]
and chose the best checkpoint model in 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 et al., 2018)
[Table]
and C&W ∞ attack (Carlini
Wagner,
[Table]
In Tabel
[Table]
we report the average results of tests with three different random seeds.
It can be found that RST-AM can achieve both higher clean test accuracies and robust test accuracies across various models on CIFAR10/CIFAR100.
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 Figure 12.
We can observe the phenomenon of robust overﬁtting (Rice et al., 2020) from these curves, which justiﬁes 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.
[Figure]
D.2. E XPERIMENTS ON TRAINING A GENERATIVE ADVERSARIAL NETWORK
We describe our setting of training a generative adversarial network (GAN) here.
Like the adversarial training, the GAN training process is also a min-max problem.
The stability of an optimizer is critical for the training process.
To demonstrate the applicability of RST-AM, we conducted experiments on a GAN which was equipped with spectral normalization (Miyato et al., 2018) (SNGAN).
The experimental setting was the same as that of AdaBelief (Zhuang et al., 2020): the dataset was CIFAR-10;
ResNets were used as the generator and the discriminator networks; the steps for optimization in the discriminator and the generator per iteration were 5 and 1, respectively; the minbatch size was 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 et 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 as the preconditioner and set the damping parameter α k = 0 . 6 .
The regularization parameters c 1 = 1 and c 2 = 1 × 10 − 7 were still kept unchanged just as the previous experiments.
[Figure]
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 is stable for this min-max optimization problem.
[Table]
Since we only tuned the damping parameter α k for RST-AM, we report the FID scores of other choices of α k in Table 9 during our experiment.
It shows that even with the suboptimal choices of α k , e.g. α k = 0 . 1 , 0 . 5 , 0 . 8 , RST-AM can still outperform the baselines.
D. A DDITIONAL 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 classiﬁcation task in ImageNet (Deng et al., 2009) and the Transformer (Vaswani et al., 2017) based neural machine translation task in the IWSTL14 DE-EN (Cettolo et al., 2014) dataset.
D.3. E XPERIMENTS ON I MAGE N ET
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
Ti GPUs were used for each test.
The hyperparameters of SGDM were set as the recommended setting in PyTorch.
For RST-AM, the hyperparameters were kept the same as those in the CIFAR experiments.
The weight-decay was 1 × 10 − 4 .
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 to an acceptable solution than SGDM, we adopted the early learning rate decay strategy recommended in (Zhang et al., 2019): decay the α k and β k for
RST-AM at the 30th, 50th and 70th epochs.
[Table]
In Table 10, we report the TOP1 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%.
It shows RST-AM needs 72 epochs while SGDM needs 88 epochs.
The curves of the training accuracy and test accuracy are shown in Figure 14.
The results suggest that RST-AM is still a competitive optimizer in training a larger model in a larger dataset.
[Figure]
D.3. T RAINING T RANSFORMER
We conducted the neural machine translation task with Transformer.
We implemented our RST-AM method and integrated it into the fairseq framework 5 .
The basic experimental setting was set as the recommended setting in (Yao et al., 2021).
We trained the model for 50 epochs and the BLEU (Papineni et al., 2002) score was calculated using the average model of the last ﬁve checkpoints.
We added a baseline optimizer RAdam (Liu et al., 2019) which was inspired by the warmup procedure in training Transformer.
[Table] [Table]
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 classiﬁcation task.
Also, the results demonstrate that RST-AM can still outperform Adam in this task.
[Table]
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
-AM over Adam and the cost