Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 10: strong accept, should be highlighted at the conference
review_content: 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.
The 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.
The authors implemented four illustrative applications of their framework ( Hyperparameter Optimization Of Multiclass SVM; Dataset Distillation; Task-Driven Dictionary Learning; Sensitivity Analysis Of Molecular Dynamics).
Code and implementation in JAX are provided along with the paper.

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.
The 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. 
Bi-level optimization problems are ubiquitous in ML (hyperparameter optimization, meta-learning, NAS) and the framework proposed allows to address these problems efficiently.

The 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.


COMMENTS:

1. 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.
2. 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.
3. 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.


MINOR TYPOS/GRAMMAR CORRECTION:
- “exiting implicit differentiation” -> “existing implicit differentiation”
- “we hope that this paper” -> “we hope this paper”
- “The derivation and implementation in these works is always case-by-case” -> “The derivation and implementation in these works are always case-by-case”
- “Often times” -> “Oftentimes”
- “this learned data set achieves small loss” -> “this learned data set achieves a small loss”

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.
review_title: null

final_score: rating: 3: reject, not good enough
review_content: 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.
main_review: ## Strengths

* 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.
* The theorem is a good, nice-to-know theorem for the literature.
* The test tasks are nice and varied and showcase the different ways the package could be applied.

## Weaknesses

* 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].
* 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.

## References

[1] https://arxiv.org/abs/2009.09346

[2] https://arxiv.org/abs/1703.00443

[3] https://arxiv.org/abs/1909.01377
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.
review_title: null

final_score: rating: 8: accept, good paper
review_content: 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. 
main_review: My detailed comments are given as below.

Strength: 

1. 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. 

2. 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. 

Weakness:

1. 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.

2. 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. 

1) K. Ji, J. Yang, and Y. Liang. Bilevel optimization: Convergence analysis and enhanced design. In Proc.
International Conference on Machine Learning (ICML), 2021.

2) 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.

3. 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. 
 





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.
review_title: null

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