Efficient and Modular Implicit Differentiation

By Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-L'opez, Fabian Pedregosa, Jean-Philippe Vert. 2021

360 citations recorded by the source corpus.

Automatic differentiation (autodiff) has revolutionized machine learning. It allows to express complex computations by composing elementary ones in creative ways and removes the burden of computing their derivatives by hand. More recently, differentiation of optimization problem solutions has attracted widespread attention with applications such as optimization layers, and in bi-level problems such as hyper-parameter optimization and meta-learning. However, so far, implicit differentiation remained difficult to use for practitioners, as it often required case-by-case tedious mathematical derivations and implementations. In this paper, we propose automatic implicit differentiation, an efficient and modular approach for implicit differentiation of optimization problems. In our approach, the user defines directly in Python a function $F$ capturing the optimality conditions of the problem to be differentiated. Once this is done, we leverage autodiff of $F$ and the implicit function theorem to automatically differentiate the optimization problem. Our approach thus combines the benefits of implicit differentiation and autodiff. It 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. We show that seemingly simple principles allow to recover many existing implicit differentiation methods and create new ones easily. We demonstrate the ease of formulating and solving bi-level optimization problems using our framework. We also showcase an application to the sensitivity analysis of molecular dynamics.

Subfields: Computational Theory and Mathematics, Numerical Analysis, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computational Theory and Mathematics

In Numerical Analysis

In Optimization & Training

Citations in the corpus

Multiset-Equivariant Set Prediction with Approximate Implicit Differentiation (2021, methodology): “This observation is inline with theoretical findings on strongly convex inner objectives (Blondel et al., 2021; Ablin et al., 2020), which indicate that implicit differentiation is better at approximating the Jacobian when the optimizer solution is inexact.”
Gradients are Not All You Need (2021, methodology): “These properties have been successfully optimized using automatic differentiation [Schoenholz and Cubuk, 2020, Blondel et al., 2021, Goodrich et al., 2021] and do not suffer from bias as in the case of truncated gradients.”