Towards Understanding Why Lookahead Generalizes Better Than SGD and Beyond

Abstract
To train networks, lookahead algorithm [1] updates its fast weights k times via an inner-loop optimizer before updating its slow weights once by using the latest fast weights. Any optimizer, e.g. SGD, can serve as the inner-loop optimizer, and the derived lookahead generally enjoys remarkable test performance improvement over the vanilla optimizer. But theoretical understandings on the test performance improvement of lookahead remain absent yet. To solve this issue, we theoretically justify the advantages of lookahead in terms of the excess risk error which measures the test performance. Specifically, we prove that lookahead using SGD as its inner-loop optimizer can better balance the optimization error and generalization error to achieve smaller excess risk error than vanilla SGD on (strongly) convex problems and nonconvex problems with Polyak-Łojasiewicz condition which has been observed/proved in neural networks. Moreover, we show the stagewise optimization strategy [2] which decays learning rate several times during training can also benefit lookahead in improving its optimization and generalization errors on strongly convex problems. Finally, we propose a stagewise locally-regularized lookahead (SLRLA) algorithm which sums up the vanilla objective and a local regularizer to minimize at each stage and provably enjoys optimization and generalization improvement over the conventional (stagewise) lookahead. Experimental results on CIFAR10/100 and ImageNet testify its advantages. Codes is available at https://github.com/sail-sg/SLRLA-optimizer.

1 Introduction
Deep neural networks (DNNs) have been successfully applied to various applications, such as image classification [3–8], speech recognition [9–11], and classic games [12, 13]. Typically, their training models can be formally formulated as a finite-sum optimization problem:
minθ FS(θ) , 1
n ∑n i=1 `(f(xi;θ),yi), (1)
where (x,y) is a sample pair from an unknown distribution D, the loss `(f(x;θ),y) measures the discrepancy between the prediction f(x;θ) parameterized by θ and the target y, and n is the training sample number. Besides DNNs, the formulation (1) also encapsulates a large body of other problems, e.g., least square regression and logistic regression. Though many algorithms, e.g., variance-reduced algorithms [14–17] and adaptive gradient algorithms [18, 19], can solve problem (1), SGD [20, 21] is one of the most preferable algorithms because of its efficiency and good generalization [22].
In this work, we are particularly interested in the lookahead algorithm [1] for solving (1). Its core idea is to maintain two kinds of network parameters, i.e. “fast weights" v and “slow weights" θ, and update them in turn. Specifically, in the inner loop, it takes the slow weights θ as warm-start
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
initialization and updates the fast weights k times to obtain vk using any vanilla optimizer; for the outer loop, it updates the slow weights as θ+ = (1− α)θ + αvk, where α ∈ (0, 1]. Any standard optimizer, e.g. SGD, Adam [18] and RAdam [23], can serve as the inner-loop optimizer, and the derived lookahead algorithm generally enjoys remarkable test performance improvement than the standard optimizer [1]. Because of its simplicity and strong compatibility, lookahead has been widely used [24–29]. However, the theoretical reasons for the superiority in test performance of lookahead are rarely investigated, though heavily desired. Moreover, in practice, to train faster and generalize better, one often uses the stagewise optimization strategy [2], namely running a large learning rate at the beginning and geometrically decaying it several times during the following training. But for lookahead, the theoretical benefits of the stagewise optimization strategy still remain unclear.
Our Contributions. In this work, we provide a theoretical viewpoint to understand the test performance improvement of lookahead, and also show the benefits of stagewise optimization strategy in lookahead. Moreover, we further propose a new stagewise locally-regularized lookahead algorithm which provably enjoys faster convergence speed, smaller generalization error, and better test performance than conventional (stagewise) lookahead. Our contributions are highlighted below.
Firstly, we prove that on convex problems, lookahead using SGD as its inner-loop optimizer has optimization error O
( 1 αηkT + η ) , where T and k respectively denote the outer- and inner-loop
iteration numbers, and η is the learning rate for the inner-loop optimizer, i.e. SGD. Then we prove that lookahead has generalization error boundO ( αηkT n ) on convex problems with n training samples. Since the excess risk error can well measure the test performance of an algorithm and is upper bounded by the sum of optimization and generalization errors, we can bound the excess risk error of lookahead with O ( 1
αηkT + η + αηkT n
) . When α = 1, lookahead degenerates to vanilla SGD and has
excess risk error O (
1 ηkT + η + ηkT n
) . Since the optimum of α in lookahead to optimally balance the
optimization and generalization errors is often not one, lookahead can enjoy a smaller excess risk error than vanilla SGD. Similarly, on strongly convex problems and a class of nonconvex problems that obey Polyak-Łojasiewicz (PŁ) condition, our results also show that lookahead can better trade-off optimization and generalization errors and achieve smaller excess risk error than vanilla SGD. For PŁ condition, it is observed /proved for deep learning models in [30–34] and our empirical results. These results well explain the better test performance of lookahead than SGD.
Secondly, we propose a Stagewise Locally-Regularized LookAhead (SLRLA) algorithm, and prove its advantages over lookahead and stagewise lookahead in terms of excess risk error. SLRLA first divides the optimization into several stages. Then at each stage, it minimizes a locally-regularized function that contains the vanilla loss FS(θ) in (1) and a local regularizer β2 ‖θ−θq‖
2 with the output θq of the previous stage. A similar stagewise strategy is commonly used in practice, but does not has the local regularizer. Our results show two advantages of SLRLA over lookahead and stagewise lookahead. (i) On strongly convex problems and weakly quasi convex problems (a class of nonconvex problems including convex problems as a special case), SLRLA achieves faster convergence rates (i.e. smaller optimization error), and enjoys smaller generalization error than lookahead and its stagewise variant. (ii) On nonconvex problems under PŁ condition, when the problems are not heavily nonconvex (the smallest eigenvalue of∇2FS(θ) is not very small), SLRLA achieves linear convergence rate, and greatly improves optimization and generalization errors in lookahead, while stagewise lookahead cannot enjoy the linear convergence rate and the improvement.
Finally, when β = 0, SLRLA degenerates to conventional stagewise lookahead, and our results show the advantages of stagewise lookahead over lookahead in terms of excess risk error on strongly convex problems, explaining the benefits of the stagewise strategy in lookahead.

