# Efficient and Modular Implicit Differentiation

arXiv: https://arxiv.org/abs/2105.15183 · HTML: https://mlsanity.com/paper/2105.15183/

## Key facts

- Cited by 203 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 90 references that are in the map.
- First posted on arXiv 31 May 2021.

## Builds on (most cited first, 30 of 37)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [CVXPY: A Python-Embedded Modeling Language for Convex Optimization](https://mlsanity.com/paper/1603.00943/) (2016)
- [Deep Equilibrium Models](https://mlsanity.com/paper/1909.01377/) (2019)
- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Meta-Learning with Implicit Gradients](https://mlsanity.com/paper/1909.04630/) (2019)
- [Differentiable Convex Optimization Layers](https://mlsanity.com/paper/1910.12430/) (2019)
- [Bilevel Programming for Hyperparameter Optimization and Meta-Learning](https://mlsanity.com/paper/1806.04910/) (2018)
- [From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification](https://mlsanity.com/paper/1602.02068/) (2016)
- [Optimizing Millions of Hyperparameters by Implicit Differentiation](https://mlsanity.com/paper/1911.02590/) (2019)
- [Forward and Reverse Gradient-Based Hyperparameter Optimization](https://mlsanity.com/paper/1703.01785/) (2017)
- [Hyperparameter optimization with approximate gradient](https://mlsanity.com/paper/1602.02355/) (2016)
- [Dataset Distillation](https://mlsanity.com/paper/1811.10959/) (2018)
- [Bilevel Optimization: Convergence Analysis and Enhanced Design](https://mlsanity.com/paper/2010.07962/) (2020)
- [Fast Differentiable Sorting and Ranking](https://mlsanity.com/paper/2002.08871/) (2020)
- [On the Iteration Complexity of Hypergradient Computation](https://mlsanity.com/paper/2006.16218/) (2020)
- [Implicit Deep Learning](https://mlsanity.com/paper/1908.06315/) (2019)
- [On Differentiating Parameterized Argmin and Argmax Problems with Application to Bi-level Optimization](https://mlsanity.com/paper/1607.05447/) (2016)
- [Structured Attention Networks](https://mlsanity.com/paper/1702.00887/) (2017)
- [Smooth and Sparse Optimal Transport](https://mlsanity.com/paper/1710.06276/) (2017)
- [Differentiating Through a Cone Program](https://mlsanity.com/paper/1904.09043/) (2019)
- [Understanding Short-Horizon Bias in Stochastic Meta-Optimization](https://mlsanity.com/paper/1803.02021/) (2018)
- [SparseMAP: Differentiable Sparse Structured Inference](https://mlsanity.com/paper/1802.04223/) (2018)
- [Neural reparameterization improves structural optimization](https://mlsanity.com/paper/1909.04240/) (2019)
- [A Regularized Framework for Sparse and Structured Neural Attention](https://mlsanity.com/paper/1705.07704/) (2017)
- [Implicit differentiation of Lasso-type models for hyperparameter optimization](https://mlsanity.com/paper/2002.08943/) (2020)
- [Variational Data Assimilation with a Learned Inverse Observation Operator](https://mlsanity.com/paper/2102.11192/) (2021)
- [Super-efficiency of automatic differentiation for functions defined as a minimum](https://mlsanity.com/paper/2002.03722/) (2020)
- [Implicit differentiation for fast hyperparameter selection in non-smooth convex learning](https://mlsanity.com/paper/2105.01637/) (2021)

## Built on by (most cited first, 30 of 203)

- [CRAFT: Concept Recursive Activation FacTorization for Explainability](https://mlsanity.com/paper/2211.10154/) (2022)
- [Sinkformers: Transformers with Doubly Stochastic Attention](https://mlsanity.com/paper/2110.11773/) (2021)
- [Equivariance with Learned Canonicalization Functions](https://mlsanity.com/paper/2211.06489/) (2022)
- [Linear Adversarial Concept Erasure](https://mlsanity.com/paper/2201.12091/) (2022)
- [Gradients are Not All You Need](https://mlsanity.com/paper/2111.05803/) (2021)
- [Tutorial on amortized optimization](https://mlsanity.com/paper/2202.00665/) (2022)
- [A Survey of Contextual Optimization Methods for Decision Making under Uncertainty](https://mlsanity.com/paper/2306.10374/) (2023)
- [Theseus: A Library for Differentiable Nonlinear Optimization](https://mlsanity.com/paper/2207.09442/) (2022)
- [Amortized Implicit Differentiation for Stochastic Bilevel Optimization](https://mlsanity.com/paper/2111.14580/) (2021)
- [A-NeSI: A Scalable Approximate Method for Probabilistic Neurosymbolic Inference](https://mlsanity.com/paper/2212.12393/) (2022)
- [Synergies between Disentanglement and Sparsity: Generalization and Identifiability in Multi-Task Learning](https://mlsanity.com/paper/2211.14666/) (2022)
- [Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods](https://mlsanity.com/paper/2502.01384/) (2025)
- [Learning differentiable solvers for systems with hard constraints](https://mlsanity.com/paper/2207.08675/) (2022)
- [Leveraging sparse and shared feature activations for disentangled representation learning](https://mlsanity.com/paper/2304.07939/) (2023)
- [On Implicit Bias in Overparameterized Bilevel Optimization](https://mlsanity.com/paper/2212.14032/) (2022)
- [Learning Model Predictive Controllers with Real-Time Attention for Real-World Navigation](https://mlsanity.com/paper/2209.10780/) (2022)
- [A Unified Framework for Implicit Sinkhorn Differentiation](https://mlsanity.com/paper/2205.06688/) (2022)
- [One-step differentiation of iterative algorithms](https://mlsanity.com/paper/2305.13768/) (2023)
- [Betty: An Automatic Differentiation Library for Multilevel Optimization](https://mlsanity.com/paper/2207.02849/) (2022)
- [Rigid Body Flows for Sampling Molecular Crystal Structures](https://mlsanity.com/paper/2301.11355/) (2023)
- [Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey](https://mlsanity.com/paper/2305.05080/) (2023)
- [Fast, Differentiable and Sparse Top-k: a Convex Analysis Perspective](https://mlsanity.com/paper/2302.01425/) (2023)
- [Automatic differentiation of nonsmooth iterative algorithms](https://mlsanity.com/paper/2206.00457/) (2022)
- [Implicit Diffusion: Efficient Optimization through Stochastic Sampling](https://mlsanity.com/paper/2402.05468/) (2024)
- [Scaling physics-informed hard constraints with mixture-of-experts](https://mlsanity.com/paper/2402.13412/) (2024)
- [Differentiable Programming for Earth System Modeling](https://mlsanity.com/paper/2208.13825/) (2022)
- [Alternating Differentiation for Optimization Layers](https://mlsanity.com/paper/2210.01802/) (2022)
- [What Can the Neural Tangent Kernel Tell Us About Adversarial Robustness?](https://mlsanity.com/paper/2210.05577/) (2022)
- [TorchDEQ: A Library for Deep Equilibrium Models](https://mlsanity.com/paper/2310.18605/) (2023)
- [FSNet: Feasibility-Seeking Neural Network for Constrained Optimization with Guarantees](https://mlsanity.com/paper/2506.00362/) (2025)

## Abstract (arXiv)

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.
