PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient and Modular Implicit Differentiation

Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-López, Fabian Pedregosa, Jean-Philippe Vert

arXiv:2105.15183v5cs.LGmath.NAstat.ML

TL;DR

Differentiating optimization problem solutions has been difficult because implicit differentiation often requires tedious, case-by-case derivations and implementations. The paper introduces automatic implicit differentiation by combining user-defined optimality conditions with autodiff and the implicit function theorem. The framework supports existing solvers, recovers many optimality-condition methods, and is demonstrated in bi-level optimization and molecular-dynamics sensitivity analysis.

  • Problem

    Implicit differentiation of optimization solutions has often required case-by-case mathematical derivations and implementations, limiting its accessibility to practitioners.

  • Method

    Users define a Python mapping F for the problem’s optimality conditions, which the framework differentiates with autodiff and the implicit function theorem while remaining applicable to existing solvers.

  • Results

    The framework covers a large catalog of optimality conditions, provides bounds for Jacobian error under approximate solves, implements four applications, and supports molecular-dynamics sensitivity analysis.

  • Takeaways & Limitations

    The approach makes implicit differentiation easier to formulate while retaining the efficiency of state-of-the-art optimization solvers.

  • Takeaways & Limitations

    The approach theoretically requires differentiability and invertibility conditions for the implicit function theorem to hold.

Abstract

from arXiv · show

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.

1 Introduction

Differentiating optimization solutions is difficult because explicit formulas are usually unavailable and existing implicit approaches often require case-by-case derivations. The paper introduces automatic implicit differentiation, which combines autodiff with the implicit function theorem through user-defined optimality conditions and existing solvers.

  • Optimization solutions typically lack explicit input formulas, preventing direct use of autodiff.
  • Unrolling optimization iterations requires reimplementing algorithms for autodiff and can incur computational costs tied to variables or iterations.Forward-mode time scales with the number of variables, while reverse-mode memory scales with the number of iterations.
  • Implicit differentiation avoids unrolling but has often required tedious, case-by-case mathematical derivations and implementations.
  • The proposed framework lets users define Python mappings that encode optimality conditions, then applies autodiff and the implicit function theorem automatically.
  • The framework integrates with existing solvers, supports a broad catalog of optimality conditions, provides approximate-solution Jacobian bounds, and implements four illustrative applications.
  • The authors present the approach as an efficient, modular blueprint that can extend beyond its JAX implementation to other frameworks.

2 Automatic implicit differentiation

Automatic implicit differentiation lets users specify optimality conditions in Python and combines autodiff with the implicit function theorem to differentiate solver solutions. The framework uses matrix-free linear solves and separates optimality-condition specifications from the differentiation mechanism, while applying under differentiability and invertibility conditions.

  • Motivation: Manual, case-by-case derivations make traditional implicit differentiation difficult to implement for constrained optimization problems.
  • Framework: Users define a mapping F that captures the problem’s optimality conditions, then attach @custom_root to an existing solver.For unconstrained differentiable objectives, F can be the gradient of the objective in its first argument.
  • Framework: The implicit function theorem reduces differentiation of x⋆(θ) to solving a linear system involving the partial derivatives of F.For one-dimensional root problems, the gradient simplifies to ∇x⋆(θ) = B^T/A.
  • Examples and scope: The approach recovers existing and new implicit differentiation methods by allowing different optimality conditions or fixed points, with computational trade-offs summarized in Table 1.Examples include stationary conditions, KKT conditions, proximal gradient fixed points, and mirror descent fixed points.
  • Computation: Jacobian-vector and vector-Jacobian products support forward- and reverse-mode autodiff without explicitly forming Jacobians, using matrix-free linear-system solvers.Conjugate gradient, GMRES, and BiCGSTAB can operate using only matrix-vector products.
  • Limitations: The framework theoretically requires the implicit function theorem’s differentiability and invertibility conditions, including in the non-smooth settings covered by its stated assumptions.The paper identifies nonsmooth implicit function theorems as a direction for extending the approach beyond these conditions.

