Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 5: Marginally below the acceptance threshold
review_content: summary: A good paper considers a critical problem, but contributions are somewhat weak.

This paper aims at a unified approach for implicit differentiation of optimization problems. The experiments results seem good, but the proposed methods lack novelty.
main_review: The high-level idea to perform the automatic implicit differentiation of optimization problems is 1) get the root and fixed-point equations which the optimal solution should satisfy; 2) differentiates these equations (by auto-diff) to and get the desired Jacobian by solving a linear equation.

1. The main problem for this framework is non-smoothness. In most cases, the optimal condition is the monotone inclusion problems, which means that the so-called root and fixed-point equations are non-differentiable. The submission lacks the details for this case, and obviously that the projector and the proximal operator are non-differentiable in most cases.

2. Another confusing thing comes from Fig.2. When the scale of the optimization problem grows, solving the large-scale equation is not guaranteed faster than unrolling iterative method, especially for the optimization algorithms with fast convergence. A more discussion is needed here.
limitations_and_societal_impact: No need to consider the Societal Impact of this submission.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This paper presents a module for implicit differentiation in JAX, which allows efficient differentiation through optimization and root-finding problems. The paper presents the general principle of combining the implicit function theorem with autodiff of optimality conditions, and discusses how this can work for a implicit functions such as a range of constrained optimization methods. Finally there are example applications to three machine learning tasks and one molecular dynamics simulation task where this framework improves on currently used methods.
main_review: This paper has clear exposition and will be useful for ML people wishing to understand and use implicit differentiation. The JAX library which the authors will open source is a very important contribution and will fuel much other work in machine learning and physical simulation. My main complaint is that I think the paper should do a better job recognizing prior work combining autodiff + implicit differentiation, and making clear how its contributions (which are important) are positioned in the broader literature. I will raise my score if some of my concerns under "originality" are addressed (or if it becomes apparent I'm mistaken about the scope of contributions).

--
Originality: 
I'm not aware of other ML papers combining implicit differentiation with autodiff of optimality conditions. It is important to introduce these to ML practitioners. The JAX library described provides functionality which wasn't previously available for ML.

My pedantic gripe: the paper almost exclusively cites other ML papers. Autodiff and implicit differentiation have a long history outside of ML, e.g. in engineering where some software packages implementing the adjoint sensitivity method combine autodiff of constraints on a solution and objective functionals on the solution with the implicit function theorem. They have also been extensively studied in the automatic differentiation literature, e.g. "Evaluating Derivatives", Griewank 2000 chapter 15, which discusses general principles for combining the implicit function theorem with autodiff of optimality criteria, and which has bounds on the error of the estimated Jacobian which are somewhat similar to Thm 1 in the submitted paper.

As such, the paper should link back to this literature. It would be good for readers to realize autodiff+implicit diff is not completely new to this paper, and also for readers to have pointers to follow back to this prior work. 

It would be helpful to understand if and how the general framework (2.1) and the specific example implicit differentiation methods (2.2) differ from previous work in the autodiff and adjoint literature.

The JAX library is a significant and wholly unique contribution. I'm not aware of other software packages which combine autodiff and implicit differentiation in a general and modular way or even which cater for these for ML problems. 

Reading the paper, I found myself wishing it was more explicitly written about the JAX library, the design decisions behind it, and how it can be used. The general principles of autodiff+implicit function theorem are out there in the literature, and while this packages / presents and applies them for an ML audience one doesn't learn so much about the principles beyond e.g. what's presented in Griewank, whereas this is the first library of its kind and an insight into the library design would be fascinating.

--
Quality, Clarity: 
The paper is well written and clear. The experiments are demonstrations of a software package rather than measurement of a model's test error. They are appropriate, seem well set up, and demonstrate clearly some advantages of the proposed method / library.

Nit:  The paragraph just before (8) on page 6 seems incorrect. If g is the commonly defined indicator function for the set, $g = 1$ in the set and $g = 0$ outside, then it will encourage (but not force) x* to lie outside the set. To recover (8) we need to force (not encourage) x* to lie inside the set. I think the definition (8) is recovered from the proximal $x^* = \text{argmin}_x f(x) + g(x)$ when $g = 0$ in the set and $g = \infty$ (or $g = \max_x f(x) - \min_x f(x)$) outside the set.

--
Significance: 
This is an important paper and important library. Implicit differentiation has become increasingly important as a method to scale many things such as bi-level optimization and meta-learning. The paper will help people understand the principles behind it, and the JAX package will help fuel many other research projects which use implicit differentiation.
limitations_and_societal_impact: Yes. This is a methods paper and I don't think there are any specific negative societal impacts.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: The authors propose a unified modular framework for implicit differentiation, which removes the need to manually derive the jacobians for each individual case by leveraging automatic differentiation of the optimality conditions. The proposed framework is shown to encompass various existing implicit differentiation schemes (such as stationary point- or KKT conditions) while also creating new ones (e.g. mirror descent fixed point). To handle the practical case in which the optimization problem can only be solved approximately, a theoretically justified bound on the error of the jacobian estimate is provided, showing that the error on the estimated jacobian is at most of the order of the deviation of the approximate solution from the optimal solution.
Additionally, a JAX implementation of the proposed framework is described, with the promise of an open-source JAX library for implicit differentiation upon acceptance.
Finally, four applications of the framework are demonstrated, highlighting its flexibility and ease of use.
main_review: This paper provides a good unification of the various existing implicit differentiation schemes. The principle idea of using the implicit function theorem for implicit differentiation of optimal solutions instead of unrolling the iterations of the solver is not new, but the simple implementation using autodiff opens up many possibilities of employing modular implicit differentiation without the need to focus on tedious calculations of jacobians. This is strongly supported by the promise of an open-source JAX library for modular implicit differentiation.

In the case of differentiating a fixed point, matrix B is defined in the equations following line 108. The derivative with respect to the second argument of F (which is $\theta$) can be confusing, as the optimal solution (depending on theta) is already plugged in as the first argument of F. The clarity could be improved by taking the derivative with x in the neighborhood of $x^*(\theta)$, and then plugging in $x^*(\theta)$ after calculating the derivative.

The authors show that part of the computation for solving the linear system on the backward pass can be reused to calculate the JVP or VJP. While in general this might be the most efficient one can get, in special cases previous work (e.g. [1]) has been able to achieve high efficiency by reusing calculations from the forward pass of a specifically designed solver, dramatically reducing the computation on the backward pass. It is unclear to me whether this could also be easily exploited using this framework. As [1] is listed as one of the schemes encompassed by the proposed framework, the potential difference in efficiency should be discussed.

The general structure of the framework allows making general statements on the error of the Jacobians when the optimization problem is only solved approximately. However, it is not entirely clear how well these conditions hold in practice, that is, how large the derived constant is in practice and whether the assumptions of the theorems hold. This could be made clearer by showing the applicability of the theorem to one of the described experiments.

The presented experiments clearly show the ease of use in various applications of bi-level optimization. The presentation of the paper is very clear. 

[1] B. Amos and J. Z. Kolter. Optnet: Differentiable optimization as a layer in neural networks. In Proc. of ICML, pages 136–145, 2017.
limitations_and_societal_impact: Societal Impact: No societal impact

Limitations: Efficiency and applicability of the theorem as mentioned in the main review.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: The paper promises extension of Google’s JAX library by automatic and modular differentitation of implicitly defined functions. Given a solver $\theta\mapsto x^*(\theta)$ one asks for its Jacobian $\partial x^*(\theta)$. In cases when $x^*$ does not have an explicit form that can be automatically differentiated, the Jacobian can be computed from the optimality conditions $F$ having the general form $F(x^*(\theta),\theta)=0$. Jacobian $\partial x^*(\theta)$ is then given by a solution of a linear system involving $\partial_1 F$ and $\partial_2 F$. The authors describe a unified implementation in JAX and show the resulting simple interface for end-users. Various use cases are presented on examples and code-snippets figures, including stationary point conditions, KKT conditions, proximal gradient, projective gradient and mirror descent fixed point and more. On the theoretical side, the authors discuss the case when the optimization problem solution is only approximate and derive Jacobian error bounds under certain assumptions on optimality conditions $F$.
main_review: The paper is well written, clearly motivated and easy to follow. It contains many illustrative examples, all notions are explained and properly defined.

The main contribution is not in the paper itself but in the code it promises. Implicit differentiation is a well-known concept that has been used many times in concrete scenarios; the paper quotes at least [6, 9, 10, 11, 19, 35, 46, 51, 52, 53]. The goal and contribution of the paper are unifying all these tedious case by case derivations and offering a tool to build new schemes. The universality and modularity are clear, strong sides of the project.

Concerning the submission is mainly about the code, no code was submitted. My biggest concern is that we are forced to believe “the organization” “will have it available in open source by the time of this paper’s publication”. From this point of view, the paper itself does not contribute and contains just an advertisement what the organization will release in December. The implicit function theorem (IFT) is a part of the second-year course of analysis. *The code* is what JAX framework users will benefit from, but not even a demo was so far provided. We can guess that the library is already functioning only from the experiments’ results presented in the paper.

As mentioned, from mathematical point of view, IFT is a classics. It is however not stated correctly (page 3, lines 94-96). Consider, for instance $F(x,\theta)=x^2-1$ and a solution $x^*(\theta)=1$ for $\theta\in(-\infty,0]$ and $x^*(\theta)=-1$ for $\theta\in(0,\infty)$. Equation (1) is clearly satisfied, $F$ is smooth, $\partial_1F(x^*(\theta),\theta)\ne 0$ for every $\theta$ but $\partial x^*(0)$ does not exist. The problem is that IFT asserts that under the assumptions, there exists a unique $C^1$ function $x^*$ in some neighbourhood of $x^*(\theta)\times\theta$ and equality (2) holds. In the setting of the paper, the solver $x^*$ is given. What assumptions do you impose on the solver? What if multiple optimal solutions or bifurcation occur ($\partial_1F$ conditionality blowup)? 

Jacobian bounds section: Definition 1 is probably not needed since it seems that it is not used anywhere else in the paper. Theorem 1 is correct; the proof is an exercise. However, I do not understand its value in this context. The global bounds on $\partial_1 F$ and $\partial_2$ (mostly the well-conditionality and boundedness) seems too restrictive, and the resulting constant in the bound (between lines 221 and 222 — consider labelling it) will be probably very large. Is it possible to demonstrate the use of Theorem 1 (or any of Corollaries 1 and 2) on an example? If not applicable, such a result does not increase the paper’s value.

Edit: After submitting my review the code was publicly released (two days before the deadline for reviews). It is impossible to treat the submission as anonymous when the authors are publicly active and attract attention on social networks, blogs etc. during the reviewing period.
limitations_and_societal_impact: Yes
review_title: null

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