2 Related Work
Optimization Algorithms. When problem (1) involves deep networks, SGD and adaptive gradient algorithms, e.g. Adam [18] and AdaGrad [19], are more preferable than other algorithms, such as variance-reduced SGD [14, 15, 35, 36], because of their high efficiency and good generalization. Moreover, SGD generally enjoys better generalization performance than adaptive gradient algorithms, since SGD tends to converge to flat minima while adaptive gradient algorithms approach sharp minima [22, 37–41]. To further boost the generalization performance of SGD and adaptive gradient algorithms, Zhang et al. [1] proposed lookahead which can employ any standard optimizer as its
inner-loop optimizer and brings remarkable generalization improvement. Recently, a few works analyze lookahed but focus on its optimization performance, e.g. small gradient noise variance on least square problems [1] and sublinear convergence rate on nonconvex problems [42]. In contrast, we analyze the intrinsic theoretical reasons for the superiority of lookahead in terms of test performance.
Generalization Analysis of Algorithms. Uniform stability [43] is a classical tool to analyze generalization error of an algorithm. For instance, Hardt et al. [44] and Zhang et al. [45] analyzed the generalization of SGD via uniform stability. We also utilize stability to analyze generalization but target at analyzing the test performance, including optimization error and generalization error, of an algorithm and its lookahead variant, which is of more practical interest especially in deep learning. Yuan et al. [46] analyzed the test performance of SGD and stagewise SGD and showed advantages of the statewised strategy. Differently, we show that lookahead can well balance the optimization and generalization errors, and thus enjoys better test performance than its vanilla inner-loop optimizer.

3 Notations and Preliminaries
Convexity, Lipschitz Continuity, and Smoothness. For analysis, we first introduce necessary definitions, i.e. convexity, Lipschitz continuity and smoothness. These definitions are commonly used in the convergence and generalization analysis of optimization algorithms, e.g. [47–51]. Definition 1 (Convexity, Lipschitz Continuity and Smoothness). We say a function f(θ) is λ-strongly convex if ∀θ1,θ2, f(θ1)≥ f(θ2)+〈∇f(θ2),θ1 − θ2〉+ λ2 ‖θ1 − θ2‖
2. If λ = 0, then we say f(θ) is convex. Moreover, we say f(θ) is G-Lipschitz continuous if ‖f(θ1)− f(θ2)‖2 ≤ G‖θ1 − θ2‖2. f(θ) is said to be L-smooth if its gradient obeys ‖∇f(θ1)−∇f(θ2)‖2 ≤ L‖θ1 − θ2‖2 (∀θ1,θ2).
Polyak-Łojasiewicz (PŁ) Condition and Weakly Quasi-Convexity. For nonconvex problems, such two conditions establish the relation between the gradient norm and the loss distance at two points. Definition 2 (PŁ Condition & Weakly Quasi-Convexity). Let θ∗∈argminθ f(θ). We say a function f(θ) satisfies µ-PŁ condition if it satisfies 2µ(f(θ) − f(θ∗))≤ ‖∇f(θ)‖2 (∀θ) with a universal constant µ. f(θ) is said to be ρ-weakly-quasi-convex if it obeys 〈∇f(θ),θ− θ∗〉≥ρ(f(θ)− f(θ∗)).
Excess Risk Error Decomposition. Given a dataset S = {(xi,yi)}ni=1 where (xi,yi) is drawn from an unknown distribution D, one often minimizes the empirical risk FS(θ) , 1n ∑n i=1 `(f(xi;θ),yi) in (1) via a randomized algorithmA, e.g. SGD, to find an estimated optimum θA,S ≈ argminθ FS(θ). However, this empirical solution θA,S differs from the desired optimum θ∗D of the population risk θ∗D ∈ argminθ F (θ) , E(x,y)∼D[`(f(x;θ),y)]. This raises a particularly important question: what performance of the estimated optimum θA,S can achieve on the test data (x,y) ∼ D? To answer this question, we analyze the test error EA,S [F (θA,S)] of θA,S via investigating the well-known excess risk error εexc defined as εexc = EA,S [F (θA,S)]−EA,S [FS(θ∗S)]= EA,S [F (θA,S)−FS(θA,S)]+EA,S [FS(θA,S)−FS(θ∗S)], (2) where θ∗S ∈ argminθ FS(θ) is the optimum of empirical risk FS . Generally, one can expect very small EA,S [FS(θ∗S)], since by selecting a powerful model, e.g., a deep network, one can well fit the data. So to bound the test loss EA,S [F (θA,S)], one only needs to upper bound the right side of Eqn. (2). The optimization error εopt ,EA,S [FS(θA,S) − FS(θ∗S)] denotes the difference between the exact optimum θ∗S and the estimated solution θA,S ; the generalization error εgen , EA,S [F (θA,S)− FS(θA,S)] measures the effects of minimizing empirical risk instead of population risk. So one often analyzes εopt and εgen to compare test performance of different algorithms.
Uniform Stability and Generalization. One popular approach to analyze generalization error εgen of a randomized algorithm A is uniform stability [44, 45]. This is because as shown in following lemma, for an algorithm A, if it is -uniformly stable, its generalization error is upper bounded by . In the following, we also analyze the uniform stability of lookahead to bound its generalization error. Lemma 1 (Uniform Stability and Generalization Error). [44] We say a randomized algorithm A is -uniformly stable if for all datasets S∼D and S ′∼D where S and S ′ differ in at most one sample,
sup(x,y)∼D EA [`(f(x;θA,S);y)− `(f(x;θA,S′);y)] ≤ . Moreover, if A is -uniformly stable, then its generalization error εgen which is defined as εgen = |EA,S [F (θA,S)− FS(θA,S)]| satisfies εgen ≤ .
Algorithm 1: Lookahead Optimization Procedure (FS(θ), η, T, α, k,θ0,A,S) Input :Objective FS(θ), dataset S, inner-loop optimizer A, inner-loop step number k and learning rate {{η(t)τ }k−1τ=0}Tt=1, outer-loop learning rate α ∈ (0, 1), initialization θ0. for t = 1, 2, ..., T do
v (t) 0 = θt−1; for τ = 1, 2, ..., k do v
(t) τ = A(FS(θ),v(t)τ−1, η (t) τ−1,S);
end θt = (1− α)θt−1 + αv(t)k .
end Output :θA,S = θT for strongly convex problem; θA,S = 1Tk ∑T t=1 ∑k−1 τ=0 v (t) τ for convex and
nonconvex problems.

