Efficient and Modular Implicit Differentiation
Authors: Mathieu Blondel
Quentin Berthet
marco cuturi
Roy Frostig
Stephan Hoyer
Felipe Llinares-López
Fabian Pedregosa
Jean-Philippe Vert
Venue: NeurIPS 2022 Accept (2022)
Keywords: implicit differentiation
bilevel optimization
autodiff
jax

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.

Efficient and Modular Implicit Differentiation
Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-Lóóez, Fabian Pedregosa, Jean-Philippe Vert" Google Research
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 differen- tiation remained difficuult to use for practitioners, it often required case-by-cass this paper, we propose efficient and modular approach for implicit
tedious mathematical derivations and implementations. automatic implicit differentiation, differentiation of optimization problems. In our approach, the user defines directly in Python function F capturing the optimality conditions of the problem be differentiated. Once this done, we leverage autodiff of F and the implicit func- tion 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 top of any state-of-the-ar solver and modular the optimality condition specification is decoupled from the implicit differentiation mechanism. We show that seemingly simple principles allow 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
Automatic differentiation (autodiff) is now an inherent part of machine learning software. It allows express complex computations by composing elementary ones in creative ways and removes the tedious burden of computing their derivatives by hand. In parallel, the differentiation of optimization problem solutions has found many applications. A classical example is bi-level optimization, which typically involves computing the derivatives of a nested optimization problem in order to solve an outer one. Examples of applications in machine learning include hyper-parameter optimization [23, 77, 70, 38, 12. 13], neural networks [59], and meta-learning [39, 72]. Another line of active research involving differentiation of optimization problem solutions are optimization layers [55, 6, 65, 31, 46), which can be used to encourage structured outputs, and implicit deep networks [8, 36, 43, 44, 73]. which have smaller memory footprint than backprop-traiee trained networks.
Since optimization problem solutions typically do not enjoy an explicit formula in terms of their inputs, autodiff cannot be used directly to differentiate these functions. In recent years, two main approaches have been developed to circumvent this problem. The first one consists of unrolling the iterations of an optimization algorithm and using the final iteration problem solution [83, 32, 29, 39, 1]. This allows to explicitly construct a computational graph relating the algorithm output the inputs, on which autodiff can then be used transparently. However, this requires a reimplementation of the algorithm using the autodiff system, and not all algorithms
proxy for the optimization
Work done while at Google Research, now at Apple and Owkin, respectively.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

are necessarily autodiff friendly. Moreover, forde autodiff has time complexity that scales linearly with the number of variables and reverse-mode autodiff has memory complexity that scales linearly with the number of algorithm iterations. In contrast, a second approach consists in implicitly relating an optimization problem solution to its inputs using optimality conditions. a machine learning context, such implicit differentiation has been used for stationarity conditions [11, 59], KKT conditions [23, 45, 6, 67, 66] and the proximal gradient fixed point [65, 12, 13). An advantage of implicit differentiation is that solver reimplementation is not needed, allowing to build upon decades of state-of-the-arr software. Although implicit differentiation has a long history in numerical analysis it required a case-by-case
[48, 10, 57, 20), so far, it remained difficult to use for practitioners, tedious mathematical derivation and implementation. CasADi [7] allows to differentiate various optimization and root finding problem algorithms provided by the library. However, it does not allow to easily add implicit differentiation on top of existing solvers from optimality conditions expressed by 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 not cover a large catalog of optimality condition mappings this issue by adding implicit differentiation on top of cvxpy [30]. This works by reducing all convex optimization problems to conic program and using conic programming 's optimality conditions to derive an implicit differentiation formula. While this approach is very generic, solving a convex optimization problem using a conic programming solver-as an ADMM-based splitting conic solver [68] in the case of evxpy is rarely state-of-ie offt--of-th--rrr for every problem instance.
do. Other work [2] attempts to address
In this work, we ambition to achieve for optimization problem solutions what autodiff did for computational graphs. We propose automatic implicit differentiation, a simple approach to add implicit differentiation on top of any existing solver. In this approach, the user defines directly in Python a mapping function F capturing the optimality conditions of the problem solved by the algorithm. Once this done, we leverage autodiff of F combined with the implicit function theorem to automatically differentiate the optimization problem solution. Our approach is generic, yet it can exploit the efficiency of state-of-the-ard solvers. It therefore combines the benefits of implicit differentiation and autodiff. To summarize, we make the following contributions.
We deseribe our framework and its JAX [21, 42] implementation (https: / com/ google/jaxopt/). Our framework significantly 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 large catalog of optimality conditions (Table 1), recovering existing schemes and obtaining new ones, such as the mirror descent fixed point based one.
On the theoretical side, we provide new bounds on the Jacobian error when the optimization problem is only solved approximately, and empirically validate them.
We implement four illustrative applications, demonstrating our framework's 's case of use.
Beyond our software implementation in JAX, we hope this paper provides self-contained blueprint for creating an efficient and modular implementation of implicit differentiation other frameworks.
We denote the gradient and Hessian of f: Rd + R evaluated at & E Rd by (f(x) E Rd and v²f(z) E Rdxd We denote the Jacobian of F:: Id + Ip evaluated at x E Rd by 0F(z) e Rexd When ff or F have several arguments, we denote the gradient, Hessian and Jacobian in the 7th argument by V. V and o, respectively. The standard probability simplex is denoted by Ad : fI E Rd: || 1, r 2 0}. For any set C c Rd. we denote the indicator function Te: Rdd R U [+oo} where (x) = 0 if E C, Ie(x) = too otherwise. For a vector or matrix A, we note || the Frobenius (or Euclidean) norm, and |||p the operator norm.
Notation.
2 Automatic implicit differentiation
2.1 General principles
Overview. Contrary to autodiff through unrolled algorithm iterations, implicit differentiation typi- cally involves a manual, sometimes complicated, mathematical derivation. For instance, numerous works [23, 45, 6, 67, 66] use Karush-Kuhn-Tuccer (KKT) conditions in order to relate constrained optimization problem's 's solution to its inputs, and manually derive formula for its derivatives. The derivation and implementation in these works are typically case-by-case.

---

