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 paper proposes a bilevel optimization method based on unrolled gradient-descent iterations of the lower-level (LL) problem, while providing a theoretical connection between the ideal bilevel optimization and the unrolled one when LL problem might not be convex. The main idea is two-fold: Regarding the initial point (i.e., z) of LL variable as another variable in the UL problem, and taking the pessimistic value from the unrolled iteration, i.e., max_k F(x, y_k(x, z) ). The theoretical derivation presents that the ideal bilevel problem (min phi(x)) is asymptotically equivalent to the unrolled version (min phi_K(x, z)). The resulting algorithm is very simple: (1) perform a few unrolled steps (with computational graph) for LL, (2) take the pessimistic estimate (y_k) from the steps, (3) and perform gradient descent for x and z for the UL problem based on the y_k(x, z). The paper also proposes a Nesterov version (for LL iterations). The proposed method shows state-of-the-art performance for a toy problem and two applications.