Efﬁcient and Modular Implicit Differentiation
Anonymous Author(s)
AfﬁliationAddress
Abstract
Automatic differentiation (autodiff) has revolutionized machine learning.
It allowsto express complex computations by composing elementary ones in creative waysand removes the burden of computing their derivatives by hand.
More recently,differentiation of optimization problem solutions has attracted widespread attentionwith applications such as optimization layers, and in bi-level problems such ashyper-parameter optimization and meta-learning.
However, so far, implicit differentiation remained difﬁcult to use for practitioners, as it often required case-by-casetedious mathematical derivations and implementations.
In this paper, we proposeautomatic implicit differentiation, an efﬁcient and modular approach for implicitdifferentiation of optimization problems.
In our approach, the user deﬁnes directlyin Python a function F capturing the optimality conditions of the problem to bedifferentiated.
Once this is done, we leverage autodiff of F and the implicit function theorem to automatically differentiate the optimization problem.
Our approachthus combines the beneﬁts of implicit differentiation and autodiff.
It is efﬁcient asit can be added on top of any state-of-the-art solver and modular as the optimalitycondition speciﬁcation is decoupled from the implicit differentiation mechanism.
We show that seemingly simple principles allow to recover many existing implicitdifferentiation methods and create new ones easily.
We demonstrate the ease offormulating and solving bi-level optimization problems using our framework.
Wealso showcase an application to the sensitivity analysis of molecular dynamics.
Introduction
Automatic differentiation (autodiff) is now an inherent part of machine learning software.
It allowsto express complex computations by composing elementary ones in creative ways and removes thetedious burden of computing their derivatives by hand.
In parallel, the differentiation of optimizationproblem solutions has found many applications.
A classical example is bi-level optimization, whichtypically involves computing the derivatives of a nested optimization problem in order to solve anouter one.
Examples of applications in machine learning include hyper-parameter optimization [23,76, 69, 38, 12, 13], neural networks [59], and meta-learning [39, 71].
Another line of active researchinvolving differentiation of optimization problem solutions are optimization layers [55, 6, 64, 31, 46],which can be used to encourage structured outputs, and implicit deep networks [8, 36, 43, 44, 72],which have a smaller memory footprint than backprop-trained networks.
Since optimization problem solutions typically do not enjoy an explicit formula in terms of theirinputs, autodiff cannot be used directly to differentiate these functions.
In recent years, two mainapproaches have been developed to circumvent this problem.
The ﬁrst one consists of unrolling theiterations of an optimization algorithm and using the ﬁnal iteration as a proxy for the optimizationproblem solution [81, 32, 29, 39, 1].
This allows to explicitly construct a computational graph this requires a reimplementation of the algorithm using the autodiff system, and not all algorithms are necessarily autodiff friendly.
Moreover, forward-mode autodiff has time complexity that scaleslinearly with the number of variables and reverse-mode autodiff has memory complexity that scaleslinearly with the number of algorithm iterations.
In contrast, a second approach consists in implicitlyrelating an optimization problem solution to its inputs using optimality conditions.
In a machinelearning context, such implicit differentiation has been used for stationarity conditions [11, 59], KKTconditions [23, 45, 6, 66, 65] and the proximal gradient ﬁxed point [64, 12, 13].
An advantage ofimplicit differentiation is that a solver reimplementation is not needed, allowing to build upon decadesof state-of-the-art software.
Although implicit differentiation has a long history in numerical analysis[48, 10, 57, 20], so far, it remained difﬁcult to use for practitioners, as it required a case-by-casetedious mathematical derivation and implementation.
CasADi [7] allows to differentiate variousoptimization and root ﬁnding problem algorithms provided by the library.
However, it does not allowto easily add implicit differentiation on top of existing solvers from optimality conditions expressedby the user, as we do.
A recent tutorial explains how to implement implicit differentiation in JAX[34].
However, the tutorial requires the user to take care of low-level technical details and does notcover a large catalog of optimality condition mappings as we do.
Other work [2] attempts to addressthis issue by adding implicit differentiation on top of cvxpy [30].
This works by reducing all convexoptimization problems to a conic program and using conic programming’s optimality conditions toderive an implicit differentiation formula.
While this approach is very generic, solving a convexoptimization problem using a conic programming solver—an ADMM-based splitting conic solver[67] in the case of cvxpy—is rarely state-of-the-art for every problem instance.
In this work, we ambition to achieve for optimization problem solutions what autodiff did forcomputational graphs.
We propose automatic implicit differentiation , a simple approach to addimplicit differentiation on top of any existing solver.
In this approach, the user deﬁnes directly in
Python a mapping function F capturing the optimality conditions of the problem solved by thealgorithm.
Once this is done, we leverage autodiff of F combined with the implicit function theoremto automatically differentiate the optimization problem solution.
Our approach is generic , yet itcan exploit the efﬁciency of state-of-the-art solvers.
It therefore combines the beneﬁts of implicitdifferentiation and autodiff.
To summarize, we make the following contributions
• We describe our framework and its JAX [21, 42] implementation ( https://github.com/google/jaxopt/ ).
Our framework signiﬁcantly lowers the barrier to use implicit differen-
tiation, thanks to the seamless integration in JAX, with low-level details all abstracted away.
• We instantiate our framework on a large catalog of optimality conditions (Table 1), recoveringexisting schemes and obtaining new ones, such as the mirror descent ﬁxed point based one.
• On the theoretical side, we provide new bounds on the Jacobian error when the optimizationproblem is only solved approximately, and empirically validate them.• We implement four illustrative applications , demonstrating our framework’s ease of use.
Beyond our software implementation in JAX, we hope this paper provides a self-contained blueprintfor creating an efﬁcient and modular implementation of implicit differentiation in other frameworks.
Notation.
We denote the gradient and Hessian of f : R d → R evaluated at x ∈ R d by ∇ f ( x ) ∈ R dand ∇ 2 f ( x ) ∈ R d × d .
We denote the Jacobian of F : R d → R p evaluated at x ∈ R d by ∂F ( x ) ∈
R p × d .
When f or F have several arguments, we denote the gradient, Hessian and Jacobian inthe i th argument by ∇ i , ∇ 2 i and ∂ i , respectively.
The standard probability simplex is denoted by(cid:52)
d := { x ∈ R d : (cid:107) x (cid:107) 1 = 1 , x ≥ 0 } .
For any set C ⊂ R d , we denote the indicator function
I C : R d → R ∪ { + ∞} where I C ( x ) =
if x ∈ C , I C ( x ) = + ∞ otherwise.
For a vector or matrix A ,we note (cid:107) A (cid:107) the Frobenius (or Euclidean) norm, and (cid:107) A (cid:107) op the operator norm.
Automatic implicit differentiation
2.1 General principles
Overview.
Contrary to autodiff through unrolled algorithm iterations, implicit differentiationtypically involves a manual, sometimes complicated, mathematical derivation.
For instance, numerousworks [23, 45, 6, 66, 65] use Karush–Kuhn–Tucker (KKT) conditions in order to relate a constrainedoptimization problem’s solution to its inputs, and to manually derive a formula for its derivatives.
The derivation and implementation in these works are typically case-by-case.
X_train, y_train = load_data() # Load features and labels def f(x, theta): # Objective function residual = jnp.dot(X_train, x) - y_
[Figure]
return (jnp.sum(residual ** 2) + theta * jnp.sum(x ** 2)) /
# Since f is differentiable and unconstrained, the optimality # condition F is simply the gradient of f in the 1st argument
F = jax.grad(f, argnums=0)
[Figure] [Figure]
(F)
[Figure]
_solver(init_x, theta):
[Figure]
Initialization
[Figure]
in this
[Figure] [Figure]
Identity
[Figure]
In this work, we propose a generic way to easily add implicit differentiation on top of existing solvers.
In our approach, the user deﬁnes directly in Python a mapping function F capturing the optimalityconditions of the problem solved by the algorithm.
We provide reusable building blocks to easilyexpress such F .
The provided F is then plugged into our Python decorator @custom_root , whichwe append on top of the solver declaration we wish to differentiate.
Under the hood, we combine theimplicit function theorem and autodiff of F to automatically differentiate the optimization problemsolution.
A simple illustrative example is given in Figure 1.
Differentiating a root.
Let F : R d × R n → R d be a user-provided mapping, capturing the optimalityconditions of a problem.
An optimal solution, denoted x (cid:63) ( θ ) , should be a root of F :
[Equation]
We can see x (cid:63) ( θ ) as an implicitly deﬁned function of θ ∈ R n , i.e., x (cid:63) : R n → R d .
More precisely,from the implicit function theorem [48, 57], we know that for ( x 0 , θ 0 ) satisfying F ( x 0 , θ 0 )
=
0with a continuously differentiable F Olive color indicates assumptions already present beforesubmission , if the Jacobian ∂ 1 F evaluated at ( x 0 , θ 0 ) is a square invertiblematrix, then there exists a function x (cid:63) ( · ) deﬁned on a neighborhood of θ 0 such that x (cid:63) ( θ 0 ) =
x 0 .
Furthermore, for all θ in this neighborhood, we have that F ( x (cid:63) ( θ ) , θ ) = 0 and ∂x (cid:63) ( θ ) exists.
Usingthe chain rule, the Jacobian ∂x (cid:63) ( θ ) satisﬁes
[Equation]
Computing ∂x (cid:63) ( θ ) therefore boils down to the resolution of the linear system of equations
[Equation]
When (1) is a one-dimensional root ﬁnding problem ( d = 1 ), (2) becomes particularly simple sincewe then have ∇ x (cid:63) ( θ )
=
B (cid:62) /A , where A is a scalar value.
We will show that existing and new implicit differentiation methods all reduce to this simple principle.
We call our approach automatic implicit differentiation as the user can freely express the optimization solution to be differentiated through the optimality conditions F .
Our approach is efﬁcientas it can be added on top of any state-of-the-art solver and modular as the optimality conditionspeciﬁcation is decoupled from the implicit differentiation mechanism.
This contrasts with existingworks, where the derivation and implementation are speciﬁc to each optimality condition.
Differentiating a ﬁxed point.
We will encounter numerous applications where x (cid:63) ( θ ) is insteadimplicitly deﬁned through a ﬁxed point :
[Equation]
where T : R d × R n → R d .
This can be seen as a particular case of (1) by deﬁning the residual
[Equation]
In this case, when T is continuously differentiable Blue color indicates newly addedassumptions or remarks , using the chain rule, we have
[Equation]
Computing JVPs and VJPs.
In most practical scenarios, it is not necessary to explicitly formthe Jacobian matrix, and instead it is sufﬁcient to left-multiply or right-multiply by ∂ 1 F and ∂ 2 F .
These are called vector-Jacobian product (VJP) and Jacobian-vector product (JVP), and are useful forintegrating x (cid:63) ( θ ) with reverse-mode and forward-mode autodiff, respectively.
Oftentimes, F will beexplicitly deﬁned.
In this case, computing the VJP or JVP can be done via autodiff.
In some cases, Fmay itself be implicitly deﬁned, for instance when F involves the solution of a variational problem.
In this case, computing the VJP or JVP will itself involve implicit differentiation.
The right-multiplication (JVP) between J = ∂x (cid:63) ( θ ) and a vector v , Jv , can be computed efﬁcientlyby solving A ( Jv ) =
Bv .
The left-multiplication (VJP) of v (cid:62) with J , v (cid:62) J , can be computed by ﬁrstsolving A (cid:62) u
= v .
Then, we can obtain v (cid:62) J by v (cid:62) J = u (cid:62) AJ = u (cid:62) B .
Note that when B changesbut A and v remain the same, we do not need to solve A (cid:62) u
= v once again.
This allows to computethe VJP w.r.t.
different variables while solving only one linear system.
To solve these linear systems, we can use the conjugate gradient method [51] when A is symmetricpositive semi-deﬁnite and GMRES [74] or BiCGSTAB [79] otherwise.
These algorithms are allmatrix-free: they only require matrix-vector products.
Thus, all we need from F is its JVPs or VJPs.
An alternative to GMRES/BiCGSTAB is to solve the normal equation AA (cid:62) u
=
Av using conjugategradient.
This can be implemented using JAX’s transpose routine jax.linear_transpose [41].
Incase of non-invertibility, a common heuristic is to solve a least squares min J (cid:107) AJ − B (cid:107) 2 instead.
Pre-processing and post-processing mappings.
Oftentimes, the goal is not to differentiate θ perse, but the parameters of a function producing θ .
One example of such pre-processing is to convertthe parameters to be differentiated from one form to another canonical form, such as a quadraticprogram [6] or a conic program [2].
Another example is when x (cid:63) ( θ ) is used as the output of a neuralnetwork layer, in which case θ is produced by the previous layer.
Likewise, x (cid:63) ( θ ) will often notbe the ﬁnal output we want to differentiate.
One example of such post-processing is when x (cid:63) ( θ )is the solution of a dual program and we apply the dual-primal mapping to recover the solution ofthe primal program.
Another example is the application of a loss function, in order to reduce x (cid:63) ( θ )to a scalar value.
We leave the differentiation of such pre/post-processing mappings to the autodiffsystem, allowing to compose functions in complex ways.
Implementation details.
When a solver function is decorated with @custom_root , we usejax.custom_jvp and jax.custom_vjp to automatically add custom JVP and VJP rules to thefunction, overriding JAX’s default behavior.
As mentioned above, we use linear system solvers basedon matrix-vector products and therefore we only need access to F through the JVP or VJP with ∂ 1 Fand ∂ 2 F .
This is done by using jax.jvp and jax.vjp , respectively.
Note that, as in Figure 1, thedeﬁnition of F will often include a gradient mapping ∇ 1 f ( x, θ ) .
Thankfully, JAX supports secondorder derivatives transparently.
For convenience, our library also provides a @custom_fixed_pointdecorator, for adding implicit differentiation on top of a solver, given a ﬁxed point iteration T ; seecode examples in Appendix B.
2.2 Examples
We now give various examples of mapping F or ﬁxed point iteration T , recovering existing implicitdifferentiation methods and creating new ones.
Each choice of F or T implies different trade-offs interms of computational oracles ; see Table 1.
Source code examples are given in Appendix B.
Stationary point condition.
The simplest example is to differentiate through the implicit function
[Equation]
where f : R d × R n → R is twice differentiable, ∇ 1 f is continuously differentiable, and ∇ 21 f isinvertible at ( x (cid:63) ( θ ) , θ ) .
In this case, F is simply the gradient mapping F (
[Equation]
f (
[Equation]
[Table]
We then have ∂ 1 F ( x, θ ) = ∇ 21 f ( x, θ ) and ∂ 2 F ( x, θ )
=
∂ 2 ∇ 1 f ( x, θ ) , the Hessian of f in its ﬁrstargument and the Jacobian in the second argument of ∇ 1 f ( x, θ ) .
In practice, we use autodiff tocompute Jacobian products automatically.
Equivalently, we can use the gradient descent ﬁxed point
T ( x, θ ) = x
− η ∇ 1 f ( x, θ ) , (5) for all η > 0 .
Using (3), it is easy to check that we obtain the same linear system since η cancels out.
KKT conditions.
As a more advanced example, we now show that the KKT conditions, manuallydifferentiated in several works [23, 45, 6, 66, 65], ﬁt our framework.
As we will see, the key will beto group the optimal primal and dual variables as our x (cid:63) ( θ ) .
Let us consider the general problemargmin z ∈ R p
f ( z, θ ) subject to G ( z, θ ) ≤ 0 , H ( z, θ )
= 0 , where z ∈ R p is the primal variable, f : R p × R n → R , G : R p × R n → R r and H :
R p × R n → R qare twice differentiable convex functions, and ∇ 1 f , ∂ 1 G and ∂ 1 H are continuously differentiable.
The stationarity, primal feasibility and complementary slackness conditions give
[Equation]
where ν ∈ R q and λ ∈ R r + are the dual variables, also known as KKT multipliers.
The primaland dual feasibility conditions can be ignored almost everywhere [34].
The system of (potentiallynonlinear) equations (6) ﬁts our framework, as we can group the primal and dual solutions asx (cid:63) ( θ ) =
( z (cid:63) ( θ ) , ν (cid:63) ( θ ) , λ (cid:63) ( θ )) to form the root of a function F ( x (cid:63) ( θ ) , θ ) , where F : R d × R n → R dand d = p + q + r .
The primal and dual solutions can be obtained from a generic solver, such asan interior point method.
In practice, the above mapping F will be deﬁned directly in Python (see
Figure 7 in Appendix B) and F will be differentiated automatically via autodiff.
Proximal gradient ﬁxed point.
Unfortunately, not all algorithms return both primal and dualsolutions.
Moreover, if the objective contains non-smooth terms, proximal gradient descent may bemore efﬁcient.
We now discuss its ﬁxed point [64, 12, 13].
Let x (cid:63) ( θ ) be implicitly deﬁned as
[Equation]
where f : R d × R n → R is twice-differentiable convex and g : R d × R n → R is convex but possiblynon-smooth.
Let us deﬁne the proximity operator associated with g by
[Equation]
To implicitly differentiate x (cid:63) ( θ ) , we use the ﬁxed point mapping [68, p.150]
[Equation]
for any step size η > 0 .
The proximity operator is 1 -Lipschitz continuous [63].
By Rademacher’stheorem, it is differentiable almost everywhere.
If, in addition, it is continuously differentiable in aneighborhood of ( x (cid:63) ( θ ) , θ ) and if I − ∂ 1 T ( x (cid:63) ( θ ) , θ ) is invertible, then our framework to differentiatex (cid:63) ( θ ) applies.
Similar assumptions are made in [13].
Many proximity operators enjoy a closed formand can easily be differentiated, as discussed in Appendix C.
An implementation is given in Figure 2.
[Figure]
Projected gradient ﬁxed point.
As a special case, when g ( x, θ ) is the indicator function I C ( θ ) ( x ) ,where C ( θ ) is a convex set depending on θ , we obtain
[Equation]
The proximity operator prox g becomes the Euclidean projection onto C ( θ )
[Equation]
and (7) becomes the projected gradient ﬁxed point
[Equation]
Compared to the KKT conditions, this ﬁxed point is particularly suitable when the projection enjoysa closed form.
We discuss how to compute the JVP / VJP for a wealth of convex sets in Appendix C.
Current limitations.
While we have not observed issues in practice, we note that the approachdeveloped in this section theoretically only applies to settings where x (cid:63) ( θ ) is a differentiable functionat θ , and where the implicit function theorem is valid, namely, where optimality conditions satisfythe differentiability and invertibility conditions stated in §2.1. While this covers a wide range ofsituations even for non-smooth optimization problems (e.g., the solution of a Lasso regression is a.e.differentiable with respect to the regularization parameter), an interesting direction for future work isto extend the framework to handle cases where the differentiability and invertibility conditions arenot satisﬁed, using, e.g., the theory of nonsmooth implicit function theorems [25, 19].
Jacobian precision guarantees
In practice, either by the limitations of ﬁnite precision arithmetic or because we perform a ﬁnitenumber of iterations, we rarely reach the exact solution x (cid:63) ( θ ) .
Instead, we reach an approximatesolution ˆ x and apply the implicit differentiation equation (2) at this approximate solution.
Thismotivates the need for precision guarantees of this approach.
We introduce the following formalism.
Deﬁnition 1.
Let F : R d × R n → R d be a continously differentiable optimality criterion mapping.
Let A := − ∂ 1 F and B := ∂ 2 F .
We deﬁne the Jacobian estimate at ( x, θ ) , when A ( x, θ ) is invertible, asthe solution to the linear equation A ( x, θ ) J ( x, θ ) =
B ( x, θ ) .
It is a function J : R d × R n → R d × n .
It holds by construction that J ( x (cid:63) ( θ ) , θ ) =
∂x (cid:63) ( θ ) .
Computing J (ˆ x, θ ) for an approximate solutionˆ x of x (cid:63) ( θ ) therefore allows to approximate the true Jacobian ∂x (cid:63) ( θ ) .
In practice, an algorithm usedto solve (1) depends on θ .
Note however that, what we compute is not the Jacobian of ˆ x ( θ ) , unlikeworks differentiating through unrolled algorithm iterations, but an estimate of ∂x (cid:63) ( θ ) .
We thereforeuse the notation ˆ x , leaving the dependence on θ implicit.
We develop bounds of the form (cid:107) J (ˆ x, θ ) − ∂x (cid:63) ( θ ) (cid:107) <
C (cid:107)
ˆ x − x (cid:63) ( θ ) (cid:107) , hence showing that the erroron the estimated Jacobian is at most of the same order as that of ˆ x as an approximation of x (cid:63) ( θ ) .
These bounds are based on the following main theorem, whose proof is included in Appendix D.
Theorem 1.
Let F : R d × R n → R d be continuously differentiable.
If there are α, β, γ, ε, R > 0 s.t. A = − ∂ 1 F and B = ∂ 2 F satisfy, for all v ∈ R d , θ ∈ R n and x s.t. (cid:107) x − x (cid:63) ( θ ) (cid:107) ≤ ε : is well-conditioned, Lipschitz: (cid:107) A ( x θ ) v (cid:107) ≥ α (cid:107) v (cid:107) , (cid:107)
A ( x θ − A ( x (cid:63) ( θ θ (cid:107) op γ (cid:107) x − bounded and Lipschitz: (cid:107) B ( x (cid:63) ( θ
, θ (cid:107) R , (cid:107) B ( x, θ ) −
(cid:63) ( θ ) , θ ) (cid:107) (cid:107) x − x (cid:63) ( (cid:107) Under these conditions, when (cid:107) ˆ x − x (cid:63) ( θ ) (cid:107) ≤ ε , we have
[Equation] This result is inspired by [52, Theorem 7.2], that is concerned with the stability of solutions to inverseproblems.
As a difference, we consider that A ( · , θ ) is uniformly well-conditioned, rather than only atx (cid:63) ( θ ) .
This does not affect the ﬁrst order in ε of this bound, and makes it valid for all ˆ x .
Our goalwith Theorem 1 is to provide a result that works for general F but can be tailored to speciﬁc cases.
In particular, for the gradient descent ﬁxed point (5), this yields [Equation] x
[Equation]
By specializing Theorem 1 for this ﬁxed point, we obtain Jacobian precision guarantees withconditions directly on f rather than F ; see Corollary 1 in Appendix D.
These guarantees holdfor instance for the dataset distillation experiment in Section 4.
Our analysis reveals in particularthat Jacobian estimation by implicit differentiation gains a factor of t compared to automaticdifferentiation , after t iterations of gradient descent in the strongly-convex setting [1, Proposition3.2].
While our guarantees concern the Jacobian of x (cid:63) ( θ ) , we note that other studies [47, 54, 13] giveguarantees on hypergradients (i.e., the gradient of an outer objective). [Equation] [Equation]
[Equation]
3.0 || x ( ) x || 2 0.
|| x ( ) J ( x , [Figure] We illustrate these results on ridge regression, wherex (cid:63) ( θ ) =
argmin x (cid:107) Φ x − y (cid:107)
+
(cid:80)
i θ i x 2 i .
This problemhas the merit that the solution x (cid:63) ( θ ) and its Jacobian∂x (cid:63) ( θ ) are available in closed form.
By running gradientdescent for t iterations, we obtain an estimate ˆ x of x (cid:63) ( θ )and an estimate J (ˆ x, θ ) of ∂x (cid:63) ( θ ) ; cf.
Deﬁnition 1.
Bydoing so for different numbers of iterations t , we cangraph the relation between the error (cid:107) x
(cid:63) ( θ ) − ˆ x (cid:107) 2 andthe error (cid:107) ∂x (cid:63) ( θ ) − J (ˆ x, θ ) (cid:107) 2 , as shown in Figure 3,empirically validating Theorem 1.
The results in Figure 3were obtained using the diabetes dataset from [35], withother datasets yielding a qualitatively similar behavior.
We derive similar guarantees in Corollary 2 in Appendix
D for proximal gradient descent.
Experiments
In this section, we demonstrate the ease of solving bi-level optimization problems with our framework.
We also present an application to the sensitivity analysis of molecular dynamics.
4.1 Hyperparameter optimization of multiclass SVMs
In this example, we consider the hyperparameter optimization of multiclass SVMs [27] trained in thedual.
Here, x (cid:63) ( θ ) is the optimal dual solution, a matrix of shape m × k , where m is the number oftraining examples and k is the number of classes, and θ ∈ R
+ is the regularization parameter.
Thechallenge in differentiating x (cid:63) ( θ ) is that each row of x (cid:63) ( θ ) is constrained to belong to the probabilitysimplex (cid:52) k .
More formally, let X tr ∈ R m × p be the training feature matrix and Y tr ∈ { 0 , 1 } m × k bethe training labels (in row-wise one-hot encoding).
Let W ( x, θ ) :
=
X (cid:62) tr
( Y tr − x ) /θ
∈
R p × k be thedual-primal mapping.
Then, we consider the following bi-level optimization problem
[Equation]
subject to
[Equation]
where C = (cid:52) k × · · · × (cid:52) k is the Cartesian product of m probability simplices.
We apply the changeof variable θ = exp( λ ) in order to guarantee that the hyper-parameter θ is positive.
The matrix
W ( x (cid:63) ( θ ) , θ ) ∈
R p × k contains the optimal primal solution, the feature weights for each class.
Theouter loss is computed against validation data X val and Y val .
While KKT conditions can be used to differentiate x (cid:63) ( θ ) , a more direct way is to use the projectedgradient ﬁxed point (9).
The projection onto C can be easily computed by row-wise projections on the [Figure] simplex.
The projection’s Jacobian enjoys a closed form (Appendix C).
Another way to differentiatex (cid:63) ( θ ) is using the mirror descent ﬁxed point (13).
Under the KL geometry, projections correspond toa row-wise softmax.
They are therefore easy to compute and differentiate.
Figure 4 compares theruntime performance of implicit differentiation vs. unrolling for the latter two ﬁxed points.
4.2 Dataset distillation
Dataset distillation [80, 59] aims to learn a small synthetic training dataset such that a model trainedon this learned data set achieves a small loss on the original training set.
Formally, let X tr ∈ R m × pand y tr ∈ [ k ] m denote the original training set.
The distilled dataset will contain one prototypeexample for each class and therefore θ ∈ R k × p .
The dataset distillation problem can then naturallybe cast as a bi-level problem, where in the inner problem we estimate a logistic regression modelx [Equation]
( θ ) ∈
R p × k trained on the distilled images θ ∈ R k × p , while in the outer problem we want tominimize the loss achieved by x (cid:63) ( θ ) over the training set:
[Equation]
subject to x
[Equation]
where f ( W, X ; y ) :
= (cid:96) ( y, XW ) , (cid:96) denotes the multiclass logistic regression loss, and ε = 10 − 3 is aregularization parameter that we found had a very positive effect on convergence.
In this problem, and unlike in the general hyperparameteroptimization setup, both the inner and outer problems are highdimensional, making it an ideal test-bed for gradient-basedbi-level optimization methods.
For this experiment, we usethe MNIST dataset.
The number of parameters in the innerproblem is p = 28 2 = 784 .
while the number of parametersof the outer loss is k × p = 7840 .
We solve this problemusing gradient descent on both the inner and outer problem,with the gradient of the outer loss computed using implicitdifferentiation, as described in §2.
This is fundamentallydifferent from the approach used in the original paper, wherethey used differentiation of the unrolled iterates instead.
Forthe same solver, we found that the implicit differentiationapproach was 4 times faster than the original one.
The obtaineddistilled images θ are visualized in Figure 5.
4.3 Task-driven dictionary learning
Task-driven dictionary learning was proposed to learn sparse codes for input data in such a way thatthe codes solve an outer learning problem [60, 77, 83].
Formally, given a data matrix X tr ∈ R m × p
[Table]
and a dictionary of k atoms θ ∈ R k × p , a sparse code is deﬁned as a matrix x (cid:63)
( θ ) ∈
R m × k thatminimizes in x a reconstruction loss f
( x, θ ) :
= (cid:96) ( X tr , xθ ) regularized by a sparsity-inducing penaltyg ( x ) .
Instead of optimizing the dictionary θ to minimize the reconstruction loss, [60] proposed tooptimize an outer problem that depends on the code.
Given a set of labels Y tr ∈ { 0 , 1 } m , we considera logistic regression problem which results in the bilevel optimization problem:
[Equation]
subject to x (cid:63)
[Equation]
argmin
[Equation]
When (cid:96) is the squared Frobenius distance between matrices, and g the elastic net penalty, [60, Eq.21] derive manually, using optimality conditions (notably the support of the codes selected at theoptimum), an explicit re-parameterization of x (cid:63) ( θ ) as a linear system involving θ .
This closedform allows for a direct computation of the Jacobian of x (cid:63) w.r.t.
θ .
Similarly, [77] derive ﬁrstorder conditions in the case where (cid:96) is a β -divergence, while [83] propose to use unrolling of ISTAiterations.
Our approach bypasses all of these manual derivations, giving the user more leisure tofocus directly on modeling (loss, regularizer) aspects.
We illustrate this on breast cancer survival prediction from gene expression data.
We frame it as abinary classiﬁcation problem to discriminate patients who survive longer than 5 years ( m 1 = 200 ) vspatients who die within 5 years of diagnosis ( m 0 = 99 ), from p = 1 , 000 gene expression values.
Asshown in Table 2, solving (11) (Task-driven DictL) reaches a classiﬁcation performance competitivewith state-of-the-art L 1 or L 2 regularized logistic regression with 100 times fewer variables.
4.4 Sensitivity analysis of molecular dynamics
with respect to increasing the diameter of the blue particles
Many physical simulations require solving optimization problems,such as energy minimization in molecular [75] and continuum [9]mechanics, structural optimization [53] and data assimilation [40].
In this experiment, we revisit an example from JAX-MD [75], theproblem of ﬁnding energy minimizing conﬁgurations to a systemof k packed particles in a 2 -dimensional box of size (cid:96)
[Equation]
where x (cid:63)
( θ ) ∈
R k × 2 are the optimal coordinates of the k particles,
U ( x i,j , θ ) is the pairwise potential energy function, with half theparticles at diameter 1 and half at diameter θ = 0 .
6 , which weoptimize with a domain-speciﬁc optimizer [15].
Here we considersensitivity of particle position with respect to diameter ∂x (cid:63) ( θ ) ,rather than sensitivity of the total energy from the original experiment.
Figure 6 shows results calculated via forward-mode implicitdifferentiation (JVP).
Whereas differentiating the unrolled optimizer happens to work for total energy,here it typically does not even converge, due to the discontinuous optimization method.
Conclusion
We proposed in this paper an approach for automatic implicit differentiation, allowing the userto freely express the optimality conditions of the optimization problem whose solutions are to bedifferentiated, directly in Python.
The applicability of our approach to a large catalog of optimalityconditions is shown in the non-exhaustive list of Table 1, and illustrated by the ease with which wecan solve bi-level and sensitivity analysis problems.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
[Table]
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
[Table] (a) Did you include the full text of instructions given to participants and screenshots, if [Table]
applicable?
[N/A]
[Table] (b) Did you describe any potential participant risks, with links to Institutional Review [Table]
Board (IRB) approvals, if applicable?
[N/A]
[Table]
(c) Did you include the estimated hourly wage paid to participants and the total amount
[Table]
spent on participant compensation?
[N/A]
[Table]
AppendixA More examples of optimality criteria and ﬁxed points
To demonstrate the generality of our approach, we describe in this section more optimality mapping
F or ﬁxed point iteration T .
Mirror descent ﬁxed point.
We again consider the case when x (cid:63) ( θ ) is implicitly deﬁned as thesolution of (8).
We now generalize the projected gradient ﬁxed point beyond Euclidean geometry.
Let the Bregman divergence D ϕ : dom( ϕ ) × relint(dom( ϕ ))
→ R + generated by ϕ be deﬁned by
[Equation]
We deﬁne the Bregman projection of y onto C ( θ ) ⊆ dom( ϕ ) by
[Equation]
Deﬁnition (12) includes the mirror map ∇ ϕ ∗ ( y ) for convenience.
It can be seen as a mapping from 576 R d to dom( ϕ ) , ensuring that (12) is well-deﬁned.
The mirror descent ﬁxed point mapping is then
[Equation]
Because T involves the composition of several functions, manually deriving its JVP/VJP is errorprone.
This shows that our approach leveraging autodiff allows to handle more advanced ﬁxed pointmappings.
A common example of ϕ is ϕ
( x ) = (cid:104) x, log x − 1 (cid:105) , where dom( ϕ ) =
R d + .
In this case,
D ϕ is the Kullback-Leibler divergence.
An advantage of the Kullback-Leibler projection is that itsometimes easier to compute than the Euclidean projection, as we detail in Appendix C.
Newton ﬁxed point.
Let x be a root of G ( · , θ ) , i.e., G ( x, θ ) = 0 .
The ﬁxed point iteration of
Newton’s method for root-ﬁnding is
[Equation]
By the chain and product rules, we have
[Equation]
Using (3), we get A = − ∂ 1 F ( x, θ ) = ηI .
Similarly,
[Equation]
Newton’s method for optimization is obtained by choosing G ( x, θ )
= ∇ 1 f ( x, θ ) , which gives
[Equation]
It is easy to check that we recover the same linear system as for the gradient descent ﬁxed point(5).
A practical implementation can pre-compute an LU decomposition of ∂ 1 G ( x, θ ) , or a Choleskydecomposition if ∂ 1 G ( x, θ ) is positive semi-deﬁnite.
Proximal block coordinate descent ﬁxed point.
We now consider the case when x (cid:63) ( θ ) is implicitlydeﬁned as the solution
[Equation]
where g 1 , . .
. , g m are possibly non-smooth functions operating on subvectors (blocks) x 1 , . . .
, x m ofx .
In this case, we can use for i ∈ [ m ] the ﬁxed point
[Equation] [Equation]
where η 1 , . . . ,
η m are block-wise step sizes.
Clearly, when the step sizes are shared, i.e., η 1 =· · · = η m = η , this ﬁxed point is equivalent to the proximal gradient ﬁxed point (7) with g ( x, θ )
= 596 (cid:80) ni
=1 g i ( x i , θ ) .
Quadratic programming.
We now show how to use the KKT conditions discussed in §2.2 todifferentiate quadratic programs, recovering Optnet [6] as a special case.
To give some intuition, letus start with a simple equality-constrained quadratic program (QP)
[Equation]
subject to H ( z, θ ) =
Ez − d = 0 , where Q ∈ R p × p , E ∈ R q × p , d ∈ R q .
We gather the differentiable parameters as θ =
( Q, E, c, d ) .
The stationarity and primal feasibility conditions give
[Equation]
In matrix notation, this can be rewritten as
[Equation]
[Equation]
We can write the solution of the linear system (16) as the root x
=
( z, ν ) of a function F ( x, θ ) .
Moregenerally, the QP can also include inequality constraints
[Equation]
subject to H
[Equation]
z, θ
[Equation]
= Ez
[Equation]
d = 0 , G ( z, θ ) =
Mz − h ≤ 0 .
where M ∈ R r × p and h ∈ R r .
We gather the differentiable parameters as θ =
( Q, E, M, c, d, h ) .
The stationarity, primal feasibility and complementary slackness conditions give
[Equation]
In matrix notation, this can be written as
[Equation]
While x = ( z, ν, λ ) is no longer the solution of a linear system, it is the root of a function F ( x, θ )and therefore ﬁts our framework.
With our framework, no derivation is needed.
We simply deﬁne f ,
H and G directly in Python.
Conic programming.
We now show that the differentiation of conic linear programs [3, 5], at theheart of differentiating through cvxpy layers [2], easily ﬁts our framework.
Consider the problem
[Equation]
subject to Ez + s =
d, s ∈ K
[Equation]
(17) where λ
=
( c, E, d ) , E ∈ R m × p , d ∈ R m , c ∈ R p and K ⊆
R m is a cone;
z and s are the primaland slack variables, respectively.
Every convex optimization problem can be reduced to the form (17).
Let us form the skew-symmetric matrix
[Equation]
where N = p + m
+ 1 .
Following [3, 2, 5], we can use the homogeneous self-dual embedding toreduce the process of solving (17) to ﬁnding a root of the residual map
[Equation]
(18) where Π = proj R p ×K ∗ × R + and K ∗ ⊆
R m is the dual cone.
The splitting conic solver [67],which is based on ADMM, outputs a solution F ( x (cid:63) ( θ ) , θ ) = 0 which is decomposed as x (cid:63) ( θ ) =( u (cid:63) ( θ ) , v (cid:63) ( θ ) , w (cid:63) ( θ )) .
We can then recover the optimal solution of (17) usingz (cid:63) λ )))
(cid:63) ( θ λ )) .
The key oracle whose JVP/VJP we need is therefore Π , which is studied in [4].
The projection onto afew cones is available in our library and can be used to express F .
Frank-Wolfe.
We now consider
[Equation]
19) where C ( θ ) is a convex polytope
, i.e., it is the convex hull of vertices v 1 ( θ ) , . .
. , v m ( θ ) .
The FrankWolfe algorithm requires a linear minimization oracle (LMO)
[Equation]
and is a popular algorithm when this LMO is easier to compute than the projection onto C ( θ ) .
However, since this LMO is piecewise constant, its Jacobian is null almost everywhere.
Inspired by
SparseMAP [66], which corresponds to the case when f is a quadratic, we rewrite (19) as
[Equation]
where V ( θ ) is a d × m matrix gathering the vertices v 1 ( θ ) , . .
. , v m ( θ ) .
We then have x (cid:63) ( θ ) =
V ( θ ) p (cid:63) ( θ ) .
Since we have reduced (19) to minimization over the simplex, we can use the projectedgradient ﬁxed point to obtain
[Equation]
We can therefore compute the derivatives of p (cid:63) ( θ ) by implicit differentiation and the derivatives ofx (cid:63) ( θ ) by product rule.
Frank-Wolfe implementations typically maintain the convex weights of thevertices, which we use to get an approximation of p (cid:63) ( θ ) .
Moreover, it is well-known that after titerations, at most t vertices are visited.
We can leverage this sparsity to solve a smaller linear system.
Moreover, in practice, we only need to compute VJPs of x (cid:63) ( θ ) .
B Code examples
B.1 Code examples for optimality conditions
Our library provides several reusable optimality condition mappings F or ﬁxed points T .
Wenevertheless demonstrate the ease of writing some of them from scratch.
KKT conditions.
As a more advanced example, we now describe how to implement the KKTconditions (6).
The stationarity, primal feasibility and complementary slackness conditions read
[Equation]
Using jax.vjp to compute vector-Jacobian products, this can be implemented asgrad = jax.grad(f) def F(x, theta): z, nu, lambd = x theta_f, theta_H, theta_G = theta _, H_vjp = jax.vjp(H, z, theta_H) stationarity =
(grad(z, theta_f) + H_vjp(nu)[0])
primal_feasability = H(z, theta_H) _, G_vjp =
jax.vjp(G, z, theta_G) stationarity += G_vjp(lambd)[0] comp_slackness = G(z, theta_G)
* lambd return stationarity, primal_feasability, comp_slackness
Similar mappings F can be written if the optimization problem contains only equality constraints oronly inequality constraints.
Mirror descent ﬁxed point.
Letting η = 1 and denoting θ =
( θ f , θ proj ) , the ﬁxed point (13) is
[Equation]
We can then implement it as follows.
[Equation]
return bregman_projection(y, theta_proj)
Although not considered in this example, the mapping ∇ ϕ could also depend on θ if necessary.
B.2 Code examples for experiments
We now sketch how to implement our experiments using our framework.
In the following, jnp isshort for jax.numpy .
In all experiments, we only show how to compute gradients with the outerobjective.
We can then use these gradients with gradient-based solvers to solve the outer objective.
Multiclass SVM experiment.
[Equation] = load_ [Equation] ) def W(x, theta): # dual-primal map return jnp.dot(X_tr.
T, Y_tr - x) / theta def f(x, theta): # inner objective return (0.5 * theta * jnp.sum(W(x, theta) ** 2) + jnp.vdot(x, Y_tr)) grad = jax.grad(f)
proj =
jax.vmap(projection_simplex)
# row-wise projections def T(x, theta): return proj(x - grad(x, theta)
) @custom_fixed_point(T) defmsvm_dual_solver(init_x, theta):
# [...] return x_star # solution of the dual objective def outer_loss(lambd): theta = jnp.exp(lambd) x_star = msvm_dual_solver(init_x, theta) # inner solution
Y_pred = jnp.dot(W(x_star, theta), X_val) return 0.5 * jnp.sum((Y_pred - Y_val) ** 2) print(jax.grad(outer_loss)(lambd))
Task-driven dictionary learning experiment.
X_tr, y_tr =
load_data() def f(x, theta): # dictionary loss residual =
X_tr - jnp.dot(x, theta) return huber_loss(residual) grad = jax.grad(f) def T(x, theta): # proximal gradient fixed point return prox_lasso(x - grad(x, theta))
@custom_fixed_point(T)
defsparse_coding(init_x, theta): # inner objective # [...] return x_star # lasso solution def outer_loss(theta, w): # task-driven loss x_star = sparse_coding(init_x, theta)
# sparse codes y_pred = jnp.dot(x_star, w) return logloss(y_tr, y_pred) print(jax.grad(outer_loss, argnums=(0,1)))
Dataset distillation experiment.
X_tr, y_tr =
load_data() logloss = jax.vmap(loss.multiclass_logistic_loss) def f(x, theta, l2reg=1e-3): # inner objective scores
=
jnp.dot(theta, x) distilled_labels = jnp.arange(10) penalty =
l2reg * jnp.sum(x * x) return jnp.mean(logloss(distilled_labels, scores))
+ penalty
F = jax.grad(f) @custom_root(F) deflogreg_solver(init_x, theta):
# [...] return x_star def outer_loss(theta): x_star =
logreg_solver(init_x, theta) # inner solution scores = jnp.dot(X_tr, x_star) return jnp.mean(logloss(y_tr, scores)) print(jax
[Figure]
Molecular dynamics experiment.energy_fn = soft_sphere_energy_fn(diameter) init_fn, apply_fn = jax_md.minimize.fire_descent( energy_fn, shift_fn) x0 = random.uniform(key, (N, 2))
R0 =
L * x # transform to physical coordinates
R = lax.fori_loop( 0,num_optimization_steps, body_fun=lambdat,state:apply_fn(state, t=t), init_val=
init_fn(R0)).position x_star =
R / L def F(x, diameter): # normalized forces energy_fn = soft_sphere_energy_fn(diameter) normalized_energy_fn =
lambda x: energy_fn(L * x) return -jax.grad(normalized_energy_fn)(x) dx = root_jvp(F, x_star, diameter, 1.0, solve=linear_solve.solve_bicgstab) print(dx)
C Jacobian products
Our library provides numerous reusable building blocks.
We describe in this section how to computetheir Jacobian products.
As a general guideline, whenever a projection enjoys a closed form, we leavethe Jacobian product to the autodiff system.
C.1 Jacobian products of projections
We describe in this section how to compute the Jacobian products of the projections (in the Euclideanand KL senses) onto various convex sets.
When the convex set does not depend on any variable, wesimply denote it C instead of C ( θ ) .
Non-negative orthant.
When C is the non-negative orthant, C = R d + , we obtain proj C
( y )
=
max( y, 0) , where the maximum is evaluated element-wise.
This is also known as the ReLu function.
The projection in the KL sense reduces to the exponential function, proj ϕ C ( y ) =
exp( y ) .
Box constraints.
When C ( θ ) is the box constraints C ( θ ) =
[ θ 1 , θ 2 ] d with θ ∈ R 2 , we obtainproj C
( y, θ )
=
clip ( y, θ 1 , θ 2 ) :
= max(min
( y, θ 2 ) , θ 1 ) .
This is trivially extended to support different boxes for each coordinate, in which case θ ∈ R d × 2 .
Probability simplex.
When C is the standard probability simplex, C = (cid:52) d , there is no analyticalsolution for proj C ( y ) .
Nevertheless, the projection can be computed exactly in O ( d ) expected time or
O ( d log d ) worst-case time [22, 62, 33, 26].
The Jacobian is given by diag ( s ) − ss (cid:62) / (cid:107) s (cid:107) 1 , wheres ∈ { 0 , 1 } d is a vector indicating the support of proj C ( y ) [61].
The projection in the KL sense, on theother hand, enjoys a closed form: it reduces to the usual softmax proj ϕ C ( y ) =
exp( y ) / (cid:80) dj
=1 exp( y j ) .
Box sections.
Consider now the Euclidean projection z (cid:63) ( θ ) =
proj C
( y, θ ) onto the set C ( θ )
=
{ z ∈ R d : α i ≤ z i ≤ β i , i ∈ [ d ]; w (cid:62)
z
= c } , where θ =
( α, β, w, c ) .
This projection is a singlyconstrained bounded quadratic program.
It is easy to check (see, e.g., [65]) that an optimal solutionsatisﬁes for all
i ∈
[ d ]z
[Equation]
where L : R × R n → R d is the dual-primal mapping and x (cid:63) ( θ ) ∈ R is the optimal dual variable ofthe linear constraint, which should be the root of
[Equation]
The root can be found, e.g., by bisection.
The gradient ∇ x (cid:63) ( θ ) is given by ∇ x (cid:63) ( θ ) =
B (cid:62) /A and the
Jacobian ∂z (cid:63) ( θ ) is obtained by application of the chain rule on L .
Norm balls.
When C ( θ ) = { x ∈ R d : (cid:107) x (cid:107) ≤ θ } , where (cid:107) · (cid:107) is a norm and θ ∈ R + ,
proj C
( y, θ )becomes the projection onto a norm ball.
The projection onto the (cid:96) 1 -ball reduces to a projection ontothe simplex, see, e.g., [33].
The projections onto the (cid:96) 2 and (cid:96) ∞ balls enjoy a closed-form, see, e.g.,[68, §6.5].
Since they rely on simple composition of functions, all three projections can therefore beautomatically differentiated.
Afﬁne sets.
When C ( θ ) = { x ∈ R d : Ax = b } , where A ∈ R p × d , b ∈ R p and θ =
( A, b ) , we getproj C
( y, θ ) =
y − A †
[Equation]
Ay − b ) =
y
[Equation]
A
[Equation]
AA
[Equation]
Ay − b ) where A † is the Moore-Penrose pseudoinverse of A .
The second equality holds if p < d and A isfull rank.
A practical implementation can pre-compute a factorization of the Gram matrix AA (cid:62) .
Alternatively, we can also use the KKT conditions.
Hyperplanes and half spaces.
When C ( θ ) = { x ∈ R d : a (cid:62) x = b } , where a ∈ R d and b ∈ R andθ = ( a, b ) , we getproj
[Equation]
When C ( θ ) = { x ∈ R d : a (cid:62) x ≤ b } , we simply replace a (cid:62) y − b in the numerator by max( a (cid:62) y − b, 0) .
Transportation and Birkhoff polytopes.
When C ( θ ) =
{ X ∈ R p ×
d : X 1 d
= θ 1 , X (cid:62) 1 p
=θ 2 , X ≥ 0 } , the so-called transportation polytope, where θ 1 ∈ (cid:52) p and θ 2 ∈ (cid:52) d are marginals, wecan compute approximately the projections, both in the Euclidean and KL senses, by switching to thedual or semi-dual [17].
Since both are unconstrained optimization problems, we can compute their
Jacobian product by implicit differentiation using the gradient descent ﬁxed point.
An advantage ofthe KL geometry here is that we can use Sinkhorn [28], which is a GPU-friendly algorithm.
The
Birkhoff polytope, the set of doubly stochastic matrices, is obtained by ﬁxing θ 1 = θ 2 = 1 d /d .
Order simplex.
When C ( θ ) = { x ∈ R d : θ 1 ≥ x 1 ≥ x 2 ≥ · · · ≥ x d ≥ θ 2 } , a so-called ordersimplex [49, 16], the projection operations, both in the Euclidean and KL sense, reduce to isotonicoptimization [58] and can be solved exactly in O
( d log d ) time using the Pool Adjacent Violatorsalgorithm [14].
The Jacobian of the projections and efﬁcient product with it are derived in [31, 18].
Polyhedra.
More generally, we can consider polyhedra, i.e., sets of the form C ( θ ) = { x ∈ 707 R d : Ax = b, Cx ≤ d } , where A ∈ R p × d , b ∈ R p , C ∈ R m × d , and d ∈ R m .
There are severalways to differentiate this projection.
The ﬁrst is to use the KKT conditions as detailed in §2.2.
Asecond way is consider the dual of the projection instead, which is the maximization of a quadraticfunction subject to non-negative constraints [68, §6.2].
That is, we can reduce the projection ona polyhedron to a problem of the form (8) with non-negative constraints, which we can in turnimplicitly differentiate easily using the projected gradient ﬁxed point, combined with the projectionon the non-negative orthant.
Finally, we apply the dual-primal mapping , which enjoys a closed formand is therefore amenable to autodiff, to obtain the primal projection.
C.2 Jacobian products of proximity operators
We provide several proximity operators, including for the lasso (soft thresholding), elastic net andgroup lasso (block soft thresholding).
All satisfy closed form expressions and can be differentiatedautomatically via autodiff.
For more advanced proximity operators, which do not enjoy a closed form,recent works have derived their Jacobians.
The Jacobians of fused lasso and OSCAR were derived in[64].
For general total variation, the Jacobians were derived in [78, 24].
D Jacobian precision proofs
Proof of Theorem 1.
To simplify notations, we note A (cid:63) :
= A ( x (cid:63) , θ ) and ˆ A :=
A (ˆ x, θ ) , and similarlyfor B and J .
We have by deﬁnition of the Jacobian estimate function A (cid:63) J (cid:63) =
B (cid:63) and ˆ
A ˆ J = ˆ B .
Therefore we have
[Equation]
For any invertible matrices M 1 , M 2 , it holds that M −
[Equation]
so
[Equation]
Therefore,
[Equation]
As a consequence, the second term in J (ˆ x, θ ) − ∂x (cid:63) ( θ ) can be upper bounded and we obtain
[Equation]
which yields the desired result.
Corollary 1 (Jacobian precision for gradient descent ﬁxed point) .
Let f be such that f ( · , θ ) istwice differentiable and α -strongly convex and ∇ 21 f ( · , θ ) is γ -Lipschitz (in the operator norm)
and∂ 2 ∇ 1 f ( x, θ ) is β -Lipschitz and bounded in norm by R .
The estimated Jacobian evaluated at ˆ x isthen given by
[Equation]
For all θ ∈ R n , and any ˆ x estimating x (cid:63) ( θ ) , we have the following bound for the approximationerror of the estimated Jacobian
[Equation]
Proof of Corollary 1.
This follows from Theorem 1, applied to this speciﬁc
A ( x, θ ) and B ( x, θ ) .
For proximal gradient descent, where T ( x, θ ) =
prox ηg ( x − η ∇ 1 f ( x, θ ) , θ ) , this yields
[Equation]
We now focus in the case of proximal gradient descent on an objective f ( x, θ ) + g ( x ) , where g issmooth and does not depend on θ .
This is the case in our experiments in §4.3.
Recent work alsoexploits local smoothness of solutions to derive similar bounds [13, Theorem 13]
Corollary 2 (Jacobian precision for proximal gradient descent ﬁxed point) .
Let f be such thatf ( · , θ ) is twice differentiable and α -strongly convex and ∇ 21 f ( · , θ ) is γ -Lipschitz (in the operatornorm) and ∂ 2 ∇ 1 f ( x, θ ) is β -Lipschitz and bounded in norm by R .
Let g : R d → R be a twicedifferentiable µ -strongly convex (with special case µ = 0 being only convex), for which the function
Γ η ( x, θ )
= ∇ 2 g
( prox ηg ( x − η ∇ 1 f ( x, θ )) is κ η -Lipschitz in it ﬁrst argument.
The estimated Jacobianevaluated at ˆ x is then given by
[Equation]
For all θ ∈ R n , and any ˆ x estimating x (cid:63) ( θ ) , we have the following bound for the approximationerror of the estimated Jacobian
[Equation]
Proof of Corollary 2.
First, let us note that prox ηg
( y, θ ) does not depend on θ , since g itself does notdepend on θ , and is therefore equal to classical proximity operator of ηg which, with a slight overloadof notations, we denote as prox ηg ( y )
(with a single argument).
In other words,
[Equation]
Regarding the ﬁrst claim (expression of the estimated Jacobian evaluated at ˆ x ), we ﬁrst have thatprox ηg
( y ) is the solution to ( x (cid:48) − y )
+ η ∇ g ( x (cid:48) )
=
in x (cid:48) - by ﬁrst-order condition for a smoothconvex function.
We therefore have that
[Equation]
ηg ( y )))
− 1 , the ﬁrst I and inverse being functional identity and inverse, and the second I d and inverse being inthe matrix sense, by inverse rule for Jacobians ∂h ( z ) =
[ ∂h − 1 ( h ( z ))] − 1 (applied to the prox).
As a consequence, we have, for Γ η ( x, θ )
= ∇ 2 g
( prox ηg ( x − η ∇ 1 f ( x, θ ))
that
[Equation]
As a consequence, for all x ∈ R d , we have that
[Equation]
In the following, we modify slightly the notation of both A and B , writing
[Equation]
With the current hypotheses, following along the proof of Theorem 1, we have that ˜
A is ( α + µ )well-conditioned, and ( γ + κ η ) -Lipschitz in its ﬁrst argument, and ˜ B is β -Lipschitz in its ﬁrstargument and bounded in norm by R .
The same reasoning yields
[Equation]
E Experimental setup and additional results
Our experiments use JAX [21], which is Apache2-licensed and scikit-learn [70], which is BSDlicensed.
[Figure] [Figure]
d a t i o
[Figure]
multiple problem sizes setting as Figure 4)
As be seen, all methods performed similarly in terms of validation loss.
This conﬁrms that the faster runtimes for implicit differentiation compared to unrolling shown in Figure 4 (CPU) and Figure(GPU) are not at the cost of worse validation loss.
[Figure]
The ground-truth solution x (cid:63) ( θ ) is computed using the liblinear solver [37] available in scikit-learn [70] with a very low tolerance of 10 − 9 .
Unlike in Figure 3, which was on ridge regression, the ground-truth Jacobian ∂x (cid:63) ( θ ) cannot be computed in closed form, in the more difﬁcult setting of multiclass SVMs.
We therefore use a ﬁnite difference to approximately compute ∂x (cid:63) ( θ ) .
Our results nevertheless conﬁrm similar trends as in Figure 3.
E.1 Hyperparameter optimization of multiclass SVMs
Experimental setup.
Synthetic datasets were generated using scikit-learn ’ssklearn.datasets.make_classification
[70], following a model adapted from [50].
All datasets consist of m = 700 training samples belonging to k = 5 distinctclasses.
To simulate problems of different sizes, the number of features is varied asp ∈ { 100 , 250 , 500 , 750 , 1000 , 2000 , 3000 , 4000 , 5000 , 7500 , 10000 } , with 10% of featuresbeing informative and the rest random noise.
In all cases, an additional m val = 200 validationsamples were generated from the same model to deﬁne the outer problem.
For the inner problem, we employed three different solvers: (i) mirror descent, (ii) (accelerated)proximal gradient descent and (iii) block coordinate descent.
Hyperparameters for all solvers wereindividually tuned manually to ensure convergence across the range of problem sizes.
For mirrordescent, a stepsize of 1 . 0 was used for the ﬁrst 100 steps, following a inverse square root decayafterwards up to a total of 2500 steps.
For proximal gradient descent, a stepsize of 5 · 10 − 4 was usedfor 2500 steps.
The block coordinate descent solver was run for 500 iterations.
All solvers used thesame initialization, namely, x init =
1 k 1 m × k , which satisﬁes the dual constraints.
For the outer problem, gradient descent was used with a stepsize of 5 · 10 − 3 for the ﬁrst 100 steps,following a inverse square root decay afterwards up to a total of 150 steps.
Conjugate gradient was used to solve the linear systems in implicit differentiation for at most 2500iterations.
All results reported pertaining CPU runtimes were obtained using an internal compute cluster.
GPUresults were obtained using a single NVIDIA P100 GPU with 16GB of memory per dataset.
For eachdataset size, we report the average runtime of an individual iteration in the outer problem, alongside a90% conﬁdence interval estimated from the corresponding 150 runtime values.
Additional results Figure 13 compares the runtime of implicit differentiation and unrolling on
GPU.
These results highlight a fundamental limitation of the unrolling approach in memory-limitedsystems such as accelerators, as the inner solver suffered from out-of-memory errors for most problemsizes ( p ≥ 2000 for mirror descent, p ≥ 750 for proximal gradient and block coordinate descent).
While it might be possible to ameliorate this limitation by reducing the maximum number of iterationsin the inner solver, doing so might lead to additional challenges [82] and require careful tuning.
Figure 14 depicts the validation loss (value of the outer problem objective function) at convergence.
Itshows that all approaches were able to solve the outer problem, with solutions produced by differentapproaches being qualitatively indistinguishable from each other across the range of problem sizesconsidered.
Figure 15 shows the Jacobian error achieved as a function of the solution error, when varying thenumber of features.
E.2 Task-driven dictionary learning
We downloaded from http://acgt.cs.tau.ac.il/multi_omic_benchmark/download.htmla set of breast cancer gene expression data together with survival information generated by the TCGA
Research Network ( https://www.cancer.gov/tcga ) and processed as explained by [73].
Thegene expression matrix contains the expression value for p=20,531 genes in m=1,212 samples, fromwhich we keep only the primary tumors (m=1,093).
From the survival information, we select thepatients who survived at least ﬁve years after diagnosis ( m 1 = 200 ), and the patients who died beforeﬁve years ( m 0 = 99 ), resulting in a cohort of m = 299 patients with gene expression and binarylabel.
Note that non-selected patients are those who are marked as alive but were not followed for 5years.
To evaluate different binary classiﬁcation methods on this cohort, we repeated 10 times a random splitof the full cohort into a training (60%), validation (20%) and test (20%) sets.
For each split and eachmethod, 1) the method is trained with different parameters on the training set, 2) the parameter thatmaximizes the classiﬁcation AUC on the validation set is selected, 3) the method is then re-trained onthe union of the training and validation sets with the selected parameter, and 4) we measure the AUC of that model on the test set.
We then report, for each method, the mean test AUC over the 10 repeats,together with a 95% conﬁdence interval deﬁned a mean ± 1.96 × standard error of the mean.
We used Scikit Learn’s implementation of logistic regression regularized by (cid:96) 1 (lasso) and (cid:96) 2 (ridge)penalty from sklearn.linear_model.
LogisticRegression , and varied the C regularizationparameter over a grid of 10 values: { 10 − 5 , 10 − 3 , . .
For the unsupervised dictionary learning experiment method, we estimated a dictionary from the gene expression data in the trainingand validation sets, using sklearn.decomposition.
DictionaryLearning(n_components=10,alpha=2.0) , which produces sparse codes in k = 10 dimensions with roughly 50% nonzero coefﬁcients by minimizing the squared Frobenius reconstruction distance with lasso regularization on thecode.
We then use sklearn.linear_model.
LogisticRegression to train a logistic regression onthe codes, varying the ridge regularization parameter C over a grid of 10 values { 10 − 1 , 10 0 , . .
Finally, we implemented the task-driven dictionary learning model (11) with our toolbox, followingthe pseudo-code in Figure 10.
Like for the unsupervised dictionary learning experiment, we set thedimension of the codes to k = 10 , and a ﬁxed elastic net regularization on the inner optimizationproblem to ensure that the codes have roughly 50% sparsity.
For the outer optimization problem, wesolve an (cid:96) 2 regularized ridge regression problem, varying again the ridge regularization parameter
C over a grid of 10 values { 10 − 1 , 10 0 , . . .
Because the outer problem is non-convex, weminimize it using the Adam optimizer [56] with default parameters.
E.3 Dataset Distillation
Experimental setup.
For the inner problem, we used gradient descent with backtracking linesearch, while for the outer problem we used gradient descent with momentum and a ﬁxed step-size.
The momentum parameter was set to 0 .
9 while the step-size was set to 1 .
Figure 5 was produced after 4000 iterations of the outer loop on CPU (Intel(R) Xeon(R)
Platinum
P-8136 CPU @ 2.00GHz), which took 1h55.
Unrolled differentiation took instead 8h:05 ( 4 timesmore) to run the same number of iterations.
As can be seen in Figure 16, the output is the same inboth approaches.
[Figure]
E.4 Molecular dynamics
Our experimental setup is adapted from the JAX-MD example notebook available at https://github.com/google/jax-md/blob/master/notebooks/meta_optimization.ipynb .
We emphasize that calculating the gradient of the total energy objective, f ( x, θ )
=
(cid:80) ij U ( x i,j , θ ) ,with respect to the diameter θ of the smaller particles, ∇ 1 f ( x, θ ) , does not require implicit differentiation or unrolling.
This is because ∇ 1 f ( x, θ ) =
at x = x (cid:63) ( θ ) :
[Equation]
This is known as Danskin’s theorem or envelope theorem.
Thus instead, we consider sensitivities ofposition ∂x (cid:63) ( θ ) directly, which does require implicit differentiation or unrolling.
Our results comparing implicit and unrolled differentiation for calculating the sensitivity of positionare shown in Figure 17.
We use BiCGSTAB [79] to perform the tangent linear solve.
Like in theoriginal JAX-MD experiment, we use k = 128 particles in m = 2 dimensions.
[Figure]