Efficient and Modular Implicit Differentiation

Abstract

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.

arXiv commentV3: added more related work and Jacobian precision figure

References in corpus (28)

• Adam: A Method for Stochastic Optimization
• Scikit-learn: Machine Learning in Python
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Discussion of "Least angle regression" by Efron et al
• Discussion of "Least angle regression" by Efron et al
• Discussion of "Least angle regression" by Efron et al
• Discussion of "Least angle regression" by Efron et al
• Discussion of "Least angle regression" by Efron et al
• Discussion of "Least angle regression" by Efron et al
• Rejoinder to "Least angle regression" by Efron et al
• Discussion of "Least angle regression" by Efron et al
• Least Angle Regression
• Discussion of "Least angle regression" by Efron et al
• Task-Driven Dictionary Learning
• From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification
• Deep Equilibrium Models
• On Differentiating Parameterized Argmin and Argmax Problems with Application to Bi-level Optimization
• Smooth and Sparse Optimal Transport
• A Regularized Framework for Sparse and Structured Neural Attention
• SparseMAP: Differentiable Sparse Structured Inference
• Neural reparameterization improves structural optimization
• Implicit differentiation of Lasso-type models for hyperparameter optimization
• Structured Prediction with Projection Oracles
• LP-SparseMAP: Differentiable Relaxed Optimization for Sparse Structured Prediction
• Variational Data Assimilation with a Learned Inverse Observation Operator
• Learning Composable Energy Surrogates for PDE Order Reduction
• Implicit differentiation for fast hyperparameter selection in non-smooth convex learning
• Decomposing reverse-mode automatic differentiation

Cited by in corpus (12)

• Differentiable quantum chemistry with PySCF for molecules and materials at the mean-field level and beyond
• Learning Pair Potentials using Differentiable Simulations
• Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey
• Inverse molecular design and parameter optimization with Hückel theory using automatic differentiation
• Sinkformers: Transformers with Doubly Stochastic Attention
• Conditional average treatment effect estimation with marginally constrained models
• A research framework for writing differentiable PDE discretizations in JAX
• Integrating Quantum Processor Device and Control Optimization in a Gradient-based Framework
• Amortized Implicit Differentiation for Stochastic Bilevel Optimization
• Multiset-Equivariant Set Prediction with Approximate Implicit Differentiation
• Gradients are Not All You Need
• Scalable Data Assimilation with Message Passing