Towards Gradient-based Bilevel Optimization with Non-convex Followers and Beyond

1 Introduction . Bi-Level Optimization ( BLO ) has been widely used to formulate problems in the field of deep learning [ 1 , 2 ] , especially for hyperparameter optimization [ 3 , 4 , 5 ] , meta learning [ 6 , 7 , 8 , 9 ] , neural architecture search [ 10 , 11 , 12 ] , adversarial learning [ 13 ] , and reinforcement learning [ 14 ] , etc . BLO aims to tackle nested optimization structures appearing in applications , which has emerged as a prevailing optimization technique for modern machine learning tasks with underlying hierarchy . In the last decade , a large number of BLO methods have been proposed to address different machine learning tasks . In fact , Gradient Methods ( GMs ) , which can effectively handle BLO problems of large scale , thus gain popularity . According to different types of strategies for gradient calculations , existing GMs can be divided into two categories , i.e. , the explicit approaches which aim to replace the Lower-Level ( LL ) problem with dynamic iterations and implicit schemes that apply the implicit function theorem to formulate the first-order optimality condition of the LL problem . Explicit Gradient Methods for BLOs . In this type , solving the LL problem is regarded as the evolution path of the dynamic system starting from a given initial point of the LL variable . The gradient of the Upper-Level ( UL ) variable can be directly calculated by automatic differentiation * Corresponding author 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . based on the trajectory of LL variable . This class of methods can be further divided into three types , namely , recurrence-based EG ( e.g. , [ 3 , 15 , 6 , 16 , 10 ] ) , initialization-based EG ( e.g. , [ 17 , 18 ] ) and proxy-based EG methods ( e.g. , [ 19 , 20 , 21 , 22 ] ) , differing from each other in the way of accessing the gradient of constructed dynamic trajectory . While most of this type of works assume the LLC and Lower-Level Singleton ( LLS ) to simplify their optimization process and theoretical analysis , cases where the LLS assumption does not hold have been tackled in the recent work [ 23 , 24 ] . In particular , to eliminate the LLS assumption which is too restrctive to be satisfied in real-world complex tasks , [ 23 ] first considers incorporating UL objective information into the dynamic iterations , but the more general cases where LLC does not hold remain unsolved . On the other hand , while most of the mentioned works focus on the asymptotic convergence , the progress on nonasymptotic convergence analysis has been recently witnessed see , e.g. , [ 25 , 26 , 27 ] . Implicit Gradient Methods for BLOs . This type , also known as implicit differentiation [ 28 , 7 , 29 ] , replaces the LL problem with its first-order optimality condition and uses the implicit function theorem to calculate the gradient of the UL problem by solving a linear system . This method decouples the calculation of UL gradient from the dynamic system of LL , resulting in a significant speed increase when the dynamic system iterates many times . However , because of the burden originated from computing a Hessian matrix and its inverse , IG methods are usually computationally expensive when linear systems are ill-conditioned . To alleviate this computational issue , there are mainly two kinds of techniques , i.e. , IG based on Linear System ( LS ) [ 28 , 7 ] and Neumann Series ( NS ) [ 29 ] . On the theoretical side , IG methods rely on the strong convexity of LL problems heavily , which is even more restrictive than the LLC and LLS together . Initialization Optimization for Learning . In deep learning , the selection of initialization scheme has a great influence on the training speed and performance [ 30 ] . As the most representative work in recent years , Model-Agnostic Meta-Learning ( MAML ) [ 31 ] applies the same initialization to all tasks , and is optimized by a loss function common to the task that evaluates the effect of the initial value , resulting in an initialization that achieves good generalization performance with only a few gradient steps on new tasks . Due to its simple form this method has been widely studied and applied [ 32 , 33 , 34 , 35 ] . [ 36 ] noticed that not all network parameters are suitable for the same initialization , and therefore proposed a strategy to apply co-initialization only on a part of parameters . In theory , [ 37 , 38 , 39 ] give comprehensive study on the convergence and convergence rate of MAML and some MAML-type approaches based on the meta objective function . However , the convergence theory of these existing results are given based on the loss function for evaluating initial values , and the convergence analysis of such type of methods from the perspective of each task is still lacked . Value-Function Approach . The value function based methods have also emerged as a promising branch to solve BLO problems [ 40 ] . Under the special case where the LL is jointly convex with respect to both the UL and LL variables , the BLO problem can be equivalently reformulated into a difference-of-convex program [ 41 ] , which is numerically solvable . Typically , by reformulating the BLO into an Inner Single Bi-level ( ISB ) optimization problem with value-function approach , a gradient-based interior-point method name BVFIM [ 42 ] is proposed to solve the BLO tasks , which effectively avoids the expensive Hessian-vector and Jacobian-vector products . Generally speaking , the value-function does not admit an explicit form , and is always nonsmooth , non-convex and with jumps . 1.1 Our Motivations and Contributions . As mentioned above , some theoretical progresses have long been witnessed in diversified learning areas , but for most existing BLO methods , extra restrictive assumptions ( e.g. , LLS , LLC and LL strong convexity ) have to be enforced . Their algorithm design and associated theoretical analysis actually are only valid for optimization with a simplified problem structure . Unfortunately , it has been well recognized that LL non-convexity frequently appears in a variety of applications , e.g. , sparse ` q regularization ( 0 < q < 1 ) for avoiding over-fitting , and learning parameters of coupled multi-layer neural networks , etc . Therefore , in challenging real-world scenarios , we are usually required to consider BLO problems where these assumptions ( e.g. , LLS , LLC and even LL strong convexity ) are naturally violated . These fundamental theoretical issues motivate us to propose a series of new techniques to address BLO with non-convex LL problems , which have been frequently appeared in various learning applications . In particular , by introducing an Initialization Auxiliary ( IA ) to the LL optimization dynamics and operating a Pessimistic Trajectory Truncation ( PTT ) strategy during the UL approximation , we construct a Gradient-based Method ( GM ) , named IAPTT-GM , to address BLO in challenging optimization scenarios ( i.e. , with non-convex follower tasks ) . We analyze the convergence behaviors of IAPTT-GM on BLOs without LLC and also investigate theoretical properties of the accelerated version of our algorithm on BLOs under LLC . Extensive experiments verify our theoretical results and demonstrate the effectiveness of IAPTT-GM on different learning applications . The main contributions of our IAPTT-GM are summarized as follows : • We propose IA and PTT , two new mechanisms to efficiently handle complex BLOs where the follower is facing with a non-convex task ( i.e. , without LLS and even LLC ) . IA actually paves the way for jointly optimizing both the UL variables and the dynamical initialization , while PTT adaptively reduces the complexity of backward recurrent propagation . • To our best knowledge , we establish the first strict convergence guarantee for gradientbased method on BLOs with non-convex follower tasks . We also justify the quality of our IAPTT-GM embedded with Nesterov ’ s accelerated dynamics under LLC . • We conduct a series of experiments to verify our theoretical findings and evaluate IAPTT-GM on various challenging BLOs , in which the follower tasks are either with non-convex loss functions ( e.g. , few-shot learning ) or coupled network structures ( e.g. , data hyper-cleaning ) . 2 The Proposed Algorithmic Framework . In this work , we consider the BLO problem in the form : min x∈X , y F ( x , y ) , s.t . y ∈ S ( x ) , ( 1 ) where x ∈ Rn , y ∈ Rm are UL and LL variables respectively , and S ( x ) denotes the set of solutions of the LL problem , i.e. , S ( x ) : = arg min y∈Y f ( x , y ) , ( 2 ) where f is differentiable w.r.t . y . To ensure the BLO model in Eq . ( 1 ) is well-defined , we assume that S ( x ) is nonempty for all x ∈ X . Observe further that the above BLO is structurally different from those in existing literature in the sense that no convexity assumption is required in the LL problem . In the following , we describe the proposed IAPTT-GM to solve the class of BLOs defined in Eq . ( 1 ) . The mechanism of a classical dynamics-embedded gradient method , approximates the LL solution via a dynamical system drawn from optimization iterations . Choosing gradient descent as the optimization dynamics for example , the approximation yK ( x ) is accessed by operations repeatedly performed by K − 1 steps parameterized by UL variable x yk+1 ( x ) = yk ( x ) − s∇yf ( x , yk ( x ) ) , k = 0 , · · · , K − 1 , ( 3 ) where s is a step size , and y0 is a fixed initial value . 2.1 Initialization Auxiliary . Embedding the dynamical iterations into the UL problem returns the approximate version F ( x , yK ( x ) ) . As long as ∇yf ( x , yK ( x ) ) uniformly converges to zero w.r.t . UL variable x varying in X , we call this a good approximation . To this end , usually restrictive LL strong convexity assumptions are imposed , thus the desired convergence of solutions of approximation problems towards those of the original BLO follows . By drawing inspiration from the classic dynamics-embedded gradient method which replaces the LL problem with certain optimization dynamics , hence resulting in an approximation of the bi-level problem , we propose a new gradient scheme to solve BLO in Eq . ( 1 ) without LLC restriction . Specifically , we let K be a prescribed positive integer and construct the following approximation yK ( x , z ) of the LL solution drawn from projected gradient descent iterations y0 ( x , z ) = z , yk+1 ( x , z ) = ProjY ( yk ( x , z ) − α k y∇yf ( x , yk ( x , z ) ) ) , k = 0 , · · · , K − 1 , ( 4 ) where { αky } is a sequence of steps sizes . We next embed the dynamical iterations yk ( x , z ) into max1≤k≤K { F ( x , yk ( x , z ) ) } , which can be regarded as a pessimistic trajectory truncation of the UL objective . The mechanism of the above scheme , in comparison , accesses the approximation parameterized by UL variable x and LL initial point y0 . Our motivation for the initialization auxiliary variable z comes from convergence theory [ 42 ] of non-convex first-order optimization methods . In fact , when the non-convex LL problem admits multiple solutions , the gradient descent steps with a “ bad ” initial point y0 can not return a desired point in the LL solution set , simultaneously optimizing the UL objective . To overcome such a difficulty , instead of using a fixed initial value , we introduce an initialization auxiliary variable z . Therefore , when it comes to solving the UL approximation problems , together with the UL variable x , the auxiliary variable z is also updated and hence optimized . As a consequence , we may search for the “ best ” initial value , starting from which the gradient descent steps approach a solution to the BLO in Eq . ( 1 ) , i.e. , a point in the LL solution set , simultaneously minimizing the UL objective . Algorithm 1 The Proposed IAPTT-GM 1 : Initialize x0 and z0 . 2 : for t = 0→ T − 1 do 3 : y0 = zt . 4 : for k = 0→ K − 1 do 5 : % LL Updating with xt and zt 6 : yk+1 = ProjY ( yk − αky∇yf ( xt , yk ) ) . 7 : end for 8 : % Pessimistic Trajectory Truncation 9 : k̄ = arg maxk { F ( x , yk ) } Kk=1 . 10 : % UL Updating with yk̄ ( x , z ) 11 : xt+1 = ProjX ( x t − αx∇xF ( xt , yk̄ ) ) . 12 : % Initialization Updating with yk̄ ( x , z ) 13 : zt+1 = ProjY ( z t − αz∇zF ( xt , yk̄ ) ) . 14 : end for

Summary: This work proposes a truncated unrolling type method for solving a bilevel optimization problem with non-convex lower-level problem and the convergence of the proposed method is also shown. Compared to the results in the existing literature, the convergence analysis in this work does not require any convexity assumptions on both upper level and lower-level objective functions. Efficiency of the proposed method is shown in the numerical experiments of artificial problems and real application problems.