3 Jacobian precision guarantees

The paper defines Jacobian estimates at approximate optimization solutions and bounds their error relative to the exact solution's Jacobian. Specialization to gradient-descent fixed points yields precision guarantees and a factor-t advantage over autodiff in the strongly-convex setting.

  • J(x, θ) is defined by solving A(x, θ)J(x, θ) = B(x, θ), when A(x, θ) is invertible.Here A = −∂1F and B = ∂2F.
  • The estimate J(ˆx, θ) approximates the true Jacobian ∂x⋆(θ), rather than the Jacobian of the solver's approximate trajectory ˆx(θ).This distinction separates implicit differentiation from differentiating through unrolled iterations.
  • ∥J(ˆx, θ) − ∂x⋆(θ)∥ < C∥ˆx − x⋆(θ)∥, so Jacobian-estimation error is at most of the same order as solution error.The theorem assumes regularity conditions on A and B near x⋆(θ), including boundedness and Lipschitz behavior of B.
  • The theorem is designed for general optimality mappings F while allowing specialization to particular fixed points and problem structures.The authors note that uniform conditioning of A makes the bound valid for all approximate solutions ˆx.
  • A factor-t precision gain over automatic differentiation is obtained for Jacobian estimation after t gradient-descent iterations in the strongly-convex setting.The result is illustrated empirically by relating solution error and Jacobian error across iteration counts.

4 Experiments

The experiments evaluate implicit differentiation across bi-level optimization, dataset distillation, task-driven dictionary learning, and molecular-dynamics sensitivity analysis. The framework supports flexible solver and fixed-point choices, accelerates dataset distillation, remains competitive with far fewer variables, and handles a sensitivity task where unrolling typically fails.

  • Hyperparameter optimization of multiclass SVMs: Implicit differentiation and unrolling are compared for multiclass SVM hyperparameter optimization across multiple problem sizes and solver choices.The experiments use mirror descent, proximal gradient, and block coordinate descent, with differentiation fixed points independently selectable from the solver.
  • Dataset distillation: Dataset distillation uses implicit differentiation for high-dimensional inner and outer problems, with the same solver running 4 times faster than the original unrolled approach.The experiment uses MNIST, with 784 inner parameters and 7,840 outer parameters.
  • Task-driven dictionary learning: Task-driven dictionary learning reaches classification performance competitive with state-of-the-art L1 or L2-regularized logistic regression using 100 times fewer variables.The evaluation concerns breast cancer survival prediction from 1,000 gene-expression values.
  • Sensitivity analysis of molecular dynamics: In molecular-dynamics sensitivity analysis, forward-mode implicit differentiation computes particle-position sensitivity, whereas unrolled differentiation typically fails to converge.The failure is attributed to the discontinuous optimization method; unrolling happens to work for total energy in the original setting.

5 Conclusion

The framework expresses optimization optimality conditions as Python mappings or fixed points, then uses autodiff to differentiate them implicitly. It supports diverse optimization structures, including KKT systems, projections, fixed-point methods, conic programs, and Frank–Wolfe formulations.

  • Framework: The framework differentiates optimization solutions by representing them as roots of user-defined optimality mappings F or fixed points T.The mapping specification is decoupled from the differentiation mechanism, so no case-specific derivative derivation is required.
  • Fixed points: Mirror descent, proximal block coordinate descent, Newton, and projected-gradient fixed points provide alternative implicit representations of optimization solutions.Mirror descent extends projected-gradient fixed points beyond Euclidean geometry, while Newton fixed points recover familiar linear systems.
  • Optimization criteria: KKT conditions for equality- and inequality-constrained quadratic programs become roots of F, including stationarity, feasibility, and complementary slackness.Users define f, H, and G directly in Python.
  • Conic programming: Conic programs fit the framework by reducing their solution process to a root of a residual map and recovering primal and slack solutions from the solver output.The required Jacobian products center on projection onto cones and their duals.
  • Frank-Wolfe: Frank–Wolfe differentiation rewrites optimization over a polytope as simplex optimization, differentiates the weights implicitly, and obtains solution derivatives by the product rule.Visited-vertex sparsity can reduce the linear system size, and implementations commonly maintain the convex weights.
  • Implementation: The library includes reusable mappings and demonstrates writing KKT and mirror-descent mappings from scratch with autodiff-based Jacobian products.The same implementation pattern supports equality-only or inequality-only constraints and parameter-dependent mirror maps when needed.

