Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 3: reject, not good enough
review_content: summary_of_the_paper: The aim of this work is to integrate the Stein variational inference methods in the existing probabilistic programming language NumPyro. The implemented methods include variants of Stein variational gradient descent algorithms with a range of kernel functions, non-linear scaling of update terms, and matrix-valued kernels. Empirical results with existing baselines for real-world problems are provided.

main_review: [Strength]

The paper is overall well-written and the method is clearly explained. The literature review is thorough.


[Weakness]
- Motivation is unclear to me. Why implementing the Stein Variational inference methods in PPL framework is preferable to non-PPL framework? And also why NumPyro is preferable to other PPLs? In the introduction, it is briefly mentioned that "Such integration would solve one of the most prominent limitations of traditional VI: its lack of flexibility in capturing rich correlations in the approximated posterior", which I consider as one of the motivations for using PPL framework. Still, theoretical or empirical evidence for this statement is missing and this work would be more appealing if the authors elaborate on it.

- The originality is low since all the Stein variational inference algorithms implemented in NumPyro in this work are proposed by prior work and the authors do not propose any new algorithm.
summary_of_the_review: My main concern is that the original contribution of this work is not significant enough for being published as a conference paper.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper proposed a new probabilistic programming framework for stein variational gradient descent and its variants.  The framework put its basis on Numpyro. Experimental results show that this framework shows faster computation for inference in wall clock time. Moreover, using this framework, the authors developed a new Stein mixture algorithm for deep Markov models, which shows better performance than existing methods.

main_review: # Strength
- Integration of SteinVI into numpyro seems very useful. Users can easily take advantage of the state-of-the-art SteinVI algorithms for their own Bayesian modelings. Moreover, experimental results show that it is faster than the existing implementation.
- The concept of ELBO-within-Stein is innovative and seems interesting to be explored.
- Extending the stein mixture method to deep Markov models is a novel application.

# Weakness
- No code example is presented. Thus I could not tell how EinSteinVI is easier or better than the standard Numpyro and PyMC3.
- Related to the above, I was not convinced by the motivation of this work in the Introduction. The authors explained that 
> Currently Stein VI methods are available in PyMC3 (Salvatier et al., 2016) and Edward (Tran et al., 2016), neither of which include the advanced methods available in EinSteinVI. $\ldots$ While algorithmic power is growing, a distinct lack of integration with a general probabilistic programming language (PPL) framework remains. Such integration would solve one of the most prominent limitations of traditional VI: its lack of flexibility in capturing rich correlations in the approximated posterior

-  As far as I understood, these phrases correspond to the motivation of this work. But after reading this paper, I thought that why PyMC3 or naive Jax implementation is not enough to resolve the above issues and why proposed EinSteinVI resolves them differently from PyMC3 or Jax implementation. Add an explanation or show additional experimental results for that.

