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: The paper studies the problem of bi-level optimization. In particular, it considers algorithms based on inexact implicit differentiation, where the inner problem and the implicit differentiation are not solved exactly. Warm-up is used when solving the inner problem and implicit differentiation. The convergence of the proposed method is analyzed by viewing the iterates of the proposed method as a dynamical system using the idea of singularly perturbed systems.