Towards Understanding Why Lookahead Generalizes Better Than SGD and Beyond

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 can not 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 .

Summary: 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).