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: This paper investigates a bi-level stochastic optimization problem in which the upper level function depends on the minimizer of the lower level objective. A new algorithm based on stochastic momentum-assisted gradient estimator is introduced, by exploiting prior iterations gradient estimates to improve the quality of the current gradient estimation (by reducing its variance). The new algorithm matches the best complexity bounds as the optimal SGD algorithms for single-level optimization.