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 2021 Submitted (2021)
Keywords: implicit differentiation bilevel optimization autodiff

Abstract: Automatic differentiation (autodiff) has revolutionized machine learning.  It allows expressing complex computations by composing elementary ones in creative ways and removes the tedious burden of computing their derivatives by hand. More recently, differentiation of optimization problem solutions has attracted a great deal of research, with applications as a layer in a neural network, and in bi-level optimization, including hyper-parameter optimization. However, the formulae for these derivatives often involves a tedious manual derivation. In this paper, we propose a unified, 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$ to automatically differentiate the optimization problem. This way, our approach combines the benefits of implicit differentiation and autodiff.  We show that seemingly simple principles allow to recover all recently proposed implicit differentiation methods and create new ones easily.  We demonstrate the ease of formulating and solving bi-level optimization problems using our framework.


Efficient and Modular Implicit Differentiation
Anonymous Author(s) Affiliation Address
email
Abstract
Automatic differentiation (autodiff) has revolutionized machine learning. It allows expressing 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 as a layer, and in bi-level problems such as hyper-parameter optimization and meta-learning. However, the formulas for these derivatives often involve case-by-case tedious mathematical derivations. In this paper, propose unified, efficient and modular approach for implicit differentia- tion of optimization problems. In our approach, the user defines (in Python in the case of our implementation) function f' capturing the optimality conditions of the problem to be differentiated. Once this is done, we leverage autodiff of F' and implicit differentiation automatically differentiate the optimization problem. Our approach thus combines the benefits of implicit differentiation and autodiff. We show that seemingly simple principles allow to recover many recently proposed im- plicit 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.
3
5
10
11
12
13
14
15
16
17
Introduction
Automatic differentiation (autodiff) is now an inherent part of machine learning software. allows expressing complex computations by composing elementary ones in creative ways and removes the tedious burden of computing their derivatives by hand. 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 order to solve outer one. Examples of applications in machine learning include hyper- aaper-pprramete optimization [19, 63, 57, 30, 10, 11]. neural networks [46], and meta-learning [31, 59]. Another line of applications layer" [42, 6, 51, 27, 36, 7]. which usually includes regularization constraints 27 in the optimization problem in order to impose desirable structure on the layer output. In addition providing well sought-for interpretability. recent research indicates that such structure could be beneficial in order to improve generalization of neural networks [20].
19
21
23
25
"optimization
29
Since optimization problem solutions typically do not enjoy explicit formula in terms of their 31 inputs, autodiff cannot used directly to differentiate these functions. In recent years, two main approaches have been developed to circumvent this problem. The first one consists in unrolling the iterations of an optimization algorithm and to use the final iteration as a proxy for the optimization problem solution [68, 28, 25, 31]. An advantage of this approach is that autodiff through the algorithm
33
35 iterates can then be used transparently. However, this requires a reimplementation of the algorithm necessarily autodiff friendly. Moreover, forward- 37 mode 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. A second
using the autodiff system, and not all algorithms
Submitted to 35th Conference on Neural Information Processing Systems (NeurIPS 2021). Do not distribute.

---

39
approach is to see optimization problem solutions as implicitly-define defined functions of certain optimality 40
conditions. Examples include stationary conditions [9, 46], KKT conditions [19, 35, 6, 53, 52] and the proximal gradient fixed point [51, 10, 111. An advantage of such implicit differentiation is that a reimplementation is not needed, allowing to build upon state-of-the-arr software. However, so far, obtaining the implicit differentiation formulas required case-by-can tedious mathematical derivation. Recent work [2] attempls to address this issue by adding implicit differentiation on top of evxpy [26]. This works by reducing all convex optimization problems to conic program and using conic programming's s optimality conditions to derive an implicit differentiation formula. While this approach is very generic, solving convex optimization problem using conic programming solver an ADMM-based splitting conic solver [54] in the case of cvxpy- is rarely the state-of-the-ar approach for each particular problem instance.
41
42
48
44
45
46
48
50 In this work, we adopt a different strategy, which allows to easily add implicit differentiation on top of existing solvers. In our approach, the user defines (in Python in the case of our implementation) a mapping function F capturing the optimality conditions of the problem solved by the algorithm. Once this is done, we leverage autodiff' of I' combined with implicit differentiation techniques to automatically differentiate the optimization problem solution. In this way, our approach is very generic and still the efficiency of state-of-the-art solvers. It therefore combines the benefits of implicit differentiation and autodiff. To summarize, we make the following contributions.
51
52
54
56
57
We delineate extremely general principles for implicitly differentiating through an optimization problem solution. Our approach can be seen as "hybrid", in the sense that it combines implicit differentiation with autodiff of the optimality conditions.
58
59
60
We show how to instantiate our framework in order to recover many recently-proposee implicit differentiation schemes, thereby providing unifying perspective. We also obtain new implicit differentiation schemes, such as the one based on the mirror descent fixed point.
61
62
63
On the theoretical side, we provide new bounds on the Jacobian error when the optimization problem is only solved approximately.
65
We describe JAX implementation and provide blueprint for implementing our approach in other frameworks. We will open-souree a full-fledged library for implicit differentiation in JAX.
66
67
We implement four illustrative applications, demonstrating our framework's ease of use.
66 In essence, our implementation significantly extends JAX's default autodiff system in context of the numerical optimization domain. From end user's perspective, autodiff with JAX simply becomes more efficient if they use solvers with implicit differentiation set up by our framework.
70
Notation. We denote the gradient and Hessian of f: Rd - R evaluated Rd and fz) Ra We denote the Jacobian of F: Rd - Rp evaluated at Rdd by 73 (F(x) E Roxd When f or F have several arguments, we denote the gradient, Hessian and Jacobian 74 in the 1th argument by V,, V? and 0, respectively. The standard probability simplex is denoted 1, 2 0}. For any set C C Rd,, we denote by Ie the function 76 Rdd (+(c)) where Tc(x) = Of C, Ic(x:) = +oo otherwise. For vector matrix A,
R by Vf(z)
75 by Ad fo E Rd: ||^|||
77 note |||l the Frohenius (or Euclidean) norm, and || llop the operator norm.
2 Proposed framework: combining implicit differentiation and autodiff
78
2.1
General principles
80 Overview.
Contrary to unrolling of algorithm iterations, implicit differentiation typically involves 81 a manual, sometimes complicated, mathematical derivation. For instance, numerous works [19, 35, 6, 53, 52] use Karush-Kuhn-Tuckee (KKT) conditions in order to relate a constrained optimization 83 problem's 's solution to its inputs, and to manually derive a formula for its derivatives. The derivation
84 and implementation is typically case-by-case.
85 In this work, we propose a general way to easily add implicit differentiation on top of existing solvers.
86 In our approach, the user defines (in Python in the case of our implementation) a mapping function
87 F capturing the optimality conditions of the problem solved by the algorithm. We provide reusable
building blocks to easily express such F. Once this is done, we leverage autodiff of F combined with