4 Excess Risk Analysis of Lookahead Algorithm
The lookahead algorithm [1] to solve problem (1) is described in Algorithm 1. It maintains (i) slow weights θ updated in the outer loop, and (ii) fast weights v updated in the inner loop. For the inner loop, one can run any standard optimizer A, e.g. SGD or Adam used in deep learning, k steps via an operator A(FS(θ),v(t)τ−1, η (t) τ−1,S) to update the fast weights v, where k is often small, e.g. k = 5 in [1]. Here the operator A(FS(θ),v(t)τ−1, η (t) τ−1,S) denotes a minibatch gradient descent step in SGD or Adam, given the loss FS(θ), current solution v (t) τ−1, learning rate η (t) τ−1, and dataset S. Next, lookahead uses the fast weights v to update the slow weights θ as θ = θ + α(v − θ) with an outer-loop learning rate α ∈ [0, 1]. Because of its effectiveness and compatibility, lookahead has been widely used to boost the performance of SGD [1], Adam [1, 52], RAdam [23, 24] and natural gradient algorithm [25], and sets new state-of-the-arts in image classification and generation, and machine translation [1, 26–29]. However, there is no theoretical analysis that explicitly justifies the performance improvement of lookahead over the vanilla inner optimizer A, hindering the development of new and more advanced optimizers in a principle way. The following sections aim to solve this problem by comparing the optimization and generalization errors of lookahead with its vanilla inner optimizer A. For analysis, we choose SGD as A, as SGD is widely used in deep learning. In this way, the inner-loop updating becomes
v(t)τ = A(FS(θ),v (t) τ−1, η (t) τ−1,S) = v (t) τ−1 − η (t) τ−1g (t) τ−1,
where g(t)τ−1 = 1 |B| ∑ (x,y)∈B∇`(f(x;v (t) τ−1);y). Here B is the sampled minibatch at the (t, τ)-th iteration. In the following, we investigate the optimization and generalization errors of lookahead, and combine these two errors to upper bound its excess risk error which measures the test performance.

4.1 Results on Strongly Convex Problems
To begin with, we first investigate the convergence performance of lookahead when its inner optimizer A is SGD. Our main results are summarized in Theorem 1. See its proof in Appendix D.1. Theorem 1. Suppose that FS(θ) is λ-strongly convex, and each individual loss `(f(x;θ),y) is G-Lipschitz and L-smooth w.r.t. θ. Let θ∗S = argminθ FS(θ). By setting the inner learning rate η (t) τ =
λ+L λL((t−1)k+τ+2) , the optimization error of the output θA,S of lookahead satisfies
εopt =EA,S [FS(θA,S)−FS(θ∗S)]≤  3L(k+2)2α 2((T+1)k+2)2αE[‖θ0−θ ∗ S‖2] + 16LG 2 λ2((T+1)k+2)2α(1−2α) , 0<α< 1 2 , 3L(k+2) 2(T+1)k+2E[‖θ0−θ ∗ S‖2] + 16LG2 log(Tk+2) λ2((T+1)k+2) , α = 1 2 , 4L(k+2)2α
((T+1)k+2)2αE[‖θ0−θ ∗ S‖2] + 90LG
2
λ2(2α−1)(Tk+2) , 1 2 <α≤1.
Theorem 1 shows that lookahead using SGD as its inner optimizer can converge on the strongly convex problems. The optimization error εopt has two terms: the first bias term characterizes the effect of initialization θ0; the second term reveals the impact of the stochastic gradient noise. For
the outer-loop learning rate α, it can be observed that with the increase of α, εopt becomes smaller. So when α=1, lookahead degenerates to vanilla SGD and achieves the smallest optimization error O (
1 T 2 + 1 λ2Tk
) . This rate matches the one O ( 1
λ2Tk
) of SGD in [53] under the same assumptions,
as k is often much smaller than iteration number T , e.g. k = 5 in [1]. So lookahead indeed does not benefit the convergence of SGD. It can be intuitively understood: for every k steps, lookahead updates θ as θt = θt−1 + α(v (t) k − θt−1) = θt−1 −α ∑k−1 τ=0 η (t) τ g (t) τ and goes forward slowly due
to α < 1, while SGD updates θ as θt = θt−1 − ∑k−1 τ=0 η (t) τ g (t) τ and runs faster, where g (t) τ denotes the stochastic gradient.
Next, to bound the excess risk error, we investigate the generalization error of lookahead by analyzing its stability, since as shown in Lemma 1, the uniform stability can upper bound generalization error. Theorem 2. Suppose the assumptions and parameter setting in Theorem 1 hold. The generalization error of the output θA,S of lookahead satisfies εgen = EA,S [F (θA,S) − FS(θA,S)] ≤ 16G2
nλ (Tk+1)α−1 ((T+1)k+2)α .
See its proof in Appendix D.2. Theorem 2 shows that when α 6= 0, the generalization error εgen of lookahead using SGD as its inner optimizer can be upper bounded by O ( G2
nλ
) . Particularly, for
α=0, εgen becomes zero. This is because α=0 means no updating, namely θt = θ0 (∀t), and thus EA,S [F (θA,S)−FS(θA,S)]=ES [F (θ0)−FS(θ0)]=0. For the effects of α on εgen, one can find that when α increases, εgen also becomes larger. This can be intuitively understood: larger α means quick updating of θt. Accordingly, the empirical risk FS(θt) quickly decreases, while the population risk F (θt) may not due to possible overfiting. Interestingly, when α=1 indicating lookahead becomes vanilla SGD, our generalization error bound matches the previous bound O ( G2
λn
) of vanilla SGD
in [44], even though Hardt et al. [44] used constant learning rate while we use decaying learning rate.
Based on Theorems 1 and 2, we can derive the excess risk error bound in Corollary 1. Corollary 1. With the same assumptions and parameter setting in Theorem 1, the excess risk error εexc in (2) of the output θA,S of lookahead obeys εexc ≤ εopt + εgen, where εopt and εgen are given in Theorems 1 and 2, respectively.
See its proof in Appendix D.3. Corollary 1 shows that the excess risk error εexc of lookahead using SGD as its inner optimizer satisfies εexc ≤ εopt +εgen, guaranteeing good test performance of the estimated solution θA,S by lookahead. In this work, we are particularly interested in the effects of α on εexc. When α∈(0, 12 ), εexc is of the order O ( L T 2α + LG2 λ2T 2αk2α + G2 nλ (1− 1 Tαkα ) ) . As in most cases, the factor LG 2
λ2 is much larger than G2
nλ , especially for ill-conditioned problems where the strongly convex parameter λ is very small, increasing α or the total training iteration number Tk will decrease εexc. When α∈( 12 , 1), εexc is of the order O ( e(α) ) where e(α)= LT 2α + LG2 λ2(2α−1)Tk+ G2 nλ (1− 1 Tαkα ). Then we have e′(α)=− 2L lnTT 2α − 2LG2 λ2(2α−1)2Tk + G2 nλ ln(Tk) Tαkα . There are two common cases that lead to e′(α)< 0: (i) the problem is large-scale but not ill-conditioned, and thus the iteration number T is not large since the problem is easy, leading to 2L lnTT 2α > G2 nλ ln(Tk) Tαkα ; (ii) the problem is heavily ill-conditioned, but iteration number T is moderate due to the moderate precision requirement, giving 2LG2
λ2(2α−1)2Tk > G2 nλ ln(Tk) Tαkα . For both cases, increasing α can decrease εexc. For other cases, in theory, by choosing a proper α∈ (0, 1), one can expect better balance between the optimization error and generalization error, and could achieve a smaller excess risk error than vanilla SGD which corresponds to α= 1. As Sec. 3 shows that the excess risk error can well measure the test performance of an algorithm, our results explain the better test performance of lookahead than SGD. Moreover, though in practice, it is hard to precisely decide whether e′(α) is positive or not, from the above discussion, at least we know that α should be selected from the range [ 12 , 1], providing some guidance to set α. Finally, our result is the first one that uses the same learning rate strategy to analyze both optimization and generalization errors, and matches the lower bounds of optimization error [53] and generalization error [45] of SGD.

