Efficient and Modular Implicit Differentiation

1 INTRODUCTION . Automatic differentiation ( autodiff ) is now an inherent part of machine learning software . It allows to express complex computations by composing elementary ones in creative ways and removes the tedious burden of computing their derivatives by hand . In parallel , the differentiation of optimization problem solutions has found many applications . A classical example is bi-level optimization , which typically involves computing the derivatives of a nested optimization problem in order to solve an outer one . Examples of applications in machine learning include hyper-parameter optimization ( Chapelle et al. , 2002 ; Seeger , 2008 ; Pedregosa , 2016 ; Franceschi et al. , 2017 ; Bertrand et al. , 2020 ; 2021 ) , neural networks ( Lorraine et al. , 2020 ) , and meta-learning ( Franceschi et al. , 2018 ; Rajeswaran et al. , 2019 ) . Another line of active research involving differentiation of optimization problem solutions are optimization layers ( Kim et al. , 2017 ; Amos & Kolter , 2017 ; Niculae & Blondel , 2017 ; Djolonga & Krause , 2017 ; Gould et al. , 2019 ) , which can be used to encourage structured outputs , and implicit deep networks ( Bai et al. , 2019 ; El Ghaoui et al. , 2019 ) , which have a smaller memory footprint than backprop-trained networks . Since optimization problem solutions typically do not enjoy an explicit formula in terms of their inputs , autodiff can not be used directly to differentiate these functions . In recent years , two main approaches have been developed to circumvent this problem . The first one consists of unrolling the iterations of an optimization algorithm and using the final iteration as a proxy for the optimization problem solution ( Wengert , 1964 ; Domke , 2012 ; Deledalle et al. , 2014 ; Franceschi et al. , 2018 ; Ablin et al. , 2020 ) . This allows to explicitly construct a computational graph relating the algorithm output to the inputs , on which autodiff can then be used transparently . However , this requires a reimplementation of the algorithm using the autodiff system , and not all algorithms are necessarily autodiff friendly . Moreover , forward-mode autodiff has time complexity that scales linearly with the number of variables and reverse-mode autodiff has memory complexity that scales linearly with the number of algorithm iterations . In contrast , a second approach consists in implicitly relating an optimization problem solution to its inputs using optimality conditions . In a machine learning context , such implicit differentiation has been used for stationarity conditions ( Bengio , 2000 ; Lorraine et al. , 2020 ) , KKT conditions ( Chapelle et al. , 2002 ; Gould et al. , 2016 ; Amos & Kolter , 2017 ; Niculae et al. , 2018 ; Niculae & Martins , 2020 ) and the proximal gradient fixed point ( Niculae & Blondel , 2017 ; Bertrand et al. , 2020 ; 2021 ) . An advantage of implicit differentiation is that a solver reimplementation is not needed , allowing to build upon decades of state-of-the-art software . Although implicit differentiation has a long history in numerical analysis ( Griewank & Walther , 2008 ; Bell & Burke , 2008 ; Krantz & Parks , 2012 ; Bonnans & Shapiro , 2013 ) , so far , it remained difficult to use for practitioners , as it required a case-by-case tedious mathematical derivation and implementation . CasADi ( Andersson et al. , 2019 ) allows to differentiate various optimization and root finding problem algorithms provided by the library . However , it does not allow to easily add implicit differentiation on top of existing solvers from optimality conditions expressed by the user , as we do . A recent tutorial explains how to implement implicit differentiation in JAX ( Duvenaud et al. , 2020 ) . However , the tutorial requires the user to take care of low-level technical details and does not cover a large catalog of optimality condition mappings as we do . Other work ( Agrawal et al. , 2019a ) attempts to address this issue by adding implicit differentiation on top of cvxpy ( Diamond & Boyd , 2016 ) . This works by reducing all convex optimization problems to a conic program and using conic programming ’ s optimality conditions to derive an implicit differentiation formula . While this approach is very generic , solving a convex optimization problem using a conic programming solver—an ADMM-based splitting conic solver ( O ’ Donoghue et al. , 2016 ) in the case of cvxpy—is rarely the state-of-the-art approach for each particular problem instance . In this work , we adopt a different strategy that makes it easy to add implicit differentiation on top of any existing solver . In our approach , the user defines directly in Python a mapping function F capturing the optimality conditions of the problem solved by the algorithm . Once this is done , we leverage autodiff of F combined with implicit differentiation to automatically differentiate the optimization problem solution . In this way , our approach is generic , yet it can exploit the efficiency of state-of-the-art solvers . It therefore combines the benefits of implicit differentiation and autodiff . To summarize , we make the following contributions . • We describe our framework and its JAX implementation ( provided in the supplementary material ) . Our framework significantly lowers the barrier to use implicit differentiation thanks to the use of autodiff of the optimality conditions and the seamless integration in JAX . Our framework significantly extends JAX for numerical optimization , with low-level details all abstracted away . • We instantiate our framework on a large catalog of optimality conditions ( Table 1 ) , recovering existing schemes and obtaining new ones , such as the mirror descent fixed point based one . • On the theoretical side , we provide new bounds on the Jacobian error when the optimization problem is only solved approximately , and empirically validate them . • We implement four illustrative applications , demonstrating our framework ’ s ease of use . Beyond our software implementation , we hope this paper provides a self-contained blueprint for creating an efficient and modular implementation of implicit differentiation . Notation . We denote the gradient and Hessian of f : Rd → R evaluated at x ∈ Rd by∇f ( x ) ∈ Rd and ∇2f ( x ) ∈ Rd×d . We denote the Jacobian of F : Rd → Rp evaluated at x ∈ Rd by ∂F ( x ) ∈ Rp×d . When f or F have several arguments , we denote the gradient , Hessian and Jacobian in the ith argument by ∇i , ∇2i and ∂i , respectively . The standard probability simplex is denoted by 4d : = { x ∈ Rd : ‖x‖1 = 1 , x ≥ 0 } . For any set C ⊂ Rd , we denote the indicator function IC : Rd → R ∪ { +∞ } where IC ( x ) = 0 if x ∈ C , IC ( x ) = +∞ otherwise . For a vector or matrix A , we note ‖A‖ the Frobenius ( or Euclidean ) norm , and ‖A‖op the operator norm . 2 COMBINING IMPLICIT DIFFERENTIATION AND AUTODIFF . 2.1 GENERAL PRINCIPLES . Overview . Contrary to autodiff through unrolled algorithm iterations , implicit differentiation typically involves a manual , sometimes complicated , mathematical derivation . For instance , numerous works ( Chapelle et al. , 2002 ; Gould et al. , 2016 ; Amos & Kolter , 2017 ; Niculae et al. , 2018 ; Niculae & Martins , 2020 ) use Karush–Kuhn–Tucker ( KKT ) conditions in order to relate a constrained optimization problem ’ s solution to its inputs , and to manually derive a formula for its derivatives . The derivation and implementation in these works are always case-by-case . In this work , we propose a generic way to easily add implicit differentiation on top of existing solvers . In our approach , the user defines directly in Python a mapping function F capturing the optimality conditions of the problem solved by the algorithm . We provide reusable building blocks to easily express such F . The provided F is then plugged into our Python decorator @ custom_root , which we append on top of the solver declaration we wish to differentiate . Under the hood , we combine implicit differentiation and autodiff of F to automatically differentiate the optimization problem solution . A simple illustrative example is given in Figure 1 . Differentiating a root . Let F : Rd×Rn → Rd be a user-provided mapping , capturing the optimality conditions of a problem . An optimal solution , denoted x ? ( θ ) , should be a root of F : F ( x ? ( θ ) , θ ) = 0 . ( 1 ) We can see x ? ( θ ) as an implicitly defined function of θ ∈ Rn , i.e. , x ? : Rn → Rd . More precisely , from the implicit function theorem ( Griewank & Walther , 2008 ; Krantz & Parks , 2012 ) , we know that for ( x0 , θ0 ) satisfying F ( x0 , θ0 ) = 0 with a continuously differentiable F , if the Jacobian ∂1F evaluated at ( x0 , θ0 ) is a square invertible matrix , then there exists a function x ? ( · ) defined on a neighborhood of θ0 such that x ? ( θ0 ) = x0 . Furthermore , for all θ in this neighborhood , we have that F ( x ? ( θ ) , θ ) = 0 and ∂x ? ( θ ) exists . Using the chain rule , the Jacobian ∂x ? ( θ ) satisfies ∂1F ( x ? ( θ ) , θ ) ∂x ? ( θ ) + ∂2F ( x ? ( θ ) , θ ) = 0 . Computing ∂x ? ( θ ) therefore boils down to the resolution of the linear system of equations −∂1F ( x ? ( θ ) , θ ) ︸ ︷︷ ︸ A∈Rd×d ∂x ? ( θ ) ︸ ︷︷ ︸ J∈Rd×n = ∂2F ( x ? ( θ ) , θ ) ︸ ︷︷ ︸ B∈Rd×n . ( 2 ) When ( 1 ) is a one-dimensional root finding problem ( d = 1 ) , ( 2 ) becomes particularly simple since we then have ∇x ? ( θ ) = B > /A , where A is a scalar value . We will show that existing and new implicit differentiation methods all reduce to this simple principle . We call our approach hybrid , since it combines implicit differentiation ( and as such requires solving a linear system ) with autodiff of the optimality conditions F . Our approach is efficient as it can be added on top of any state-of-the-art solver and modular as the optimality condition specification is decoupled from the implicit differentiation mechanism . This contrasts with existing works , where the mathematical derivation and implementation are specific to each optimality condition . Differentiating a fixed point . We will encounter numerous applications where x ? ( θ ) is implicitly defined through a fixed point : x ? ( θ ) = T ( x ? ( θ ) , θ ) , where T : Rd × Rn → Rd . This can be seen as a particular case of ( 1 ) by defining the residual F ( x , θ ) = T ( x , θ ) − x . ( 3 ) In this case , using the chain rule , we have

Summary: The paper promises extension of Google’s JAX library by automatic and modular differentitation of implicitly defined functions. Given a solver $\theta\mapsto x^*(\theta)$ one asks for its Jacobian $\partial x^*(\theta)$. In cases when $x^*$ does not have an explicit form that can be automatically differentiated, the Jacobian can be computed from the optimality conditions $F$ having the general form $F(x^*(\theta),\theta)=0$. Jacobian $\partial x^*(\theta)$ is then given by a solution of a linear system involving $\partial_1 F$ and $\partial_2 F$. The authors describe a unified implementation in JAX and show the resulting simple interface for end-users. Various use cases are presented on examples and code-snippets figures, including stationary point conditions, KKT conditions, proximal gradient, projective gradient and mirror descent fixed point and more. On the theoretical side, the authors discuss the case when the optimization problem solution is only approximate and derive Jacobian error bounds under certain assumptions on optimality conditions $F$.