[Paper Review] Efficient and Modular Implicit Differentiation

Mathieu Blondel, Quentin Berthet+6 more|arXiv (Cornell University)|May 31, 2021

Advanced Optimization Algorithms Research70 references37 citations

TL;DR

This paper introduces automatic implicit differentiation, a Python/JAX framework that differentiates optimization problem solutions by specifying the optimality conditions F, enabling modular, solver-agnostic differentiation on top of existing solvers, with Jacobian error guarantees and multiple applications.

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.

Motivation & Objective

• Lower the barrier to using implicit differentiation by letting users specify optimality conditions directly in Python.
• Combine implicit differentiation with automatic differentiation to differentiate optimization solutions without reimplementing solvers.
• Provide a modular framework that works with state-of-the-art solvers and various optimality conditions.
• Offer theoretical Jacobian precision guarantees for approximate solutions.
• Demonstrate practical bi-level optimization and sensitivity analysis through diverse applications.

Proposed method

• User defines a mapping F capturing the optimality conditions of the problem solved by the algorithm.
• Use a Python decorator (@custom_root) to attach implicit differentiation on top of the solver based on F.
• Apply the implicit function theorem to relate x*(theta) to theta via the linear system -∂1F(x*(theta),theta) ∂x*(theta) = ∂2F(x*(theta),theta).
• Compute Jacobian-vector and vector-Jacobian products (JVP/VJP) efficiently using matrix-free linear solvers (CG, GMRES, BiCGSTAB).
• Support pre/post-processing mappings so differentiation can be composed with other neural network or loss operations.
• Provide practical implementations of various optimality condition mappings (stationary, KKT, proximal gradient fixed point, projected gradient fixed point, etc.).

Experimental results

Research questions

• RQ1 Can automatic implicit differentiation cover a wide catalog of optimality conditions by using a user-defined F?
• RQ2 What are the Jacobian precision guarantees when the inner optimization is solved approximately?
• RQ3 How does the proposed framework compare in efficiency and flexibility to unrolling methods across different solvers and fixed-point representations?
• RQ4 Can the approach readily handle bi-level problems such as hyperparameter optimization, dataset distillation, and dictionary learning?
• RQ5 What are practical guidelines for implementing and differentiating through common optimization schemes (proximal, projection, mirror descent) within this framework?

Key findings

• The framework enables differentiating optimization problem solutions by differentiating the residual map F with autodiff and applying the implicit function theorem.
• Jacobian errors from approximate solutions are bounded and scale with the residual of the inner solve (Theorem 1).
• The approach recovers existing implicit differentiation methods and enables new ones within a single, modular system.
• Experiments show efficient differentiation for hyperparameter optimization of multiclass SVMs, dataset distillation, task-driven dictionary learning, and molecular dynamics sensitivity analysis.
• Using fixed-point formulations (mirror descent, proximal gradient, projected gradient) with implicit differentiation yields practical, scalable bi-level optimization workflows.
• The method achieves favorable runtime compared to unrolling in several bi-level tasks, demonstrating practical gains in speed and simplicity.