B.2 Code examples for experiments

The experiments are implemented by composing reusable projection and optimization building blocks with autodiff-based Jacobian products. The examples cover common convex sets and proximity operators used by the framework.

  • Experiment setup: Experiments compute outer-objective gradients with the framework and use those gradients with gradient-based solvers for the outer optimization problem.The code examples use jnp as shorthand for jax.numpy.
  • Projection building blocks: Closed-form projections, including norm-ball, non-negative-orthant, and box projections, can be differentiated automatically through their function compositions.The ℓ1-ball reduces to simplex projection, while ℓ2- and ℓ∞-ball projections have closed forms.
  • Projection building blocks: Simplex projection has exact algorithms with O(d) expected or O(d log d) worst-case time, while its KL counterpart has a closed form related to softmax.The Euclidean projection Jacobian uses the support indicator s.
  • Projection building blocks: Affine, hyperplane, half-space, transportation, Birkhoff, order-simplex, and polyhedral projections are handled through pseudoinverses, KKT systems, isotonic optimization, or implicit differentiation.Order-simplex projections can be solved exactly in O(d log d) time using the Pool Adjacent Violators algorithm.
  • Advanced operators: More advanced projections and proximity operators are differentiated using KKT conditions, dual formulations, implicit fixed points, or known Jacobians when closed forms are unavailable.The library includes lasso, elastic-net, and group-lasso proximity operators with closed-form expressions.

D Jacobian precision proofs

The paper derives Jacobian estimates for approximate optimization solutions and establishes error bounds under smoothness, convexity, and conditioning assumptions. The analysis specializes the general result to gradient and proximal-gradient fixed points.

  • General precision result: Theorem 1 analyzes the difference between Jacobian estimates evaluated at an approximate solution and at the exact solution.The proof compares the linear systems A⋆J⋆ = B⋆ and ÂĴ = B̂.
  • Gradient descent: For gradient descent fixed points, the estimated Jacobian uses the inverse Hessian-like term and parameter derivative of the gradient evaluated at the approximate solution.The corollary assumes twice differentiability, strong convexity, Lipschitz Hessians, and bounded parameter-gradient derivatives.
  • Proximal gradient: For proximal-gradient fixed points, the analysis assumes smooth strongly convex f, convex or strongly convex g, and regularity of the proximal-gradient components.The special case considered has g independent of θ and includes µ = 0 for merely convex g.
  • Proximal gradient: The proximal-gradient proof derives the Jacobian estimate through the smooth optimality condition of the proximity operator and the inverse rule for Jacobians.This yields a bound on the approximation error between the estimated Jacobian and the true solution Jacobian.

E The Lasso case

The Lasso analysis addresses the nonsmoothness of the soft-thresholding fixed point and identifies when the smooth implicit function theorem applies. Away from solution-path kinks, the required regularity holds with probability one under random-design assumptions.

  • Motivation: Nonsmooth proximity operators can invalidate direct use of the smooth implicit function theorem, requiring analysis of F near the solution.The paper studies this issue for Lasso differentiation with respect to the regularization parameter.
  • Lasso formulation: The Lasso solution uses f(x, θ) = (1/2)∥Φx − b∥2 and g(x, θ) = e^θ∥x∥1, whose proximal operator is soft-thresholding.The resulting root equation is a proximal-gradient fixed point parameterized by η > 0.
  • Solution path: The Lasso solution path is continuous and piecewise linear, with finitely many nondifferentiable kinks where the solution is not differentiable.At those kinks, the smooth implicit function theorem using the root equation does not apply.
  • Regularity result: With Φ entries drawn from a continuous distribution, the smooth implicit function theorem holds with probability one at every non-kink point of the Lasso solution path.The proof establishes continuous differentiability near the solution and invertibility of the solution-variable Jacobian.