- Some important Stein VI methods seem lacking, e.g,
    - Function space particle optimization for bayesian neural network (https://arxiv.org/abs/1902.09754)
    - Understanding and Accelerating Particle-Based Variational Inference (https://arxiv.org/abs/1807.01750)
    - Learning Equivariant Energy Based Models with Equivariant Stein Variational Gradient Descent (https://arxiv.org/abs/2106.07832)
- The presentation of the paper is poor. I think the concept of ELBO-within-Stein and its application to deep Markov models are the technical contributions of this work, but they are poorly explained. 
- No experiments to support the usefulness of EinSteinVI for Non-linear Stein VI, Matrix-valued kernel stein VI, and message passing stein VI.


# The reason for the score
Although the developed library seems useful for the community and the concept of ELBO-within-Stein seems interesting, I think the contribution is not enough experimentally and theoretically as I described in Weakness above. 


# Comments
- The bandwidth tuning is critical for Stein VI. I recommend adding the tuning method other than the median method, e.g, the heat kernel approach introduced in Understanding and Accelerating Particle-Based Variational Inference (https://arxiv.org/abs/1807.01750)
- I think the presentation can be improved to make this work more attractive.  The technical contribution, that is, the concept of ELBO-within-Stein and its application to deep Markov models, should be more emphasized. The authors should explain how novel they are in more detail. Also, the authors should add more experiments to support the usefulness of the concept of ELBO-within-Stein and stein mixtures in deep Markov models.
summary_of_the_review: The contribution is not enough experimentally and theoretically to be published. Although the proposed framework seems beneficial for the community, it lacks a detailed comparison regarding the usability or performance compared to existing frameworks, e.g., PyMC3. Regarding the technical contribution, the extension of stein mixtures seems promising, but it lacks the explanation, numerical experiments, and theoretical justification.
review_title: null

final_score: rating: 3: reject, not good enough
review_content: summary_of_the_paper: The paper presents the implementation of a framework to work with different "Stein VI" methods in NumPyro.
A quick introduction to the different algorithms is given as well as an evaluation of a series of examples.
main_review: Although I am fully supportive of having source code presented in academia, I do not think this paper brings anything of particular interest, in particular for this venue.
Although the introduction to the different Stein-VI methods is nicely written and very understandable, this paper does not aim to be a review which is already done in Anastasiou 21' as the authors point out.
What the paper could bring, would be novel techniques or implementations to improve speed, but unfortunately only using `vmap` to solve things seems a bit light.
Another major weakness is that no code is available, for a submission on a software this is unacceptable.

The positive point I found was the generalization of these different algorithms through a unified algorithm, but this alone is not enough to make the work sufficiently significant to be accepted.

summary_of_the_review: The paper does not bring anything interesting enough research-wise and such a paper should be submitted to a software venue/journal like [JOSS](https://joss.theoj.org/).
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper:   
This paper describes an implementation of EinsteinVI, an improvement over Stein VI, a method that can be used to apply VI in prob. models assuming a flexible distribution represented in terms of particles. The paper describes how this method has been implemented in NumPyro, a software for prob. programming. The method is evaluated and compare to other alternative prob. languages such as PyMC3 that uses also methods based on Stein VI.

main_review: It is not clear what is the difference between EinSteinVI and the Stein VI methods included in Edward and PyMC3.

Th experiments are extensive, considering different prob. models and problems. The computational time of each method is recorded comparing favorably for the implemented one.

The paper compares favorably with competing software for prob. programing languages such as PyMC3. However, it does not compare results with Edward. This questions the significance of the results.

The novelty of the paper is low since it does not describe any new method. It only consists in an implementation of an already known method to be used within NumPyro for prob. programming.

My main concern with the paper hence the lack of novelty, although some practitioners may consider the implementation interesting and useful.


summary_of_the_review: A nice implementation of an already known method that could be used for machine learners using prob. programs.
review_title: null

Meta-review: The paper aims to integrate Stein variational inference methods into the existing probabilistic programming language NumPyro. The implemented methods include variantions of Stein variational gradient descent with different types of kernel functions, non-linear scaling of update terms, and matrix-valued kernels. The paper includes empirical results with a comparsion with existing baselines in real-world problems. Using this framework, the authors developed a new Stein mixture algorithm for deep Markov models, which shows better performance than existing methods.

Strengths:

- The paper is overall well-written and the method is clearly explained. 
- The literature review is thorough.
- Integration of SteinVI into numpyro seems useful. Users can easily take advantage of the state-of-the-art SteinVI algorithms for their own Bayesian modelings.
- Extending the stein mixture method to deep Markov models is a novel application.

Weaknesses:

- The originality is low the authors propose algorithms that are very similar to previous work and there is a lack of experiments to verify the usefulness of the proposed method, for example,
  to verify the decreased variance of the gradient estimates claimed by the authors.
- Efforts are required to illustrate why ELBO-within-Stein is preferred over the existing work.
- Some important Stein VI methods seem lacking.
- No experiments to support the usefulness of EinSteinVI for Non-linear Stein VI, Matrix-valued kernel stein VI, and message passing stein VI.

All reviewers vote for rejection. I recommend the authors to addrss the limitatoins mentioned above and improve the paper before its resubmission to another venue.
Decision: Reject