Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: EDIT: After the discussion with the authors and the revisions they submitted, I modified my overall score from 3 to 6, the presentation score from 3 to 4, and the soundness score from 1 to 3.

The paper presents a blueprint for automatic implicit differentiation of solutions to optimization problems, along with its implementation in the JAX library. They claim that their blueprint is widely applicable by listing several common optimization problem templates (e.g., those solvable by mirror descent, proximal gradient descent, Newton's method, conic programming, etc) that they claim one can apply their blueprint to. The main idea of the blueprint is to find a fixed point equation representing the optimality conditions and then apply the implicit function theorem to this fixed point equation. They also prove a theorem regarding the precision of the estimated implicit Jacobian in terms of the precision of the estimated fixed point, which they give numerical support for. Finally, they report numerical experiments for four problems, primarily comparing their method to unrolling.
strengths_and_weaknesses: One strength of the paper is the software implementation in JAX, which is user-friendly, designed to be broadly applicable, and seemingly efficient; later theoretical flaws limit the justified use of this software and thus make this contribution less significant.

The paper is ambiguous when defining functions, not specifying if a function is continuous, differentiable, twice differentiable, etc, at key points, like in the blueprint and the theorem about Jacobian precision. For example, around line 98 when differentiating the root is introduced as the main principle which the rest of the paper relies on, there is no specification on the regularity of the function F besides that it should be "a user-provided mapping, capturing the optimality conditions of a problem." However, the argument that follows relies on the smooth implicit function theorem, which in this context requires the function F to be C1 (continuously differentiable) in a neighborhood. This continues in line 118 where the chain rule is used without specifying that F is smooth, in line 168 when the regularity of f, G, and H are not specified, in line 199 for the definition of Jacobian estimate, in line 210 for the main/only theorem, etc.

This ambiguity becomes important for applying the results. Many of the functions involved in Table 1 on line 161, for instance, are frequently nonsmooth in machine learning contexts (e.g., the prox, the projection operator, etc), and thus the results developed in the paper are not applicable to them. There is mention in line 183 that, because the prox is differentiable a.e., that the smooth implicit function theorem can be applied a.e.. This is false - stronger assumptions than differentiability at a point are needed to apply the smooth implicit function theorem, i.e., F must be C1 on an open neighborhood of the solution (along with an invertibility condition). While nonsmooth implicit function theorems exist, e.g., [Clarke 1990], [Bolte et al 2022], they are nowhere mentioned and their requirements to be applied are different than those of the smooth implicit function since the generalized gradients involved are set-valued. This issue comes up again in line 202 where it's assumed that the solution xstar is differentiable, which is not true in general (e.g., the Lasso solution can have kinks as a function of the l1 weight).

Some numerical experiments are also affected by this, for example the euclidean projection onto the simplex isn't smooth (line 257) and cannot be treated by the blueprint. While it's true that sometimes a smooth fixed point equation can be associated to a nonsmooth problem, such as in section 4.1 where the mirror descent formulation was used with the smooth KL projection on the simplex, there is no discussion of this phenomenon in the paper, nor a way to systematically construct such smooth fixed equations in general. The nonsmoothness also raises questions for the numerical comparisons regarding unrolling, which is mentioned in the questions section.

There is also a concern about the originality of the ideas in the paper. The main idea of the blueprint, to associate an equation modeling optimality to the problem instead of using the set-valued optimality conditions, is trivial in the smooth case since the optimality conditions will no longer be set-valued. Since little beyond line 183 is said about implicitly differentiating nonsmooth functions, the contribution of the blueprint here reduces to the fact that we can associate a fixed point equation to optimality conditions, which is well-known, e.g., for maximal monotone inclusions in [Bauschke, Combettes 2011].

Because the results are ultimately relegated to the smooth setting, I find the paper to be lacking in originality and significance. The main result that remains is the Jacobian precision theorem for smooth functions which, while compelling, is not sufficient on its own nor in conjunction the numerical experiments for smooth problems to deserve publication here. I think an integration of the nonsmooth case is necessary to make the contribution significant enough to worthy of publication, especially due to the prevalence of nonsmooth optimization problems in machine learning.

"Optimization and nonsmooth analysis" - FH Clarke 1990

"Nonsmooth Implicit Differentiation for Machine Learning and Optimization" - Jérôme Bolte, Tam Le, Edouard Pauwels, Antonio Silveti-Falls 2021

"Convex analysis and monotone operator theory in Hilbert spaces" - HH Bauschke, PL Combettes, 2011
questions: In line 262, what is meant by unrolling if the algorithm is nonsmooth?
limitations: A major limitation of the work is that it does not apply to the nonsmooth setting in which the fixed point equation associated to the optimality conditions is not continuously differentiable. Many of the optimization problems coming from machine learning are nonsmooth and so this severely limits the impact of the proposed blueprint.
review_title: null

final_score: rating: 9: Very Strong Accept: Technically flawless paper with groundbreaking impact on at least one area of AI/ML and excellent impact on multiple areas of AI/ML, with flawless evaluation, resources, and reproducibility, and no unaddressed ethical considerations.
review_content: summary: Implicit differentiation can be used to through the solution to an optimization problem without having to backpropagate through the method by which this solution was determined. While this approach has been used in many different contexts, implicit differentiation must often be tailored towards the specific problem in question. This article introduces a library that implements automatic implicit differentiation, addressing this limitation. It goes on to provide precision estimates if an optimization problem has not been exactly solved, and concludes with several examples showcasing the library.
strengths_and_weaknesses: This submission fills an important gap by providing a library implementing automatic implicit differentiation. In my view, this promises to make the use of implicitly defined layers (regardless of their specific form) more accessible for practitioners, in particular making it easier to try out a particular idea. The library implementation is well documented and seems to integrate well with other Jax packages. The article itself is well written, providing a nice balance of motivation, code, theory, and examples. Theorem 1 is closely related to existing results about inverse stability (as the authors note), but it is helpful to have it stated explicitly in the context of implicit differentiation. The proof seems to be correct.

As a minor issue, I was surprised to not see any example applications in the context of deep networks. Reference [55] implements dataset distillation on deep networks, yet the authors seem to be focused on a logistic regression case (see Questions). The code provided in the supplementary material also seems to include deep learning examples, but the authors make no note of this in the appendix. I don't see the inclusion of deep learning examples in the article as a necessity at all, but was a little surprised about the lack thereof.
questions: - Returning to the dataset distillation example, reference [55] only uses a few gradient steps for the inner loop to make their algorithm scalable. Do the authors think their approach towards dataset distillation for logistic regression would scale towards deeper networks?
- l. 224: What do the authors mean by "implicit differentiation gains a factor of t compared to automatic differentiation"? Moreover, does automatic differentiation in this context refer to automatic differentiation through the t iterations rather than implicit automatic differentiation?
- l. 227: What is missing to apply these results to hypergradients? Would that simply be a matter of applying the chain rule to compute the gradients with respect to the objective and obtain corresponding precision guarantees?
- I did not understand the illustration in figure 6; since $\theta\in\mathbb{R}^k$ specifies the diameter of each particle, isn’t $\partial x^{\ast}(\theta)\in\mathbb{R}^{k\times(k\times 2)}$? Is the figure depicting the diagonal elements of that Jacobian?
- Where is digit ‘9’ in figure 5?

One typo: l. 324: 'converge, due to the'
limitations: I believe the authors have adequately addressed the limitations of their work.
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: This paper introduces a Jax package for automatic implicit differentiation. Specifically, the authors propose an efficient and modular approach for implicit differentiation of optimization problems. Their approach combines the benefits of autodiff and implicit differentiation. The authors also provide precision guarantees of the proposed approach. Moreover, the author demonstrate the effectiveness of the proposed approach in bi-level optimization and the sensitivity analysis.
strengths_and_weaknesses: My detailed comments are given as below.

Strength:

1 The motivation of this paper is clear. I believe that this paper makes an important contribution to very relevant topics, e.g., the bi-level optimization, implicit deep neural networks.

2 This paper is well written and the proposed approach is easy to follow. The authors provide some illustrative figures to demonstrate their approach. I feel this is good.

3 The authors implemented four illustrative examples to demonstrate the effectiveness of their approach.

Weakness:

1 The main contribution of this paper is about the software, but the theoretical contribution is overstated. The proof of the theorem is quite standard and I do not get some new insight from it.

2 Direct runtime comparisons with existing methods are missing. The proposed approach is based on implicit differentiation which usually requires additional computational costs. Thus, the direct runtime comparison is necessary to demonstrate the efficiency of the proposed approach.

3 Recently, implicit deep learning has attracted many attentions, which is very relevant to the topic of this paper. An implementation example of implicit deep neural networks should be included. Moreover,  many Jacobian-free methods e.g., [1-3] have been proposed to reduce the computational cost. The comparisons (runtime and accuracy) with these methods are preferred.

[1] Fung, Samy Wu, et al. "Fixed point networks: Implicit depth models with Jacobian-free backprop." (2021).

[2] Geng, Zhengyang, et al. "On training implicit models." Advances in Neural Information Processing Systems 34 (2021): 24247-24260.

[3] Ramzi, Zaccharie, et al. "SHINE: SHaring the INverse Estimate from the forward pass for bi-level optimization and implicit models." arXiv preprint arXiv:2106.00553 (2021).
questions: Please see weakness 2 and 3. If the authors can address these two concerns, I will improve my scores.
limitations: Yes.
review_title: null

Meta-review: metareview: The reviewers have discussed the paper at length and have reached a consensus after the authors have clarified the applicability and limitations of their proposed method. I recommend that the authors continue to polish their manuscript with the points they raised in their summary to the Area Chairs and congratulate them on the acceptance of their submission.
confidence: Certain
Decision: Accept