F Experimental setup and additional results

The experiments use JAX and scikit-learn, under Apache-2 and BSD licenses respectively.

  • The experiments use JAX, which is Apache-2 licensed.
  • The experiments use scikit-learn, which is BSD licensed.
  • The software stack combines JAX with scikit-learn.

F.1 Hyperparameter optimization of multiclass SVMs

The multiclass SVM experiments vary problem size and compare several inner solvers, implicit differentiation, and unrolling. Implicit differentiation maintains comparable validation loss while avoiding many GPU memory failures observed for unrolling.

  • Experimental setup: The experiments vary the number of features from 100 to 10,000 across synthetic five-class datasets with 700 training and 200 validation samples.
  • Experimental setup: Mirror descent, accelerated proximal gradient descent, and block coordinate descent serve as inner problem solvers with tuned hyperparameters.
  • Runtime comparison: GPU runtime comparisons use a single 16GB NVIDIA P100 and report per-outer-iteration averages with 90% confidence intervals.
  • Runtime comparison: Unrolling encounters out-of-memory errors for most larger GPU problems: p ≥2000 for mirror descent and p ≥750 for proximal and block coordinate descent.
  • Validation loss: All approaches achieve qualitatively indistinguishable validation-loss solutions across the considered problem sizes.
  • Jacobian accuracy: The Jacobian-error experiment varies feature count and evaluates estimates against a liblinear solution computed with tolerance 10^-9.

F.2 Task-driven dictionary learning

The task-driven dictionary-learning experiment uses breast-cancer gene-expression and survival data, repeated train-validation-test splits, and parameter selection by validation AUC.

  • Data and evaluation: The dataset contains gene-expression measurements for 20,531 genes across 1,093 retained primary-tumor samples.
  • Data and evaluation: Ten random splits divide each cohort into 60% training, 20% validation, and 20% test sets.
  • Data and evaluation: Parameters are selected by validation AUC before retraining each method on the combined training and validation data.
  • Compared methods: The comparison includes l1- and l2-regularized logistic regression, unsupervised dictionary learning, and task-driven dictionary learning.

F.3 Dataset Distillation

The dataset-distillation experiment compares implicit and unrolled differentiation using gradient-based inner and outer optimization. Implicit differentiation produces the same output in one quarter of the runtime.

  • Optimization setup: The inner problem uses gradient descent with backtracking line search, while the outer problem uses momentum with a fixed step size.
  • Runtime comparison: After 4,000 outer-loop iterations, implicit differentiation takes 1h55 on CPU versus 8h05 for unrolled differentiation.
  • Result: Implicit differentiation is 4 times faster while producing the same distilled MNIST output as unrolled differentiation.

F.4 Molecular dynamics

The molecular-dynamics experiment compares implicit and unrolled differentiation for position sensitivities, using a BiCGSTAB tangent linear solve on 128 particles in two dimensions. Across 40 random initial conditions, implicit differentiation converges whereas gradients through the unrolled FIRE optimizer generally do not.

  • Sensitivity target: Position sensitivities require implicit differentiation or unrolling, unlike the total-energy gradient, which vanishes at the optimum by Danskin’s theorem.The experiment therefore targets ∂x⋆(θ) directly rather than differentiating the objective with respect to particle diameter.
  • Experimental setup: The setup uses k = 128 particles in m = 2 dimensions, adapted from the JAX-MD meta-optimization example.The comparison uses BiCGSTAB to perform the tangent linear solve.
  • Results: For 40 random initial conditions, gradients through the unrolled FIRE optimizer do not converge, in contrast to implicit differentiation.Figure 17 plots the L1 norm of position sensitivities, with different colored lines representing the initial conditions.
Loading 2105.15183v5…