---

X_tr, y_tr load_data()
f(x, theta):
residual jap. dot x) y_tr
return (jnp. sum (residuul ** 2) theta * jum ** 2)) /
F jax. jax.grad(f)
@custom_root((F)
def ridge_aolver (theta):
XX jup dot (X_tr.1 X_tr)
Xy = jup (ott y_tr) I jup.eye(X_ sye(p_)r.shappe())
return jup. linalg. sXlve theta Xy)
(10.0
Figure 1: Example: adding implicit differentiation on top of ridge regression solver. The function f(z, 0) defines the objective function and the mapping F, here simply equation (4), captures the optimality conditions. The decorator @custom_roo (provided by our library) automatically adds implicit differentiation to the solver for the user. The last line evaluates the Jacobian at A 10.
89 implicit differentiation to automatically differentiate the optimization problem solution. simple illustrative example is given in Figure 1.
90
Differentiating root. Let F; IRd x R" + Rd be user-provided mapping, capturing the opti- mality conditions of a problem. An optimal solution, denoted x (0), should be root of
81
92
F (0), 0.
33 can see (0) as an implicitly defined function of 0 R", i.e., x²: R + Rdd Our goal
is to differentiate x² (0) w.r.t. 0. From the implicit function theorem [44], if F is continuously
85 differentiable and the Jacobian 0 evaluated x (0) x a square invertible matrix, then Ox
exists. Using the chain rule, we know that the Jacobian Ox* (0) satisfies
0₂ F Fz 0xxx (0) + 32 F x (0), 0) 0.
97
Computing 0x*(0) boils down to the resolution of the linear system of equations
0 F(x Ox (8) O2F(x"
(2)
AERexa
JER
BERaxa
When (1) is one-dimensional root finding problem (d = 1), (2) becomes particularly simple since
Vxr (0)= = BT /A, where A is a scalar value.
89
We will show that existing and new implicit differentiation methods all reduce this simple principle. We call our approach hybrid, since it combines implicit differentiation (it involves the resolution of a linear system) with the autodiff of the optimality conditions F. Our approach is efficient can be added on top of any state-of- the-art solver and modular as the optimality conditon specification is decoupled from the implicit differentiation mechanism.
101
103
105
Differentiating a fixed point. We will encounter numerous applications where x²' x0) is implicitly defined through fixed point iteration:
106
r²(0) T(x((())))
where T: Rd R Rdd. This can be seen as particular case of (1) with
107
F(x" (()) 0) ( (z)) x² (0).
(3)
In this case, using the chain rule, have
A = -0 (zT and B 0zF(2 (0). = 0,T(xx (0).0)).
3

---

109
mappings to Computing JVPs and VJPs. In practice, all we need to know from F is how to left-multiply or right-multply 0, F and 0₂F with a vector of appropriate size. These are called vector-Jacobian product (VJP) and Jacobian-vector product (JVP), and are useful for integrating x² (0) with reverse- mode and swaard-mode autodiff, respectively. Often times, will be explicitly defined. In this case, computing the VJP or JVP can be done via autodiff. Other times, F may itself be implicitly defined, for instance when F' involves the solution of a variational problem. In this case, computing the VJP or JVP will itself involve implicit differentiation.
110
112
113
114
116
The
(JVP) between J= J=0z² 00) and vector v, Je, can be computed efficiently by solving A(Jv) Bv. The left-multiplicationn (VJP) of with J, v² J, can be computed by first Then, we can obtain v J by 0' J = u AJ u. B. Note that when B changes but and a remain the same, we do not need to solve A Au = v once again. This allows to compute the VJP w.r.t. different variables while solving only one linear system.
117
solving A u
119
121
To solve these linear systems, we can use the conjugate gradient method [39] when A is positive semi definte and GMRES [61] or BiCGSTAB [66] when A is not. All algorithms are matrix-free free, i.., they only require matrix-vector products (linear maps). Thus, all we need from is its JVPs or VJPs. An alternative to GMRES/BiCGSTAN is to solve the normal equation AA T Av using
122
123
124
125
conjugate gradient, which we find faster in some scenarios.
126
Pre- Pre-processing and post-processing mappings. Often times, the goal is not to differentiate 0 per 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 x (0) is used as the output of a neural network layer, in which case 0 is produced by the previous layer. Likewise, x² x²(() (0) will often not
127
se, but the parameters of function producing 0. One example of such pre-processing
128
129
130
131
be the final output we want to differentiate. One example of such post-procsssing is when =0)) is the solution of a dual program and we apply the dual-primal mapping to recover the solution of the primal program. Another example is the application of loss function, in order to reduce (0) to a
132
133
134
scalar value. In all these cases, we leave the differentiation of the the autodiff system, allowing us to compose functions in complex ways.
135
136
Usage and implementation details. Our implementation is based on JAX [17, 34]. JAX's autodiff features enter the picture in at least two ways: (i) we lean heavily JAX within our implementation, and (ii) we integrate the differentiation routines introduced by our framework into JAX's existing autodiff system. In doing the latter, we override JAX's default autodiff behavior (e.g. of differentiating transparently through an iterative solver's unrolled iterations).
137
138
139
140
We delincate here what features are needed from autodiff system to implement our proposed framework. As mentioned, we only need access to F' through the JVP or VJP of 0 and FzF.. Since the definition of F will often include gradient mapping V (if 8) (see examples in $2.2), -ocond-order derivatives need also be supported. Our library provides two decorators, custom_root and custom_ wauttomfiieeeeeee fixed__ppppee _pixee,_ppp,,..... for adding implicit differentiation top of a solver, given optimality conditions F or fixed point iteration T. This functionality requires the ability to add custom JVP and/or VJP to function. All these features are supported by recent autodiff systems, including JAX [17], TensorFlow [1] and PyTTorch [56]. Our implementation, made in JAX, also uses JAX-specific features. We make extensive
142
143
144
146
147
148
149
of automatic batching with jax.vmap. JAX's vectorizing map transformation. In order to solve the normal equation AA u = Av, we also use JAX's ability to automatically transpose linear map using jax. linear ttteeerrttraspsseeee [33].
150
151
152
2.2 Examples
153
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 F or T implies different trade-offs in terms of computational oracles; see Table Source code examples are given in Appendix A.
154
155
156
Stationary point condition.
The simplest example is to differentiate through the implicit function
x² (0) argmin (xx, 0),
zEr
where f: RRd x R + R is twice differentiable. In this case, F is simply the gradient mapping F(x, 0) vif(z,0) 0).
157
(4)
4

---

