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: This paper studies the evolution strategies (ES) based method to estimate the Jacobian matrix used in bilevel optimization. The main idea is to use first-order information to estimate the second-order information in the Jacobian matrix. Some convergence results have been provided, companied with some simulations.