4.2 Results on Convex Problems
Now we analyze lookahead using SGD as its inner optimizer A on the convex problems. Our main results are summarized in Theorem 3 with proof in Appendix D.4.
Theorem 3. Suppose that FS(θ) is convex, and each loss `(f(x;θ),y) is G-Lipschitz w.r.t. θ. By setting the inner learning rate η(t)τ = η, we have following properties. (1) The optimization error εopt of the output θA,S of lookahead satisfies εopt = EA,S [FS(θA,S) − FS(θ ∗ S)] ≤ ∆2αηkT + ηG2 2 , where θ ∗ S ∈ argminθ FS(θ) and ∆ = E [ ‖θ0 − θ∗S‖2 ] . (2) The generalization error εgen of the output θA,S obeys εgen = EA,S [F (θA,S)−FS(θA,S)] ≤ αηG2kT
n . (3) The excess risk error εexc of the output θA,S of lookahead satisfies εexc ≤ εopt + εgen.
To begin with, Theorem 3 guarantees the convergence of lookahead on the convex problems. By setting η = ∆ 0.5 α0.5k0.5T 0.5G , lookahead achieves the optimization error O ( G∆0.5 α0.5k0.5T 0.5 ) . Moreover, when α increases, the optimization error εopt decreases. This accords with the analysis results on the strongly convex problems that large α can reduce εopt. When α=1, lookahead degenerates to vanilla
SGD and its optimization error matches the one of vanilla SGD in [54] under the same assumptions. The second part of Theorem 3 shows that the generalization error εgen of lookahead is bounded by O ( αηG2kT
n
) . When α = 1, this error bound is consistent with the lower bound O ( kT n ) of SGD
with a constant learning rate η in [45]. Moreover, one can find that smaller α can lead to a smaller generalization error, which also accords with the results on the strongly convex problems.
Finally, by combining the optimization error εopt and generalization error εgen, we can bound the excess risk error εexc of lookahead by O
( ∆
2αηkT + ηG2 2 + αηG2kT n
) . So when fixing the learning rate,
tuning α ∈ (0, 1] can yield smaller excess risk error. It means that a proper α can benefit lookahead in terms of excess risk error, explaining the better test performance of lookahead than SGD.

4.3 Results on Nonconvex Problems
For general nonconvex problems, one often uses the gradient norm E[‖∇FS(θ)‖2] instead of the loss distance E[FS(θ)− FS(θ∗S)] to measure whether θ is a stationary point. This is because many stationary points may exist in a nonconvex problem. But as shown in Sec. 3, to bound the excess risk error, one needs to bound the loss distance. To solve this issue, we follow [46] and are particularly interested in nonconvex problems under PŁ condition which allows us to bound the loss distance, since PŁ condition in Definition 2 establishes the relation between gradient norm and loss distance. Moreover, as observed/proved in [30–34] and our empirical results in Sec. 6.2, deep learning models often satisfy PŁ condition. We summarize our main results in Theorem 4 with proof in Appendix D.5. Theorem 4. Assume each loss `(f(x;θ),y) is G-Lipschitz and L-smooth w.r.t. θ. Suppose that FS(θ) obeys the µ-PŁ condition. By setting η (t) τ = 1 tk+τ+1 and α > 1 2 , we have following properties. (1) The optimization error of the output θA,S produced by lookahead satisfies
εopt = EA,S [FS(θA,S)− FS(θ∗S)] ≤ 4∆′
(Tk + 1)2α + 2αLG2 (α+ 2(1− α)(k − 1)) µ2(Tk + 1)2α−1 ,
where θ∗S ∈ argminθ FS(θ) and ∆′ = E [FS(θ0)− FS(θ∗S)]. (2) The generalization error of the output θA,S produced by lookahead satisfies
εgen = EA,S [F (θA,S)− FS(θA,S)] ≤ ξ
n− 1 α
1 1+γ (Tk) γ γ+1 ,
where γ = (1− 1n ) αL µ and ξ = `
γ 1+γ max
[ 2G2
µ
] 1 1+γ
in which `max = maxθ,(x,y) `(f(x;θ),y). (3) The excess risk error εexc of the output θA,S of lookahead satisfies εexc ≤ εopt + εgen.
For the optimization error εopt of lookahead on nonconvex problems, it is of the order O ( 1 (Tk)2α + 1 µ2(Tk)2α−1 ) . The same as (strongly) convex problems, larger α benefits the convergence of lookahead, for which we have discussed the reasons in Sec. 4.1. When α=1, lookahead degenerates to SGD and achieves the smallest optimization error O ( 1
(Tk)2 + 1 µ2Tk
) which matches the one of SGD in [55].
For the generalization error εgen, it is of the order O ( 1 nα 1 1+γ (Tk) γ γ+1 )
which accords with the one of SGD in [44]. The sublinear dependence on kT and the inverse linear dependence on n also match the
Algorithm 2: Stagewise Locally-Regularized LookAhead (SLRLA)
Input :Loss FS(θ), constant {βq}Qq=1, inner-loop iteration number {kq} Q q=1 and learning rate
{ηq}Qq=1,outer-loop learning rate {αq} Q q=1, inner optimizer A, dataset S, initialization
θ0. for q = 1, 2, ..., Q do
Fq(θ) = FS(θ) + βq 2 ‖θ − θq−1‖ 2; θq = Look-ahead(Fq(θ), ηq, Tq, αq, kq,θq−1,A,S).
end Output :θA,S = θQ.
lower bound in [45]. Similarly, to achieve smaller generalization error, one should use small α. For the excess risk error εexc, it is bounded by O (εopt + εgen). So similar to (strongly) convex problems, when α is well chosen, the optimization error εopt and generalization error εgen can be balanced well, giving smaller excess risk error and better test performance than SGD which corresponds to α = 1.
Regarding the lookahead method with inner optimizers other than SGD, we believe that one could still expect similar performance trade-off between optimization and generation with respect to the choice of α. Intuitively, for any inner optimizer, let gtτ denote the “gradient" (or any descent direction) at the (t, τ)-th iteration. After the k inner-steps, lookahead updates parameter θ as θt = θt−1 + α(v (t) k − θt−1) = θt−1 −α ∑k−1 τ=0 η (t) τ g (t) τ . Obviously, α ≈ 1 is preferable for preserving the optimization speed of the inner optimizer. When it comes to the generalization error, obviously the best possible performance occurs at the initialization point as it is not dependent on the training data. Along with more training iterations, the network will gradually fit the training data and thus could give larger and larger prediction discrepancy between training data and test data. Therefore, it is desirable to have α 1 as opposed to α ≈ 1 for generalization. Overall, for generic inner-loop optimizers, lookahead is still expected to be able to balance the optimization and generalization performances with proper choices of α.

5 Stagewise Locally-Regularized Lookahead
We introduce the Stagewise Locally-Regularized LookAhead (SLRLA) algorithm in Algorithm 2. SLRLA divides the optimization into Q stages. For the q-th stage, it first combines the vanilla loss FS(θ) and a local regularizer βq 2 ‖θ−θq−1‖ 2 2 to construct a locally regularized loss Fq(θ)=FS(θ)+ βq 2 ‖θ−θq−1‖ 2 2. Here θq−1 is the output of the (q−1)-th stage, and βq is a constant. The intuition behind this local regularizer is that (i) it improves the convexity of the loss, e.g. converting an illconditioned loss to a well-conditioned one, accelerating convergence; (ii) it may avoid overfitting by preventing current solution θ having extreme values and far from θq−1. Theoretically, for λ-stronglyconvex problems, as shown in Theorems 1 & 2 that both optimization and generalization errors depend on O(1/λ). For nonconvex problems under µ-PŁ condition, Theorem 4 also shows that both optimization and generalization error scale with O(1/µ). So on strongly convex problems, adding a regularization on the vanilla loss can enhance the convexity and thus reduces both optimization and generalization errors; on nonconvex problems, regularization can also increase the PŁ condition parameter µ and thus helps optimization and generalization. We use ‖θ−θq−1‖22 instead of ‖θ− 0‖2 or ‖θ−θ0‖22. This is because compared with 0 and θ0, θq−1 is closer to the optimum θ∗S of FS(θ) and thus allows to use larger βq , improving the convexity of a loss more and benefiting convergence and generalization more (see discussion below). Next, SLRLA uses lookahead, i.e. Algorithm 1, to minimize Fq(θ), where constant inner- and outer-loop learning rates ηq and αq are used. A similar conventional stagewise optimization strategy which however does not regularize FS(θ) is widely used in SGD. But it is theoretically unclear whether conventional stagewise strategy benefits lookahead.
In the following sections, we investigate two questions: (i) whether the conventional stagewise strategy improves lookahead; (ii) what advantages SLRLA has over the conventional (stagewise) lookahead. For (i), we show the advantages of stagewise lookahead over vanilla lookahead in terms of the optimization error. For (ii), we prove that SLRLA can improve the optimization and generalization of conventional (stagewise) lookahead because of the local regularizer in SLRLA.

5.1 Results on Strongly Convex Problems
We analyze the optimization and generalization errors of SLRLA in Theorem 5 with proof in Appendix E.1, and then analyze the aforementioned two questions.
Theorem 5. Suppose that FS(θ) is λ-strongly convex, and each loss `(f(x;θ),y) is G-Lipschitz and L-smooth w.r.t. θ. By setting εq = ∆2q , βq = β ≤ 1 6λ, αq =α∈ [0, 1], ηq ≤ εq 3G2 , ηqkqTq ≥ 6 λαq
, ∆′=E[FS(θ0)−FS(θ∗S)] and θ∗S ∈ argminθ FS(θ), the following properties hold. (1) The optimization error εopt of the output θA,S of SLRLA satisfies εopt≤ ∆ ′
2Q . Moreover, to achieve
εopt≤ , Q satisfies Q ≥ log ∆ ′ and the stochastic gradient complexity is ∑Q q=1Tqkq= 36G2 λα .
(2) The generalization error εgen of the output θA,S of SLRLA satisfies εgen ≤ γ1n ( β α + λL λ+L )−1( 1−
exp ( − 6QLλ+L )) , where γ1 = 2G2/ ( 1− exp ( − 6Lλ+L )) . (3) The excess risk error εexc of the output θA,S of lookahead satisfies εexc ≤ εopt + εgen.
By Theorem 5, we here show the advantages of the Stagewise LookAhead (SLA for short) [1] over lookahead, and also discuss the superiority of SLRLA over SLA. When βq=0 and the learning rate is geometrically decayed as SLRLA after each stage, SLRLA degenerates to SLA, and Theorem 5 still holds. Theorem 5 shows the linear convergence of both SLRLA and SLA w.r.t. the stage number Q. For both SLRLA and SLA, to achieve optimization error , i.e. εopt =E[FS(θA,S)−FS(θ∗S)]≤ , Q is of the orderO ( log 1 ) and the stochastic gradient complexity (evaluation number) is ∑Q q=1 Tqkq=
O ( G2
λα
) . With the optimization error , Theorem 1 shows that vanilla lookahead needs stochastic
gradient complexity O (( L ) 1 2α + ( LG2 (1−2α)λ2 ) 1 2α ) for α ∈ (0, 12 ), O ( LG2 log 1 λ2 ) for α = 12 , and
O ( LG2
(2α−1)λ2 ) for α∈( 12 , 1]. By comparison, both SLA and SLRLA improve the dependences on two important factors, i.e. λ and α, in vanilla lookahead. Specifically, for factor λ, SLA and SLRLA rely on O ( 1 λ ) , while lookahead depends on at least O ( 1 λ2 ) . For factor α, SLA and SLRLA only
linearly depend on O (
1 α ) . In contrast, when α ∈ (0, 12 ), lookahead exponentially depends on 1 α due
to the term O (( 1 ) 1 2α ) . For α = 12 , SLA and SLRLA improve lookahead by removing the logarithm factor log 1 . When α ∈ ( 1 2 , 1], SLA and SLRLA also enjoy smaller dependence on α than lookahead, as the factor 1α (≤ 2) in SLA and SLRLA is often smaller than the factor 1 2α−1 in lookahead.
For generalization error, by comparing Theorems 5 and 2, SLA and lookahead enjoy the same generalization error boundO
( G2
nλ
) , while SLRLA has superior oneO ( G2
n(β/α+λ)
) especially for small
α. This is because β can be at the same order of λ in Theorem 5. By combining the optimization and generalization errors together, SLRLA enjoys smaller excess risk error than SLA which however outperforms lookahead, when the computational budget (stochastic gradient complexity) is the same.