I_train, y_train load_ # Load features and labels
def f(x, theta):
Objectsve functson
residual jap jap.dot( (X_train, x) y_train
return (jup.aun(rasidual 2) thata jap. *x 2)) /
Ssace f s8 differentiable unconstrassed, optsmalsty condétion F is ssepty the gradsent of f in the 1st argument F jux grad(E, argnuns-0)
Scuston_root
def ridge.colvar(intt. theta):
del init_x # Instsalisation not ueed this solver XX jap. dot(X_ttriai I, X_train)
Xy jap.eyyeX_ttaal (b)) Identsty natrie
y.train)
Finds the ridge solution by solvsng Isnear system return sapplinalg solve(XX theta + I, Xy)
init_x None
(rant (ride argnuns-1) (inst_x, 10.0))
Figure I: Adding implicit differentiation on top of ridge regression solver. The function f(x, 0) defines the objective function and the mapping F, here simply equation (4), captures the optimality conditions. Our decorator @custom__ corotoooooooooooooooooooooooooto automatically adds implicit differentiation the solver for the user, overriding JAX's s default behavior. The last line evaluates the Jacobian at 8 = 10.
In this work, we propose a generic way to easily add implicit differentiation on top of existing solvers. In our approach, the user defines directly in Python mapping function F capturing the optimality conditions of the problem solved by the algorithm. We provide reusable building blocks to easily express such F. The provided is then plugged into our Python decorator rcustom_rood which we append on top of the solver declaration we wish to differentiate. Under the hood, we combine the implicit function theorem and autodiff of to automatically differentiate the optimization problem solution. A simple illustrative example is given in Figure 1.
Differentiating root. Let F: R x R" Rd be user-provided mapping, capturing the opti- mality conditions of a problem. An optimal solution, denoted x (0), should be a root of
(O)) 0.
We can see x² (0) an implicitly defined function of 0 IR" i.e., x²: RR" Rdd More precisely, from the implicit function theorem [48, 57J, we know that for (zo, 0o) satisfying F(zo, 0o) 0 with continuously differentiable F, if the Jacobian F, evaluated at (zo. fo) is square invertible matrix, then there exists a function x= (-) defined on neighborhood of 6g such that xt (0o) zo- Furthermore, for all in this neighborhood. we have that ((xx()) O) and òx (0) exists. Using the chain rule, the Jacobian Or"(0) satisfies
0 Fa (0), 0)0x" (0) + F2F (0), =
Computing Oz (0) therefore boils down the resolution of the linear system of equations
F (1), (0), Ox (0) F2F(x*()),
(2)
AERdxa
JERUXN
BERaXn
When (1) is one-dimensional root finding problem (d 1), (2) becomes particularly simple since we then have Vx+ (0) = B /A, where 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 optimiza- tion solution to be differentiated through the optimality conditions F. Our approach is efficient as it can be added on top any state-of-the-art solver and modular as the optimality condition specification is decoupled from the implicit differentiation mechanism. This contrasts with existing works, where the derivation and implementation are specific to each optimality condition.
Differentiating fixed point. We will encounter numerous applications where x (0) is instead implicitly defined through fixed point:
x (0) (0),
3

---

