A Near-Optimal Algorithm for Stochastic Bilevel Optimization via Double-Momentum

1 Introduction . Many learning and inference problems take a “ hierarchical '' form , wherein the optimal solution of one problem affects the objective function of others [ 27 ] . Bilevel optimization is often used to model problems of this kind with two levels of hierarchy [ 27 , 8 ] , where the variables of an upper level problem depend on the optimizer of certain lower level problem . In this work , we consider unconstrained bilevel optimization problems of the form : minx2Rdup ` ( x ) = f ( x , y ⇤ ( x ) ) : = E⇠ [ f ( x , y⇤ ( x ) ; ⇠ ) ] s.t . y⇤ ( x ) = argminy2Rdlo g ( x , y ) : = E⇣ [ g ( x , y ; ⇣ ) ] , ( 1 ) where f , g : Rdup ⇥ Rdlo ! R with x 2 Rdup and y 2 Rdlo ; f ( x , y ; ⇠ ) with ⇠ ⇠ ⇡f ( resp . g ( x , y ; ⇣ ) with ⇣ ⇠ ⇡g ) represents a stochastic sample of the upper level objective ( resp . lower level objective ) . Note here that the upper level objective f depends on the minimizer of the lower level objective g , and we refer to ` ( x ) as the outer function . Throughout this paper , g ( x , y ) is assumed to be strongly-convex in y , which implies that ` ( x ) is smooth but possibly non-convex . The applications of ( 1 ) include many machine learning problems that have a hierarchical structure . Examples are meta learning [ 13 , 31 ] , data hyper-cleaning [ 35 ] , hyper-parameter optimization [ 12 , ⇤Corresponding Author : Mingyi Hong . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 13 , 29 ] , and reinforcement learning [ 22 ] , etc .. To better contextualize our study , below we describe examples on meta-learning problem and data hyper-cleaning problem : Example 1 : Meta learning . The meta learning problem aims to learn task specific parameters that generalize to a diverse set of tasks [ 30 ] . Suppose we have M tasks { Ti , i = 1 , . . . , M } and each task has a corresponding loss function L ( x , yi ; ⇠i ) with ⇠i representing a data sample for task Ti , x 2 Rdup the model parameters shared among tasks , and yi 2 Rd i lo the task specific parameters . The goal of meta learning is then to solve the following problem : minx2Rdup Lts ( x , ȳ⇤ ( x ) ) : = 1 M PM i=1 E⇠i⇠Di [ L ( x , y⇤i ( x ) ; ⇠i ) ] s.t . ȳ⇤ ( x ) 2 argmin ȳ2R PM i=1 d i lo Ltr ( x , ȳ ) : = 1 M PM i=1 E⇣i⇠Si [ L ( x , yi ; ⇣i ) ] +R ( yi ) , ( 2 ) where ȳ = [ yT1 , . . . , yTM ] T , R ( · ) is a strongly convex regularizer while Si and Di are the training and testing datasets for task Ti . Compared to the number of tasks , the dataset sizes are usually small for meta-learning problems , so the stochasticity in tackling ( 2 ) results from the fact that at each iteration we can only sample a subset m out of M tasks . Note that this problem is a special case of ( 1 ) . ⇤ Example 2 : Data hyper-cleaning . The data hyper-cleaning is a hyperparameter optimization problem that aims to train a classifier model with a dataset of randomly corrupted labels [ 35 ] . The optimization problem is formulated below : minx2Rdup ` ( x ) : = P i2Dval L ( a > i y ⇤ ( x ) , bi ) ( 3 ) s.t . y⇤ ( x ) = argminy2Rdlo ckyk 2 + P i2Dtr ( xi ) L ( a > i y , bi ) . In this problem , we have dup = |Dtr| and dlo is the dimension of the classifier . Moreover , ( ai , bi ) is the ith data point ; L ( · ) is the loss function , with y being the model parameter ; xi is the parameter that determines the weight for the ith data sample , and : R ! R+ is the weight function ; c > 0 is a regularization parameter ; Dval and Dtr are validation and training sets , respectively . Clearly , ( 3 ) is a special case of ( 1 ) where the lower level problem finds the classifier y⇤ ( x ) with the training set Dtr , and the upper level problem finds the best weights x with respect to the validation set Dval . ⇤ A natural approach to tackling ( 1 ) is to apply alternating stochastic gradient ( SG ) updates . Let , ↵ > 0 be some step sizes , one performs the recursion y + y r̂yg ( x , y ) , x + x ↵r̂x ˆ̀ ( x ; y ) ( 4 ) such that r̂yg ( x , y ) , r̂x ˆ̀ ( x ; y ) are stochastic estimates of ryg ( x , y ) , r ` ( x ) , respectively . Notice that ( 4 ) is significantly different from the standard alternating primal-dual gradient algorithm for saddle point problems . Particularly , the design of r̂x ˆ̀ ( x ; y ) is crucial to the SG scheme in ( 4 ) . Observe that r ` ( x ) can be computed using the implicit function theorem , and its evaluation requires f ( · , · ) and y ? ( x ) , the minimizer of g ( x , y ) given x ( cf . ( 5 ) ) . This gives rise to a unique challenge to bilevel optimization , where y ? ( x ) can only be approximated by y obtained in the first relation of ( 4 ) . In light of the above observations , previous endeavors have considered two approaches to improve the estimate of y ? ( x ) while r̂x ˆ̀ ( x ; y ) is used as a biased approximation of r ` ( x ) . The first approach is to apply the double-loop algorithms . For example , [ 14 ] proposed to repeat the y+ update for multiple times to obtain a better estimate of y ? ( x ) before performing the x+ update , [ 19 ] proposed to take a large batch size to estimate ryg ( x , y ) . While simple to analyze , these algorithms may suffer from a poor sample complexity for the inner problem . The second approach is to apply single-loop algorithms where the y+-updates are performed simultaneously with the x+-updates . Instead , advanced techniques are utilized that allows y+ to accurately track y ? ( x ) . For example , [ 18 ] suggested to tune the step size schedule with ↵ , [ 5 , 17 ] proposed single-timescale algorithms with advanced variance reduction techniques . However , the latter two algorithms require Hessian projections onto a compact set along with Hessian matrices inversion which scales poorly with dimension ( i.e. , in O ( d3lo ) ) . We summarize and compare the complexity results of the state-of-the-art algorithms in Table 1 . A careful inspection on the above results reveals a gap in the iteration/sample complexity compared to single-level stochastic optimization . For instance , an optimal stochastic gradient algorithm finds an ✏-stationary solution [ cf . Definition 1.1 ] to minx E⇠ [ ` ( x ; ⇠ ) ] in O ( ✏ 3/2 ) iterations [ 10 , 7 , 37 , 42 ] . For bilevel optimization , the fastest rate available is only O ( ✏ 2 ) to the best of the authors ’ knowledge . In comparison , the proposed algorithm achieves a rate of O ( ✏ 3/2 ) . During the preparation of the current paper , a preprint [ 17 ] has appeared which extended [ 5 ] , and achieves an improved rate of O ( ✏ 3/2 ) . We remark that the latter work follows a different design philosophy from ours and maybe less efficient ; see the detailed discussion at the end of Sec . 3 . Contributions . In this paper , we depart from the prior developments which focused on finding better inner solutions y⇤ ( x ) to approximate r̂x ˆ̀ ( x ; y ) ⇡ r ` ( x ) . Our idea is to exploit the gradient estimates from prior iterations to improve the quality of the current gradient estimation . This leads to momentum-assisted stochastic gradient estimators for both ryg ( x , y ) and r ` ( x ) using similar techniques in [ 7 , 37 ] for single-level stochastic optimization . The resultant algorithm only requires O ( 1 ) samples at each update , and updates x and y using step sizes of the same order , hence the name single-timescale double-momentum stochastic approximation ( SUSTAIN ) algorithm . Additionally , it is worth noting that our algorithm has a O ( d2lo ) per iteration complexity , compared to the O ( d 3 lo ) complexity of STABLE [ 5 ] and SVRB [ 17 ] . That is , the SUSTAIN algorithm is both sample and computation efficient . Our specific contributions are : • We propose the SUSTAIN algorithm for bilevel problems which matches the best complexity bounds as the optimal SGD algorithms for single-level stochastic optimization . That is , it requires O ( ✏ 3/2 ) [ resp . O ( ✏ 1 ) ] samples to find an ✏-stationary solution for non-convex ( resp . stronglyconvex ) bilevel problems ; see Table 1 . Furthermore , the algorithm utilizes a single-loop update with step sizes of the same order for both upper and lower level problems . Such complexity bounds match the optimal sample complexity of stochastic gradient algorithms for single-level problems . • By developing the Lipschitz continuous property of the ( biased ) stochastic estimates of r ` ( x ) , we show that obtaining a good estimate of r ` ( x ) does not require explicit ( sampled ) Hessian inversion . This key result ensures that our algorithm depends favorably on the problem dimension . • Comparing with prior works such as TTSA [ 18 ] , BSA [ 14 ] , STABLE [ 5 ] and SVRB [ 17 ] , our analysis reveals that improving the gradient estimation quality for bothryg ( x , y ) and r ` ( x ) is the key to obtain a sample and computation efficient stochastic algorithm for bilevel optimization . Related works . The study of the bilevel problem ( 1 ) can be traced to that of game theory [ 36 ] and was formally introduced in [ 2–4 ] . It is also related to the broader class of problems of Mathematical Programming with Equilibrium Constraints [ 26 ] . Related algorithms include approximate descent [ 9 , 38 ] , and penalty-based methods [ 40 ] ; see [ 6 ] and [ 25 ] for a comprehensive survey . In addition to the works cited in Table 1 , recent works on bilevel optimization have focused on algorithms with provable convergence rates . In [ 34 ] , the authors proposed BigSAM algorithm for solving simple bilevel problems ( with a single variable ) with convex lower level problem . Subsequently , the works [ 24 , 23 ] utilized BigSAM and developed algorithms for a general bilevel problem for the cases when the solution of the lower level problem is not a singleton . Note that all the aforementioned works [ 34 , 24 , 23 ] assumed the upper level problem to be strongly-convex with convex lower level problem . In a separate line of work , backpropagation based algorithms have been proposed to approximately solve bilevel problems [ 12 , 35 , 16 , 15 ] . However , the major focus of these works was to develop efficient gradient estimators rather than on developing efficient optimization algorithms . Notation . For any x 2 Rd , we denote kxk as the standard Euclidean norm ; as for X 2 Rn⇥d , kXk is induced by the Euclidean norm . For a multivariate function f ( x , y ) , the notationrxf ( x , y ) [ resp.ryf ( x , y ) ] refers to the partial gradient taken with respect to ( w.r.t . ) x [ resp . y ] . For some µ > 0 , a function f ( x , y ) is said to be µ-strongly-convex in x if f ( x , y ) µ2 kxk 2 is convex in x . For some L > 0 , the map A : Rd ! Rm is said to be L-Lipschitz continuous if kA ( x ) A ( y ) k Lkx yk for any x , y 2 Rd . A function f : Rd ! R is said to be L-smooth if its gradient is L-Lipschitz continuous . Uniform distribution over a discrete set { 1 , . . . , T } is represented by U { 1 , . . . , T } . Finally , we state the following definitions for the optimality criteria of ( 1 ) . Definition 1.1 ( ✏-Stationary Point ) . A point x is called ✏-stationary if kr ` ( x ) k2 ✏ . A stochastic algorithm is said to achieve an ✏-stationary point in t iterations if E [ kr ` ( xt ) k2 ] ✏ , where the expectation is over the stochasticity of the algorithm until time instant t. Definition 1.2 ( ✏-Optimal Point ) . A point x is called ✏-optimal if ` ( x ) ` ⇤ ✏ , where ` ⇤ : = minx2Rdup ` ( x ) . A stochastic algorithm is said to achieve an ✏-optimal point in t iterations if E [ ` ( xt ) ` ⇤ ] ✏ , where the expectation is over the stochasticity of the algorithm until time instant t .

Summary: The work developed a new algorithm for unconstrained bilevel optimization with strongly convex lower level subproblems. The proposed algorithm executes on a single-timescale, without the need to use either two-timescale updates, large batch gradients, or double-loop algorithm. Authors showed that SUSTAIN is both sample and computation efficient, because it matches the best-known sample complexity guarantees on single-level problems with non-convex and strongly convex objective functions, while matching the best-known per-iteration computational complexity for the same class of bi-level problems.