Table Summary of optimality condition mappings. Oracles are accessed through their JVP or VJP.
Name
Equation
Solution needed
Oracles needed
Stationary
(4), (5)
Primal
Vif V.J. H. G. d HH
KKT
(6)
Primal and dual
Proximal gradient
(7)
Primal
V.S. prox
Projected gradient
(9)
Primal
ViS. proje
Mirror descent
(11)
Primal
V.f. proje Vip
Newton
(15)
Primal
[ 0)] Vif(z,0)
Block proximal gradient
(16)
Primal
[V [prox projar
Conic programming
(20)
Residual map root
158
We then have F,F(2,, 0) - viS(2, 0) and - 0). the Hessian of f in its first
O2Vif(x; 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
160
Tzz, 0) 17 if 0),
(5)
which holds for all step sizes 1) > 0. Using (3), it is easy verify that we end up with the same linear system since cancels out.
182
KKT conditions. We now show that the KKT conditions, manually differentiated in several works
184
[19, 35, 6, 53, 52], fit our framework. As we will see, the key will be to group the optimal primal and dual variables our x' (0). Let consider the general problem
185
argmin f(z 0) subject to G(z, 0) - 0, H(z, 0) 0,
SERY
where Rº is the primal variable, f: R x R" - R. G: x R' and H: R" + Ra. The stationarity, primal feasibility and complementary slackness conditions give
187
Vif(z, (0 [0, A ( 0)] v 0
H(z,0) =0 0
A G(2 =0
(6)
188 where v E RD and À E RT are the dual variables, also known as KKT multipliers. The system of can group the primal and dual solutions as x² (0) = (0), v" (0), (0)))((0))))))))))))))))))))(()) (f)) to form the root of a function (0), 0), where F: Rd x R Rd and = p +r. The primal and dual solutions can obtained from a generic solver, such as
189
(potentially nonlinear) equations (6) fits our framework,
170
171
an interior point method. In practice, the above mapping F will be defined directly in Python (see mappinn Figure in Appendix A) and F will be differentiated automatically via autodiff.
173
174
Proximal gradient fixed point. Unfortunately, not all algorithms return both primal and dual solutions. Morcover, if the objective contains non-smooth terms, proximal gradient descent may be more efficient. We now discuss its fixed point. Let x=* (0) be implicitly defined as
175
(0) == argmin (ax 0) g(x, 0),
#ER*
177 where f: Rdd x R + R is twice-differentiable convex and g: Rdd x R -> R is convex but possibly non-smooth. Let us define the proximity operator associated with g by
178
prox g (y,
8) argmin
_ || 2112 + 0).
=ERd
179 To implicitly differentiate through x (0), we use the fixed point mapping [55, p50]
T 0) = prox (x - 1f 0),
(7)
for any step size >> O. The proximity operator is 1 I-Lipschitz continuous [50]. By Rademacher's 181 theorem, differentiable almost everywhere. Many proximity operators enjoy a closed form and can casily be differentiated, as discussed in Appendix B.
5

---

183
Projected gradient fixed point. As special case, when g(z, 0) is the indicator function Ic(0)(2). (z),
184
where C(O) is convex set depending on 0, we obtain
x' (0) =
argmin ((, 0). =EC(()
(8)
185
The proximity operator , becomes the Euclidean projection onto C(0)
prox (v, 0) = (v, 0) == argmin || - -vl
=EC(O)
186
and (7) becomes the projected gradient fixed point
Tæx,
(V
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 Appendix B.
188
Mirror descent fixed point. We again consider the case when x (0) is implicitly defined as the solution of (8). We now generalize the projected gradient fixed point beyond Euclidean geometry. Let the Bregman divergence Dip dom((²) relint(dom(s)) R+ generated by yp be defined by
190
191
D1 v) p(x) ( (Vp(v), 2).
192
We define the Bregman projection of onto C()) C do by
proje (y. 0) = argmin (x (y)). aEC(+)
(10)
193
Definition (10) incudes the mirror map (y) for convenience. can be seen as a mapping from
Rdd (o). ensuring that (10) well-defined. The mirror descent fixed point mapping is then
194
â Vp(x)
y = 1 - VI)(<,,)
(11) Because T involves the composition of several functions, manually deriving its JVP/VJP is error prone. This shows that our approach leveraging autodiff allows to handle more advanced fixed point mappings. A common example of &o is p(x) = (z, log x - 1), where dom(pp) I2. In this case, Di is the Kullback-Leilerr divergence. An advantage of the Kullback-Leibler projection is that it sometimes easier compute than the Euclidean projection, as we detail in Appendix B.
T(x, 0) = proje (y. 0).
196
197
198
199
200
Other fixed points. More fixed points are described in Appendix C.
201
2.3
Jacobian bounds
202
In practice, either by the limitations of finite precision arithmetic or because we perform a finite number of iterations, we rarely reach the exact solution x* (0), but instead only reach an approximate solution & and apply the implicit differentiation equation (2) at this approximate solution. This
203
204
205
motivates the need for approximation guarantees of this approach
Definition 1. Let F IRd x R + Rdd be an optimality criterion mapping. Let A= - and B 0zF. We define the Jacobian estimate at (z, 0) the solution the following linear equation A(x, (z, 0) = B(x, 0). It is function J IRdd x IR"" + Rdxn.
206
208
209
holds by construction that (xx (x), 0) = dx* (0). Computing (2, 0) for an approximate solution & of xx (0) therefore allows approximate the true Jacobian (0). In practice, an algorithm used to solve (1) depends on A. Note however that, what we compute is not the Jacobian of (()), unlike works unrolling an algorithm that differentiate through its iterations, but estimate of (0). We therefore use the notation अ, leaving the dependence on 9 implicit.
211
213
We develop bounds of the form || 0) Or (0)l < ClI x² (0)ll. hence showing that the ervor 215 on the estimated Jacobian at most of the same order as that of as an approximation x= (0).
These bounds are based on the following main theorem, whose proof is included in Appendix D.

---

Mirrar descent (MD)
Preximal gradlent (PG)
Block coordinate descent (BCD)
Runtime por 2t0p (secanda)
Unrelinn
500
Unro
Unroling
impic dir 3Di
500
ID
Fand
2.50
102
300
200
230
100
100
Number features
10000
2000
Number of features
4000
6000
8030
10000
2000
Number features
4.00
5000
5000
10000
(a)
(b)
(C)
Figure 2: 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 solver, with mirror descent fixed point for implicit differentiation. (b) Proximal gradient solver, with proximal gradient fixed point for implicit differentiation. (c) Block coordinate descent solver; for implicit differentiation we obtain z²( by BCD but perform differentiation with the mirror descent and proximal gradient fixed points. This showcases that the solver and fixed point can be independently chosen.
Theorem (Jacobian estimate). Let F IR x R" - Rd. Assume that there exist a, 7,
217
such that A = F F and B 02F satisfy, for all Rd, E IR" and such that ||z -z-xx (0)] <
219 A cell-conditioned, Lipschitz: ||, 0) a al|b|| |l [|,(x, 0) A(x* 0) lop S x² (0)
220
B is bounded and Lipschitz: [B (8). || v R, [ 0) (x), - B x² (0)]
Under these conditions, when | x=' (0)] - s, we have
[| 0) Ox (0)]l < (Ba- 1 + yRa || x0)1l (0)
This result is similar [40, Theorem 7.2], that is concerned with the stability of solutions inverse problems. Here we consider that A(, is uniformly well-conditioned, rather than only at x= (0). This does not affect the first order in of this bound, and makes it valid for all ft. It is also more tailored to applications to equation-specific cases.
223
225
Indeed, Theorem can be applied to specific functions F or T for some root and fixed-point equations. In particular, for gradient descent fixed point, where T(x, 0) = x Vi,(E, 0), this yields
227
A(x, 0) = fz, 0) and Bax, 0) = -122 0).
This guarantees precision on the estimated Jacobian under regularity conditions on f directly; see Corollary in Appendix D.
229
230
For proximal gradient descent, where T'I, 0) = prox (x - 1/ if 0), 0), this yields
A(x, 0) = 0, a, = - (I V prox (x v V 0), prox B(x, 0) = a2prox ng (x ( - 1 0 (x
- v, f (2x
An important special case that of function to minimize in the form f(x, 0) + gz)) where the prox function is smooth, and does not depend on 0, as is the case in our experiments in $3.2. For this setting, we derive similar guarantees in Corollary 2 in Appendix D. Recent work also exploits local smoothness of solutions to derive similar bounds [11, Theorem 13].
232
233
234
235
3
Experiments
To conclude this work, demonstrate the ease of formulating and solving bi-level optimization problems with our modular framework. We also present an application to the sensitivity analysis of molecular dynamics.
237
238
3.1
Hyperparameter optimization of multiclass SVMs
240
In this example, we consider the hyperparamer optimization of multiclass SVMs [23] trained in the dual. Here, x² (0) is the optimal dual solution, matrix of shape m k, where m is the number of
7

---

Table 2: Mean AUC (and 95% confidence interval) for the cancer survival prediction problem.
Method
L1 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
training examples and k is the number of classes, and 0 E 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 Xtr RRxx be the training feature matrix and Y+r E {0, 1]m* be the training labels (in row-wise one-hot encoding). Let W (z, == (Y z)/0 R be the dual-primal mapping. Then, we consider the following bi-level optimization problem
243
244
245
248
argmin 1 0=exp(A]
51l 21l ((x Yoa1 11?
subject to
x (0) = argmin [(x, #EC
2
W O)?.
outer peoblem
inner problem
(12)
where C = A
x A is the Cartesian product of probability simplices. We apply the change of variable = exp(A) in order to guarantee that the hyper-parameter 0 is positíve. The matrix W (0), 0) E P contains the optimal primal solution, the feature weights for each class. The outer loss is computed against validation data Xva and Yva1-
248
249
251
In order to differentiate x² (6), several ways are possible using our framework. The first one would be to map (12) to a quadratic program form (18) and use the KKT conditions to form a mapping F(x, f). more direct way is to use proximal gradient fixed point (9). Since C is a Cartesian product, the projection can be easily computed by row-wise projections on the simplex, which we map over rows $B.1, this projection's Jacobian enjoys a closed form. A third way to differentiate x (0) is using the mirror descent fixed point (1I). Under the KL geometry, proje (y, 0) corresponds to a row-wise softmax. It is therefore easy to compute and differentiate. Figure 2 compares the runtime performance of implicit differentiation vs. unrolling for the latter two fixed points. code example is included in Figure 8 in the Appendix.
252
253
255
(using vectorized operations) via jax. vmap. As we explained
257
259
260
3.2
Task-driven dictionary learning
-riked dictionary learning was proposed to learn sparse codes for input data in such a way that 262
the codes solve an outer learning problem [47, 64, 70]. Formally, given a data matrix X E mxx and a dictionary of & atoms 0 E RRp sparse code is defined as a matrix x² x0) E mg that a reconstruction loss fxx 6) f(xur, .f) regularized by a sparsity-inducing penalty g(z). Instead of optimizing the dictionary 0 to minimize the reconstruction loss, [47] proposed to 286
264
minimizes
optimize an outer problem that depends on the code. For example, given a set of labels Yux E {0, 1)m we consider logistic regression problem which results in the bilevel optimization problem:
min
(xz"(0)w + b; Vor) subject to x (0) E argmin f(z, 0) + g(z) rER=*
(13)
DERAXP DER
outer problem
inner problem
When é is the squared Frobenius distance between matrices, and g the elastic net penalty, [47, Eq. 269 21) derive manually, using optimality conditions (notably the support of the codes selected at the optimum), an explicit e-paraaeterizattin of x* x*((), (0) as a linear system involving 0. This closed- form allows for a direct computation of the Jacobian of xt w.r.t. Similarly, [64] derive first order conditions in the case where is B-divergence, while [70] propose to unrolling of ISTA iterations. Our approach bypasses all of these computations, giving the user more leisure to focus directly on modeling (loss, regularizer) aspects (see code snippet in Figure 9 in Appendix).
271
272
273
275
We illustrate this on a problem of breast cancer survival prediction from gene expression data, framed as a binary classification problem to discriminate patients who survive longer than 5 years
276
277
(ml 200) vs patients who die within 5 years of diagnosis (mi = 99), from p = 1, 000 gene expression values. As shown in Table 2, solving (13) (Task-driven DictL.) reaches classification performance competitive with sfate-of-the-art L, or Toz regularized logistic regression with 100 times fewer variables. See Appendix E.2 for more details.
279
8

---

Dataset: Distillation (MNIST), Generalization Accuracy: 0.8556
Figure 3: Distilled MNIST dataset 0 Rxxp obtained by solving (14). We learn one image per class such that a logistic regression model trained on 0 achieves the lowest logistic loss on the MNIST training set. Implicit differentiation was times faster than unrolling.
281
3.3
Dataset distillation
Dataset distillation [67, 46] aims to learn a small synthetic training dataset such that a model trained 282
on this learned data set achieves small loss on the original training set. Formally, let Xtr E mxx and Ytr E [k]" denote the original training set. The distilled dataset will contain one prototype example for each class and therefore 0 E R The dataset distillation problem can then naturally be cast as bi-level problem, where in the inner problem we estimate logistic regression model x (0) E IP
284
285
trained on the distilled images A E RKX,, while in the outer problem we want to minimize the loss achieved by x (0) over the training set:
287
288
argmin [a (0), Xeri Utr) subject to x= (0) argmin (zr,0; (b:) + e||x|| BERtxy
(14)
zER
outer problem
inner problem
where (;)) 6(Xx,, y), l denotes the multiclass logistic regression loss, and & 10-3 3 is regularization parameter that we found had very positive effect on convergence.
289
290
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 gradient- -hased bi-level optimization methods. For this experiment, we use the MNIST dataset. The number of parameters in the inner problem is p = 28 = 784. while the number of parameters of the outer loss is k x p 7840. We solve this problem using gradient descent on both the inner and outer problem, with the gradient of the outer loss computed using implicit differentiation, as described in $2 This is fundamentally different from the approach used in the original paper, where they used differentiation of the unrolled
292
293
294
295
296
297
298
iterates instead. For the same solver, we found that the implicit differentiation approach was 4 times faster than the original one. The obtained distilled images 0 are visualized in Figure and code example is given in Figure 10 in the Appendix.
299
301
3.4
Sensitivity analysis of molecular dynamics
302
Many applications of physical simulations require solving optimiza- tion problems, such as energy minimization in molecular [62] and continuum [8] mechanics, structural optimization [41] and data assimilation [32]. However, even fully differentiable simulators
303
304
305
306
may not have efficient or accurate derivatives. We revisit an exam- ple from JAX-MD [62], the problem of finding energy minimizing configurations to a system of packed particles in an mi-dimensional box of width E,
308
x" (0) argmin
Σ Utay mod L, mm
Figure 4: Particle positions and position sensitivity vectors, with respect to increasing the diameter of the blue particles.
310
where U (z.j, 0) is the pairwise potential energy function, with half the particles at diameter and half at diameter 0 = 0.6, which we optimize with a domain-specific optimizer [13]. Here we consider sensitivity of particle position with respect to diameter Ox (0),
312
313
314
rather than sensitivity of the total energy from the original experi- ment. Figure 4 shows results calculated via forward-mode implicit differentiation (JVP). Whereas
316
differentiating the unrolled optimizer happens to work for total energy, here typically does not even
converge (see Appendix Fig. 15), due the discontinuous optimization method.

---

References
318
319
[1] M. Abadi, P. Barham, J. Chen, Z. Chen, A. Davis, J. Dean, M. Devin, S. Ghemawat, G. Irving, M. Isard, et al. Tensorflow: A system for large-scale machine learning. In 12th {USENIX] symposium on operating systems design and implementation ((OSDI) 16), pages 263-283, 2016.
321
323
[2] A. Agrawal, B. Amos, S. Barratt, S. Boyd, S. Diamond, and Z. Kolter. Differentiable convex optimization layers. arXiv preprint arXiv: 12430, 2019.
324
[3] A. Agrawal, S. Barratt, S. Boyd, E. Busseti, and W. M. Moursi. Differentiating through cone program. arXiv preprint arXiv: 1e0e: 09043, 2019.
326
[4] A. Ali, E. Wong, and J. Z. Kolter. A semismooth newton method for fast, generic convex programming. In International Conference on Machine Learning, pages 70-79. PMLR, 2017.
328
329
[5] B. Amos. Disserentiable optimization-baxeee modeling for machine learning. PhD thesis, PhD thesis. Carnegie Mellon University, 2019.
330
331
[6] B. Amos and J. Z. Kolter. Optnet: Differentiable optimization as layer in neural networks. In Proc. of ICML. pages 136-145, 2017.
333
[7] S. Bai, Z. Kolter, and V. Koltun. Deep equilibrium models. arXiv preprint arXiv:190..01377 2019.
334
[8] A. Beatson, J. Ash, G. Roeder, T. 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 in Neural Information Processing Systems, volume 33, pages 338-348. Curran Associates, Inc., 2020.
336
338
339
[9] Y. Bengio. Gradient-based optimization of hyperparameters. Neural computation, 12(8):1889- 1900, 2000.
340
341
[10] Q. Bertrand, Q. Klopfenstein, M. Blondel, S. Vaiter, A. Gramfort, and J. Salmon. Implicit differentiation of lasso-type models for hyperparameter optimization. In Proc. of ICML, pages 310-821, 2020.
343
344
[11] Q. Bertrand, Q. Klopfenstein, M. Massias, M. Blondel, S. Vaiter, A. Gramfort, and J. Salmon. Implicit differentiation for fast hyperparameter selection in non-smooth convex learning. arXiv preprint arXiv:2105.01637, 2021.
345
346
[12] M. J. Best, N. Chakravarti, and V. A. Ubhaya. Minimizing separable convex functions subject to simple chain constraints. SIAM Journal on Optimization, 10(3):658-672 2000.
348
[13] F. Bitzek, P. Koskinen, F. Gähler, M. Moseler, and P. Gumbsch. Structural relaxation made simple. Phys. Rev. Lett., 97:1 Oct 2006.
350
351
[14] M. Blondel. Structured prediction with projection oracles. In Proc. of NeurlPS, 2019.
352
[15] M. Blondel, V. Seguy, and A. Rolet. Smooth and sparse optimal transport. In Proc. of AISTATS, pages 880-88 180-889.11 PMLR, 2018
353
354
[16] M. Blondel, O. Teboul, Q. Berthet, and J. Djolonga. Fast differentiable sorting and ranking. In Proc. of ICML, pages 950 909,959991 2020.
356
[17] Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, and Wanderman- Milne. Jax: composable transformations of python+ numpy programs, 2018. URL http://github. com/google/jax, 4: 4:16, 2020.
357
358
359
[18] P. Brucker. An (n) algorithm for quadratic knapsack problems. Operations Research Letters, 363)1633- 1984.
380
361
[19] O. Chapelle, Vapnik, O. Bousquet, and S. Mukherjee. Choosing multiple parameters for support vector machines. Machine learning, 46(1):131-159, 2002.
10

---

363
[20] X. Chen, Y. Zhang, C. Reisinger. and L. Song. Understanding deep architecture with reasoning layer. Advances in Neural Information Processing Systems, 33, 2020.
384
385
[21] H. Cherkaoui, J. Sulam, and T. Moreau. Learning to solve tv regularised problems with unrolled algorithms, Advances Neural Information Processing Systems, 33, 2020.
386
367
[22] L. Condat. Fast projection onto the simplex and the & ball. Mathematical Programming, 158(12)):55-558, 2016.
388
369
[23] K. Crammer and Y. Singer. On the algorithmic implementation of multiclass kernel-based vector machines. Journal of machine learning research, 2(Dec): 292, 2001.
371
[24] M. Cuturi. Sinkhorn distances: lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems, volume 2, 2013.
373
[25] C.-A. Deledalle, S. Vaiter, Fadili, and G. Peyré. Stein unbiased gradient estimator of the risk (sugar) for multiple parameter selection. SIAM Journal on Imaging Sciences, 7(4):2448-2487, 2014.
375
376
[26] S. Diamond and S. Boyd. Cvxpy: A python-embeddeed modeling language for convex optimiza- tion. The Journal of Machine Learning Research, 17(1):2909-29133 2016.
377
378
[27] J. Djolonga and A. Krause. Differentiable learning of submodular models. Proc. of NeurlPS, 30:1013-100 2017.
379
380
[28] J. Domke. Generic methods for opumization-based modeling, In Artificial Intelligence and Statistics, pages 318-326. PMLR, 2012.
381
382
[29] J. C. Duchi, S. Shalev-Shwartz, Y. Singer, and T. Chandra. Efficient projections onto the &i-ball for learning high dimensions. In Proc. of ICML, 2008.
383
384
[30] L. Franceschi, M. Donini, P. Frasconi, and M. Pontil. Forward and reverse gradient-basee based hyperparameter optimization. In International Conference on Machine Learning, pages 1165- 1173. PMLR, 2017.
386
[31] 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.
386
390
[32] T. Frerix, D. Kochkov, J. A. Smith, D. Cremers, M. Brenner, and S. Hoyer. Variational data assimilation with a learned inverse observation operator. 2021.
382
[33] R. Frostig, M. Johnson, Maclaurin, Paszke, and A. Radul. Decomposing roverse-mode automatic differentiation. In LAFI 2021 workshop at POPL, 2021.
383
394
[34] R. Frostig, M. Johnson, and C. Leary. Compiling machine learning programs via high-level tracing. Machine Learning and Systems (MLSys), 2018.
395
396
[35] S. Gould, B. Fernando, A. Cherian, P. Anderson, R. Cruz, and E. Guo. On differentiating parameterized argmin and argmax problems with application to bi-level optimization. arXiv preprint arXiv: (60v:1660.0000 05447, 2016
387
398
389
[36] S. Gould, R. Hartley, and D. Camphell. Deep declarative networks: new hope. arXiv preprint 1rXiv:190..0455 (999900466, 2019.
400
401
[37] S. Grotzinger and C. Witzgall. Projections onto order simplexes. Applied mathematics and Optimization, 12(1):247-27)) 1984.
403
[38] 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.
405
[39] M. R. Hestenes, E. Stiefel, et al. Methods of conjugate gradients for solving linear systems, volume 49. NBS Washington, DC, 1952.

---

407
[40] N. J. Higham. Accuracy and Stability of Numerical Algorithms. Society for Industrial and Applied Mathematics, second edition, 2002.
408
409
[41] Hoyer, J. Sohl-Dicksteen, and S. Greydanus. Neural reparameterizztioon improves structural optimization. 2019.
[42] Y. Kim, C. Denton, L. Hoang, and A. M. Rush, Structured attention networks. arXiv preprint arXiv:1702.0083 00887, 2017.
412
413
[43] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization.
arXiv preprint
arXiv:1412.6980, 2014.
[44] S. G. Krantz and H. R. Parks. The implicit function theorem: history, theory, and applications. Springer Science & Business Media, 2012.
416
417
[45] 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.
[46] Lorraine, P. Vicol, and D. Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In International Conference on Artificial Intelligence and Statisties, pages 1540-1552 PMLR, 2020.
420
[47] J. Mairal, F. Bach, and J. Ponce. Task-driven dictionary learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 34(4):791-804, 2012.
423
424
[48] A. F. Martins and R. F. Astudillo. From softmax to sparsemax: A sparse model of attention and
425
multi-label classification. In Proc. of ICML, 2016.
426
[49] C. Michelot. A finite algorithm for finding the projection of point onto the canonical simplex of R". Journal of Optimization Theory and Applications, 50(1):195-200) 1986.
427
[50] J.-J Moreau, Proximité et dualité dans un espace hilbertien. Bulletin de la S.M.F,, 93:273-299,
429
1 1965.
430
[1]] Niculae and M. Blondel. regularized framework for sparse and structured neural attention. In Proc. of NeurIPS, 2017.
[52] V. Niculae and A. Martins. Lp-sparsemap: Differentiable relaxed optimization for sparse structured prediction. In International Conference on Machine Learning, pages 7348-7359, 2020.
435
[53] V. Niculae, A. Martins, M. Blondel, and C. Cardie. Sparsemap: Differentiable sparse structured inference. In International Conference on Machine Learning, pages 3799-3808 PMLR, 2018.
436
437
[54] 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, 1642)) 2016.
439
440
[55] N. Parikh and Boyd. Proximal algorithms. Foundations and Trends in optimization, 1(3):127- 239, 2014.
442
[56] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. arXiv preprint arXiv:1912.01703, 2019.
443
444
445
[57] Pedregosa.
Hyperparameter optimization with approximate gradient. In International conference on machine learning. PMLR, 2016.
446
447 [58] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12:2825-288 2830, 2011.
448
449
450
451
[59] A. Rajeswaran, C. Finn, S. Kakade, and Levine. Meta-learning with implicit gradients. arXiv preprint 1901vv:100.04660 2019.
12

---

453
[60] N. Rappoport and R. Shamir. Multi-omis and multi-view clustering algorithms: review and cancer benchmark. Nucleic Acids Res., 46: 16:10056-115555 10562,, 2018.
454
455
[61] Y. Saad and M. H. Schultz. Gmres: A generalized minimal residual algorithm for solving nonsymmetric linear systems. SIAM Journal scienlific and statistical computing, 7(3):856- 869, 1986.
457
458
[62]
S. Schoenholz and E. 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 11488-11441 Curran Associates, Inc., 2020.
480
462
[63] M. W. Seeger. Cross-validation optimization for large scale structured classification kernel methods. Journal of Machine Learning Research, 9(6), 2008.
463
[64] 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.
465
467
488
[65] Vaiter, C.-A. Deledalle, G. Peyré, C. Dossal, and Fadili. Local behavior of sparse analysis regularization: Applications to risk estimation, Applied and Computational Harmonic Analysis, (5(3):3334411 2013.
470
471
[66]
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 nonsymmetric linear systems. SIAM Journal on Scientific and Statistical Computing, 13(2):631-644, 1992.
472
473
[67] T. Wang, J.-Y. Zhu, A. Torralba, and A. A. Efros. Dataset distillation. arXiv: 1811. 10959, 2018.
arXiv preprint
475
476
[68] R. E. Wengert. A simple automatic derivative evaluation program. Communications of the ACM, 7(8):463 464, 1964.
[69]
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 2018, Conference Track Proceedings. OpenRe- view.net, 2018.
479
480
481
482
[70] J. Zarka, L. Thiry, T. Angles, and S. Mallat. Deep network classification by scattering and homotopy dictionary learning. arXiv preprint arXhv:1910.0356/ 2019.
483
13

---

Checklist
484
485
1. For all authors...
486
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes] Our paper builds on the premise that reinstantiating the implicit function theorem for every optimization problem a user may encounter is cumbersome. We make the case that a modular approach is needed to bypass that issue. This approach raises several challenges, notably in the way these implicit solvers can be automatically instantiated, consistently, across the large corpus of optimization approaches favored by users.
488
490
402
498
(b) Did you describe the limitations of your work? [Yes], We discuss several limitations in our work. For instance, implicit differentiation requires è to be sufficiently close to x" to be meaningful. This is the main topic of $2.3
485
486
(c) Did you discuss any potential negative societal impacts of your work? [N/A] As a purely methodological paper, we do not foresee negative societal impacts of our work. (d) Have you read the ethies review guidelines and ensured that your paper conforms to them? [Yes] We confirm our paper conforms to those guidelines.
487
498
489
500
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes], The paper contains one theoretical section, $2.3.
502
503
(b) Did you include complete proofs of all theoretical results? [Yes] All proofs are included in the Appendix D
505
3. If you ran experiments..
506
(a) Did you include the code, data, and instructions needed to reproduce the main exper- imental results (either in the supplemental material or as a URL.)? [No] At the time of submission, we are in the course of an approval process for open-source release required by our organization. We believe that the library itself comprises a contribution, and will have available in open source by the time of this paper's publication (at the latest).
507
508
509
510
511
512
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]. Experiments were mostly run with minimal parameter tuning to reflect the simplicity of the approach we advocate. This is reflected in $3 and Appendix
513
514
515
516
(c) Did you report error bars (e.g., with respect to the random seed after running ex- periments multiple times)? [Yes], see Figure 2 and std for the dictionary learning task.
517
518
519
(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], see Appendix E
520
521
4. If you are using existing assets (c.g., code, data, models) or curating/rellasingg new assets... (a) If your work uses existing assets, did you cite the creators? [Yes], see $3 and Appendix
523
524
(b) Did you mention the license of the assets? [Yes], see Appendix E
525
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you's 're using/curating? [N/A]
528
529
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
531
5. If you used crowdsourcing or conducted research with human subjects...
532
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
533
534
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
535
536
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
537
14


Reviews:
[{"review_id": "ydfA8rwcQ4Y", "reviewer": "Reviewer_u8Y1", "summary": "A good paper considers a critical problem, but contributions are somewhat weak.\n\nThis paper aims at a unified approach for implicit differentiation of optimization problems. The experiments results seem good, but the proposed methods lack novelty.", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "The high-level idea to perform the automatic implicit differentiation of optimization problems is 1) get the root and fixed-point equations which the optimal solution should satisfy; 2) differentiates these equations (by auto-diff) to and get the desired Jacobian by solving a linear equation.\n\n1. The main problem for this framework is non-smoothness. In most cases, the optimal condition is the monotone inclusion problems, which means that the so-called root and fixed-point equations are non-differentiable. The submission lacks the details for this case, and obviously that the projector and the proximal operator are non-differentiable in most cases.\n\n2. Another confusing thing comes from Fig.2. When the scale of the optimization problem grows, solving the large-scale equation is not guaranteed faster than unrolling iterative method, especially for the optimization algorithms with fast convergence. A more discussion is needed here.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No need to consider the Societal Impact of this submission.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4h"}, "extra_text": {}}, {"review_id": "gEyUnIgUNqQ", "reviewer": "Reviewer_oNfk", "summary": "The authors propose a unified modular framework for implicit differentiation, which removes the need to manually derive the jacobians for each individual case by leveraging automatic differentiation of the optimality conditions. The proposed framework is shown to encompass various existing implicit differentiation schemes (such as stationary point- or KKT conditions) while also creating new ones (e.g. mirror descent fixed point). To handle the practical case in which the optimization problem can only be solved approximately, a theoretically justified bound on the error of the jacobian estimate is provided, showing that the error on the estimated jacobian is at most of the order of the deviation of the approximate solution from the optimal solution.\nAdditionally, a JAX implementation of the proposed framework is described, with the promise of an open-source JAX library for implicit differentiation upon acceptance.\nFinally, four applications of the framework are demonstrated, highlighting its flexibility and ease of use.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "This paper provides a good unification of the various existing implicit differentiation schemes. The principle idea of using the implicit function theorem for implicit differentiation of optimal solutions instead of unrolling the iterations of the solver is not new, but the simple implementation using autodiff opens up many possibilities of employing modular implicit differentiation without the need to focus on tedious calculations of jacobians. This is strongly supported by the promise of an open-source JAX library for modular implicit differentiation.\n\nIn the case of differentiating a fixed point, matrix B is defined in the equations following line 108. The derivative with respect to the second argument of F (which is $\\theta$) can be confusing, as the optimal solution (depending on theta) is already plugged in as the first argument of F. The clarity could be improved by taking the derivative with x in the neighborhood of $x^*(\\theta)$, and then plugging in $x^*(\\theta)$ after calculating the derivative.\n\nThe authors show that part of the computation for solving the linear system on the backward pass can be reused to calculate the JVP or VJP. While in general this might be the most efficient one can get, in special cases previous work (e.g. [1]) has been able to achieve high efficiency by reusing calculations from the forward pass of a specifically designed solver, dramatically reducing the computation on the backward pass. It is unclear to me whether this could also be easily exploited using this framework. As [1] is listed as one of the schemes encompassed by the proposed framework, the potential difference in efficiency should be discussed.\n\nThe general structure of the framework allows making general statements on the error of the Jacobians when the optimization problem is only solved approximately. However, it is not entirely clear how well these conditions hold in practice, that is, how large the derived constant is in practice and whether the assumptions of the theorems hold. This could be made clearer by showing the applicability of the theorem to one of the described experiments.\n\nThe presented experiments clearly show the ease of use in various applications of bi-level optimization. The presentation of the paper is very clear. \n\n[1] B. Amos and J. Z. Kolter. Optnet: Differentiable optimization as a layer in neural networks. In Proc. of ICML, pages 136–145, 2017.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Societal Impact: No societal impact\n\nLimitations: Efficiency and applicability of the theorem as mentioned in the main review.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7"}, "extra_text": {}}, {"review_id": "LRT86TMSUp", "reviewer": "Reviewer_8dYU", "summary": "The paper promises extension of Google’s JAX library by automatic and modular differentitation of implicitly defined functions. Given a solver $\\theta\\mapsto x^*(\\theta)$ one asks for its Jacobian $\\partial x^*(\\theta)$. In cases when $x^*$ does not have an explicit form that can be automatically differentiated, the Jacobian can be computed from the optimality conditions $F$ having the general form $F(x^*(\\theta),\\theta)=0$. Jacobian $\\partial x^*(\\theta)$ is then given by a solution of a linear system involving $\\partial_1 F$ and $\\partial_2 F$. The authors describe a unified implementation in JAX and show the resulting simple interface for end-users. Various use cases are presented on examples and code-snippets figures, including stationary point conditions, KKT conditions, proximal gradient, projective gradient and mirror descent fixed point and more. On the theoretical side, the authors discuss the case when the optimization problem solution is only approximate and derive Jacobian error bounds under certain assumptions on optimality conditions $F$.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "The paper is well written, clearly motivated and easy to follow. It contains many illustrative examples, all notions are explained and properly defined.\n\nThe main contribution is not in the paper itself but in the code it promises. Implicit differentiation is a well-known concept that has been used many times in concrete scenarios; the paper quotes at least [6, 9, 10, 11, 19, 35, 46, 51, 52, 53]. The goal and contribution of the paper are unifying all these tedious case by case derivations and offering a tool to build new schemes. The universality and modularity are clear, strong sides of the project.\n\nConcerning the submission is mainly about the code, no code was submitted. My biggest concern is that we are forced to believe “the organization” “will have it available in open source by the time of this paper’s publication”. From this point of view, the paper itself does not contribute and contains just an advertisement what the organization will release in December. The implicit function theorem (IFT) is a part of the second-year course of analysis. *The code* is what JAX framework users will benefit from, but not even a demo was so far provided. We can guess that the library is already functioning only from the experiments’ results presented in the paper.\n\nAs mentioned, from mathematical point of view, IFT is a classics. It is however not stated correctly (page 3, lines 94-96). Consider, for instance $F(x,\\theta)=x^2-1$ and a solution $x^*(\\theta)=1$ for $\\theta\\in(-\\infty,0]$ and $x^*(\\theta)=-1$ for $\\theta\\in(0,\\infty)$. Equation (1) is clearly satisfied, $F$ is smooth, $\\partial_1F(x^*(\\theta),\\theta)\\ne 0$ for every $\\theta$ but $\\partial x^*(0)$ does not exist. The problem is that IFT asserts that under the assumptions, there exists a unique $C^1$ function $x^*$ in some neighbourhood of $x^*(\\theta)\\times\\theta$ and equality (2) holds. In the setting of the paper, the solver $x^*$ is given. What assumptions do you impose on the solver? What if multiple optimal solutions or bifurcation occur ($\\partial_1F$ conditionality blowup)? \n\nJacobian bounds section: Definition 1 is probably not needed since it seems that it is not used anywhere else in the paper. Theorem 1 is correct; the proof is an exercise. However, I do not understand its value in this context. The global bounds on $\\partial_1 F$ and $\\partial_2$ (mostly the well-conditionality and boundedness) seems too restrictive, and the resulting constant in the bound (between lines 221 and 222 — consider labelling it) will be probably very large. Is it possible to demonstrate the use of Theorem 1 (or any of Corollaries 1 and 2) on an example? If not applicable, such a result does not increase the paper’s value.\n\nEdit: After submitting my review the code was publicly released (two days before the deadline for reviews). It is impossible to treat the submission as anonymous when the authors are publicly active and attract attention on social networks, blogs etc. during the reviewing period.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7"}, "extra_text": {}}, {"review_id": "GMwHO53Vt-", "reviewer": "Reviewer_jPkD", "summary": "This paper presents a module for implicit differentiation in JAX, which allows efficient differentiation through optimization and root-finding problems. The paper presents the general principle of combining the implicit function theorem with autodiff of optimality conditions, and discusses how this can work for a implicit functions such as a range of constrained optimization methods. Finally there are example applications to three machine learning tasks and one molecular dynamics simulation task where this framework improves on currently used methods.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "This paper has clear exposition and will be useful for ML people wishing to understand and use implicit differentiation. The JAX library which the authors will open source is a very important contribution and will fuel much other work in machine learning and physical simulation. My main complaint is that I think the paper should do a better job recognizing prior work combining autodiff + implicit differentiation, and making clear how its contributions (which are important) are positioned in the broader literature. I will raise my score if some of my concerns under \"originality\" are addressed (or if it becomes apparent I'm mistaken about the scope of contributions).\n\n--\nOriginality: \nI'm not aware of other ML papers combining implicit differentiation with autodiff of optimality conditions. It is important to introduce these to ML practitioners. The JAX library described provides functionality which wasn't previously available for ML.\n\nMy pedantic gripe: the paper almost exclusively cites other ML papers. Autodiff and implicit differentiation have a long history outside of ML, e.g. in engineering where some software packages implementing the adjoint sensitivity method combine autodiff of constraints on a solution and objective functionals on the solution with the implicit function theorem. They have also been extensively studied in the automatic differentiation literature, e.g. \"Evaluating Derivatives\", Griewank 2000 chapter 15, which discusses general principles for combining the implicit function theorem with autodiff of optimality criteria, and which has bounds on the error of the estimated Jacobian which are somewhat similar to Thm 1 in the submitted paper.\n\nAs such, the paper should link back to this literature. It would be good for readers to realize autodiff+implicit diff is not completely new to this paper, and also for readers to have pointers to follow back to this prior work. \n\nIt would be helpful to understand if and how the general framework (2.1) and the specific example implicit differentiation methods (2.2) differ from previous work in the autodiff and adjoint literature.\n\nThe JAX library is a significant and wholly unique contribution. I'm not aware of other software packages which combine autodiff and implicit differentiation in a general and modular way or even which cater for these for ML problems. \n\nReading the paper, I found myself wishing it was more explicitly written about the JAX library, the design decisions behind it, and how it can be used. The general principles of autodiff+implicit function theorem are out there in the literature, and while this packages / presents and applies them for an ML audience one doesn't learn so much about the principles beyond e.g. what's presented in Griewank, whereas this is the first library of its kind and an insight into the library design would be fascinating.\n\n--\nQuality, Clarity: \nThe paper is well written and clear. The experiments are demonstrations of a software package rather than measurement of a model's test error. They are appropriate, seem well set up, and demonstrate clearly some advantages of the proposed method / library.\n\nNit:  The paragraph just before (8) on page 6 seems incorrect. If g is the commonly defined indicator function for the set, $g = 1$ in the set and $g = 0$ outside, then it will encourage (but not force) x* to lie outside the set. To recover (8) we need to force (not encourage) x* to lie inside the set. I think the definition (8) is recovered from the proximal $x^* = \\text{argmin}_x f(x) + g(x)$ when $g = 0$ in the set and $g = \\infty$ (or $g = \\max_x f(x) - \\min_x f(x)$) outside the set.\n\n--\nSignificance: \nThis is an important paper and important library. Implicit differentiation has become increasingly important as a method to scale many things such as bi-level optimization and meta-learning. The paper will help people understand the principles behind it, and the JAX package will help fuel many other research projects which use implicit differentiation.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes. This is a methods paper and I don't think there are any specific negative societal impacts. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper got 1 "Marginally below the acceptance threshold", 2 "Marginally above the acceptance threshold"s and 1 "Good paper, accept". The average score, 6, is slightly above the threshold. The major challenges include that the focus of the paper should be the integration with JAX (three reviewers all pointed out this), whether the implicit differentiation for several operators presented in the paper is of sufficient contribution, and leaving out the whole literature on implicit differentiation and autodiff outside machine learning. The author rebuttals were somewhat successful, and the authors promised to revise the paper accordingly. The AC agreed with the reviewers that how to integrate with JAX should be much more important and deducing the implicit differentiation for several operators used in machine learning is just an exercise level work. Since the paper needs $\mathbf{substantial}$ revision in order to change the focus of the paper, adding related work on implicit differentiation and autodiff outside machine learning, and clarify some details or ambiguities, and the decision should be made based on the current exposition rather than the promised unseen version, the AC had to recommended rejection. The SAC confirmed this decision.
Decision: Reject