5.2 Results on Nonconvex Problems
Now we analyze SLRLA on two classes of nonconvex problems. The first one requires the smallest
eigenvalue (−σ) of the Hessian ∇2FS(θ) to satisfy σ ≤ 16µ, where µ is the parameter in µ-PŁ condition in Definition 2. It actually means the function is not heavily nonconvex. The second one satisfies the weakly-quasi-convex assumption which guarantees that any local minimizer of the loss is also a global minimizer. Indeed, any convex function is 1-weakly-quasi-convex, but the converse is generally not true. We use these two classes of nonconvex problems as examples to investigate the aforementioned two problems. Theorem 6 with proof in Appendix E.2 summarizes the main results.
Theorem 6. Assume each `(f(x;θ),y) is G-Lipschitz and L-smooth w.r.t. θ, and FS(θ) satisfies µ-PŁ condition. Let ∆′ = E[F (θ0)− F (θ∗S)] and εq = ∆ ′ 2q , where θ ∗ S ∈ argminθ FS(θ). (1) When σ ≤ 16µ where∇ 2FS(θ) −σI(∀θ), by setting σ ≤ βq = β ≤ 16µ, ηq ≤ εq 3G2 , ηqkqTq ≥
6 µαq , the output θA,S of SLRLA satisfies
εopt ≤ ∆′
2Q , εgen ≤ γ2 n (β − σ α + µL µ+ L )−1( 1− exp ( − 6QL λ+ L )) , εexc ≤ εopt + εgen,
where γ2 = 2G2/
(
1 − exp
(
− 6Lµ+L
))
. Moreover, to achieve εopt ≤ , the total stage number Q
satisfies Q ≥ log ∆ and the total stochastic gradient complexity is ∑Q q=1 Tqkq = 36G2
µα . (2) When Fq(θ) is ρ-weakly-quasi-convex, by setting βq = β ≤ 16µ, ηq ≤ ρεq 3G2 , ηqkqTq ≥ 6 µραq , α ≤ βL , the output θA,S of SLRLA satisfies
εopt ≤ ∆
2Q , εgen ≤ γ2 n (β α − L )−1( 1− exp ( − 6S(β − αL) µρα )) , εexc ≤ εopt + εgen,
where γ2 = 2G2/ ( 1− exp ( − 6(β−αL)µρα ))) . Moreover, to achieve εopt ≤ , the total stage number Q
satisfies Q ≥ log ∆ and the total stochastic gradient complexity is ∑Q q=1 Tqkq = 36G2 µρ2α .
Here we discuss the advantages of SLRLA over SLA and lookahead. Theorem 6 shows that for the nonconvex loss FS(θ) which obeys∇2FS(θ) − 16µI , SLRLA enjoys linear convergence rate and has stochastic gradient complexity O ( G2
µα
) to achieve εopt ≤ . Compared with lookahead
which has complexity O (( LG2
µ2
)1/α) , SLRLA improves the factor 1
µ2/α in lookahead to 1µ , and also
reduces its exponential dependence on 1α to linear dependence. These improvements accord with the ones on the strongly convex problems. For generalization error εgen, SLRLA has an upper bound O ( 1 n/ ( β−σ α +µ )) that does not reply on the total iteration number Tk. In contrast, Theorem 4 shows
that lookahead has generalization error bound O (
1 n (Tk)
γ γ+1 )
sublinearly relying on Tk. So SLRLA enjoys smaller excess risk error than lookahead. For SLA, Theorem 6 cannot guarantee its linear convergence and small generalization error, as Theorem 6 requires βq≥σ>0 but SLA needs βs=0. It is because the regularizer βq2 ‖θ−θq−1‖
2 in SLRLA transforms the nonconvex loss FS(θ) into a strongly convex one and greatly accelerates the convergence. This shows the advantage of SLRLA over SLA. On weakly-quasi-convex problems, SLRLA has stochastic gradient complexity O ( G2
µρ2α
) to achieve
εopt ≤ , and has generalization error O ( γ2 n / ( β α − L )) . When viewing ρ as a constant, then same as the above case, SLRLA makes improvements on vanilla lookahead in terms of both stochastic complexity and generalization error. Moreover, the results on SLRLA also do not hold for SLA, since Theorem 6 needs βq ≥ αL > 0 while SLA sets βq = 0. Thus, all these results show the advantages of SLRLA over SLA and lookahead on achieving smaller excess risk error and thus enjoying better test performance. Limitation Discussion. As explained in Sec. 4.3, one cannot bound the loss distance E[FS(θ) − FS(θ ∗ S)] for general nonconvex probems (GNP). So our main limitation is that our analysis is not applicable to GNP. But we analyze lookahead on GNP under PŁ condition which is observed/proved for deep networks [30–34] and our empirical results. See more discussions in Appendix A.
Societal Impact. This work analyzes the intrinsic theoretical reasons for the superiority of lookahead in terms of test performance, and further proposes a general and more advanced deep learning optimizer with provable improvement over lookahead which could advance deep network training. For the negative social impact of this work, it is mainly determined by which applications the developed optimizer is applied to.
6 Experiments

6.1 Results on Strongly Convex Problems
Here we investigate the effects of α on the performance of lookahead, stagewise lookahead [1] (SLA) and SLRLA on a regularized softmax problem with MNIST [56]. The regularized softmax problem is strongly convex, as its regularization constant is set to 5×10−6. Following our theory, we use a linearly decayed learning rate (LR) for lookahead, and multi-step decayed LRs for SLA/SLRLA. See more details in Appendix B. Fig. 1 (a) shows that when α increases, for lookahead & SLA, (i) their training loss reflecting the optimization error εopt decreases faster; (ii) the distance between their training and test losses reflecting the generalization error εgen becomes larger; (iii) their test loss first decreases and then increases which indicates the balance of α. Fig. 1 (b) shows the balance impact of α on test accuracy of SLRLA, and also testifies the superiority of SLRLA over lookahead and SLA.
6.2 Results on Nonconvex Problems
Assumption Investigation. We investigate the key assumption, PŁ assumption on the nonconvex problems, for networks. We train ResNet18 [4] and wide-ReseNet-16-10 (WRN) [57] via SLA/SLARA, and report µ , ‖∇FS(θt)‖2/[2(FS(θt)− FS(θ
∗))]. We estimate the optimum θ∗ of FS(θ) as the solution found by SLA/SLARA after 200 epochs which gives a small objective value (≈10−3) and is almost an optimum. Fig. 2 shows that µ in SLA & SLARA is larger than 5×10−4. So on networks, PŁ condition holds at least along the optimization trajectory. This accords with the observations/theories in [30–34].
Classification Results. We evaluate SLA and SLRLA on CIFAR10/100 [58] and ImageNet [59] using different network architectures, i.e. ResNet18 [4], VGG16 [60] and WRN-16-10 [57]. For all experiments, SLRLA and SLA set k=5, a momentum of 0.9, and a multi-stage learning rate (LR) decay at the {0.3S, 0.6S, 0.8S}-th epoch with total epoch number S. On CIFAR10/100, we train 200 epochs with α= 0.8, a weight decay of 10−3, and set LR decay rate as 0.2. On Imagenet, we run 100 epochs using α=0.5, a weight decay of 10−4 and an LR decay rate of 0.1. These settings follow [1, 61, 62]. For regularization constant βq, SLRLA selects it from {0.02, 0.2, 2.0, 20} via cross validation, and finally sets it as 0.2 on CIFAR10/100 and 20 on ImageNet.
Table 1 reports the average accuracy and variance of 5 random seeds. SLRLA achieves the highest accuracy on CIFAR10/100 and ImageNet. This because (i) our theories in Sec. 5 show the advantages of SLRLA over SGD/SLA to achieve small excess risk error, indicating better test performance; (ii) as observed/proved in [22, 37–41], compared with adaptive algorithms, e.g. Adam and its variants, SGD-alike algorithms, e.g. SLA/SLARA, often converge to flatter minima and thus generalize better. The results in Appendix B show the stable performance of SLARA on ImageNet when tuning the regularization constant βq in a large range, testifying the robustness of SLARA.

7 Conclusion
In this work, for the first time we theoretically show the advantages of lookahead in terms of the excess risk error, explaining its better test performance than its vanilla inner optimizer. Moreover, we prove that the stagewise optimization strategy can benefit lookahead in improving its excess risk error. Finally, we propose SLRLA which locally regularizes the vanilla objective to further improve the excess risk error of stagewise lookahead. Experimental results validated the advantages of SLRLA.

Acknowledgements
The authors sincerely thank the anonymous reviewers for their constructive comments on this work.
1. Funding. Pan Zhou, Hanshu Yan, Jiashi Feng and Shuicheng Yan are supported by Sea AI Lab, mainly for their GPU resource support. Xiao-Tong Yuan is supported in part by the National Key Research and Development Program of China under Grant No. 2018AAA0100400 and in part by the Natural Science Foundation of China (NSFC) under Grant No.61876090 and No.61936005.
2. Competing Interests. Pan Zhou, Jiashi Feng and Shuicheng Yan are staffs in Sea AI Lab. Hanshu Yan is an intern in Sea AI Lab, and is also a Ph.D. student in National University of Singapore. Xiao-Tong Yuan works as a professor in Nanjing University of Information Science & Technology, Nanjing, China.

Limitation
Limitation Discussion. As explained in Sec. 4.3, one cannot bound the loss distance E[FS(θ) − FS(θ ∗ S)] for general nonconvex probems (GNP). So our main limitation is that our analysis is not applicable to GNP. But we analyze lookahead on GNP under PŁ condition which is observed/proved for deep networks [30–34] and our empirical results. See more discussions in Appendix A. Societal Impact. This work analyzes the intrinsic theoretical reasons for the superiority of lookahead in terms of test performance, and further proposes a general and more advanced deep learning optimizer with provable improvement over lookahead which could advance deep network training. For the negative social impact of this work, it is mainly determined by which applications the developed optimizer is applied to. 6 Experiments

Limitations (refined)

Reviewer Summary
Reviewer_1: The authors analyze the Lookahead optimizer (Zhang et al.) algorithm, showing that it has nice bounds on the excess risk error compared to SGD on strongly convex problems and noncovex problems which satisfy the PL condition. They also propose a locally-regularized version of the Lookahead algorithm which has some better theoretical properties and performs well on empirical evaluations.

Reviewer_2: Lookahead is a relatively recent algorithm that has shown great performances in both single-objective and two-objective training, but it is not well understood theoretically why it improves training particularly in the former case. This paper aims to elaborate why Lookahead improves test accuracy using excess risk error which is decomposed into optimization and generalization terms. Interestingly, the authors show that the extra hyperparameters in Lookahead when combined with SGD allow for a better trade-off between the two terms relative to SGD. Moreover, they propose a new "stagewise" variant algorithm that: (i) uses constant step sizes for Lookahead, but changes/decreases these at each stage, and (ii) adds regularization term that enforces doing more "conservative" steps (the term penalizes high distance between predicted iterate and last iterate of the previous stage, see Alg.2).

Reviewer_3: The paper provides the excessive risk rounds for the lookahead algorithms on (strongly) convex and non-convex problems with PL condition. The theoretical results are obtained via (1) convergence analysis (2) the uniform stability analysis. From the excessive risk bounds, the authors provide justification for the performance of the lookahead algorithms. Finally, the paper proposes a stagewise locally-regularized lookahead algorithm called SLRLA, which adds a L2 regularizer and decays the learning rate for each stage.

Reviewer Comments
Reviewer_1: Strengths
The theoretical results are novel and useful for understanding the lookahead algorithm and how it compares to SGD. Theorem 1 provides bounds on the optimization error of lookahead(SGD); Theorem 2 bounds the generalization error.
The authors do a good job of offering intuition for the their theoretical results and discussing implications.
Stagewise locally-regularized lookahead is well-motivated and the proof is technically sound. It is impressive to improve on an existing, popular algorithm.
While most of the results are theoretical, the empirical evaluations are strong and performance on CIFAR and ImageNet exceeds the performance on the original Lookahead algorithm across 5 seeds.
Weaknesses
Writing can be a bit choppy at points and would benefit from careful proofreading.
Questions
Does stagewise lookahead refer to conventional lookahead (Zhang et al.) with a stagewise learning rate (geometric) decay? If so, I believe it would be better to make this point more explicit and mention that it is the learning rate schedule typically used.
Could you say anything about lookahead with other inner optimizers?
Minor
83: the sharp minima story is a bit more nuanced (as mentioned in some of your references) or "Sharp Minima Can Generalize For Deep Nets" (Dinh et al.)
86: typo in lookahead
In algorithm 1, it would improve clarity if the order of the description of the inputs matches that of the function arguments
Post-rebuttal: I have read the other reviewer comments and responses and am maintaining my score.
Limitations And Societal Impact:
Limitations are sufficiently discussed throughout the paper and in the appendix. The authors do not directly address societal impact ("This is not applicable to a theoretical work"), though I believe it is worth some further thought e.g. see the AdaBelief paper.
Needs Ethics Review: No
Time Spent Reviewing: 9

Reviewer_2: In my opinion, this paper does a nice balance between theoretical discussions and experiments, and I find it very well written. Regarding the latter, the authors introduce the necessary definitions, list theoretical results, but it is nice that they also discuss the results.
While I enjoyed reading the first part discussing the original Lookahead (LA), I have several questions/concerns for the newly proposed algorithm called SLRLA.
--- Questions ---
Following your analysis on LA, could you elaborate on how SLRLA is motivated? While I understand that given its formulation you obtain a small excess risk error, my question is towards understanding how your LA analysis guided finding new algorithms, and it might be useful adding such a comment to also connect the two parts.
Is introducing the stage-wise step-size strategy necessary for SLRLA, or alternatively, is the excess risk error worse in that case? In contrast, if I understand well, for LA you used decaying learning rate in the analysis. I am not sure I understand, could a stage-wise step-size strategy provide better excess risk error for LA than the one provided?
While SLRLA has better excess risk error, it also introduces two additional hyper-parameters making it harder to tune (
Q
and
β
q
). Could you comment from the analysis perspective whether the obtained results require optimal hyperparameters? Also, is
β
q
easy to select in your experiments, and do the performances vary across the values you used in the cross-validation?
Regarding the empirical setup, seems like only the
β
q
hyperparameter of SLRLA is cross-validated, and the rest are fixed. This might give an advantage to SLRLA.
Figure 1. Surprisingly, LA in (b) performs poorly (for MNIST). How are the hyperparameters selected for it, and what are the performances of SGD? Also, as test accuracy is not used in both a.3 and b, it is harder to compare the performances of SLRLA and SLA with varying hyperparameters.
As a comment, in my opinion, the nested version of Lookahead proposed in [21] -- which applies LA recursively -- is related to SLRLA as both push the current iterate toward a past "anchored" iterate, but they are yet different methods. It would be interesting to explore how the two differ (in terms of the error trade-off).
--- Minor/typos ---
ln.67: has->have
Fig.2: maybe ensure legend text does not overlap much, and improve caption
--- Recommendation ---
In summary, I find the result/analysis on the original Lookahead interesting and intuitive, and in my opinion, the paper is well-written. In terms of novelty, (i) the analysis used in this paper is not novel, (ii) while the paper shows how LA hyperparameters influence the trade-off between the two errors terms, it is not clear how these can be selected -- thus several questions still remain open for LA (how to select the hyperparameters of LA in practice and deriving tighter rates for it). However, I would be glad to raise my score if the authors address some of the above questions.
--- Update after discussion ---
The authors addressed most of my concerns, and I increased my score.
Limitations And Societal Impact:
Yes
Needs Ethics Review: No
Time Spent Reviewing: 7

Reviewer_3: Originality
The convergence analysis and the uniform stability analysis of the lookahead algorithms are new and interesting.
Quality
The theoretical part of the paper is sound.
The authors provide convergence analysis for lookahead algorithms under different settings, e.g., (strongly) convex and PL-condition. For the generalization part, the authors derive the uniform stability bounds for lookahead algorithms. By combing two parts, the authors show that there is an optimal choice for
α
which balances optimization error and generalization error.
The empirical result is strong.
The authors test their proposed algorithm in three commonly used datasets. They show that SLRLA is better than other baseline algorithms.
Some concerns for the theory part:
One caveat of the uniform stability analysis is that for general convex problem (not strongly convex), the generalization depends badly on
T
. Therefore if we calculate the optimal
α
we found that
α
scales with
1
/
T
, which means for small excessive risk we might need to restrict the effective learning rate.
Line 163: In [41], actually the optimal rate for SGD is
O
(
1
/
λ
T
)
. Also in Line 291: The dependency for
λ
is optimal for SLA and SLRLA, but I think the
O
(
1
/
λ
2
)
bound for lookahead algorithms is due to your analysis techniques. The comparison may be unfair.
Clarity
The paper is well-written.
Limitations And Societal Impact:
Yes.
Needs Ethics Review: No
Time Spent Reviewing: 1.5