ES-Based Jacobian Enables Faster Bilevel Optimization

1 INTRODUCTION . Bilevel optimization has recently arisen as a powerful tool to capture various modern machine learning problems , including meta-learning ( Bertinetto et al. , 2018 ; Franceschi et al. , 2018 ; Rajeswaran et al. , 2019 ; Ji et al. , 2020a ; Liu et al. , 2021a ) , hyperparamater optimization ( Franceschi et al. , 2018 ; Shaban et al. , 2019 ) , neural architecture search ( Liu et al. , 2018a ; Zhang et al. , 2021 ) , etc . Bilevel optimization generally takes the following mathematical form : min x∈Rp Φ ( x ) : = f ( x , y∗ ( x ) ) s.t . y∗ ( x ) = arg min y∈Rd g ( x , y ) , ( 1 ) where the outer and inner objectives f : Rp×Rd → R and g : Rp×Rd → R are both continuously differentiable with respect to ( w.r.t . ) the inner and outer variables x ∈ Rp and y ∈ Rd . In many machine learning problems , the loss functions f and g take a finite-sum form over a set of given data Dn , m = { ξi , ζj , i = 1 , ... , n , j = 1 , ... , m } as follows : f ( x , y ) = 1n ∑n i=1 F ( x , y ; ξi ) , g ( x , y ) = 1 m ∑m i=1G ( x , y ; ζi ) ( 2 ) where the sample sizes n and m are typically very large . In this work , we consider a popular setting where the total objective function Φ ( x ) in eq . ( 1 ) is possibly nonconvex w.r.t . x and the inner function g is strongly convex w.r.t . y. Gradient-based methods have served as a popular tool for solving bilevel optimization problems . Two types of approaches have been widely used : the iterative differentiation ( ITD ) method ( Domke , 2012 ; Maclaurin et al. , 2015 ; Finn et al. , 2017 ; Franceschi et al. , 2017 ; Shaban et al. , 2019 ; Rajeswaran et al. , 2019 ; Liu et al. , 2020a ) and the approximate iterative differentiation ( AID ) method ( Domke , 2012 ; Pedregosa , 2016 ; Gould et al. , 2016 ; Liao et al. , 2018 ; Lorraine et al. , 2020 ) . In the ITD method , the objective value f ( x , ŷ ( x ) ) is first computed using an approximation point ŷ ( x ) close to y∗ ( x ) , and the hypergradient∇Φ ( x ) = ∂f ( x , y ∗ ( x ) ) ∂x is then approximated by ∂f ( x , ŷ ( x ) ) ∂x using either reverse or forward mode automatic differentiation . Alternatively , the AID method leverages the Implicit Function Theorem ( IFT ) to establish an implicit expression for the hypergradient . Due to the bilevel structure of the problem , these gradient-based ITD and AID approaches typically involve second-order matrix computations . Existing efficient implementations of these methods often adopt Jacobian- or/and Hessian-vector computations , which can still be very costly in practice for high-dimensional problems with deep neural networks . To overcome the computational challenge of current gradient-based methods , Evolution Strategies ( ES ) can be a good candidate to eliminate the second-order computations , which uses function values for estimating its gradient in blackbox optimization ( Nesterov & Spokoiny , 2017 ) . An ES-based bilevel optimizer has recently been proposed in Gu et al . ( 2021 ) for hyperparameter optimization , which uses the ES method to approximate the hypergradient of Φ ( x ) , treating Φ ( x ) fully as a blackbox objective . However , the hypergradient in bilevel optimization has a much more involved structure ( see eq . ( 3 ) ) than the gradient in standard minimization problems , and hence direct use of the ES method without exploiting the structure of the hypergradient can encounter a large estimation bias and result in poor performance in practice , especially in deep learning as demonstrated in our experiments in Section 4 . This hence motivates the following intriguing question : Can we design a better hypergradient estimator based on the ES method by leveraging the analytical structure of the hypergradient , and then use such an estimator to construct more efficient ES-based bilevel optimizers , which are scalable for high-dimensional problems ? Furthermore , since the hypergradient in bilevel optimization is much more challenging to estimate than the gradient in standard minimization problems , it is unclear whether the ES-based estimator can yield a provably convergent algorithm for bilevel optimization . So far the ES-based bilevel optimizer in Gu et al . ( 2021 ) did not come with a polynomial-time complexity guarantee . Our empirical experiments in Section 4 show that such an algorithm fail to converge even with shallow neural networks . This thus motivates the second question we ask : Can the ES method lead to bilevel optimizers with theoretical finite-time convergence guarantee ? This paper provides the affirmative answers to both questions . 1.1 MAIN CONTRIBUTIONS . Novel ES-based Jacobian bilevel optimizer . We propose a novel bilevel optimizer with the ESbased Jacobian estimation , which we call as ESJ . In contrast to the existing ES-based optimizer in Gu et al . ( 2021 ) , which estimates the hypergradient by treating the outer-level objective Φ ( x ) fully as a blackbox , ESJ estimates only the respone Jacobian matrix ( i.e. , gradient of y∗ ( x ) with respect to x ) , which is the major computational bottleneck , and then leverages the analytical structure of the hypergradient to construct a much more accurate estimator . Further , ESJ has only gradient computations , and is computationally much more efficient than existing AID and ITD bilevel optimizers that require Hessian- and/or Jacobian-vector product computations . We further propose a stochastic bilevel optimizer ESJ-S to allow the algorithm scalable over large dataset . Convergence guarantee . Theoretically , we characterize the convergence rate and their computational complexity for both ESJ and ESJ-S to achieve an -accurate solution . Technically , in contrast to the standard analysis of ES-based methods on the smoothed blackbox function values , we develop tools to analyze the ES estimator on the output of the execution of an inner optimizer . Such an analysis can be of independent interest for ES-based bilevel optimization . Superior performance . Experimentally , our algorithms ESJ and ESJ-S achieve superior performance compared to the current state of the art bilevel optimizers . Our approach is not only efficient and scalable , but also stable and robust across various bilevel problems . In particular , on the few-shot meta-learning problem , we meta-learn the twelve millions weights of a ResNet-12 network . Previous applications of other bilevel algorithms to meta-learning were limited only to settings where the backbone network has only up to few hundreds of thousands of parameters . 1.2 RELATED WORK . Bilevel optimization . Bilevel optimization has been studied for decades since Bracken & McGill ( 1973 ) . A variety of bilevel optimization algorithms were then developed , including constraintbased approaches ( Hansen et al. , 1992 ; Shi et al. , 2005 ; Moore , 2010 ) , approximate implicit differ- entiation ( AID ) approaches ( Domke , 2012 ; Pedregosa , 2016 ; Gould et al. , 2016 ; Liao et al. , 2018 ; Lorraine et al. , 2020 ) and iterative differentiation ( ITD ) approaches ( Domke , 2012 ; Maclaurin et al. , 2015 ; Franceschi et al. , 2017 ; Finn et al. , 2017 ; Shaban et al. , 2019 ; Rajeswaran et al. , 2019 ; Liu et al. , 2020a ) . These methods often suffer from expensive computations of second-order information ( e.g. , Hessian-vector products ) . Hessian-free algorithms were recently proposed based on interiorpoint method ( Liu et al. , 2021b ) and Evolution Strategies ( ES ) ( Gu et al. , 2021 ) . This paper proposes a more efficient ES-based approach via exploiting the benign structure of the hypergradient . Recently , the convergence rate has been established for gradient-based ( Hessian involved ) bilevel algorithms ( Grazzi et al. , 2020 ; Ji et al. , 2021 ; Rajeswaran et al. , 2019 ; Ji & Liang , 2021 ) . This paper provides the convergence analysis for the proposed Hessian-free ES-based approach . Stochastic bilevel optimization . Ghadimi & Wang ( 2018 ) ; Ji et al . ( 2021 ) ; Hong et al . ( 2020 ) proposed stochastic gradient descent ( SGD ) type bilevel optimization algorithms by employing Neumann series for Hessian-inverse-vector approximation . Recent works ( Khanduri et al. , 2021a ; b ; Chen et al. , 2021 ; Guo et al. , 2021 ; Guo & Yang , 2021 ; Yang et al. , 2021 ) then leveraged momentumbased variance reduction to further reduce the computational complexity of existing SGD-type bilevel optimizers . In this paper , we propose a stochastic ES-based method , which eliminate the computation of second-order information required by all aforementioned stochastic methods . Please see Appendix B for more works about applications of bilevel optimization and ES methods . 2 PROPOSED ALGORITHMS . We describe our proposed deterministic and stochastic algorithms for bilevel optimization . 2.1 HYPERGRADIENTS . The key step in popular gradient-based bilevel optimizers is the estimation of the hypergradient ( i.e. , the gradient of the objective with respect to the outer variable x ) , which takes the following form : ∇Φ ( x ) = ∇xf ( x , y∗ ( x ) ) + J∗ ( x ) > ∇yf ( x , y∗ ( x ) ) ( 3 ) where the Jacobian matrix J∗ ( x ) = ∂y ∗ ( x ) ∂x ∈ R d×p . Following Lorraine et al . ( 2020 ) , it can be seen that∇Φ ( x ) contains two components : the direct gradient∇xf ( x , y∗ ( x ) ) and the indirect gradient J∗ ( x ) > ∇yf ( x , y∗ ( x ) ) . The direct component can be efficiently computed using the existing automatic differentiation techniques . The indirect component , however , is computationally much more complex , because J∗ ( x ) takes the form of J∗ ( x ) = − [ ∇2yg ( x , y∗ ( x ) ) ] −1∇x∇yg ( x , y∗ ( x ) ) ( if ∇2yg ( x , y∗ ( x ) ) is invertible ) , which contains the Hessian inverse and the second-order mixed derivative . Some approaches mitigate the issue by designing Jacobian-vector and Hessian-vector products ( Pedregosa , 2016 ; Franceschi et al. , 2017 ; Grazzi et al. , 2020 ) to replace second-order computations . But these algorithms still have poor scalability to modern large-scale bilevel problems with high-dimensional variables such as neural network parameters . We next introduce the ES approach which is at the core of our idea for designing scalable Hessian-free bilevel optimizers . 2.2 EVOLUTION STRATEGIES ( ES ) METHOD . Evolution Strategies ( ES ) is a powerful technique to estimate the gradient of a function based on function values , when it is not feasible ( such as in black-box problems ) or computationally costly to evaluate the gradient . The idea of the ES method in Nesterov & Spokoiny ( 2017 ) is to approximate the gradient of a general black-box function h : Rn → R using the following oracle based only on the function values ( i.e. , the zeroth-order information ) ∇̂h ( x ; u ) = h ( x+ µu ) − h ( x ) µ u ( 4 ) where u ∈ Rn is a Gaussian random vector and µ > 0 is the smoothing parameter . Such an oracle can be shown to be an unbiased estimator of the gradient of the smoothed function Eu [ h ( x+ µu ) ] .

Summary: The paper aims to reduce the computational cost in solving a bilevel problem. In bilevel optimization, the outer gradient can be decomposed into direct and indirect gradients. The indirect gradient requires computing the response Jacobian, which involves computing & inverting the Hessian and computing mixed derivative. To avoid expensive second-order approximations, the authors propose to use the evolution strategies (ES) based method to approximate the response Jacobian. The authors provide the convergence guarantee of the algorithm and empirically demonstrate that the proposed algorithm performs well on large-scale bilevel problems.