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: This paper presents a module for implicit differentiation in JAX, which allows efficient differentiation through optimization and root-finding problems. The paper presents the general principle of combining the implicit function theorem with autodiff of optimality conditions, and discusses how this can work for a implicit functions such as a range of constrained optimization methods. Finally there are example applications to three machine learning tasks and one molecular dynamics simulation task where this framework improves on currently used methods.