where T: Rd x R + R. This can be seen as particular case of (1) by defining the residual F(z, = Txx 0) x
(3)
In this case, when T is continuously differentiable, using the chain rule, we have
F (6). 8) 0x (0), and B (2F((= = (t).
Computing JVPs and VJPs. In most practical scenarios, it is not necessary to explicitly form the Jacobian matrix, and instead it is sufficient to left-multipl or right-multiply by 8 F and F.F. These are called vector-Jacobian product (VJP) and aecobian-vector product (JVP), and are useful for integrating x² (0) with reverse-mode and forward-mode autodiff, respectively. Oftentimes, F will be explicitly defined. this case, computing the VJP or JVP can be done via autodiff. In some cases, F may itself be implicitly defined, for instance when F involves the solution of variational problem. In this case, computing the VJP or JVP will itself involve implicit differentiation.
The
(JVP) between J = 0x= and a vector v, Ju. can be computed efficiently by solving A(Jv) = Be. The left-multillication (VJP) of v with J, v J, can be computed by first solving A Tu- = v. Then, we can obtain v by J = u AJ = B. Note that when B changes but A and remain the same, we do not need to solve A the VJP w.r.t. different variables while solving only one linear system.
once again. This allows to compute
To solve these linear systems, we can use the conjugate gradient method [51] when A is symmetric positive semi-definite and GMRES [75] or BiCGSTAB [81] otherwise. These algorithms are all matrix-freet they only require matrix- 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 T u = Av using conjugate gradient. This can be implemented using JAX's transpose routine j jax max.linew www.einnrr..... _ _transposeeeeeee [41]. In case of non-invertibility. a common heuristic is to solve least squares min, || - B|2 instead.
Pre-processing and post-processing mappings. Oftentimes, the goal is not to differentiate 0 per se, but the parameters of function producing 0. One example of such pre-procee pre-processing to convert the parameters to be differentiated from one form to another canonical form, such as quadratic program [6] or conic program [2]. Another example is when =0()) is used as the output of a neural network layer, in which case 0 is produced by the previous layer. Likewise, x (0) will often not be the final output we want to differentiate. One example of such post-proot poss-processing is when x*(() (0) is the solution of a dual program and we apply the qual-prima mapping to recover the solution of the primal program. Another example the application of a loss function. in order to reduce x*(0) (0) to scalar value. We leave the differentiation of such system, allowing to compose functions complex ways.
mappings to the autodiff
Implementation details. When a solver function is decorated with @custom_ root, we use jax. custom_jvp and jax. custom_vjj vjj to automatically add custom JVP and VJP rules to the function, overriding JAX's default behavior. As mentioned above, we use linear system solvers based on matrix-vector products and therefore we only need access to F through the JVP or VJP with 0 F and F. This is done by using jax. jvp and jax.vjp, respectively. Note that, as in Figure 1, the definition of F' will often include a gradient mapping (If 0). Thankfully. JAX supports second- order derivatives transparently. For convenience, our library also provides a @custom cicttooooieeieeeeeeeeeeeeeeeeeee fixed_point decorator, for adding implicit differentiation on top of solver, given a fixed point iteration T;; see code examples in Appendix B.
2.2 Examples
We now give various examples of mapping F or fixed point iteration T, recovering existing implicit differentiation methods and creating new ones. Each choice of or T implies different trade-offs in terms of computational oracles; see Table 1. Source code examples are given in Appendix B.
Stationary point condition.
The simplest example is differentiate through the implicit function x² (0) = argmin [(x 0),
ZERd
where f: Rd x _= + IR is twice differentiable, Vif is continuously differentiable,, and V is invertible at (x" (0),0) 0). In this case, F' is simply the gradient mapping
F(x,0 0) vif(2,0)
(4)

---

Table Summary of optimality condition mappings. Oracles are accessed through their JVP VJP.
Name
Equation
Solution needed
Oracle
Stationary
(4), (5)
Primal
Vif
KKT
(6)
Primal and dual
V.S. S. H. G. &.H.&&&
Proximal gradient
(7)
Primal
ViS. proNng
Projected gradient
(9)
Primal
V.f. proje
Mirror descent
(13)
Primal
Vif. proje. Vi
Newton
(14)
Primal
f 01] [Vi [prox proj xR+
V Vif(z,
Block proximal gradient
(15)
Primal
Conic programming
(18)
Residual map root
We then have FfF(z, Vif(z, and 02F(z,0) V2Vif 8), the Hessian of f in its first argument and the Jacobian in the second argument of V (z, 0). In practice, we use autodiff to compute Jacobian products automatically. Equivalently, we can use the gradient descent fixed point
T 8)
1f. 0),
(5)
for all 1) > 0. Using (3), it is easy to check that we obtain the same linear system since r) cancels out.
KKT conditions. As differentiated in several works [23, 45, 6, 67, 66). fit our framework. As we will see, the key will be group the optimal primal and dual variables
advanced example, now show that the KKT conditions, manually
x: (0). Let consider the general problem
argmin f(z, 0) subject to G(z, 0) < O, H(2, 0) 0,
SER*
where Iº is the primal variable, f: IP x R. + R. G: RP x R + K and H: F x R + R are twice differentiable convex functions, and Vif. ò G and θ H are continuously differentiable. The stationarity, primal feasibility and complementary slackness conditions give
Vif(z, G(2 T 0 ( (2, 0)]
H(2, 0) = 0 À G(2.0)
where E R9 and À E RC are the dual variables, also known as KKT multipliers. The primal and dual feasibility conditions can be ignored almost everywhere [34]. The system of (potentially framework,
group the primal and dual solutions as x x0(() = (0), x*))), (0)) to form the root of a function F(x"(0). 0). where F: Rd x 12" + Rd and d = + The primal and dual solutions can be obtained from a generic solver, such as an interior point method. In practice, the above mapping will be defined directly in Python (see Figure in Appendix B) and F will be differentiated automatically via autodiff.
nonlinear) equations (6) fits
Proximal gradient fixed point. Unfortunately, not all algorithms return both primal and dual solutions. Moreover, if the objective contains non-smooth terms, proximal gradient descent may be more efficient. We now discuss its fixed point [65, 12, 13). Let x² (0) be implicitly defined as
argmin [ax 0) g(x, 0),
æER
where f: Rd x R + R is non-smooth. Let us define the proximity operator associated with g by
convex and g: Rd x R - R is convex but possibly
prox (v, 0) = argmin
7 ||
v g(x, 0).
sERd
To implicitly differentiate x² (0),, we use the fixed point mapping [69, p. p.10
T(x, prox,ng (z vV 0),
(7)
for any step size > O. The proximity operator is 1-Lipschitz continuous [64]. By Rademacher's theorem, differentiable almost everywhere. If, in addition, it is continuously differentiable in neighborhood of (x* (0), 0) and if 7 - 0 T(x [((="(( (0), 0) is invertible, then our framework to differentiate xr (0) applies. Similar assumptions are made in [13]. Many proximity operators enjoy closed form and can easily be differentiated, as discussed in Appendix C. An implementation is given in Figure 2.
5

---

grad jax. (axd(() # Pre-compile the gradsent,
def T(x, theta):
Unpack the parameters of f and theta_f, theta_g theta
Resurn the fized point condétion evaluated at
return prox(x grad(x, theta_f), theta_g)
Figure 2: Implementation of the proximal gradient fixed point (7) with step size r) = 1.
Projected gradient fixed point. As special case, when g(z, f) is the indicator function Ic(o) (z), where C(0) is convex set depending on 0, we obtain
x' = argmin f(x, 0).
(8)
=EC(()
The proximity operator prox, becomes the Euclidean projection onto C(6)
prox (v. 8) = proje (v. 0) = argmin || - yl
xEC(Ø)
and (7) becomes the projected gradient fixed point
T 0) = projc - v (x,
(9)
Compared to the KKT conditions, this fixed point is particularly suitable when the projection enjoys a 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 approach developed in this section theoretically only applies to settings where the implicit function theorem is valid, namely, where optimality conditions satisfy the differentiability and invertibility conditions stated in $2.1. While this covers a wide range of situations even for non-smooth optimization problems (c.g., under mild assumptions the solution of a Lasso regression can be differentiated a.c. with respect to the regularization parameter, see Appendix E), an interesting direction for future work to extend the framework handle cases where the differentiability and invertibility conditions are not satisfied, using, e.g., the theory of nonsmooth implicit function theorems [25, 19].
3 Jacobian precision guarantees
In practice, either by the limitations finite precision arithmetic or because we perform a finite number of iterations, we rarely reach the exact solution x= (0). Instead, we reach approximate solution & and apply the implicit differentiation equation (2) at this approximate solution. This motivates the need for precision guarantees of this approach. We introduce the following formalism.
Definition 1. Let F x R. + Ra be continously differentiable optimality criterion mapping, Let A = = and B O2F. We define the Jacobian estimate (z, f). when A(x, B) is invertible, as the solution the linear equation A(x, (xr, B(x, 0). It is function J Rd R Rdxn,
holds by construction that [x (x)) = Oxx (0). Computing J(e. 0) for an approximate solution & of x (0) therefore allows approximate the true Jacobian Ox (9). In practice, an algorithm used to solve (1) depends on 0. Note however that, what we compute is not the Jacobian of (0), unlike works differentiating through unrolled algorithm iterations, but an estimate of Oz (0). We therefore use the notation æ,, leaving the dependence on 0 implicit.
We develop bounds of the form [[ 9) (f)l - x² (f)ll. hence showing that the error on the estimated Jacobian at most of the same order as that of as an approximation of =+)) These bounds are based on the following main theorem, whose proof is included in Appendix D.
Theorem 1. Let F : Rd x R + Rd be continuously differentiable. If there are B,y,,, R > S.t.
A - F and B O,F satisfy. for all E Rdd. E R and s.t. || xx x0)]l $ E:
A is well-conditioned, Lipschitz: |x, 0) 2 a |oll. [|,(x, 0) A(x (0), 1)pp S >|4 x² (0)]l-
B is bounded and Lipschitz: [x((x* (0), 0) < R, [B 0) B(x (0), - B| x' (0)]

---

Under these conditions, when | 1--===00 x0)0]] S we have
[] 0) - (0)l & (Ba-) + +Ro l| ==(0]]
This result is inspired by [52, Theorem 7.2], that is concerned with the stability of solutions inverse problems. As a difference, we consider that A(-, 0) is uniformly well-conditioned, rather than only at x" (0). This does not affect the first order in e of this bound, and makes it valid for all é. Our goal with Theorem is to provide result that works for general F but can be tailored to specific cases.
In particular, for the gradient descent fixed point (5), this yields
A(x, fx, 0) and Bax, 0) = n V (1)
By specializing Theorem for this fixed point, we obtain Jacobian precision guarantees with conditions directly on f rather than F; Corollary in Appendix D. These guarantees hold for instance for the dataset distillation experiment in Section 4. Our analysis reveals in particular that Jacobian estimation by implicit differentiation gains factor of compared to automatic differentiation. after t iterations of gradient descent in the strongly-come
setting [1, Proposition 3.2]. While our guarantees concern the Jacobian of x= (0), we note that other studies [47, 54, 13] give guarantees on hypergradients (i.e., the gradient of an outer objective).
We illustrate these results ridge regression, where x² (0) = argmin, |x - y/12 + >. 0,-7 This problem has the merit that the solution x= (0) and its Jacobian Ox- (0) are available in closed form. By running gradient descent for / iterations, we obtain estimate of x" (0) and an estimate J (ē, 0) of 0x (0); cf. Definition 1. By doing so for different numbers of iterations t, we can graph the relation between the error (0) 100--11 -112 and the error |0x. (0) [ f)]l,, as shown in Figure 3, empirically validating Theorem 1. The results in Figure 3 were obtained using the diabetes dataset from [35], with other datasets yielding a qualitatively similar behavior. We derive similar guarantees in Corollary in Appendix D for proximal gradient descent.
Implictt oiff
Unrolling
(B11l
Theoretical bound
(0) 8112
Figure 3: Jacobian estimate errors. Em- pirical error of implicit differentiation fol- lows closely the theoretical upper bound. Unrolling achieves a much worse error for comparable iterate error.
4 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 the dual. Here, x² (0) is the optimal dual solution, a matrix of shape m x k, where m is the number of training examples and i is the number of classes, and 9 R4+ is the regularization parameter. The challenge in differentiating x= (0) is that each row of x (0) is constrained to belong to the probability simplex A More formally, let Xer E mR be the training feature matrix and Ytr E {0, m he the training labels (in row-wise one-hot encoding). Let W(x, 0) == (Y= x)/8 E RR be the Aual-primal mapping. Then, consider the following bi-level optimization problem
0 x x') argmin fzx 0) = / | (x. 0) + (z, Yue),
min 0 mmxxp(A) 2
[Xeai W ((=" (0), Youili
subject to
sec
outer problem
inner problem
where C = A of variable = exp(A) in order to guarantee that the hyper-paramett 0 is positive. The matrix (0), f) E R contains the optimal primal solution, the feature weights for each class. The outer loss is computed against validation data X and Y.ai-
X A is the Cartesian product of probability simplices. We apply the change
While KKT conditions can be used to differentiate x (0), more direct way is to use the projected gradient fixed point (9). The projection onto C can be easily computed by row-wise projections on the
7

---

Mirrar descent (MD)
Preximal gradient (PG)
Block coordinate descent (BCD)
Runtime por 210p (secanda)
Unrelinn
500
Unro
Unrollng
impic dir (Di
500
ID
Fand
602
300
200
230
100
100
Number features
2000
Number of features
4000
6000
8030
10000
2000
Number features
4.00
5000
5000
10000
(a)
(b)
(C)
Figure 4: CPU runtime comparison of implicit differentiation and unrolling for hyperparameter optimization of multiclass SVMs for multiple problem sizes. Error bars represent 90% confidence intervals. (a) Mirror descent (MD) solver, with MD fixed point for differentiation. (b) Proximal gradient (PG) solver, with PG fixed point for differentiation. (c) Block coordinate descent solver; for implicit differentiation we obtain x (0) by BCD but perform differentiation with the MD and PG fixed points. This shows that the solver and fixed point can be independently chosen.
simplex. The projection's 's Jacobian enjoys closed form (Appendix C). Another way to differentiate x² (0) is using the mirror descent fixed point (13). Under the KL geometry, projections correspond to a row-wise softmax. They are therefore easy to compute and differentiate. Figure compares the runtime performance of implicit differentiation vs. unrolling for the latter two fixed points.
4.2 Dataset distillation
Dataset distillation [82, 59] aims learn small synthetic training dataset such that model trained on this learned data set achieves a small loss on the original training set. Formally, let Xxr E mxx and Ver [k]" denote the original training set. The distilled dataset will contain one prototype R The dataset distillation problem be cast as a bi-level problem, where in the inner problem we estimate a logistic regression model xx (0) E Raxk trained on the distilled images 0 e R*xp, while in the outer problem we want to minimize the loss achieved by x² (0) over the training set:
example for each class and therefore
then naturally
min MEREXP
[x (0), Xeri Wer) subject to x x0)))= E argmin [(x 0; [A]) + =||=||? ZERPxE
(10)
outer problem
peablem
where f(W,X;)) := regularization parameter that we found had very positive effect on convergence.
(ly, XW), denotes the multiclass logistic regression loss, and 10-3 is
In this problem, and unlike in the general hyperparameter optimization setup, both the inner and outer problems are high- dimensional, making it an ideal test-bed for tradient-based bi-level optimization methods. For this experiment, we use the MNIST dataset. The number of parameters in the inner problem p = 28 784. while the number of parameters 7840. We solve this problem
of the outer loss is & using gradient descent on both the inner and outer problem, with the gradient of the outer loss computed using implicit differentiation, as described $2. This is fundamentally different from the approach used in the original paper, where they used differentiation the unrolled iterates instead. For the same solver, approach was 4 times faster than the original one. The obtained distilled images & are visualized in Figure 5.
found that the implicit differentiation
Figure Distilled dataset E R obtained by solving (10).
4.3 Task-driven dictionary learning
Task-driven dictionary learning was proposed to learn sparse codes for input data such a way that the codes solve outer learning problem [60, 78, 85]. Formally, given a data matrix Xxr mxx

---

Table 2: Mean AUC (and 95% confidence interval) for the cancer survival prediction problem.
Method
Li logreg
L2 logreg
DictL + L2 logreg
Task-driven DictL
AUC (%)
71.6 + 2.0
72.4 2.8
68.3 + 2.3
73.2 + 2.1
and a dictionary of k atoms 0 E R sparse code is defined as a matrix (0) E Tgnx that minimizes a reconstruction loss f(z 8) ((X.r. .zff) regularized by a sparsity-inducing penalty gzx). Instead of optimizing the dictionary 0 to minimize the reconstruction loss, [60] proposed to
optimize an outer problem that depends on the code. Given a set of labels Ytr E {0, 1 we consider
logistic regression problem which results in the bilevel optimization problem:
min DERAXP WER bER
o (x" (0)w + b; yor) subject to x (0) E argmin fzz 0) + g(x)
(1I)
zERm
outer problem
inner prublem
When is the squared Frobenius distance between matrices, and the elastic net penalty, [60, Eq. 21] derive manually, using optimality conditions (notably the support of the codes selected at the optimum), an explicit re-parameterization of (0) as a linear system involving 0. This closed- form allows for direct computation of the Jacobian of x- w.r.t 0. Similarly, [78] derive first order conditions in the case where is 3-divergence, while [85] propose to use unrolling of ISTA iterations. Our approach bypasses all of these manual derivations, giving the user more leisure to focus directly on modeling (loss, regularizer) aspects.
We illustrate this on breast cancer survival prediction from gene expression data. We frame it as binary classification problem to discriminate patients who survive longer than years (m1 = 200) vs patients who die within 5 years of diagnosis (mo 99), from p 1, 000 gene expression values. As shown in Table 2, solving (1I) (Task-driven DictL) reaches classification performance competitive with state-of-the-arr art L1 L2 regularized logistic regression with 100 times fewer variables.
4.4 Sensitivity analysis of molecular dynamics
Many physical simulations require solving optimization problems, such energy minimization in molecular [76] and continuum [9] mechanics, structural optimization [53] and data assimilation [40].
In this experiment, we revisit example from JAX-MD [76], the problem of finding energy minimizing configurations to a system of k packed particles in a 2- 2-dimensional box of size
(0) = argmin (xx 0) = Σ Ou zERtx2
where x² (8) Rx2 are the optimal coordinates of the k particles, U(ra,,, 0) is the pairwise potential energy function, with half the 0.6, which optimize with a domain-specific optimizer [15]. Here we consider sensitivity of particle position with respect diameter Oz* (0), rather than sensitivity of the total energy from the original exper- iment. Figure shows results calculated via forward-mode implicit differentiation (JVP). Whereas differentiating the unrolled optimizer happens to work for total energy, here it typically does not even converge, due the discontinuous optimization method.
Figure 6: Particle positions and position sensitivity vectors, with respect to increasing the diameter of the blue particles.
particles at diameter and half at diameter
5 Conclusion
We proposed in this paper approach for automatic implicit differentiation, allowing the user to freely express the optimality conditions the optimization problem whose solutions are to be differentiated, directly in Python. The applicability of our approach to large catalog of optimality conditions is shown in the non-exhaustive list of Table I, and illustrated by the ease with which we can solve bi-level and sensitivity analysis problems.

---

References
[1] Ablin, G. Peyré, and T. Moreau, Super-efficiency of automatic differentiation for functions defined as a minimum. In Proc. of ICML, pages 32-41, 2020.
[2] A. Agrawal, B. Amos, S. Barratt, S. Boyd, S. Diamond, and Z. Kolter. Differentiable convex optimization layers. arXiv preprint arXiv: 12400, 2019.
[3] A. Agrawal, S. Barratt, S. Boyd, E. Busseti, and W. M. Moursi. Differentiating through a cone program. arXiv preprint arXiv: 1904.09433 2019.
[4] A. Ali, E. Wong, and J. 7.. Kolter. A semismooth newton method for fast, generic convex programming. In International Conference on Machine Learning, pages 70-79. PMLR, 2017.
[5] B. Amos. Differentiable optimization- baased modeling for machine learning. PhD thesis, PhD thesis. Carnegie Mellon University, 2019.
[6] B. Amos and J. Z. Kolter. Optnet: Differentiable optimization layer in neural networks, In Proc. of ICML, pages 135-145,20 2017.
[7] J. A. Andersson, J. Gillis, G. Horn, J. B. Rawlings, and M. Diehl, Casadi: software framework for nonlinear optimization and optimal control. Mathematical Programming Computation, 11(1):1-36, 2019.
[8] S. Bai, J. Z. Kolter, and V. Koltun. Deep equilibrium models. arXiv preprint arXiv: 2019.
[9] A. Beatson, J. Ash, G. Roeder, Xue, and R. P. Adams. Learning composable energy surrogates for pde order reduction. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances Neural Information Processing Systems, volume 33, pages 338-348. Curran Associates, Inc., 2020.
[10] B. M. Bell and J. V. Burke. Algorithmic differentiation of implicit functions and optimal values. In Advances in Automatic Differentiation, pages 67-77. Springer, 2008.
[11] Y. Bengio. Gradient-based optimization of hyperparameters. Neural computation, 12(8): 1900, 2000.
[12] Q. Bertrand, Klopfenstein, M. Blondel, S. Vaiter, A. Gramfort, and J. Salmon. Implicit differentiation of lasso-type models for hyperparameter optimization. In Proc. of ICML, pages 810-88 821, 2020.
[13] Q. Bertrand, Q. Klopfenstein, M. Massias, M. Blondel, S. Vaiter, A. Gramfort, and Salmon. Implicit differentiation for fast byperparameter selection in non-smooth convex learning, arXiv preprint arXiv:2105.0166 2021.
[14] M. J. Best, Chakravarti, and A. Ubhaya. Minimizing separable convex functions subject to simple chain constraints. SIAM Journal on Optimization, 10(3):658-67) 2000.
[15] E. Bitzek, P. Koskinen, F. Gähler, M. Moseler, and P. Gumbsch. Structural relaxation made simple. Phys. Rev. Lett., 97:17020 Oct 2006.
[16] M. Blondel. Structured prediction with projection oracles. In Proc. of NeurlPS, 2019.
[17] M. Blondel, V. Seguy, and A. Rolet. Smooth and sparse optimal transport. In Proc. of AISTATS, pages 880-889. PMLR, 2018.
[18] M. Blondel, O. Teboul, Q. Berthet, and J. Djolonga. Fast differentiable sorting and ranking. In Proc. of ICML, pages 950-959, 2020.
[19] J. Bolte, T. Le, E. Pauwels, and T. Silveti-Falls. Nonsmooth implicit differentiation for machine- learning and optimization. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 13537-13549. Curran Associates, Inc., 2021.
[20] F. Bonnans and A. Shapiro. Perturbation analysis of optimization problems. Springer Science & Business Media, 2013.
[21] J. Bradbury, R. Frostig, P. Hawkins, M. Johnson, C. Leary, D. M.clauun,,, G. Necula, A. Paszke, VanderPlas, Wanderman-MMine, and Q Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
10

---

[22] P. Brucker. An (n) algorithm for quadratic knapsack problems. Operations Research Letters, 3(3):163-166, 1984.
[23] O. Chapelle, V. Vapnik, O. Bousquet, and S. Mukherjee. Choosing multiple parameters for support vector machines. Machine learning, 46 46(1):131-159 2002.
[24] H. Cherkaoui, J. Sulam, and T. Moreau. Learning to solve tv regularised problems with unrolled algorithms. Advances in Neural Information Processing Systems, 33, 2020.
[25] F. Clarke. Optimization and Nonsmooth Analysis. Wiley New York, 1983.
[26] L. Condat. Fast projection onto the simplex and the & ball. Mathematical Programming, 158(1-2):555-585 2016.
[27] K. Crammer and Y. Singer. On the algorithmic implementation of multiclass kernel-based vector machines. Journal of machine learning research, 2(Dec)):26-22) 292,, 2001.
[28] M. Cuturi. Sinkhorn distances: lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems, volume 2, 2013.
[29] C.-A. Deledalle, S. Vaiter, J. Fadili, and G. Peyré. Stein unbiased gradient estimator of the risk (sugar) for multiple parameter selection. SIAM Journal on Imaging Sciences, 7(4):2448-244 2014.
[30] S. Diamond and S. Boyd. Cvxpy: A python- modeling language for convex optimiza- tion. The Journal of Machine Learning Research, 17(1):9009-2913, 2016.
[31] J. Djolonga and A. Krause. Differentiable learning of submodular models. Proc. of NeurlPS, 30:1013-1023, 2017.
[32] Domke. Generic methods for optimization- based modeling. In Artificial Intelligence and Statistics, pages 118-326. PMLR, 2012.
[33] J. C. Duchi, S. Shalev-Shwartz, Y. Singer, and T. Chandra. Efficient projections onto the Ei-llll for learning in high dimensions. In Proc. of ICML, 2008.
[34] D. Duvenaud, J. Z. Kolter, and M. Johnson. Deep implicit layers tutorial neural ODEs, deep equilibirum models, and beyond. Neural Information Processing Systems Tutorial, 2020.
[35] B. Efron, T. Hastie, Johnstone, and Tibshirani. Least angle regression. The Annals of slutistics, 32(2):407-499, 2004.
[36] L. El Ghaoui, F. Gu, B. Travacca, A. Askari, and A. Y. Tsai. Implicit deep learning. arXiv preprint arXiv: 06315, 2, 2019.
[37] R.-E. Fan, K.-W. Chang. C.-J. Hsich, X.-R. Wang, and C.-J. Lin. Liblinear: A library for large linear classification. JMLR, 9: 1:1871-1874, 2008.
[38] L. Franceschi, M. Donini, Frasconi, and M. Pontil. Forward and reverse gradient-based hyperparameter optimization. In International Conference on Machine Learning, pages 1 1165- 1173. PMLR,, 2017.
[39] L. Franceschi, P. Frasconi, S. Salzo, R. Grazzi, and M. Pontil. Bilevel programming for hyperparameter optimization and meta-learning. In International Conference on Machine Learning, pages 1568-1577. PMLR, 2018.
[40] T. Frerix, D. Kochkov, J. A. Smith, D. Cremers, M. Brenner, and S. Hoyer. Variational data assimilation with a learned inverse observation operator. 2021.
[41] R. Frostig, M. Johnson, D. Maclaurin, A. Paszke, and A. Radul. Decomposing reverse-mode automatic differentiation. In LAFI 2021 workshop at POPL, 2021.
[42] R. Frostig, M. J. Johnson, and C. Leary. Compiling machine learning programs via high-level tracing. Machine Learning and Systems (MLSyss) 2018.
[43] S. W. Fung, H. Heaton, Q. Li, D. McKenzie, S. Osher, and W. Yin. Fixed point networks: Implicit depth models with jacobian-free backprop. arXiv preprint arXiv:2103.12803 2021.
[44] Z. Geng, X.-Y. Zhang, S. Bai, Y. Wang, and Z. Lin. On training implicit models. Advances in Neural Information Processing Systems, 34:24247-24222
2021.
[45] S. Gould, B. Fernando, A. Cherian, P. Anderson, R. S. Cruz, and E. Guo. On differentiating parameterized argmin and argmax problems with application bi-level optimization. arXiv preprint arXiv: 1107.:1660..054440 05447, 2016.

---

[46] S. Gould, R. Hartley, and D. Campbell. Deep declarative networks: A new hope. arXiv preprint (rXiiv:909.04500 2019.
[47] R. Grazzi, L. Franceschi, M. Pontil, and S. Salzo. On the iteration complexity of hypergradient computation. In International Conference on Machine Learning, pages 7748-3758. PMLR, 2020.
[48] A. Griewank and A. Walther. Evaluating derivatives: principles and techniques of algorithmic differentiation. SIAM, 2008.
[49] S. Grotzinger and C. Witzgall. Projections onto order simplexes. Applied mathematics and Optimization, 12(1): ):247-270 1984.
[50] I. Guyon. Design of experiments of the nips 2003 variable selection benchmark. In NIPS 2003 workshop on feature extraction and feature selection, volume 253, 2003.
[51] M. R. Hestenes, F. Stiefel, et al. Methods of conjugate gradients for solving linear systems, volume 49. NBS Washington, DC, 1952.
[52] N. J. Higham. Accuracy and Stability of Numerical Algorithums. Society for Industrial and Applied Mathematics, second edition, 2002.
[53] Hoyer, Sohl-Dickstein, and S. Greydanus. Neural reparametcrizzation improves structural optimization. 2019.
[54] K. Ji, J. Yang, and Y. Liang. Bilevel optimization: Convergence analysis and enhanced design.
In International Conference on Machine Learning, pages 4882-4892 PMLR, 2021.
[55] Y. Kim, C. Denton, L. Hoang, and A. M. Rush. Structured attention networks. arXiv preprint arXiv:1702. 00887, 2017.
[56] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[57] S. G. Krantz and H. R. Parks. The implicit function theorem: history, theory, and applications.
Springer Science Business Media, 2012.
[58] C. H. Lim and S. J. Wright. Efficient bregman projections onto the permutahedron and related polytopes. In Proc. of AISTATS, pages 1205-1213, PMLR, 2016.
[59] J. Lorraine, P. Vicol, and D. Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In International Conference on Artificial Intelligence and Statistics, pages I 1540-1552. PMLR, 2020.
[60] J. Mairal, F. Bach, and J. Ponce. Task-driven dictionary learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 34(4):791 2012.
[61] J. Mairal and B. Yu. Complexity analysis of the lasso regularization path. In Proceedings of the 29th International Conference on Machine Learning, 2012, Edinburgh, Scotland, UK June 26 July 1, 2012. ieml.ec Omnipress, 2012.
[62] A. F. Martins and R. F. Astudillo. From softmax to sparsemax: A sparse model of attention and multi-label classification. Proc. of ICML, 2016.
[63] C. Michelot. A finite algorithm for finding the projection of a point onto the canonical simplex of R". Journal of Optimization Theory and Applications, 50(1):195-25 -9000, 1986.
[64] J.-J Moreau. Proximité et dualité dans un espace hilbertien. Bulletin de la S.M.F., 93:273-299, 1965.
[65] V. Niculae and M. Blondel. regularized framework for sparse and structured neural attention. In Proc. of NeurIPS, 2017.
[66] V. Niculae and A. Martins. Lp-sparsemap: Differentiable relaxed optimization for sparse structured prediction. In International Conference Machine Learning, pages 7348-7359) 2020.
[67] V. Niculae, A. Martins, M. Blondel, and C. Cardie. Sparsemap: Differentiable sparse structured inference. In International Conference Machine Learning, pages 3799-3808. 3808. PMLR, 2018.
[68] B. O'Donoghue, E. Chu, N. Parikh, and S. Boyd. Conic optimization via operator splitting and homogeneous self-dual embedding. Journal of Optimization Theory and Applications, 169(3):1042-1066 2016.
12

---

[69] N. Parikh and S. Boyd. Proximal algorithms. Foundations and Trends in optimization, 1(3):127-
239, 2014.
[70] F. Pedregosa. Hyperparameter optimization with approximate gradient. In International conference on machine learning. PMLR, 2016.
[71] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel,
P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine
Learning Research, 12:2825-2830. 2011.
[72] A. Rajeswaran, C. Finn, S. Kakade, and S. Levine. Meta-learning with implicit gradients. arXiv preprint arXiv: 9999.046993 2019.
[73] Z. Ramzi, F. Mannel, S. Bai, J.-L.. Starck, P. Ciuciu, and T. Moreau. Shine: Sharing the inverse estimate from the forward pass for bi-level optimization and implicit models. arXiv preprint 1:Xiv:2106.0055: 2021.
[74] N. Rappoport and R. Shamir. Multi-omic and multi-view clustering algorithms: review and cancer benchmark. Nucleic Acids Res., 46: 00546-10562, 2018.
[75] Y. Saad and M. H. Schultz. Gmres: A generalized minimal residual algorithm for solving nonsymmetric linear systems. SIAM Journal on scientific and statistical computing, 7(3):866- 869, 1986.
[76] S. Schoenholz and D. Cubuk. Jax md: A framework for differentiable physics. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1 2020.
Curran Associates, Inc.,
[77] M. W. Seeger. Cross-validation optimization for large scale structured classification kernel methods. Journal of Machine Learning Research, 9(6), 2008.
[78] P. Sprechmann, A. M. Bronstein, and G. Sapiro. Supervised non-euclidean sparse nmf via bilevel optimization with applications to speech enhancement. In 2014 4th Joint Workshop on Hands-free Speech Communication and Microphone Arrays (HSCMA), pages 11-15. IEEE, 2014.
[79] R. J. Tibshirani. The lasso problem and uniqueness. Electronic Journal of Statisties, 7(none): 1490, 2013.
[80] S. Vaiter, C.-A. Deledalle, G. Peyré, C. Dossal, and J. Fadili. Local behavior of sparse analysis regularization: Applications to risk estimation. Applied and Computational Harmonic Analysis, 35(3):433-451 2013.
[81] H. A. v. d. Vorst and H. A. van der Vorst. Bi-CGSTAB: A fast and smoothly converging variant of Bi-CG for the solution of nonsymmetric linear systems. SIAM Journal on Scientific and Statistical Computing, 13(2):631-644, 1992.
arXiv preprint [82] T. Wang, J.-Y. Zhu, A. Torralba, and A. A. Efros. Dataset distillation. arXiv:1811 10959, 2018.
[83] R. E. Wengert. A simple automatic derivative evaluation program. Communications of the ACM, 7(8):463 1964.
[84] Y. Wu, M. Ren, R. Liao, and R. B. Grosse. Understanding short-horizon bias in stochastic meta-optimization. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30- May 3, 2018, Conference Track Proceedings. OpenRe- view.net, 2018
[85] J. Zarka, L. Thiry, T. Angles, and S. Mallat. Deep network classification by scattering and homotopy dictionary learning, arXiv preprint arXiv: 00.61100000.111110 2019.
13


Reviews:
[{"review_id": "VUl0iP3eZSH", "reviewer": "Reviewer_sKqm", "summary": "EDIT: After the discussion with the authors and the revisions they submitted, I modified my overall score from 3 to 6, the presentation score from 3 to 4, and the soundness score from 1 to 3.\n\nThe paper presents a blueprint for automatic implicit differentiation of solutions to optimization problems, along with its implementation in the JAX library. They claim that their blueprint is widely applicable by listing several common optimization problem templates (e.g., those solvable by mirror descent, proximal gradient descent, Newton's method, conic programming, etc) that they claim one can apply their blueprint to. The main idea of the blueprint is to find a fixed point equation representing the optimality conditions and then apply the implicit function theorem to this fixed point equation. They also prove a theorem regarding the precision of the estimated implicit Jacobian in terms of the precision of the estimated fixed point, which they give numerical support for. Finally, they report numerical experiments for four problems, primarily comparing their method to unrolling.", "questions": "In line 262, what is meant by unrolling if the algorithm is nonsmooth?", "limitations": "A major limitation of the work is that it does not apply to the nonsmooth setting in which the fixed point equation associated to the optimality conditions is not continuously differentiable. Many of the optimization problems coming from machine learning are nonsmooth and so this severely limits the impact of the proposed blueprint.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "One strength of the paper is the software implementation in JAX, which is user-friendly, designed to be broadly applicable, and seemingly efficient; later theoretical flaws limit the justified use of this software and thus make this contribution less significant.\n\nThe paper is ambiguous when defining functions, not specifying if a function is continuous, differentiable, twice differentiable, etc, at key points, like in the blueprint and the theorem about Jacobian precision. For example, around line 98 when differentiating the root is introduced as the main principle which the rest of the paper relies on, there is no specification on the regularity of the function F besides that it should be \"a user-provided mapping, capturing the optimality conditions of a problem.\" However, the argument that follows relies on the smooth implicit function theorem, which in this context requires the function F to be C1 (continuously differentiable) in a neighborhood. This continues in line 118 where the chain rule is used without specifying that F is smooth, in line 168 when the regularity of f, G, and H are not specified, in line 199 for the definition of Jacobian estimate, in line 210 for the main/only theorem, etc.\n\nThis ambiguity becomes important for applying the results. Many of the functions involved in Table 1 on line 161, for instance, are frequently nonsmooth in machine learning contexts (e.g., the prox, the projection operator, etc), and thus the results developed in the paper are not applicable to them. There is mention in line 183 that, because the prox is differentiable a.e., that the smooth implicit function theorem can be applied a.e.. This is false - stronger assumptions than differentiability at a point are needed to apply the smooth implicit function theorem, i.e., F must be C1 on an open neighborhood of the solution (along with an invertibility condition). While nonsmooth implicit function theorems exist, e.g., [Clarke 1990], [Bolte et al 2022], they are nowhere mentioned and their requirements to be applied are different than those of the smooth implicit function since the generalized gradients involved are set-valued. This issue comes up again in line 202 where it's assumed that the solution xstar is differentiable, which is not true in general (e.g., the Lasso solution can have kinks as a function of the l1 weight).\n\nSome numerical experiments are also affected by this, for example the euclidean projection onto the simplex isn't smooth (line 257) and cannot be treated by the blueprint. While it's true that sometimes a smooth fixed point equation can be associated to a nonsmooth problem, such as in section 4.1 where the mirror descent formulation was used with the smooth KL projection on the simplex, there is no discussion of this phenomenon in the paper, nor a way to systematically construct such smooth fixed equations in general. The nonsmoothness also raises questions for the numerical comparisons regarding unrolling, which is mentioned in the questions section.\n\nThere is also a concern about the originality of the ideas in the paper. The main idea of the blueprint, to associate an equation modeling optimality to the problem instead of using the set-valued optimality conditions, is trivial in the smooth case since the optimality conditions will no longer be set-valued. Since little beyond line 183 is said about implicitly differentiating nonsmooth functions, the contribution of the blueprint here reduces to the fact that we can associate a fixed point equation to optimality conditions, which is well-known, e.g., for maximal monotone inclusions in [Bauschke, Combettes 2011].\n\nBecause the results are ultimately relegated to the smooth setting, I find the paper to be lacking in originality and significance. The main result that remains is the Jacobian precision theorem for smooth functions which, while compelling, is not sufficient on its own nor in conjunction the numerical experiments for smooth problems to deserve publication here. I think an integration of the nonsmooth case is necessary to make the contribution significant enough to worthy of publication, especially due to the prevalence of nonsmooth optimization problems in machine learning.\n\n\"Optimization and nonsmooth analysis\" - FH Clarke 1990\n\n\"Nonsmooth Implicit Differentiation for Machine Learning and Optimization\" - Jérôme Bolte, Tam Le, Edouard Pauwels, Antonio Silveti-Falls 2021\n\n\"Convex analysis and monotone operator theory in Hilbert spaces\" - HH Bauschke, PL Combettes, 2011", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "JOFhKKzifSJ", "reviewer": "Reviewer_Mi39", "summary": "This paper introduces a Jax package for automatic implicit differentiation. Specifically, the authors propose an efficient and modular approach for implicit differentiation of optimization problems. Their approach combines the benefits of autodiff and implicit differentiation. The authors also provide precision guarantees of the proposed approach. Moreover, the author demonstrate the effectiveness of the proposed approach in bi-level optimization and the sensitivity analysis.", "questions": "Please see weakness 2 and 3. If the authors can address these two concerns, I will improve my scores.", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "My detailed comments are given as below.\n\nStrength:\n\n1 The motivation of this paper is clear. I believe that this paper makes an important contribution to very relevant topics, e.g., the bi-level optimization, implicit deep neural networks.\n\n2 This paper is well written and the proposed approach is easy to follow. The authors provide some illustrative figures to demonstrate their approach. I feel this is good.\n\n3 The authors implemented four illustrative examples to demonstrate the effectiveness of their approach.\n\nWeakness:\n\n1 The main contribution of this paper is about the software, but the theoretical contribution is overstated. The proof of the theorem is quite standard and I do not get some new insight from it.\n\n2 Direct runtime comparisons with existing methods are missing. The proposed approach is based on implicit differentiation which usually requires additional computational costs. Thus, the direct runtime comparison is necessary to demonstrate the efficiency of the proposed approach.\n\n3 Recently, implicit deep learning has attracted many attentions, which is very relevant to the topic of this paper. An implementation example of implicit deep neural networks should be included. Moreover,  many Jacobian-free methods e.g., [1-3] have been proposed to reduce the computational cost. The comparisons (runtime and accuracy) with these methods are preferred.\n\n[1] Fung, Samy Wu, et al. \"Fixed point networks: Implicit depth models with Jacobian-free backprop.\" (2021).\n\n[2] Geng, Zhengyang, et al. \"On training implicit models.\" Advances in Neural Information Processing Systems 34 (2021): 24247-24260.\n\n[3] Ramzi, Zaccharie, et al. \"SHINE: SHaring the INverse Estimate from the forward pass for bi-level optimization and implicit models.\" arXiv preprint arXiv:2106.00553 (2021).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "F7mteGJtY50", "reviewer": "Reviewer_SAnA", "summary": "Implicit differentiation can be used to through the solution to an optimization problem without having to backpropagate through the method by which this solution was determined. While this approach has been used in many different contexts, implicit differentiation must often be tailored towards the specific problem in question. This article introduces a library that implements automatic implicit differentiation, addressing this limitation. It goes on to provide precision estimates if an optimization problem has not been exactly solved, and concludes with several examples showcasing the library.", "questions": "- Returning to the dataset distillation example, reference [55] only uses a few gradient steps for the inner loop to make their algorithm scalable. Do the authors think their approach towards dataset distillation for logistic regression would scale towards deeper networks?\n- l. 224: What do the authors mean by \"implicit differentiation gains a factor of t compared to automatic differentiation\"? Moreover, does automatic differentiation in this context refer to automatic differentiation through the t iterations rather than implicit automatic differentiation?\n- l. 227: What is missing to apply these results to hypergradients? Would that simply be a matter of applying the chain rule to compute the gradients with respect to the objective and obtain corresponding precision guarantees?\n- I did not understand the illustration in figure 6; since $\\theta\\in\\mathbb{R}^k$ specifies the diameter of each particle, isn’t $\\partial x^{\\ast}(\\theta)\\in\\mathbb{R}^{k\\times(k\\times 2)}$? Is the figure depicting the diagonal elements of that Jacobian?\n- Where is digit ‘9’ in figure 5?\n\nOne typo: l. 324: 'converge, due to the'", "limitations": "I believe the authors have adequately addressed the limitations of their work. ", "rating": 9, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This submission fills an important gap by providing a library implementing automatic implicit differentiation. In my view, this promises to make the use of implicitly defined layers (regardless of their specific form) more accessible for practitioners, in particular making it easier to try out a particular idea. The library implementation is well documented and seems to integrate well with other Jax packages. The article itself is well written, providing a nice balance of motivation, code, theory, and examples. Theorem 1 is closely related to existing results about inverse stability (as the authors note), but it is helpful to have it stated explicitly in the context of implicit differentiation. The proof seems to be correct.\n\nAs a minor issue, I was surprised to not see any example applications in the context of deep networks. Reference [55] implements dataset distillation on deep networks, yet the authors seem to be focused on a logistic regression case (see Questions). The code provided in the supplementary material also seems to include deep learning examples, but the authors make no note of this in the appendix. I don't see the inclusion of deep learning examples in the article as a necessity at all, but was a little surprised about the lack thereof.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept