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

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 and implementation. 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 describe in details a JAX implementation of our framework and demonstrate the ease of differentiating through optimization problems thanks to it on four diverse tasks: hyperparameter optimization of multiclass SVMs, dataset distillation, task-driven dictionary learning and sensitivity analysis of molecular dynamics.

Under review as a conference paper at ICLR 2022
EFFICIENT AND MODULAR IMPLICIT DIFFERENTIATION
Anonymous authors
Paper under double-blind review
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 layers, and in bi-level problems such as hyper-parameter optimization and meta- meta-learning. However, so far, implicit differen- tiation remained difficult to use for practitioners, as it often required case- case-by-cue by-case tedious mathematical derivations and implementations. In this paper, we propose an efficient and modular approach for implicit differentiation of optimization prob- lems. 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 implicit differentiation to automatically differentiate the optimization problem. Our approach thus combines the benefits of implicit differ- entiation and autodiff. It efficient as it can be added on top of any state-of- 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 diffcrentiation methods and create new ones easily. We demonstrate the ease of formulating and solving bi-level optimization problems using framework. We also showease application to the sensitivity analysis of molecular dynamics.
INTRODUCTION
Automatic differentiation (autodiff) is now an inherent part of machine learning software. It allows to 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 papee--prrmeter optimization (Chapelle et al., 2002; Seeger, 2008; Pedregosa, 2016; Franceschi et al., 2017: Bertrand et al., 2020; 2021), neural networks (Lorraine et al., 2020), and meta-learning (Franceschi et al., 2018; Rajeswaran et al., 2019). Another line of active research involving differentiation of optimization problem solutions are optimization layers (Kim et al., 2017; Amos & Kolter, 2017; Niculae & Blondel, 2017; Djolonga & Krause, 2017; Gould et al., 2019), which can be used to encourage structured outputs, and implicit deep networks (Bai et al., 2019; El Ghaoui et al., 2019), which have a smaller memory footprint than tackprop-trainee 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 proxy for the optimization problem solution (Wengert, 1964; Domke, 2012; Deledalle et al.. 2014; Franceschi et al., 2018; Ablin al., 2020). This allows to explicitly construct computational graph relating the algorithm output to the inputs, on which autodiff can then be used transparently. However, this requires a reimplementation the algorithm using the autodiff system, and not all algorithms are necessarily autodiff friendly. Moreover, forward-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. In contrast, second approach consists in implicitly relating an optimization problem solution to its inputs using optimality conditions. In a machine learning context, such implicit differentiation has been used for stationarity conditions (Bengio, 2000;

---

Under review as a conference paper at ICLR 2022
Lorraine et al., 2020), KKT conditions (Chapelle et al., 2002: Gould al., 2016; Amos & Kolter, 2017: Niculae et al., 2018; Niculae & Martins. 2020) and the proximal gradient fixed point (Niculae & Blondel, 2017; Bertrand et al., 2020; 2021). An advantage of implicit differentiation is that solver reimplementation is not needed, allowing to build upon decades of state-of-the-ard software. Although implicit differentiation has a long history in numerical analysis (Griewank & Walther, 2008; Bell & Burke, 2008; Krantz& & Parks, 2012; Bonnans & Shapiro, 2013), so far, it remained difficult to use for practitioners, as required a case-by-case tedious mathematical derivation and implementation. CasADi (Andersson et al., 2019) allows to differentiate various optimization and root finding problem algorithms provided by the library. However, it does not allow to easily add implicit differentiation
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 (Duvenaud et al., 2020). 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 as we do. Other work (Agrawal et al., 2019a) attempts to address this issue by adding implicit differentiation on top of evxpy (Diamond & Boyd, 2016). This works by reducing all convex optimization problems to a 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-an ADMM-based splitting conic solver (O' Donoghue al., 2016) the case of cvxpy-is rarely the state-of-the offthee-ar art approach for each particular problem instance.
In this work, we adopt different strategy that makes it easy to add implicit differentiation on top of any existing solver. In our approach, the user defines directly in Python a mapping function F capturing the optimality conditions of the problem solved by the algorithm. Once this is done, we leverage autodiff of F combined with implicit differentiation to automatically differentiate the optimization problem solution. In this way, our approach is generic, yet it can exploit 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.
We describe our framework and its JAX implementation (provided in the supplementary material). Our framework significantly lowers the barrier to use implicit diffcrentiation thanks to the use of autodiff of the optimality conditions and the seamless integration in JAX. Our framework significantly extends JAX for numerical optimization, with low-level details all abstracted away.
We instantiate our framework on a large catalog of optimality conditions (Table 1), recovering existing schemes and obtaining new ones, such 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 ease of use.
Beyond our software implementation, we hope this paper provides a self-contained blueprint for creating an efficient and modular implementation of implicit differentiation.
Notation. We denote the gradient and Hessian of f: Rd + RR evaluated at æ E Rdd by (s(x) E Rd and f(x( E Radd We denote the Jacobian of F: Rd IR evaluated at x Rd by 0F(x)) Rexd When I or F have several arguments, we denote the gradient, Hessian and Jacobian in the 7th argument by V,, V and òg respectively. The standard probability simplex is denoted by
Ad : {I E Rdd || 1, 0}. For any set C C Rd. denote the indicator function le: Rdd + R [+ where lx(x) 0 i E C, Lc(x) = too otherwise. For a vector or matrix A, we note ||A|| the Frobenius (or Euclidean) norm, and || the operator norm.
2
COMBINING IMPLICIT DIFFERENTIATION AND AUTODIFF
2.1 GENERAL PRINCIPLES
Overview. Contrary to autodiff through unrolled algorithm iterations, implicit differentiation typically involves a manual, sometimes complicated, mathematical derivation. For instance, numerous works (Chapelle et al., 2002; Gould et al., 2016; Amos & Kolter, 2017; Niculae et al., 2018; Niculae Karush Karush-Kuhn-Th
& Martins, 2020) optimization problem's 's solution to its inputs, and to manually derive a formula for its derivatives. The derivation and implementation in these works are always case-by-case.
(KKT) conditions in order to relate constrained
2

---

Under review as a conference paper at ICLR 2022
Xrrrii, y_train loed_data() # Load features and tabela
def i(x, theta):
reeidual jap Objective function
x) y_trein
return (jnp. cunesiiuull 2) theta jup.eun 2))
Since sa differentsable and usconstrasned, the optimalsty condstion samply the gradient of f in the 1st argument jax grad(f, argnans-0)
Bcueton_root
thata):
init_x Instsalization not thés solver
XX = (xp_.dttX_ttainn. X_train)
Xy jap dot T, y_train)
(inp-eye(X_ (reetrrrainn.. tta.rrae...eeee]]...
Identsty natrie
Finds the ridge reg solatson by solving lssear system
return sop.linaig.amm solvee. theta I, xy)
init_x
None
print (jax. jacobian argnuns-1) (imit_x, 10.0))
Figure I: Adding implicit differentiation 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. Our decorator @custom_rt _root automatically adds implicit differentiation to the solver for the user, overriding JAX's default behavior. The last line evaluates the Jacobian at 0 = 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 the problem solved by the algorithm. We provide reusable building blocks to easily express such F. The provided F is then plugged into our Python decorator we append on top of the solver declaration we wish to differentiate. Under the hood, we combine implicit differentiation and autodiff of F to automatically differentiate the optimization problem solution. A simple illustrative example is given in Figure
root which
Differentiating root. Let F; R Rdd be user- user-provvded mapping, capturing the optimality conditions of a problem. An optimal solution, denoted x² (0), should be a root of
F(x* 0.
(I)
We can see x² (0) an implicitly defined function of 0 IR" i.c., x²: RR + Rd. More precisely, from the implicit function theorem (Griewank & Walther, 2008; Krantz & Parks, 2012), we know that for (xg, fu) satisfying F( (0) = 0 with continuously differentiable F, if the Jacobian 0 F evaluated at (Iro, 0o) is a square invertible matrix, then there exists function x²(-) (-) defined on neighborhood of Og such that (0p) = Tp-- Furthermore, for all 0 in this neighborhood, we have that F(x* (0), 0) = 0 and Os+ (0) exists. Using the chain rule, the Jacobian Ox (0) satisfies
F,F(x*(0) 0)rr" (0) + F,F (0), 0) = 0.
Computing (x)()) therefore boils down to the resolution of the linear system of equations
F (0),0) Ox² (0) F2F
AgRaxa
JERaXA
BERAXA
When (1) is a one-dimensional root finding problem (d = 1), (2) becomes particularly simple since we then have (0) = /A, where a scalar value.
We will show that existing and new implicit differentiation methods all reduce this simple principle. call our approach hybrid, since it combines implicit differentiation (and as such requires solving a linear system) with autodiff of the optimality conditions F. Our approach is efficient as it can be added on top of any statc-of-the-arr solver and modular as the optimality condition specification is decoupled from the implicit differentiation mechanism. This contrasts with existing works, where the mathematical derivation and implementation are specific to each optimality condition.
Differentiating fixed point. We will encounter numerous applications where z=(() is implicitly defined through a fixed point:
x x0) T(x))0),,)
3

---

Under review as a conference paper at ICLR 2022
where T: T. R" + R This can be seen particular case of (1) by defining the residual
F(x, 0) (2x,0) 0)
(3)
In this case, using the chain rule, we have
A= - (0), 0) = -- 0,7 (0),0) and B - 02f((= -27))=)))=0(0000000 =-00 O,T(2"
Computing JVPs and VJPs. In most practical scenarios, it is not necessary to explicitly form the Jacobian matrix, and instead it sufficient to meft- or right-multiply by F,F and 0.F. These are called vector-Jacobian product (VJP) and Jacobian-vector product (JVP), and are useful for integrating x² (0) with reverse-mode and borwar--mode autodiff, respectively. Oftentimes, will be explicitly defined. In this case, computing the VJP JVP can be done via autodiff.. In some cases, F may itself be implicitly defined, for instance when involves the solution a variational problem. In this case, computing the VJP JVP will itself involve implicit differentiation.
The hight-multiplicn by solving A(Jv) Bv. The left- fett-luliilliaation (VJP) of with J, v. J, can be computed by first + solving A 1u= = v. Then, we can obtain v by 0' J = u AJ = u B. Note that when B changes but A and remain the same, do need solve A u the VJP w.r.t. different variables while solving only one linear system.
(JVP) between J = Ox (0) and a vector v, Ju, can be computed efficiently
once again. This allows compute
To solve these linear systems, A is symmetric positive semi-definite and GMRES (Saad & Schultz, 1986) or BiCGSTAB (Vorst & van der Vorst, 1992) otherwise. These algorithms are all matrix-free: they only require matrix-vector products. Thus, all we need from F is its JVPs or VJPs. An alternative to GMRES/BiCGSTAB is to solve the normal equation AA u Av using conjugate gradient. We implement this using JAX's automatic transpose routine jax. 1 mea transpose (Frostig et al., 2021).
the conjugate gradient method (Hestenes et al., 1952) when
Pre-processing and post-processing mappings. Oftentimes, the goal is not to differentiate per se, but the parameters of function producing 6. One example of such pre- pre-processing! is to convert the parameters to be differentiated from one form to another canonical form, such as a quadratic program (Amos& & Kolter, 2017) is used as the output of a neural network layer, in which case is produced by the previous layer. Likewise, (0) will often not be the final output want to differentiate. One example of such post-processing is when x² (0) is the solution of dual program and we apply the dual- dual--rimal mapping to recover the solution of the primal program. Another example is the application of a loss function, in order to reduce x" (0) to a scalar value. We leave the differentiation of such pre/post-processing mappings to the autodiff system, allowing compose functions in complex ways.
conic program (Agrawal et al., 201 2019a). Another example is when x (0)
Implementation details. When a solver function is decorated with @customm_ree root, we use jax. cust on_jo and jax. custom_ye jip 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-vectes products and therefore we only need access F through the JVP or VJP with 0, F and O.F. This is done by using jax. jvp and jax. vax.vip... vip, respectively. Note that, as in Figure 1, the definition of F will often include gradient mapping V Vix 0). Thankfully, JAX supports second-order derivatives transparently. For convenience, our library also provides ecusoon i xed pedepoiiie_eeeeeeeeeeeeeee point in t decorator, for adding implicit differentiation on top of solver, given a fixed point iteration T; see code examples in Appendix A.
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 F T implies different trade-offs in terms of computational oracles; see Table 1. Source code examples are given in Appendix A.
Stationary point condition.
The simplest example is to differentiate through the implicit function
I' (0) argmin f(z, 0),
2ERd
where f: Rdd x R
R is twice differentiable. In this case, F is simply the gradient mapping F(x, Vif(x,
4

---

Under review as a conference paper at ICLR 2022
Table Summary of optimality condition mappings. Oracles are accessed through their JVP VJP.
Name
Equation
Solution needed
Oracles necded
Stationary
(4), (5)
Primal
ViS
KKT
(6)
Primal and dual
Vif. H. G. &.H.
Proximal gradient
(7)
Primal
Vif. v.f.proxps
Projected gradient
(9)
Primal
V.S. proje
Mirror descent
(11)
Primal
V V.J Vi..proiz. Ve
Newton
(14)
Primal
[ar 0)] VIf(n, [i1f,s pprox
Block proximal gradient
(15)
Primal
Conic programming
(18)
Residual map root
proj xxx xR+
We then have F1F(2,0) = (zf(x, 0) and argument and the Jacobian in the second argument of f 0). In practice, compute Jacobian products automatically. Equivalently, we can use the gradient descent fixed point
0). the Hessian of f in its first autodiff to
T(x, 8) = x Vf 8),
(5)
for all r) > 0. Using (3), it is easy to check that we obtain the same linear system since 1) cancels out.
KKT conditions. We now show that the KKT conditions, manually differentiated in several works (Chapelle et al., 2002: Gould et al., 2016; Amos Kolter, 2017; Niculae et al., 2018; Niculae & Martins, 2020), fit our framework. As will see, the key will be to group the optimal primal and dual variables as our (0). Let us consider the general problem
argmin f(z, 0) subject to G(z, 0) S O, H(z, 0) 0, zER>
where = E R is the primal variable, f: R" x R" R, G: R x R + R" and H: R" - Rd. The stationarity, primal feasibility and complementary slackness conditions give
V ViJ(z, [0, (8 H (2, 0) =
H 0) = 0
A ((2,0) O,
(6)
where v E R29 and E R are the dual variables, also known as KKT multipliers. The primal and dual feasibility conditions can be ignored almost everywhere. The system of (potentially nonlinear) equations (6) fits our framework, as we can group the primal and dual solutions as (0) = (0), n" (0), (0), (0)) to form the root of a function F(x* (0).0) 0), where F: I2 x R - Rd and d = p + + r. The primal and dual solutions can obtained from a generic solver, such as an interior point method. In practice, the above mapping F will be defined directly in Python (see Figure 7 in Appendix A) and F will be differentiated automatically via autodiff.
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 (Niculae & Blondel, 2017; Bertrand al., 2020; 2021). Let x" (0) be implicitly defined as
x (0) = argmin f(x, 0) + g(x, (),
xER
where f: Idd x R + R is differiiiffe convex and g: Rd R" + R is convex but possibly non-smooth. Let us define the proximity operator associated with g by
prox, (y, 0) = argmin
- ||æ - v112 + + 0).
.egd
To implicitly differentiate x² (0), we use the fixed point mapping (Parikh & Boyd, 2014, p. p500
T(x, prox (ar v/ 0),
for any step size 7) 0. The proximity operator is 1-Lipschitz continuous (Moreau, 1965). By Rademacher's 's theorem, it is differentiable almost everywhere. Many proximity operators enjoy closed form and can easily be differentiated, discussed in Appendix B.
5

---

Under review as a conference paper at ICLR 2022
Projected gradient fixed point. As special case, when g(2; 0) is the indicator function Ic(o) (z), where ((0) is convex set depending on 0, we obtain
x'
argmin ((z 0). =EC(0)
(8)
The proximity operator prox, becomes the Euclidean projection onto C(0)
prox, (v. 8) proje (v. 8) argmin || =ll
xEC(Ø)
and (7) becomes the projected gradient fixed point
Txx, 0) proje(2 V (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 in Appendix B.
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({p) x relint(dom(()) (y=) + R generated by yp be defined by
Dz v) ppx)) ((y) (Vp(y), v).
We define the Bregman projection of onto (()) C dom(pp) by
proje (y, 0) argmin D (z, V =EC(#)
(10)
Definition (10) includes the mirror map Vy² for convenience. It can be seen as mapping from Rd to dom( ensuring that (10) is well-defined. The mirror descent fixed point mapping is then
à V
y = â - n (V,
T(x, 0) proje (y, 0).
(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 tp is <px::) (Ir, log x 1), where dom(pp Rd In this case, De is the Kullback-eeibe divergence. An advantage of the Kullback-Leibler projection is that it sometimes easier compute than the Euclidean projection, as we detail in Appendix B.
Other fixed points. More fixed points are described in Appendix C.
JACOBIAN PRECISION GUARANTEES
In practice, either by the limitations of finite precision arithmetic because perform 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 motivates the need for precision guarantees of this approach. We introduce the following formalism.
Definition 1. Let : Rdd x R + Rd be an optimality criterion mapping. Let A = - and B = Fz We define the Jacobian estimate at (x, 0) as the solution to the linear equation A(x, O)J 0) B(x,0) 0). It is a function TRd x TR. - Rdxn.
holds by construction that (0), 0) Oxx (0). Computing J(t 0) for an approximate solution a of x (0) therefore allows approximate the true Jacobian dx*(0)) In practice, an algorithm used to solve (1) depends on 0. Note however that, what we compute is not the Jacobian ((),, unlike works differentiating through unrolled algorithm iterations, but estimate of Oz (0). We therefore use the notation ă, leaving the dependence on # implicit.
We develop bounds of the form [| 9) (f)) C|| x² (I) hence showing that the error on the estimated Jacobian at most of the same order as that of as an approximation of (0). These bounds are based the following main theorem, whose proof is included in Appendix D.

---

Under review as a conference paper at ICLR 2022
Theorem (Jacobian estimate). Let F Rd x R" - Rd. Assume that there exist ",E, R 0
such that A = F and B = 02F satisfy. for all u R E IR" and such that || - x [0) - E:
A is well-conditioned, Lipschitz: ||A(z, 0) a | ||A(z, 0) A(x 0) S <1|l= =² (0)
B is bounded and Lipschitz: [[ (0), 0) - R, [B 0) (zx=(())))) = 8 (0)ll.
Under these conditions, when || x=' (0)ll - we have
[2 0) - (0)]l < (Ba -1 + +yRa 3) || x²
This result is inspired by (Higham, 2002, Theorem 7.2). that is concerned with the stability of solutions to inverse problems. Here we consider that A(,, 0) is uniformly cell-conditioned, rather than only xx (0). This does not affect the first order in of this bound, and makes it valid for all é. is also more tailored to applications to equation-speciie
cases.
be applied to specific cases. In particular, for the gradient descent fixed point, where Tzz 0) = r v/ fir (x, and Fxz, 0) = T(x, 0) I, this yields
Indeed, Theorem
A(x, = ( 0) and B(x, 8) V1
The guarantees on Jacobian precision under regularity conditions rely on I directly; see Corollary 1 in Appendix D. This reveals in particular that Jacobian estimation by implicit differentiation gains factor of t compared to automatic differentiation, after iterations of gradient descent in the strongly-convex setting (Ablin et al., 2020, Proposition 3.2). While our guarantees concern the Jacobian of x" (f), note that other studies (Grazzi et al., 2020; Ji al., 2021; Bertrand et al., 2021) give guarantees on hypergradients (i.e., the gradient of an outer objective).
We illustrate these results ridge regression, where x (0) = argming 0 - y + Σ. 0,a,. This problem has the merit that the solution (0) and its Jacobian Ox*(0) are available in closed form. By running gradient descent for t iterations, we obtain an estimate of x (0) and an estimate [ f) of Ox* (f); ef. Definition 1. By doing so for different numbers of iterations t, we can graph the relation between the error | 2 and the error |0x, (0) ( 0)112, as shown Figure 2, empirically validating Theorem 1. The results in Figure 2 were obtained using the diabetes dataset from Efron et al. (2004), with other datasets yielding qualitatively similar behavior. We derive similar guarantees Corollary in Appendix D for proximal gradient descent.
Implici off Unvolling
(B1ll
Theoretical bound
||x
A112
Figure 2: 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.
EXPERIMENTS
To conclude this work, 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 (Crammer & Singer, 2001) trained in the dual. Here, x² x0) is the optimal dual solution, matrix of shape m x k. where m is the number of training examples and is the number of classes, and 0 E R,, the regularization parameter. The challenge in differentiating xx (0) is that each row of x (0) is constrained to belong to the probability simplex A/ More formally, let Xtr E Rmxp be the training feature matrix and Yur E {0, 1] mm be the training labels (in row-wise one-hot encoding). Let W W(x, 0) = [Ytr - / E Rx be the dual-primal mapping. Then, we consider the following bi-level optimization problem
min 1 |
1xxx ((= ())) Y 11
subject to
x x²(0) argmin f(z, 0) I 0)1 + (z. Yu), rEC inner problem
xxexp(X)
2
outer problem

---

Under review as a conference paper at ICLR 2022
Mirrr descent (MD)
Proximal gradient (PG)
Block coordinate descent (BCD)
lnnelling
1.50
Imp
320
400
Runtima stop
100
130
200
200
100
2000
4000
6010
8000
10000
2000
6000
6000
8000
10000
2003
4.00
6000
6000
15000
features
Number of features
Number features
(b)
(c)
Figure CPU runtime comparison of implicit differentiation and unrolling for hyperparameter optimization of multiclass SVMs for multiple problem sizes. Error hars 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.
where C= C=
x & is the Cartesian product of probability simplices. We apply the change of variable exp(A) in order to guarantee that the (0), 0) E xa contains the optimal primal solution, the feature weights for each class. The outer loss is computed against validation data X and Yya1-
0 is positive. The matrix
While KKT conditions can be used to differentiate z" 20), more direct way is use the projected gradient fixed point (9). The projection onto C can be easily computed by row-wise projections on the simplex. The projection's Jacobian enjoys closed form (Appendix B). Another way to differentiate x (0) is using the mirror descent fixed point (11), Under the KI. geometry, proje (y, 0) corresponds to row-wise softmax. is therefore easy to compute and differentiate. Figure 3 compares the runtime performance of implicit differentiation unrolling for the latter two fixed points.
4.2 DATASET DISTILLATION
Dataset distillation (Wang et al., 2018; Lorraine et al., 2020) aims to learn a small synthetic training dataset such that model trained on this learned data set achieves small loss on the original training set. Formally, let Xtx E gxxx and E [K]* denote the original training set. The distilled dataset will contain one prototype example for each class and therefore 0 kxp, The dataset distillation problem can then naturally be cast as a bi-level problem, where in the inner problem we estimate logistic regression model x= =0)) B2 trained on the distilled images k while in the outer problem we want to minimize the loss achieved by x= (0) over the training set:
fxx (0), X Ver) subject to x² (0) E argmin f(x, 0 [b]) + =||=|12, BEREXy
min
(12)
ZERPxè
outer problem
inner problem
where f(W, X: := Ely, XW), denotes the multiclass logistic regression loss, and 10-1 is regularization parameter that we found had very positive effect on convergence.
In this problem, and unlike the general hyperparameter optimiza- tion setup, both the inner and outer problems are high-dimensional, making an ideal test-bed for gradient- bassdd 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 7840.
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 iterates instead. For the same solver, we found that the implicit differentiation approach 4
Figure 4: Distilled dataset
k obtained by solving (12).
8

---

Under review as a conference paper at ICLR 2022
Table 2: Mean AUC (and 95% confidence interval) for the cancer survival prediction problem.
Method
Li logreg
L2 logreg
DictL La logreg
Task-driven DictI.
AUC (%)
71.6 = 2.0
72.4 1: 2.8
68.3 1 2.3
73.2 1 2.1
times faster than the original one. The obtained distilled images visualized in Figure 4.
4.3 TASK-DRIVEN DICTIONARY LEARNING
Task-driven dictionary learning was proposed to learn sparse codes
for input data in such way that the codes solve outer learning problem (Mairal et al., 2012;
Sprechmann et al., 2014; Zarka et al., 2019). Formally, given data matrix Xtr E TR=xp and a dictionary of atoms kxp sparse code is defined matrix x (0) Rmxk that minimizes in r reconstruction loss f(x, 0) = 8(Xtr, I0) regularized by a sparsity-inducing penalty g(r). Instead of optimizing the dictionary 0 to minimize the reconstruction loss, Mairal et al. (2012) proposed to optimize an outer problem that depends on the code. Given a set of labels Y E (0, 1)m,
we consider logistic regression problem which results in the bilevel optimization problem:
min
(x (f)w + b; Utr) subject to x (0) argmin fxz 0) + g(x)
(13)
DERKXP,EERR bER
vERmxx
outer problem
inner problem
When é the squared Frobenius distance between matrices, and the elastic net penalty, Mairal al. (2012, Eq. 21) derive manually, using optimality conditions (notably the support of the codes selected at the optimum), an explicit re-parameterization of x= (0) as a linear system involving 0. This closed-form allows for direct computation of the Jacobian x" w.r.t. 8. Similarly, (Sprechmann et al., 2014) derive first order conditions in the case where l is 3-divergence, while (Zarka et al., 2019) 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 5 years (m 200) vs patients who die within 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 state-of the art L1 or L2 regularized logistic regression with 100 times fewer variables.
4.4 SENSITIVITY ANALYSIS OF MOLECULAR DYNAMICS
Many physical simulations require solving optimization problems, such as energy minimization in molecular (Schoenholz & Cubuk, 2020) and continuum (Beatson et al., 2020) mechanics, structural optimization (Hoyer et al., 2019) and data assimilation (Frerix al., 2021). In this experiment, we revisit an example from JAX- MD (Schoenholz & Cubuk, 2020), the problem of finding energy minimizing configurations 2- 2-dimensional box of size É
system of packed particles in a
(0) argmin f(x, 0) = EU(Zo) =cRtx2
2.J
where (0) E kx2 are the optimal coordinates of the k particles, U(ra,, f)) 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-speciff optimizer (Bitzek et al., 2006). Here we consider sensitivity of particle position with respect to diameter (0), rather than sensitivity of the total energy from the original experiment. Figure shows results calculated via forward-modh implicit differentiation (JVP). Whereas differentiating the unrolled optimizer happens to work for total energy, here typically does not even converge (see Appendix Figure 17), due the discontinuous optimization method.
Figure 5: Particle positions and position sensitivity vectors, with respect to increasing the diameter of the blue particles.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Pierre Ablin, Gabriel Peyré, and Thomas Moreau. Super-efficieney of automatic differentiation for functions defined as a minimum. In Proc. of ICML, pp. 32-41, 2020.
Akshay Agrawal, Brandon Amos, Shane Barratt, Stephen Boyd. Steven Diamond, and Zico Kolter. Differentiable convex optimization layers. arXiv preprint arXiv:1910.124300 2019a.
Akshay Agrawal, Shane Barratt, Stephen Boyd, Enzo Busseti, and Walaa M Moursi. Differentiating through a cone program. arXiv preprint arXiv:1904.09043 2019b.
Alnur Ali, Eric Wong, and J Zico Kolter. A semismooth newton method for fast, generic convex programming. International Conference on Machine Learning, pp. 70-79. PMLR, 2017.
Brandon Amos. Differentiable optimization-basee based modeling for machine learning. PhD thesis, PhD
thesis. Careegie Mellon University, 2019.
Brandon Amos and Zico Kolter. Optnet: Differentiable optimization as layer in neural networks. In Proc. of ICML, pp. 136- 145, 2017.
Joel AE Andersson, Joris Gillis, Greg Horn, James B Rawlings, and Moritz Dichl. Casadi: software framework for nonlinear optimization and optimal control. Mathematical Programming Computaaiion 11(1):1-36, 2019.
Shaojie Bai, J Zico Kolter, and Vladlen Koltun. Deep equilibrium models. arXiv preprint 2019.
Alex Beatson, Jordan Ash, Geoffrey Roeder, Tianju Xue, and Ryan P Adams. Learning composable energy surrogates for pde order reduction. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 338- 348. Curran Associates, Inc., 2020. URL https::/proceed:: http///poceeeng....oomo//o/ 2020
Paper.pde
Bradley M Bell and James V Burke. Algorithmic differentiation of implicit functions and optimal values. In Advances Automatic Differentiation, pp. 67-77. Springer, 2008
Yoshua Bengio. Gradient-basee optimization hyperparameters. Neural computation, 12(8): 1889-1900, 2000.
Quentin Bertrand, Quentin Klopfenstein, Mathieu Blondel, Samuel Vaiter, Alexandre Gramfort, and Joseph Salmon. Implicit differentiation of lasso-type models for hyperparameter optimization. In Proc. of ICML, pp. 810-821, 2020.
Quentin Bertrand, Quentin Klopfenstein, Mathurin Massias, Mathieu Blondel, Samuel Vaiter, Alexan- dre Gramfort, and Joseph Salmon. Implicit differentiation for fast hyperparameter selection in non-smooth convex learning. arXiv preprint 1rXiv:2105.01637 2021.
Michael J Best, Nilotpal Chakravarti, and Vasant A Ubhaya. Minimizing separable convex functions subject to simple chain constraints. SIAM Journal on Optimization, 10(3):658 662,66772000 2000.
Erik Bitzek, Pekka Koskinen, Franz Gähler, Michael Moseler, and Peter Gumbsch. Structural relaxation made simple. Phys. Rev. Letr., 97:170201, Oct 2006. doi: 10.1103//hysRevLett..97 170201. URL https:tt......... / www.liinnppp........mmmm.mm org /ooi / / 103 wwt.sseysssettttt.eeeee 97 170201.
Mathieu Blondel. Structured prediction with projection oracles. In Proc. of NeurlPS, 2019.
Mathieu Blondel, Vivien Seguy, and Antoine Rolet. Smooth and sparse optimal transport. In Proc. of AISTATS, pp. 880-889. PMLR, 2018.
Mathieu Blondel, Olivier Teboul, Quentin Berthet, and Josip Djolonga. Fast differentiable sorting and ranking. In Proc. of ICML., pp. 950-959, 2020.
J Frédéric Bonnans and Alexander Shapiro. Perturbation analysis of optimization problems. Springer Science & Business Media, 2013.
10

---

Under review as a conference paper at ICLR 2022
James Bradbury. Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http:/ http//gitttt tuu... com hom/gooolee Le e/ jax.
Peter Brucker: An ((n) algorithm for quadratic knapsack problems. Operations Research Letters, (3):163-16,, 1984.
Olivier Chapelle, Vladimir Vapnik, Olivier Bousquet, and Sayan Mukherjee. Choosing multiple parameters for support vector machines. Machine learning, 46(1): 2002.
Hamza Cherkaoui, Jeremias Sulam, and Thomas Moreau. Learning to solve tv regularised problems with unrolled algorithms. Advances in Neural Information Processing Systems, 33, 2020.
Laurent Condat. Fast projection onto the simplex and the & ball. Mathematicaa Programming. 158 (1-2))575-588, 2016.
Koby Crammer and Yoram Singer. On the algorithmic implementation of multiclass kernel-base vector machines. Journal of machine learning research, 2(Dec): 2001.
Marco Cuturi. Sinkhorn distances: lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems, volume 2, 2013.
Charles-Alban Deledalle, Samuel Vaiter, Jalal Fadili, and Gabriel Peyré. Stein unbiased gradient estimator of the risk (sugar) for multiple parameter selection. SIAM Journal on Imaging Sciences, 7(4):2448- 2487, 2014.
Steven Diamond and Stephen Boyd. Cvxpy: A python-embeddee modeling language for convex optimization. The Journal of Machine Learning Research, 17(1):2909-299 -9133, 2016.
Josip Djolonga and Andreas Krause, Differentiable learning of submodular models. Proc. of NeurIPS, 30:1013-1023, 2017.
Justin Domke. Generic methods for optimization-based modeling. In Artificial Intelligence and Statistics, pp. 318-326. PMLR, 2012.
John Duchi, Shai Shalev-Shwartz, Yoram Singer, and Tushar Chandra. Efficient projections onto the billl for learning in high dimensions. In Proc. of ICMI., 2008
David Duvenaud, J. Zico Kolter, and Matthew Johnson. Deep implicit layers tutorial neural ODEs, deep equilibirum models, and beyond. Neural Information Processing Systems Tutorial, 2020.
Bradley Efron, Trevor Hastie, lain Johnstone, and Robert Tibshirani. Least angle regression. The Annals of statistics, 32(2):407-499, 2004.
Laurent El Ghaoui, Fangda Gu. Bertrand Travacca, Armin Askari, and Alicia Y Tsai. Implicit deep learning. arXiv preprint arXiv: 06315, 2, 2019.
Rong-Ei En Fan, Kai-Wa Rei Chang, Cho-Jui Hsieh, Xiang-Rui Rui Wang, and Chih-Jen Lin. Liblinear: A library for large linear classification. JMLR, 9:1871-1874 2008.
Luca Franceschi, Michele Donini, Paolo Frasconi, and Massimiliano Pontil. Forward and reverse based hyperparameter optimization. In International Conference on Machine Learning. pp. 1165-1111 173. PMLR, 2017.
Luca Franceschi, Paolo Frasconi, Saverio Salzo, Riccardo Grazzi, and Massimiliano Pontil. Bilevel programming for hyperparameter optimization and meta-leamingg In International Conference Machine Learning, pp. 1568- PMLR, 2018
Thomas Frerix, Dmitri Kochkov, Jamie A. Smith, Daniel Cremers, Michael P. Brenner, and Stephan Hoyer. Variational data assimilation with learned inverse observation operator. 2021.
Roy Frostig, Matthew Johnson, Dougal Maclaurin, Adam Paszke. and Alexey Radul. Decomposing reverse-mon moverse-mode automatic differentiation. In LAFI 2021 workshop at POPL, 2021. URL https: / rxiv. www...oog/abba....
09469.

---

Under review as a conference paper at ICLR 2022
Stephen Gould, Basura Fernando, Anoop Cherian, Peter Anderson, Rodrigo Santa Cruz, and Edison Guo. On differentiating parameterized argmin and argmax problems with application to bi-level optimization. arXiv preprint 1eX.vv:607.054000 2016.
Stephen Gould, Richard Hartley, and Dylan Campbell. Deep declarative networks: A new hope. arXiv preprinl urXiv:1909 099900466, 2019.
Riccardo Grazzi, Luca Franceschi, Massimiliano Pontil, and Saverio Salzo. On the iteration com- plexity of hypergradient computation. In International Conference on Machine Learning, pp- 3748 3758. PMLR, 2020.
Andreas Griewank and Andrea Walther. Evaluating derivatives: principles and techniques of algorithmic differentiation. SIAM, 2008
SJ Grotzinger and C Witzgall. Projections onto order simplexes. Applied mathematics and Optimiza- tion, 12(1):247-270, 1984.
Isabelle Guyon. Design of experiments of the nips 2003 variable selection benchmark. In NIPS 2003 workshop on feature extraction and feature selection, volume 253, 2003.
Magnus Rudolph Hestenes, Eduard Stiefel, et al. Methods of conjugate gradients for solving linear systems, volume 49, NBS Washington, DC, 1952.
Nicholas J. Higham. Accuracy and Stability of Numerical Algorithms. Society for Industrial and Applied Mathematics, second edition, 2002. doi: 10. 1 11111111197099999 18027. URL https: / //epubs. siam. org/ doi /aaa//aa//// /10 1137/1 97808 7780888788027
Stephan Hoyer, Jascha Sohl-Dickstein, and Sam Greydanus. Neural reparameterization improves structural optimization. 2019.
Kaiyi Ji, Junjie Yang, and Yingbin Liang. Bilevel optimization: Convergence analysis and enhanced design. In International Conference on Machine Learning, pp. 4882-4892. PMLR, 2021.
Yoon Kim, Carl Denton, Luong Hoang, and Alexander M Rush. Structured attention networks. arXiv 008877, 2017.
preprint
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980 2014.
Steven G Krantz and Harold R Parks. The functionn theorem: history, theory, and applications. Springer Science Business Media, 2012.
Cong Han Lim and Stephen Wright. Efficient bregman projections onto the permutahedron and related polytopes. In Proc. of AISTATS, pp. 1205-1213. PMLR, 2016.
Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparamcters by implicit differentiation. In International Conference on Artificial Intelligence and Statistics, pp. 1540 1550-1552. PMLR, 2020.
Julien Mairal, Francis Bach, and Jean Ponce. Task-driven dictionary learning. IEEE Transactions Pattern Analysis and Machine Intelligence, 34(4):791-804, 2012. doi: 10.109//PAM..201.115
André FT Martins and Ramón Fernandez Astudillo. From softmax to sparsemax: A sparse model of attention and
classification. In Proc. of ICML, 2016.
Christian Michelot. A finite algorithm for finding the projection of a point onto the canonical simplex of R Journal of Optimization Theory and Applications, 19(1):195-200 1986.
J.-J. Moreau. Proximité dualité dans un espace hilbertien. Bulletin de la S.M.F., 93:273-299, 1965.
Vlad Niculae and Mathieu Blondel. A regularized framework for sparse and structured neural attention. In Proc. of NeurlPS, 2017.
Vlad Niculae and Andre Martins. Lp-sparsemap: Differentiable relaxed optimization for sparse structured prediction. In International Conference Machine Learning, pp. 7348 7358,7359, 2020.
12

---

Under review as a conference paper at ICLR 2022
Vlad Niculae, Andre Martins, Mathieu Blondel, and Claire Cardie. Sparsemap: Differentiable sparse structured inference. In International Conference on Machine Learning, pp. 7799-3808. PMLR, 2018.
Brendan O' DDnooguue, Eric Chu, Neal Parikh, and Stephen Boyd. Conic optimization via operator splitting and homogeneous self-dual embedding. Journal of Optimization Theory and Applications, 169(3):1422-1068 2016.
Neal Parikh and Stephen Boyd. Proximal algorithms. Foundations and Trends in optimization, 1(3):
127-239, 2014.
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Pretten- hofer, 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-2830, 2011.
Fabian Pedregosa.
Hyperparameter optimization with approximate gradient. In International conference on machine learning. PMLR, 2016.
Aravind Rajeswaran, Chelsea Finn, Sham Kakade, and Sergey Levine. Meta-learning with implicit gradients. arXiv preprint 1rXiv:1909004630 2019.
Nimrod Rappoport and Ron Shamir. Multi-omic and multi-view clustering algorithms: review and cancer benchmark. Nucleic Acids Res., 46: 10546-10562, 2018.
Youcef Saad and Martin H Schultz. Gmres: A generalized minimal residual algorithm for solving nonsymmetric linear systems. SIAM Journal on scientific and statistical computing, 7(3):856-88) 869, 1986.
Samuel Schoenholz and Ekin Dogus Cubuk. Jax md: A framework for diffcrentiable physics. In H. Larochelle, M. Ranzato, R. Hadsell, M. E. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 11428-11444 Curran Asso- ciates, Inc., 2020. URL https:: stttpa://proceedings nwwrips httppppp.cc/ppppe/////eeee// 83d3d4b6c 95 957511 5e1 67 (ee77acaccccc0033 Papeee pdf.
Matthias W Seeger. Cross-validation optimization for large scale structured classification kernel methods. Journal of Machine Learning Research, 9(6). 2008.
Pablo Sprechmann, Alex M Bronstein, and Guillermo 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), pp. 11-15. IEEE, 2014.
Samuel Vaiter, Charles-Alban Deledalle, Gabriel Peyré, Charles Dossal, and Jalal Fadili. Local behav- ior of sparse analysis regularization: Applications to risk estimation. Applied and Computational Harmonic Analysis, 35(3):433-451, 2013.
H A van der Vorst and van der Vorst. Bi-CGSTAB: A fast and smoothly converging variant of Bi-CG for the solution of nonsymmetric linear systems. SIAM Journal Scientific and Statistical Computing, 13(2):631 644, 1992. URL http: dx doi org/ 10. 1137 0913035.
Tongzhou Wang. Jun-Yan Zhu, Antonio Torralba, and Alexci A Efros. Dataset distillation. arXiv preprint arXiv: 2018.
Robert Edwin Wengert. A simple automatic derivative evaluation program. Communications the ACM, ((8):4334644 1964.
Yuhuai Wu, Mengye Ren, Renjie Liao, and Roger B. Grosse. Understanding short-horizon bias in stochastic meta-optimization. 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30- May 3, 2018, Conference Truck Proceedings. OpenReview.net, 2018. URL https:ttp/..c.....t.
net. forum? d=H1 MczcMM-cg---..
John Zarka, Louis Thiry, Tomás Angles, and Stéphane Mallat. Deep network classification by scattering and homotopy dictionary learning, arXiv preprint 1r.00.1910.03561, 2019.
13

---

Under review as a conference paper at ICLR 2022
Appendix
A CODE EXAMPLES
A.1 CODE EXAMPLES FOR OPTIMALITY CONDITIONS
Our library provides several reusable optimality condition mappings F or fixed points T. We nevertheless demonstrate the ease of writing some of them from scratch.
Proximal gradient fixed point. The proximal gradient fixed point (7) with step size = 1 is T(x, () prox (z V Vif(x,0f) 0g). can be implemented as follows.
grad jax. (f)
def Txx, theta): theta_f, theta_g theta return prox grad theta_f), theta_g)
Figure 6: Proximal gradient fixed point T(x, 0)
We recall that when the proximity operator is a projection, recover the projected gradient fixed point as a special case. Therefore, this fixed point can also be used for constrained optimization. We provide numerous proximal and projection operators in the library.
KKT conditions. As a more advanced example, we now describe how to implement the KKT conditions (6). The stationarity, primal feasibility and complementary slackness conditions read
V (1 A [8 (2,, 0x)] V 0
H (2200n) 0
À G(2, (c) = 0. Using tax.vip compute vector-Jacobian products, this can be implemented
grad jax (f)
def F(x, theta): 2, nu, lambd x theta_f, theta_H, theta_G theta
H_jjp tax.vjp(H; z, theta_H)
stationarity = (grad(z, theta_f) + (Lvip(nu) [o])
primal.feasabilig = Hzz, theta_H)
G_vjp vG 2, theta_G) stationarity (_amb(1ambd) [0] comp_slackness G(z, theta_G) lambd
return stationarity, frimal_feasabilityy,, comp_slackness
Figure KKT conditions F(x,,)
Similar mappings F can be written if the optimization problem contains only equality constraints or only inequality constraints.
14

---

Under review as a conference paper at ICLR 2022
Mirror descent fixed point.
Letting and denoting (0 (pro)) the fixed point (11) is
f Vi((z)) = i V vif(2.01 0y)
T(x, 0) = proje (y. ppes)
We can then implement it as follows.
grad jax. (f)
def Txx, theta):
theta_f, theta_proj = params
x hat phi_ (mi__apping(x)
y x_hat grad (x, theta_f)
return bregman brogman_rrjectionnse theta_proj)
Figure 8: Mirror descent fixed point T(x,0)
Although not considered in this example, the mapping Vip could also depend on if necessary.
A.2 CODE EXAMPLES FOR EXPERIMENTS
We now sketch how implement experiments using our framework. In the following, inp is short for jax numpy. In all experiments, we only show how to compute gradients with the outer objective. We can then use these gradients with gradient-based solvers to solve the outer objective.
Multiclass SVM experiment.
X_tr, Y_tr, X_val, Y_val load_data()
def w(x, theta): # dual-prima map
return jnp.dot dX_ tr.T, Y_tr x) theta
def f(x, theta): inner objective
return (0.5 theta jnp. cop.aum(((.. theta) ** 2) (x, Y_tr))
grad jax gf) proj = jax wwwppppoojccaiiii (proooeccioonnsippleee row-wise projections
def T(x, theta):
return proj (x grad(x, theta))
ccustom_fixed poinnt(p (T)
def dovm_dual_.m www.aaalolveeeinit....oam theta): [...]
return x_star solution of the dual objective
def outer llss (lambd):
theta e (lambd)
x_star = msvm_dual suual_solver (init_x, theta) # inner solution
dot dot theta), X_val) Y_pred jnp.dot
return 0.5 jnp. sum sum((Y_preee Y_val) ** 2)
print (jax grad (outer___ood (1ambd))
Figure 9: Code example for the multiclass SVM experiment.
15

---

Under review as a conference paper at ICLR 2022
Task-driven dictionary learning experiment.
x_tr, y-tr load_data
def f(x, theta): :
dictionary loss residual X_tr (x, theta) return huber loss (residual)
grad = baax.gra(f)) def T(x, theta): prozimal gradient fixed point return prox_lasso grad(x, theta))
@custom_fixed_poot def sparse_coding (init_x, theta): # inner objective [[..]
return x_star # lasso solution
def outer looe (theta, w): task-driven loss
x_star sparse_ caaase__odinginiitt.. theta) sporse codes y-pred jup. dot (x_star, w) return logloss(y (y_tr, y-pred)
print (jax (ruterutr__oos,, argnums=(0,, (0,1))))
Figure Code example for the task-driven dictionary learning experiment.
Dataset distillation experiment.
X_tr, y_tr 1oad_datal
logloss jax. vmap (loss mutiiclasss_oggtsiicc....
def fxx, theta, 12reg=1e-3): inner objective
scores = jup dot (theta, x)
distilled_labels = jnp. arange (10)
penalty 12reg jnp.sum(( x)
return jnp mean (logloss (distilled_labels, scores)) penalty
jax. fax.grad(f))
@custom_root
def logreg_solveriinitt (init_x, theta):
[[..] return x_star
def outer_ onter_los@ll (theta):
x_star logreg_solved (init_x, theta) # inner solution
scores = jnp. dot (X_tr, x_star)
return jnp mean (logloss (y_tr, scores))
print (jax, grad (theta))
Figure 11: Code example for the dataset distillation experiment.
16

---

Under review as a conference paper at ICLR 2022
Molecular dynamics experiment.
energy_fn
eeteereeeeerr__fl(ddeee init_fn, apply_fn jax_md minimize fire_descent( energy_fn, shift_fn)
x0 random. uniform(eyy, (N, 2))
RO = L x0 # transform to physical coordinates
R = lax. fori_loop(
0, num_ www__ppimmzaaiionsteet steps,
body_fun-lambda t, state: apply_fn(state, t=t),
init stt__va==init_fn(R0)) position x_star /
def F(x, diameter): normalized forces
energy_fn = wtw.spphrrooonrryyniiiiiiaaaaaae normalized_energe lambda x: energy * x)
return -jax gaa...aaddaaaa.............a..a.
(x)
root_jvp(F, x_star, diameter, 1.0, solve=linears sewveeeerrreeeeaaeeeeeeeeeeeeeee solve. solve bicgstab)
print (dx)
Figure 12: Code for the molecular dynamics experiment.
B JACOBIAN PRODUCTS
Our library provides numerous reusable building blocks. We describe in this section how compute their Jacobian products. As general guideline, whenever a projection enjoys closed form, we leave the Jacobian product to the autodiff system.
B. JACOBIAN PRODUCTS OF PROJECTIONS
We describe in this section how to compute the Jacobian products of the projections (in the Euclidean and KI. senses) onto various convex sets. When the convex set does not depend on any variable, simply denote it C instead of C(0).
Non-negative orthant When C is the non-negative orthant, C = Rd. we obtain proje = max (y, 0), where the maximum evaluated weise. This is also known as the ReLu function. The projection in the KL sense reduces to the exponential function, proje (y) exp(y).
Box constraints.
When C(0) is the box constraints C(f) = 0,],00] with 0 E R.2, we obtain proje(v. 0) clip(y, 01, 02) = max(min(y, 01),
This is trivially extended to support different boxes for each coordinate, which case RRdx2
Probability simplex. When C is the standard probability simplex, C = Ad, there is no analytical solution for proje(y). Nevertheless, the projection be computed exactly in O(d) expected time or O(d log d) worst-case time (Brucker, 1984; Michelot, 1986; Duchi et al., 2008; Condat, 2016). /|[s 1. where s {0, 1)d is vector indicating the support
The Jacobian given by diag(s) of proje (2) (Martins & Astudillo, 2016). The projection in the KL sense, on the other hand, enjoys closed form: it reduces to the usual softmax projž (v) exp(y)/ Σ5= exp(y;)-
17

---

Under review as a conference paper at ICLR 2022
Box sections.
Consider the Euclidean projection (0) proje 0) onto the set C(0)
{= E R2d: a; S 2 I B,, E [d]: c). where 0 = (a, B, w, c). This projection is a singly- constrained bounded quadratic program. is easy to check (sec, e.g., (Niculae & Martins, 2020)) that optimal solution satisfies for all i [d]
(0) - [x (0), 0)]; = clip(w;=" (0) + yi an 3.)
where L: R x R" + Rd is the dual-pring -primal mapping and x: x=(() R is the optimal dual variable of the linear constraint, which should be the root of
F(x" () = L(x (0), 0) Tw c.
The root can be found, e.g., by bisection. The gradient Vx* (0) is given by Vx* (B))=B//AA /A and the Jacobian 0z" (0) is obtained by application of the chain rule on L.
Norm balls. When C(0) f- E R |||||| S 8). where || is norm and R,, (r 0) becomes the projection onto norm ball. The projection onto the fi-ball reduces to projection onto the simplex, see, e.g. (Duchi et al., 2008). The projections onto the & and E. balls enjoy closed-form, see, e.g., (Parikh & Boyd, 2014, $6.5). Since they rely on simple composition of functions, all three projections can therefore be automatically differentiated.
Affine sets.
When C(0) {x E =d; Ar b}, where A E R b R' and 0 = (A, b), get proic 0) y A (Ay - b) y A (AAT)-
where At is the Moore-Penrose pseudoinverse of A. The second equality holds if < d and A is full rank. A practical implementation can pre-compute a factorization of the Gram matrix AAT. Alternatvell,, we can also use the KKT conditions.
Hyperplanes and half spaces. 0 = (a, b). we get
When C(O) = {ar E Rdd Jxx = b), where a E Rd and b E IR and
(r, 0) a -
||||
When C(0) far Rd: a b). we simply replace y b in the numerator by max(a' Ty 0).
Transportation and Birkhoff polytopes. When C(Ø) = {X E Rexd; X1, = X =
02,X X 2 0}. the so-called transportation polytope, where 0, E AP and 02 E A are marginals, we can compute approximately the projections, both in the Euclidean and KL senses, by switching to the dual or semi-dual (Blondel et al., 2018). Since both are unconstrained optimization problems, we can compute their Jacobian product by implicit differentiation using the gradient descent fixed point. An advantage of the KI. geometry here is that we can use Sinkhorn (Cuturi, 2013), which is GPU-friendly algorithm. The Birkhoff polytope, the set of doubly stochastic matrices, is obtained by fixing 0, 02 = 1a/d.
xz 2 02). a so-called order simplex (Grotzinger & Witzgall, 1984; Blondel, 2019), the projection operations, both in the Euclidean and KI. sense, reduce to isotonic optimization (Lim & Wright, 2016) and can be solved exactly in (d log d) time using the Pool Adjacent Violators algorithm (Best et al., 2000). The Jacobian of the projections and efficient product with it are derived in (Djolonga & Krause, 2017; Blondel al., 2020).
Order simplex. When C(0) = {i E 2 2
Polyhedra. More generally, we can consider polyhedra, i.c., sets of the form C(0) = {o e Rd: Ar Ca S d}, where E Rexd b E RP, C Rmxd and d E R. There are several ways to differentiate this projection. The first is to use the KKT conditions as detailed in $2.2. A second way consider the dual of the projection instead, which is the maximization of a quadratic function subject to non-negative constraints (Parikh & Boyd, 2014, $6.2). That is, we can reduce the projection on polyhedron to a problem of the form (8) with neg-ntgetive constraints, which we can in turn implicitly differentiate easily using the projected gradient fixed point, combined with the projection on the non-negative orthant. Finally, we apply the dual duimml mapping. which enjoys a closed form and is therefore amenable to autodiff, to obtain the primal projection.
18

---

Under review as a conference paper at ICLR 2022
B.2 JACOBIAN PRODUCTS OF PROXIMITY OPERATORS
We provide several proximity operators, including for the lasso (soft thresholding), elastic net and group lasso (block soft thresholding). All satisfy closed form expressions and can be differentiated automatically via autodiff. For more advanced proximity operators, which do not enjoy closed form, recent works have derived their Jacobians. The Jacobians of fused lasso and OSCAR were derived in (Niculae & Blondel, 2017). For general total variation, the Jacobians were derived in (Vaiter et al., 2013; Cherkaoui et al., 2020).
C MORE EXAMPLES OF OPTIMALITY CRITERIA AND FIXED POINTS
To demonstrate the generality of our approach, we describe in this section more optimality mapping F or fixed point iteration T.
Newton fixed point. Let be root of G(:,0), i.c. G(z, 0) 0. The fixed point iteration of Newton's 's method for root-finding is
T(x, 0) = x = G((,
By the chain and product rules, we have
0 Txx, 0) I Using (3), we get A 0,F((z,)) 0) = nI. Similarly,
(1 T.)..
B (,T(2, 0) = F,F( = -
Newton's 's method for optimization is obtained by choosing G(z, 0) = Vif(x, (), which gives
T(x, 0) = x (
(14)
It is easy to check that we recover the same linear system as for the gradient descent fixed point (5). A practical implementation can pre-compute an LU decomposition of Gxx 8), or a Cholesky decomposition if 0 G(z, , 0) is positive semi-definite.
Proximal block coordinate descent fixed point. defined as the solution
We now consider the case when x (0) is implicitly
x (0) argmin f(z, + Σ 0),
ZERd
where -1..... Jm possibly non-smooth functions operating subvectors (blocks) Ti,,.., Fmm Fm of x. In this case, we can use for E [m]] the fixed point
[z, 0)]: = prox n.g. (I, - [V If 0))
(15)
where ',,,.., Tm are biock-wise step sizes. Clearly, when the step sizes are shared, i.e., 1)1 = 7m 17, this fixed point is equivalent to the proximal gradient fixed point (7) with g( 9: (x, 0).
Quadratic programming. We now show how to use the KKT conditions discussed in $2.2 to differentiate quadratic programs, recovering Opinet (Amos & Kolter, 2017) special case. To give some intuition, let us start with a simple equality-constrainn
quadratic program (QP)
1 argmin f(2, 2 +==T=++ c subject H(z,0) Ez O,
"ER*
where Q E RPXP, E E ROXP, d E Ri, We gather the differentiable parameters as 0 = (Q, E, The stationarity and primal feasibility conditions give
V.f(2,0) () - [0, H2, 8)8 ==
c- E F H(z E 0.
19

---

Under review as a conference paper at ICLR 2022
In matrix notation, this can be rewritten
Q F
c
(16)
0
d
We can write the solution of the linear system (16) as the root generally, the QP can also include inequality constraints
(z, v) of function F(x, 0). More
argmin f(z, 0) =ER"
Q= c subject to = - 0),
5 0.
where M E Rr* and R We gather the differentiable parameters as = EQ h).
The stationarity, primal feasibility and complementary slackness conditions give
Vif(zz (0,H( 0)] v [[,(2, 0)] X = Q2 + c E' v M X = 0
H(z, 0) Ez d 0
Ao G(2, 0) = diag(A)
h) - 0
In matrix notation, this can be written as
E
M
2 L
0
d
d
0
A o h
While x (z, u, X) is no longer the solution of linear system, it is the root of function F Far, 0) and therefore fits our framework. With our framework, no derivation is needed. We simply define f. H and G' directly in Python.
Conic programming. We now show that the differentiation of conic linear programs (Agrawal al., 2019 Amos, 2019), the heart of differentiating through evxpy layers (Agrawal et al., 2019a), easily fits our framework. Consider the problem
z (A),
= argmii c² subject to SERP.SERM
Ez + E
(17)
where A (c, E Rmxp, d E RT, C RP and K. S R is a cone; 2 and s are the primal and slack variables, respectively. Every convex optimization problem be reduced the form (17). Let us form the skew-symmetric matrix
E
c
8(A)
- E
0
d
E
NKN
d
0
where N = -p+++L Following (Agrawal et al., 2019b;a; Amos, 2019), we can use the homogeneous self-dual embedding to reduce the process of solving (17) to finding root of the residual map
F(x, 0) 6TIr + I = ((0 I)TI + T)ær
(18)
where 1l = projaexx: XRR, and KC* C R is the dual cone. The splitting conic solver (O' Donoonuue et al., 2016), which is based on ADMM, outputs a solution F(x* (0), 0) = O which is decomposed as x= (0) = (0), u (0), w (0)). We can then recover the optimal solution of (17) using
2 (A) n² (0(A)) and s*()) = projxc (v* (0(A)) o" (B())).
The key oracle whose JVP/VJP need is therefore which is studied in (Ali et al., 2017). The projection onto few cones is available in our library and can be used to express F.
Frank-Wolfe.
We now consider
x= (0) argmin f(z, 0),
(19)
=EC())CR4
where C(0) is convex polytope, i.e., it is the convex hull of vertices (9),..., Um (0). The Frank- Wolfe algorithm requires a linear minimization oracle (LMO)
ト argmin (s,z)
xEC(O)
20

---

Under review as a conference paper at ICLR 2022
and is popular algorithm when this LMO is easier to compute than the projection onto C(O). However, since this LMO is piecewise constant, its Jacobian is null almost everywhere. Inspired by SparseMAP (Niculae et al., 2018), which corresponds to the case when f is quadratic, rewrite (19) as
p (0) argmin g(p, 0) ((V(O)p, 0),
PÊΛ."
where V(0) is d x m matrix gathering the vertices vi (0). ..... Um (0). We then have x²(() V(0)p" (0). Since we have reduced (19) to minimization over the simplex, we can use the projected gradient fixed point to obtain
Tpp 0) = proj (0) V Vig(p" (1)(p"(0),)) 0)).
We can therefore compute the derivatives of (f) by implicit differentiation and the derivatives of x² (#) by product rule. Frank- Wolfe implementations typically maintain the convex weights of the vertices, which we use to get an approximation of p" (0). Moreover, it is well-known that after t iterations, at most vertices are visited. We can leverage this sparsity solve smaller linear system. Moreover, in practice, we only need to compute VJPs of x² (0).
D JACOBIAN PRECISION PROOFS
Proof of Theorem I. To simplify notations, we note A, = A(x",)x)) and A= A(2, 0), and similarly for B and J. We have by definition of the Jacobian estimate function A.J, B. and ÁJ B. Therefore we have
((2 (0) = - - AT'
= - B A- B, + A - A. B+
=A- - B.) + (Å- -A.
For any invertible matrices M1, M2, it holds that M, M²¹ -M M M2)M2
||M MM 1lap
||
1lap |M M1 M1llopp1.Mn llap
Therefore,
- 1, 1lup S
1
11 A-
7 - (f)1l S
a
a
As consequence, the second term J( 0) Or (0) can be upper bounded and we obtain
[| 0) Ox (0) _ | ( B.)l || + I A.
z || Å - ||opl| /3 B.ll + 2/2 |1 x² (0)] |B
which yields the desired result.
Corollary 1 (Jacobian precision for gradient descent fixed point). Let f be such that f(,0) is twice differentiable and a-strongly convex and f- 0) is --Lisschitz (in the operator norm) and V 0) is B-Lipschitz and bounded in norm by The estimated Jacobian evaluated at is then given by
J(ž 0) = (2 ( (2,
For all 0 E R", and any å estimating x= (8), we have the following bound for the approximation error of the estimated Jacobian
R a²
(xx,,))-a-'(0 <
11 x0)ll
Proof of Corollary 1. This follows from Theorem 1, applied to this specific A(z, 0) and B(x, 0).
For proximal gradient descent, where T(x, 0) = prox ng (a V2,(( 8), this yields
A(z, 8) = 1 - 0,T(2,0) = 1 - 2 (z - (
B 0) = 02prox (x V 0) 1a
21

---

Under review as a conference paper at ICLR 2022
We now focus in the case of proximal gradient descent objective f(z, 0) + g(z), where is smooth and does not depend on This is the case in our experiments in $4.3. Recent work also exploits local smoothness of solutions to derive similar bounds (Bertrand et al., 2021, Theorem 13)
Corollary 2 (Jacobian precision for proximal gradient descent fixed point). Let f be such that I(;,0) is twice differentiable and o-strongly convex and (i (,,(), is y-Lipschitz (in the operator norm) and 02 Jzr 0) is B-Lipschitz and bounded in norm by R. Let g Rd - lR. be a twice- differentiable u-strongly convex (with special case μ = 0 being only convex), for which the function (z, 0) (proxn
(g((== (x- (z-1 1f,, 0)) is Kny Lipschitz in it first argument. The estimated Jacobian evaluated at ã is then given by
J(i, 0) - (Vi (2) -.
For all 0 E IR" and any estimating x² (0), we have the following bound for the approximation error of the estimated Jacobian
8 + Kn
7R (a n1)2
[ Ox (0)( S
+
|| x x0)1l-
Proof of Corollary 2. First, let us note that prox (v. 0) does not depend on A since g itself does not depend on and is therefore equal classical proximity operator of ny which, with a slight overload of notations, we denote as prox (v) (with a single argument). In other words,
prox og (v,, 0)
= prox, ng (3)
Ôn proxx (v, 0) Öprox (v) @aprox ng (v, O.
Regarding the first claim (expression of the estimated Jacobian evaluated at .)), we first have that prox,ng (1) the solution to (a' y) + ///g(>') = in x' by first-order condition for a smooth convex function. We therefore have that
prox ng (v) = (1 + V))
Oprox (v) = (Ia v/ 9 (v)))
the first / and inverse being functional identity and inverse, and the second la and inverse being in the matrix sense, by inverse rule for Jacobians ah(2) [ah (h(2))]-1 (applied to the prox).
As a consequence, we have, for I, (2, 0) = 9 (x - V if 0)) that
A(x, 0) = Ta - (1a+ (4-+
(La 144--+00((66000)11----1111000100 (aa--a
B(x, 0) (11 nI' (2, (,)) (x,
+ (T=
As consequence, for all R.C. we have that
[(x 0) - (V +1,
In the following, we modify slightly the notation of both A and B, writing
Å(r. 0) = fi, 0) + (x,
B(x, 0) v
With the current hypotheses, following along the proof of Theorem 1, we have that Á is (a p)) well-conditioned, and (r + k Lipschitz in its first argument, and B S-Lipschitz its first argument and bounded in norm by R. The same reasoning yields
+ Ky +
RR catp p (a+ H)2 (a+#w)) || - x² x'(A)]l-
[ 0) Ox (0)( s
22

---

Under review as a conference paper at ICLR 2022
Mirror descent (MD)
Proximal gradient (PG)
(BCD)
Runtime per 2150 (2800ndd))
2.00
0.6
1.75
ID wi fixed point
fraed paint
60
1.0
1.00
40
0.6
Implich 9D)
0.73
Implice
Numberr features
Number
5000
5000
2000
Number of features
4000
6000
8000
10000
(a)
(b)
(c)
Figure 13: GPU runtime comparison of implicit differentiation and unrolling for hyperarrameter optimization of multiclass SVMs for multiple problem sizes (same setting as Figure 3). Error bars represent 90% confidence intervals. Absent data points were due out-of-memory errors (16 GB maximum).
Mirror descent (MD)
Proximal gradient (PG)
Block coordinate descent (BCD)
96-
96
80
90
90
Unrolling
Unroling in
paint
86-
86
2300 Number features
4000
2000
4000
6000
2000
10000
2000
4000
6000
8000 10000
Num
fealuures res
Number features
(a)
(b)
Figure 14: Value of the outer problem objective function (validation loss) for hyperparameter optimization multiclass SVMs for multiple problem sizes (same setting as Figure 3). As can be seen, all methods performed similarly in terms of validation loss. This confirms that the faster runtimes for implicit differentiation compared to unrolling shown in Figure (CPU) and Figure 13 (GPU) are the cost of worse validation loss.
Maaag (o Jik. aillz
tokid Unrolng
Inplik
Uno Kng
[|x]te| -ib
[1x18)) Alla
[|x -x
(a) 250 features
(b) 500 features
(c) 1000 features
Figure 15: Jacobian error (0)x [ā,, f)1l (see also Definition 1) evaluated with regularization parameter of 0 = 1, as a function of solution error || (0) - 2/1 when varying the number of features, on the multielass SVM task (see Appendix E. for detailed description of the experimental setup). The ground-truit solution x0) is computed using the liblinear solver (Fan et al., 2008) available in scikit-learn (Pedregosa al., 2011) with very low tolerance of 10 Unlike in Figure 2, which was on ridge regression, the ground-truth Jacobian 0x (0) cannot be computed in closed form, in the more difficult setting of multiclass SVMs. We therefore use finite difference to approximately compute Or (f). Our results nevertheless confirm similar trends as in Figure 2.
23

---

Under review as a conference paper at ICLR 2022
E EXPERIMENTAL SETUP AND ADDITIONAL RESULTS
Our experiments use JAX (Bradbury et al., 2018), which is Apache2-licensed and scikit-learn (Pedregosa al., 2011), which is BSD-licensed.
E.1 HYPERPARAMETER OPTIMIZATION OF MULTICLASS SVMs
Experimental setup. Synthetic datasets were generated using scikit-learts sklearn. datasets w.raasttsmakeeee....mm www._aaaeessaaaeaaaaae www.caasssficaaiin
(Pedregosa et al., 2011), following model adapted from (Guyon, 2003). All datasets consist of m 700 training samples belonging to k = 5 distinct classes. To simulate problems of different sizes, the number of features is varied as p E {100, 250, 500, 750, 1000, 2000, 3000, 4000, 5000, 7500, 10000]. with 10% of features being informative and the rest random noise. In all cases, an additional myal 200 validation samples were generated from the same model to define the outer problem.
For the inner problem, we employed three different solvers: (i) mirror descent, (ii) (accelerated) proximal gradient descent and (iii) block coordinate descent. Hyperparameters for solvers were individually tuned manually to ensure convergence across the range of problem sizes. For mirror descent, a stepsize afterwards up to total of 2500 steps. For proximal gradient descent, stepsize of 5 10 was used for 2500 steps. The block coordinate descent solver was run for 500 iterations. All solvers used the same initialization, namely, Tinit & mxk, which satisfies the dual constraints.
1.0 was used for the first 100 steps, following a inverse square root decay
For the outer problem, gradient descent used with stepsize of 5 10-3 for the first 100 steps, following inverse square root decay afterwards up to a total of 150 steps.
Conjugate gradient used to solve the linear systems in implicit differentiation for at most 2500 iterations.
All results reported pertaining CPU runtimes were obtained using an internal compute cluster. GPU results were obtained using single NVIDIA P100 GPU with 16GB of memory per dataset. For each dataset size, we report the average runtime of an individual iteration in the outer problem, alongside 90% confidence interval estimated from the corresponding 150 runtime values.
Additional results Figure 13 compares the runtime implicit differentiation and unrolling on GPU. These results highlight fundamental limitation of the unrolling approach in memory- -imited systems such as accelerators, as the inner solver suffered from mut-of-memory errors for most problem sizes (p 2000 for mirror descent, p 750 for proximal gradient and block coordinate descent). While it might be possible to ameliorate this limitation by reducing the maximum number of iterations in the inner solver, doing so might lead to additional challenges (Wu al., 2018) and require careful tuning.
Figure 14 depicts the validation loss (value of the outer problem objective function) at convergence, shows that all approaches were able to solve the outer problem, with solutions produced by different approaches being qualitatively indistinguishable from each other across the range of problem sizes considered.
Figure 15 shows the Jacobian error achieved as a function of the solution error, when varying the number of features.
E.2 TRIVED DICTIONARY LEARNING
We downloaded from http::/aaae cs tau.acc. at..aa..............a.. iilmmltiiii. omic benchmark/ download html a set of breast cancer gene expression data together with survival information generated by the TCGA Research Network (ht /www. cancer. gov (ov/tcga) and processed as explained by (Rappoport & Shamir, 2018). The gene expression matrix contains the expression value for p=20,531 genes in m=1,212 samples, from which we keep only the primary tumors (m=1,093). From the survival information, we select the patients who survived at least five years after diagnosis (m1 200), and the patients who died before five years (mo 99), resulting in cohort of m = 299 patients with gene expression and binary label. Note that non-selected patients those who are marked alive but not followed for 5 years.
24

---

Under review as a conference paper at ICLR 2022
To evaluate different binary classification methods on this cohort, we repeated 10 times random split of the full cohort into a training (60%), validation (20%) and test (20%) sets. For each split and each method, 1) the method is trained with different parameters the training set, 2) the parameter that maximizes the classification AUC on the validation set is selected, 3) the method is then re-trained on the union of the training and validation sets with the selected parameter, and 4) we measure the AUC of that model on the test set. We then report, for each method, the mean test AUC over the 10 repeats, together with a 95% confidence interval defined mean + 1.96 x standard error of the mean.
We used Scikit Leam's implementation of logistic regression regularized by & (lasso) and lg (ridge) penalty from sklearn. wii.einmnnemmelee model..el...... Soadee.Looistiitene Iwggsstcceezzsssis.............e
and varied the C regularization parameter over a grid of 10 values: {10- ...00"".....".............. 10"). For the unsupervised dictionary learning experiment method, we estimated a dic- tionary from the gene expression data in the training and validation sets, us- ing
sklearn.. www.mpos itionn www.iiiaar.eeea......... alpha-2.0r which produces sparse codes in k 10 dimensions with roughly 50% nonzero coefficients by minimizing the squared Frobenius reconstruction distance with lasso regularization on the code. We then use sklearn. Linear model to train a logistic regression on the codes, varying the ridge regularization parameter over a grid of 10 values ( (10-1, 10" 10*) }.
in__component=110,
iearrmodeel Logist wet.eeeeeeessssee
Finally, we implemented the task-driven dictionary learning model (13) with our toolbox, following the pseudo-code in Figure 10. Like for the unsupervised dictionary learning experiment, set the dimension of the codes to k = 10, and a fixed elastic net regularization on the inner optimization problem to ensure that the codes have roughly 50% sparsity. For the outer optimization problem, solve an by regularized ridge regression problem, varying again the ridge regularization parameter over a grid of 10 values {10-1, 100 ...., 10 Because the outer problem is non-convex, we minimize it using the Adam optimizer (Kingma & Ba, 2014) with default parameters.
E.3 DATASET DISTILLATION
Experimental setup. For the inner problem, we used gradient descent with backtracking line- search, while for the outer problem we used gradient descent with momentum and a fixed step-size. The momentum parameter was set to 0.9 while the step-size was set to 1.
Figure was produced after 4000 iterations of the outer loop on CPU (Intel(R) Xeon(R) Platinum P-8136 CPU @ 2.00GHz), which took 1h55. Unrolled differentiation took instead 8h:05 (4 times more) to run the same number of iterations. As can be seen in Figure 16, the output the same in both approaches.
Dataset Distillation on (MNIST). Generalization Accuracy: 0.8556
Figure 16: Distilled MNIST dataset 0 E Rx obtained by solving (12) through unrolled differen- tiation. Although there is no qualitative difference, the implicit differentiation approach is 4 times faster.
E.4 MOLECULAR DYNAMICS
Our experimental setup is adapted from the JAX-MD example notebook available at https://github:
optimi
We emphasize that calculating the gradient of the total energy objective, f(x, Eis U(r, 0), with respect to the diameter 0 of the smaller particles, fz 0), does not require implicit differenti- ation or unrolling. This is because V (zf 0) 0 at (0):
Vzf(x" (0).0) dx (0) V (i (0),0) V2I(=* v=f(x*(0),0)
25

---

Under review as a conference paper at ICLR 2022
This is known Danskin's 's theorem envelope theorem. Thus instead, consider sensitivities of position Ox (0) directly, which does require implicit diffcrentiation or unrolling.
Our results comparing implicit and unrolled diffcrentiation for calculating the sensitivity of position are shown in Figure 17. We use BiCGSTAB (Vorst & van der Vorst, 1992) to perform the tangent linear solve. Like in the original JAX-MD experiment, use k 128 particles in dimensions.
Unrolled FIRE optimizer
Implicit differentiation
10
10a
Gradient norm
10a
10a
100
10-4
2500
Number of optimization steps
Number of optimization steps
Figure 17: L1 norm of position sensitivities in the molecular dynamics simulations, for 40 dif- ferent random initial conditions (different colored lines). Gradients through the unrolled FIRE optimizer (Bitzek et al., 2006) for many initial conditions do not converge, in contrast to implicit differentiation.
26


Reviews:
[{"review_id": "wQTyRy-gPH", "reviewer": "Reviewer_b1rc", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "My detailed comments are given as below.\n\nStrength: \n\n1. The presentation of this paper is good and most parts are easy to follow. The motivation of this paper is clear, i.e., to provide a unified and easy-to-use implementation for implicit differentiation by leveraging the tool from automatic differentiation. The tool developed here seems to be useful and cover many existing designs of interest. \n\n2. Since implicit differentiation is useful and has critical use in practical applications now, e.g., meta-learning and automated machine learning, I feel an easy-to-use implementation can further promote the developments of more advanced algorithms. I feel this is good. \n\nWeakness:\n\n1. One of my concerns is the novelty of this paper. Since all things are not new, e.g., implicit differentiation, automatic differentiation, matrix-vector product computation in deep learning package, I feel this paper does not contribute a new algorithm design, but an easy-to-use package for many practical algorithms.\n\n2. The analysis is not new as well, since there are already many works on studying the iteration complexity of the response Jacobian and the hypergradient, which I find are missing in this paper. I list some of them and highly encourage the authors to add them (and some related works therein) and provide a short discussion. \n\n1) K. Ji, J. Yang, and Y. Liang. Bilevel optimization: Convergence analysis and enhanced design. In Proc.\nInternational Conference on Machine Learning (ICML), 2021.\n\n2) Grazzi, R., Franceschi, L., Pontil, M., and Salzo, S. On the iteration complexity of hypergradient computation. In Proc. International Conference on Machine Learning (ICML), 2020.\n\n3. Experiments are conducted over only small datasets and models, e.g., synthetic data and no DNN involved. I am wondering whether this package is generalizable enough for such more practical settings. Therefore, I suggest the authors can provide some experiments on larger datasets and models. \n \n\n\n\n\n", "summary_of_the_paper": "This paper provides a unified tool for combining the implicit differentiation technique and the  automatic differentiation method widely used in existing deep learning packages such as PyTorch and TensorFlow. The proposed implementation is easy to use for numeric optimization such as bilevel optimization, meta-learning and hyperparameter optimization because it covers many existing schemes such as fixed point, KKT point, projected method. In the experiments, the authors illustrate how their tool can be useful to simply the implementation. ", "summary_of_the_review": "Overall, I feel this work provides a useful tool for implementing the implicit differentiation in various scenarios. Although I feel the novelty of this work is not that high, I am still slightly positive about it. I am open to increase my score if the authors can address my concerns and add missing related works. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "TWJeqkh0P3W", "reviewer": "Reviewer_h9WU", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "## Strengths\n\n* The proposed package is a nice unification of the core implicit differentiation procedures. I believe that this construction has the potential to speed up development and application of implicit differentiation methods.\n* The theorem is a good, nice-to-know theorem for the literature.\n* The test tasks are nice and varied and showcase the different ways the package could be applied.\n\n## Weaknesses\n\n* The paper can effectively be split into the \"code\" portion, the \"theorem\" portion, and the \"examples\" portion. None of the portions, by themselves, meet the threshold for ICLR. In particular, the code is a package of preexisting work, and I am unaware of any similar types of packages being published to major conferences (see for example, the baseline TorchDyn library [1]). The theorem also has a straightforward, routine proof, so not much theoretical insight is gained here. Finally, the examples are smaller scale and are more toy than previous work such as [2, 3].\n* The core contribution, the software, is not presented in a digestible manner. For example, the only code in the main paper is given in Figure 1, and that code is confusing. In particular, the functions $f$ and $F$ are not good names for variables, the jax.grad call operates on two arguments so should have argnums=0, and the choice to write out the functions in pure jax code clutters up the entire block. In addition, I would hope that Section 2.2 could be rewritten with more code examples. Currently, this section feels like a rehash of preexisting methods, and code would be helpful in showing how one can directly implement these in practice.\n\n## References\n\n[1] https://arxiv.org/abs/2009.09346\n\n[2] https://arxiv.org/abs/1703.00443\n\n[3] https://arxiv.org/abs/1909.01377", "summary_of_the_paper": "This paper introduces a Jax package for implicitly differentiating various numerical solvers. Concretely, the authors develop a systemic methodology for producing gradients for a variety of optimization problems. Then, the authors prove that the Jacobian solution to the approximate numerical solution produces close enough gradients. Finally, the authors show the power of their framework on four test tasks.", "summary_of_the_review": "I enjoyed the paper and think that the package can contribute much to popularize implicit differentiation tools. However, I feel that the paper is not sufficiently novel on any front to warrant publication. Furthermore, I would encourage the authors to present the main contribution, the package, in a more digestible manner.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "IwKMNAl6nic", "reviewer": "Reviewer_Q3Lr", "summary": "", "questions": "", "limitations": "", "rating": 10, "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": "The paper proposes a modular and efficient framework along with its JAX implementation for the implicit differentiation of optimization problems. Firstly, the user defines the function F capturing the optimality conditions of the problem to be differentiated. Then the method frames the differentiation problem as a resolution of the linear system of equations (eq. 2) and combines implicit differentiation and autodiff of F to automatically differentiate the optimization problem.\nThe main novelty of the paper resides in the efficiency and modularity of the framework proposed to solve bi-level optimization problems. This framework allows to abstract away low-level details and significantly lowers the barrier to use implicit differentiation. \nBi-level optimization problems are ubiquitous in ML (hyperparameter optimization, meta-learning, NAS) and the framework proposed allows to address these problems efficiently.\n\nThe paper is well written and touches upon very relevant topics to the community. The framework, its implementation, and the theoretical insights of the Jacobian error are novel and useful contributions. Overall, the proposed method can positively impact the community since abstracting away low-level details of implicit differentiation allows opening new research directions. Therefore I strongly encourage accepting the paper (modulo minor comments below) and I think it should be highlighted at the conference.\n\n\nCOMMENTS:\n\n1. Figure 1 proposes to the reader an illustrative example of the framework. Although the mathematical notation in the paper is well defined in a specific paragraph, the interpretation of the code functions/notation is left to the reader. For example, when you define “X_tr, y_tr = load_data()” the meaning of X_tr is left to the reader (might be X_transpose as well as X_train). Furthermore, an inexperienced reader with the numpy/jax notation might misunderstand the example proposed (functions jnp.eye and jnp.linalg.solve are not defined). I suggest you add line comments or a paragraph with the code notation.\n2. In section 2.1 General Principle, the inner working of the method is proposed; however, it is difficult for me to untangle the novelties presented with respect to what is already known in the literature. Is the procedure of differentiating a root as a linear system part of the novelty? Or the main novelty resides in how to solve the linear system? Please clarify your contribution in this section.\n3. The paper doesn’t contain an explicit definition of “Implicit Differentiation”. I think it would be beneficial to the reader to have a brief explicit definition of Implicit Differentiation.\n\n\nMINOR TYPOS/GRAMMAR CORRECTION:\n- “exiting implicit differentiation” -> “existing implicit differentiation”\n- “we hope that this paper” -> “we hope this paper”\n- “The derivation and implementation in these works is always case-by-case” -> “The derivation and implementation in these works are always case-by-case”\n- “Often times” -> “Oftentimes”\n- “this learned data set achieves small loss” -> “this learned data set achieves a small loss”\n", "summary_of_the_paper": "The paper proposes a modular and efficient framework along with its JAX implementation for the implicit differentiation of optimization problems. The user defines the function F capturing the optimality conditions of the problem to be differentiated; then the framework combines implicit differentiation and autodiff of F to automatically differentiate the optimization problem. The proposed framework is labeled as efficient, since it doesn’t have to unroll the computational graph like in autodiff, and modular since it doesn’t require case-by-case mathematical derivation like in implicit differentiation.\nThe authors show that existing implicit differentiation methods can be instantiated in their framework. They provide and empirically validate new bounds on the Jacobian error when the optimization problem is only solved approximately.\nThe authors implemented four illustrative applications of their framework ( Hyperparameter Optimization Of Multiclass SVM; Dataset Distillation; Task-Driven Dictionary Learning; Sensitivity Analysis Of Molecular Dynamics).\nCode and implementation in JAX are provided along with the paper.\n", "summary_of_the_review": "The paper proposes a modular and efficient framework along with its JAX implementation for the implicit differentiation of optimization problems. This framework allows to abstract away low-level details and significantly lowers the barrier to use implicit differentiation. Bi-level optimization problems are ubiquitous in ML (hyperparameter optimization, meta-learning, NAS) and a procedure to automatically and efficiently tackle them is much needed by the community. I believe that this paper should be accepted and highlighted at the conference.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper was particularly discussed between the reviewers, the AC and SAC. A last minute reviewer was also called to clarify some issues raised, as one of the reviews never got into the system.

The paper was overall perceived as well written and well presented, and that the software contribution of implicit differentiation techniques is a nice asset for the community, especially its modularity.
The stability guaranty constitutes a nice (though straightforward) result providing a theoretical ground for the proposed approach.
Yet, the paper is often loose on mathematical justifications, in particular on minimal validity assumptions. Details on when the proposed framework could fail would be of interest, both on theoretical and practical parts. A discussion on the minimal assumptions required for validity of the approach should be highlighted more in the text.

Furthermore, the paper lacks discussions and comparisons with concurrent works,
for instance how would the framework compare with existing estimates for implicit differentiation or for unrolling. This could be improved along with providing more analysis on the implementation efficiency.
On the practical part, a high level description the software details would also be much beneficial.
A core discussion focused around what should be expected of this type of paper (i.e., "implementation issues, parallelization, software platforms, hardware" papers as suggested by Q3Lr)

A point of concern was the novelty aspects in the discussion phase was the novelty of the proposed framework: even if the contribution is the framework introduced, this is not new per se (the literature on implicit differentiation now contains a considerable amount of results and implementation examples).
The relevance of the work, both on theoretical and computational aspects, beyond the development of a computational library was found difficult to assess by several reviewers.
Overall, the reviewers judged the novelty and the paper's contribution more on the software side. Hence, a core discussion could focus on aspects expected for code oriented papers (i.e., implementation issues, parallelization, hardware, etc.).

Following the long discussion phase (more than 30 posts on OpenReview) and the aforementioned comments, the paper was rejected.

We encourage the authors to submit a revised version in a future conference or possibly to a software oriented journal, such as JMLR MLOSS or JOSS for instance.
Decision: Reject