Amortized Implicit Differentiation for Stochastic Bilevel Optimization

1 INTRODUCTION . Bilevel optimization refers to a class of algorithms for solving problems with a hierarchical structure involving two levels : an inner and an outer level . The inner-level problem seeks a solution y ? ( x ) minimizing a cost g ( x , y ) over a set Y given a fixed outer variable x in a set X . The outer-level problem minimizes an objective of the form L ( x ) =f ( x , y ? ( x ) ) over X for some upper-level cost f . When the solution y ? ( x ) is unique , the bilevel optimization problem takes the following form : min x∈X L ( x ) : = f ( x , y ? ( x ) ) , such that y ? ( x ) = arg min y∈Y g ( x , y ) . ( 1 ) First introduced in the field of economic game theory by Stackelberg ( 1934 ) , this problem has recently received increasing attention in the machine learning community ( Domke , 2012 ; Gould et al. , 2016 ; Liao et al. , 2018 ; Blondel et al. , 2021 ; Liu et al. , 2021 ; Shaban et al. , 2019 ) . Indeed , many machine learning applications can be reduced to ( 1 ) including hyper-parameter optimization ( Feurer and Hutter , 2019 ) , meta-learning ( Bertinetto et al. , 2018 ) , reinforcement learning ( Hong et al. , 2020b ; Liu et al. , 2021 ) or dictionary learning ( Mairal et al. , 2011 ; Lecouat et al. , 2020a ; b ) . The hierarchical nature of ( 1 ) introduces additional challenges compared to standard optimization problems , such as finding a suitable trade-off between the computational budget for approximating the inner and outer level problems ( Ghadimi and Wang , 2018 ; Dempe and Zemkoho , 2020 ) . These considerations are exacerbated in machine learning applications , where the costs f and g often come as an average of functions over a large or infinite number of data points ( Franceschi et al. , 2018 ) . All these challenges highlight the need for methods that are able to control the computational costs inherent to ( 1 ) while dealing with the large-scale setting encountered in machine learning . Gradient-based bilevel optimization methods appear to be viable approaches for solving ( 1 ) in largescale settings ( Lorraine et al. , 2020 ) . They can be divided into two categories : Iterative differentiation ( ITD ) and Approximate implicit differentiation ( AID ) . ITD approaches approximate the map y ? ( x ) by a differentiable optimization algorithm A ( x ) viewed as a function of x . The resulting surrogate loss L̃ ( x ) = f ( x , A ( x ) ) is optimized instead of L ( x ) using reverse-mode automatic differentiation ( see Baydin et al. , 2018 ) . AID approaches ( Pedregosa , 2016 ) rely on an expression of the gradient ∇L resulting from the implicit function theorem ( Lang , 2012 , Theorem 5.9 ) . Unlike ITD , AID avoids differentiating the algorithm approximating y ? ( x ) and , instead , approximately solves a linear system using only Hessian and Jacobian-vector products to estimate the gradient∇L ( Rajeswaran et al. , 2019 ) . These methods can also rely on stochastic approximation to increase scalability ( Franceschi et al. , 2018 ; Grazzi et al. , 2020 ; 2021 ) . In the context of machine-learning , Ghadimi and Wang ( 2018 ) provided one of the first comprehensive studies of the computational complexity for a class of bilevel algorithms based on AID approaches . Subsequently , Hong et al . ( 2020b ) ; Ji et al . ( 2021 ) ; Ji and Liang ( 2021 ) ; Yang et al . ( 2021 ) proposed different algorithms for solving ( 1 ) and obtained improved overall complexity by achieving a better trade-off between the cost of the inner and outer level problems . Still , the question of whether these complexities can be improved by better exploiting the structure of ( 1 ) through heuristics such as warm-start remains open ( Grazzi et al. , 2020 ) . Moreover , these studies proposed separate analysis of their algorithms depending on the convexity of the loss L and whether a stochastic or deterministic setting is considered . This points out to a lack of unified and systematic theoretical framework for analyzing bilevel problems , which is what the present work addresses . We consider the Amortized Implicit Gradient Optimization ( AmIGO ) algorithm , a bilevel optimization algorithm based on Approximate Implicit Differentiation ( AID ) approaches that exploits a warm-start strategy when estimating the gradient of L. We then propose a unified theoretical framework for analyzing the convergence of AmIGO when the inner-level problem is strongly convex in both stochastic and deterministic settings . The proposed framework is inspired from the early work of Habets ( 1974 ) on singularly perturbed systems and analyzes the effect of warm start by viewing the iterates of AmIGO algorithm as a dynamical system . The evolution of such system is described by a total energy function which allows to recover the convergence rates of unbiased oracle methods which have access to an unbiased estimate of ∇L ( c.f . Table 1 ) . To the best of our knowledge , this is the first time a bilevel optimization algorithm based on a warm-start strategy provably recovers the rates of unbiased oracle methods across a wide range of settings including the stochastic ones . 2 RELATED WORK . Singularly perturbed systems ( SPS ) are continuous-time deterministic dynamical systems of coupled variables ( x ( t ) , y ( t ) ) with two time-scales where y ( t ) evolves much faster than x ( t ) . As such , they exhibit a hierarchical structure similar to ( 1 ) . The early work of Habets ( 1974 ) ; Saberi and Khalil ( 1984 ) provided convergence rates for SPS towards equilibria by studying the evolution of a single scalar energy function summarizing these systems . The present work takes inspiration from these works to analyze the convergence of AmIGO which involves three time-scales . Two time-scale Stochastic Approximation ( TTSA ) can be viewed as a discrete-time stochastic version of SPS . ( Kaledin et al. , 2020 ) showed that TTSA achieves a finite-time complexity of O ( −1 ) for linear systems while Doan ( 2020 ) obtained a complexity of O ( −3/2 ) for general non-linear systems by extending the analysis for SPS . Hong et al . ( 2020b ) further adapted the non-linear TTSA for solving ( 1 ) . In the present work , we obtain faster rates by taking into account the dynamics of a third variable zk appearing in AmIGO , thus resulting in a three time-scale dynamics . Warm-start in bilevel optimization . Ji et al . ( 2021 ) ; Ji and Liang ( 2021 ) used a warm-start for the inner-level algorithm to obtain an improved computational complexity over algorithms without warm-start . In the deterministic non-convex setting , Ji et al . ( 2021 ) used a warm-start strategy when solving the linear system appearing in AID approaches to obtain improved convergence rates . However , it remained open whether using a warm-start when solving both inner-level problem and linear system arising in AID approaches can yield faster algorithms in the more challenging stochastic setting ( Grazzi et al. , 2020 ) . In the present work , we provide a positive answer to this question . 3 AMORTIZED IMPLICIT GRADIENT OPTIMIZATION . 3.1 GENERAL SETTING AND MAIN ASSUMPTIONS . Notations . In all what follows , X and Y are Euclidean spaces . For a differentiable function h ( x , y ) : X ×Y → R , we denote by∇h its gradient w.r.t . ( x , y ) , by ∂xh and ∂yh its partial derivatives w.r.t . x and y and by ∂xyh and ∂yyh the partial derivatives of ∂yh w.r.t x and y , respectively . Geometries Setting Algorithms Complexity BA ( Ghadimi and Wang , 2018 ) O ( κ2L ∨ κ2g log 2 −1 ) AccBio ( Ji and Liang , 2021 ) O ( κ1/2L κ 1/2 g log 2 −1 ) ( D ) AmIGO ( Corollary 1 ) O ( κLκg log −1 ) BSA ( Ghadimi and Wang , 2018 ) O ( κ4L −2 ) TTSA ( Hong et al. , 2020b ) O ( κ0.5L ( κ 8.5 g + κ 3 L ) −3/2 log −1 ) ( SC ) ( S ) AmIGO ( Corollary 2 ) O ( κ2Lκ 3 g −1 log −1 ) BA ( Ghadimi and Wang , 2018 ) O ( κ5g −5/4 ) AID-BiO ( Ji et al. , 2021 ) O ( κ4g −1 ) ( D ) AmIGO ( Corollary 3 ) O ( κ4g −1 ) BSA ( Ghadimi and Wang , 2018 ) O ( κ9g −3 + κ6g −2 ) TTSA ( Hong et al. , 2020b ) O ( κ16g −5/2 log −1 ) stocBiO ( Ji et al. , 2021 ) O ( κ9g −2+κ6g −2 log −1 ) MRBO/VRBO ? ( Yang et al. , 2021 ) O ( poly ( κg ) −3/2 log −1 ) ( NC ) ( S ) AmIGO ( Corollary 4 ) O ( κ9g −2 ) Table 1 : Cost of finding an -accurate solution as measured by E [ L ( xk ) −L ? ] ∧2−1µE [ ‖xk−x ? ‖2 ] when L is µ-strongly-convex ( SC ) and 1k ∑k i=1 E [ ‖∇L ( xi ) ‖2 ] when L is non-convex ( NC ) . The settings ( D ) and ( S ) stand for the deterministic and stochastic settings . The cost corresponds to the total number of gradients , Jacobian and Hessian-vector products used by the algorithm . κL and κg are the conditioning numbers of L and g whenever applicable . The dependence on κL and κg for TTSA and AccBio are derived in Proposition 11 of Appendix A.4 . The rate of MRBO/VRBO is obtained under the additional mean-squared smoothness assumption ( Arjevani et al. , 2019 ) . To ensure that ( 1 ) is well-defined , we consider the setting where the inner-level problem is strongly convex so that the solution y ? ( x ) is unique as stated by the following assumption : Assumption 1 . For any x ∈ X , the function y 7→ g ( x , y ) is Lg-smooth and µg-strongly convex . Assumption 1 holds in the context of hyper-parameter selection when the inner-level is a kernel regression problem ( Franceschi et al. , 2018 ) , or when the variable y represents the last linear layer of a neural network as in many meta-learning tasks ( Ji et al. , 2021 ) . Under Assumption 1 and additional smoothness assumptions on f and g , the next proposition shows that L is differentiable : Proposition 1 . Let g be a twice differentiable function satisfying Assumption 1 . Assume that f is differentiable and consider the quadratic problem : min z∈Rdy Q ( x , y , z ) : = 1 2 z > ( ∂yyg ( x , y ) ) z + z > ∂yf ( x , y ) . ( 2 ) Then , ( 2 ) admits a unique minimizer z ? ( x , y ) for any ( x , y ) in X × Y . Moreover , y ? ( x ) is unique and well-defined for any x in X and L is differentiable with gradient given by : ∇L ( x ) = ∂xf ( x , y ? ( x ) ) + ∂xyg ( x , y ? ( x ) ) z ? ( x , y ? ( x ) ) . ( 3 ) Proposition 1 follows by application of the implicit function theorem ( Lang , 2012 , Theorem 5.9 ) and provides an expression for ∇L solely in terms of partial derivatives of f and g evaluated at ( x , y ? ( x ) ) . Following Ghadimi and Wang ( 2018 ) , we further make two smoothness assumptions on f and g : Assumption 2 . There exist positive constants Lf and B such that for all x , x′ ∈ X and y , y′ ∈ Y : ‖∇f ( x , y ) −∇f ( x′ , y′ ) ‖ ≤ Lf‖ ( x , y ) − ( x′ , y′ ) ‖ , ‖∂yf ( x , y ) ‖ ≤ B . Assumption 3 . There exit positive constants L′g , Mg such that for any x , x′ ∈ X and y , y′ ∈ Y : max { ‖∂xyg ( x , y ) − ∂xyg ( x′ , y′ ) ‖ , ‖∂yyg ( x , y ) − ∂yyg ( x′ , y′ ) ‖ } ≤Mg‖ ( x , y ) − ( x′ , y′ ) ‖ ‖∂yg ( x , y ) − ∂yg ( x′ , y ) ‖ ≤ L′g‖x− x′‖ . Assumptions 1 to 3 allow a control of the variations of y ? and z ? and ensure L is L-smooth for some positive constant L as shown in Proposition 6 of Appendix B.2 . As an L-smooth function , L is necessarily weakly convex ( Davis et al. , 2018 ) , meaning thatL satisfies the inequalityL ( x ) −L ( y ) ≤ ∇L ( x ) > ( x− y ) − µ2 ‖x− y‖ 2 for some fixed µ ∈ R with |µ| ≤ L. In particular , L is convex when µ ≥ 0 , strongly convex when µ > 0 and generally non-convex when µ < 0 . We thus consider two cases for L , the strongly convex case ( µ > 0 ) and the non-convex case ( µ < 0 ) . When L is convex , we denote by L ? its minimum value achieved at a point x ? and define κL=L/µ when µ > 0 . Stochastic/deterministic settings . We consider the general setting where f ( x , y ) and g ( x , y ) are expressed as an expectation of stochastic functions f̂ ( x , y , ξ ) and ĝ ( x , y , ξ ) over a noise variable ξ . We recover the deterministic setting as a particular case when the variable ξ has zero variance , thus allowing us to treat both stochastic ( S ) and deterministic ( D ) settings in a unified framework . As often in machine-learning , we assume we can always draw a new batch D of i.i.d . samples of the noise variable ξ with size |D| ≥ 1 and use it to compute stochastic approximations of f and g defined by abuse of notation as f̂ ( x , y , D ) : = 1|D| ∑ ξ∈D f̂ ( x , y , ξ ) and ĝ ( x , y , D ) : = 1 |D| ∑ ξ∈D ĝ ( x , y , ξ ) . We make the following noise assumptions which are implied by those in Ghadimi and Wang ( 2018 ) : Assumption 4 . For any batchD , ∇f̂ ( x , y , D ) and ∂y ĝ ( x , y , D ) are unbiased estimator of∇f ( x , y ) and ∂yg ( x , y ) with a uniformly bounded variance , i.e . for all x , y ∈ X × Y : E [ ∥∥∥∇f̂ ( x , y , D ) −∇f ( x , y ) ∥∥∥2 ] ≤ σ̃2f |D|−1 , E [ ‖∂y ĝ ( x , y , D ) − ∂yg ( x , y ) ‖2 ] ≤ σ̃2g |D|−1 . Assumption 5 . For any batch D , the matrices F1 ( x , y , D ) : = ∂xy ĝ ( x , y , D ) − ∂xyg ( x , y ) and F2 ( x , y , D ) : = ∂yy ĝ ( x , y , D ) − ∂yyg ( x , y ) have zero mean and satisfy for all x , y ∈ X × Y : ∥∥E [ F1 ( x , y , D ) > F1 ( x , y , D ) ] ∥∥op ≤ σ̃2gxy |D|−1 , ∥∥E [ F2 ( x , y , D ) > F2 ( x , y , D ) ] ∥∥op ≤ σ̃2gyy |D|−1 . For conciseness , we will use the notations σ2f : =σ̃ 2 f |D| −1 , σ2g : =σ̃ 2 g |D| −1 , σ2gxy : =σ̃ 2 gxy |D| −1 and σ2gyy : =σ̃ 2 gyy |D| −1 , without explicit reference to the batch D. Next , we describe the algorithm .

Summary: This work focuses on bilevel optimization. The key innovation here is the warm start, which enables improved complexity bounds (under some settings). The analysis nicely builds on singularly perturbed systems (SPS). Numerical tests are also provided on both synthetic and real problems, where the merit of warm start is demonstrated.