Einstein VI:   General and Integrated Stein Variational Inference in NumPyro
Authors: Ahmad Salim Al-Sibahi
Ola Rønning
Christophe Ley
Thomas Wim Hamelryck
Venue:  (2021)
Keywords: Stein variational inference
variational inference
probabilistic programming
Pyro
deep probabilistic programming
deep learning

Abstract: Stein Variational Inference is a technique for approximate Bayesian inferencethat is recently gaining popularity since it combines the scalability of traditionalVariational Inference (VI) with the flexibility of non-parametric particle basedinference methods.  While there has been considerable progress in developmentof algorithms, integration in existing probabilistic programming languages (PPLs)with an easy-to-use interface is currently lacking.  EinStein VI is a lightweightcomposable library that integrates the latest Stein Variational Inference methodswith the NumPyro PPL. Inference with EinStein VI relies on ELBO-within-Stein tosupport use of custom inference programs (guides), non-linear scaling of repulsionforce, second-order gradient updates using matrix-valued kernels and parametertransforms. We demonstrate the achieved synergy of the different Stein techniquesand the versatility of EinStein VI library by applying it on examples. Comparedto traditional Stochastic VI, EinStein VI is better at capturing uncertainty andrepresenting richer posteriors. We use several applications to show how one canuse Neural Transforms (NeuTra) and second-order optimization to provide betterinference using EinStein VI. We show how EinStein VI can be used to infer theparameters of a Latent Dirichlet Allocation model with a neural guide. The resultsindicate that Einstein VI can be combined with NumPyro’s support for automaticmarginalization to do inference over models with discrete latent variables. Finally,we introduce an example with a novel extension to Deep Markov Models, calledthe Stein Mixture Deep Markov Model (SM-DMM), which shows that EinStein VIcan be scaled to reasonably large models with over 500.000 parameters

Under review as a conference paper at ICLR 2020
EINSTEIN VI: GENERAL AND INTEGRATED STEIN VARIATIONAL INFERENCE IN NUMPYRO
Anonymous authors Paper under double-bling review
ABSTRACT
Stein Variational Inference is technique for approximate Bayesian inference that is gaining popularity since it combines the scalability of traditional Variational In- ference (VI) with the flexibility of non-parametric particle based inference methods. While there has been considerable progress in development of algorithms, integra- tion in existing probabilistic programming languages (PPLs) with an easy-to-use interface is currently lacking. EinStein VI is a lightweight composable library that integrates the latest Stein Variational Inference methods with the NumPyro PPL. Inference with EinStein VI relies on ELBO-within-SSi
support use of custom inference programs (guides), non-linear scaling of repulsion force, second-order gradient updates using matrix-valued kernels and parameter transforms. We demon- strate the achieved synergy the different Stein techniques and the versatility of the EinStein VI library by applying it on examples. Compared to traditional Stochastic VI, EinStein VI better at capturing uncertainty and representing richer posteriors. We use several applications (NeuTra) and second-order optimization to provide better inference using EinStein VI. We show how EinStein VI can be used to infer novel Stein Mixture versions of realistic models. We infer the parameters of Stein Mixture Latent Dirichlet Allocation model (SM-LDA) with neural guide. The results indicate that Einstein VI can be combined with NumPyro's 's support for automatic marginalization to do inference over models with discrete latent variables. Finally, we introduce an example with novel Stein Mixture extension to Deep Markov Models, called the Stein Mixture Deep Markov Model (SM-DMM), which shows that EinStein VI can be scaled to reasonably large models with over 500.000 parameters.
show how one can use Neural Transforms
INTRODUCTION
Interest in Bayesian deep learning has surged due to the need for quantifying the uncertainty of predictions provided by machine learning algorithms. The idea behind Bayesian learning is to describe observed data x using model with latent variable z (representing model parameters and nuisance variables, latent variables given a model describing the joint distribution p(z, x) = p(x pz)( following the rules of Bayesian inference:
e.g., Fig. 4a). The goal is then to infer posterior distribution p(z|x) over
p(z]x) = Z - -='p(x/2))p222
where the normalization constant Z = [. (r(xx//)((zdc is intractable for most practical models including deep neural networks: an analytic solution is lacking or may require an infeasible number of calculations.
Variational Inference (VI) techniques (Blei et al. 2017 Hoffman et al. 2013 Ranganath et al. 2014) provide way to find an approximation of the posterior distribution. VI poses a family of distributions over latent variables qzz) E Q (c.g., Fig. Fbg.4b) and chooses the one that minimizes chosen divergened") D(g(z) || p(z]x)) (e.g., Kullback-Leibler) to the true posterior distribution. VI often provides good approximations that can capture uncertainty, scaling millions of data points using mini-batch training.
Asymmetric distance

---

Under review as a conference paper at ICLR 2020
def model(X, y=None):
def guide y=None):
W1 np- zeros((4, 3))
W1 = param WW np. zeros 3)))
Ws np. one 3))
Ws param ('Ws', np. (p,, 3)),
constraint-positivee
bl np. z3)os(3)
b1 param np. np.zeros(3))
bs (3)
bs param ('bs', np on)
constraint-positivee W sample ('W', Normal (W1, Ws))
W sample ('W1', Normal (W1, Ws))
b sample (''', Normal (bl, bs))
sample('b' Normal bs))
probs softmax(X @ W b)
with (late('data' X. shape[0]]
sample('y', Categorical (probs),
obs-y)
(a) Probabilistic
linear
regression
model
(b) Inference Program (guide) q( g(W,b) with pa- rameters We W., by and b,
p pW,
data = foad_iriS()
X_train, testt, y_train, y_test train_test__plit(data data, data.targer
stein Steinnmodel, VrappeddGuidd(guiide)) Adam(0. 1), ELBOO)
RBFKernel(), num wamm_paatccees10001
state, loss stein. train key, 15000, X train, y_train, callbacks- [Progbar ()]
(c) Running EinStein VI to learn parameters using Stein particles (Wes, Wa,s be 1
y-pred stein. stein.predict print print(f" :cccuru:::: [np my == y_test) 1100:.2. 2f} %")
X_test)
(d) Using learned parameters from EinStein VI to predict values of test data
Figure 1: Linear regression example model in NumPyro with EinStein VI for inference
Stein Variational Inference (Liu and Wang. 2016 a recent non-parametric approach to VI which uses set of particles { as the approximating distribution ((2) to provide better flexibility in capturing correlations between latent variables. The technique preserves the scalability of traditional VI approaches while offering the flexibility and modelling scope of techniques such as Markov Chain Monte Carlo (MCMC). Stein VI has been shown to be good at capturing multi multi-modality (Liu and Wang 2016 Wang and Liu 2019a), and has useful theoretical interpretations particles following gradient flow LLu 2017) and as moment matching optimization system (Liu and Wang 2018).
Many advanced inference methods based Stein VI have been recently developed, including Stein mixtures (Nalisnick) 2019), non-linear Stein (Wang and Liu 2019b), factorized graphical models (Zhuo al. 20188 Wang al. [2188), matrix-valued kernels (Wang et al. (2019) and support for higher-order gradient-baset optimization (Detommaso et al.] 2018). These techniques have been shown significantly extend the power of Stein VI, allowing more flexible and effective approximations of the true posterior distribution. While algorithmic power is growing, there remains distinct lack of integration of these techniques into general probabilistic programming language (PPL.) framework. Such an integration would solve one of the most prominent limitations of traditional VI, which lacks flexibility to capture rich correlations the approximated posterior.
This paper presents the EinStein VI library that extends the NumPyro PPL (Bingham et al. 2019 Phan et al. 2019) with support for the recent developments Stein Variational Inference in an integrated and compositional fashion (see Fig. Ic and Fig. Id]. The library takes advantage of the capabilities of NumPyro- including universal probabilistic programming (van de Meent et al., (2018), integration with deep learning using JAX (Frostig et al.] 2018), , and automatic optimization and marginalization of discrete latent variables (Obermeyer et al. (2019) to provide capabilities that work synergetically with the Stein algorithms. Concretely, contributions
2

---

Under review as a conference paper at ICLR 2020
Einstein VI, a general library that extends NumPyro and allows Stein Variational Inference to work with custom guide programs based on ELBO-within-See Steiin optimization (Nalisnick) 2 The library is compositional with NumPyro features, including supported deep learning and automatic marginalization, loss functions (ELBO, Rényi ELBO, custom losses) and optimization methods, allowing it to grow organically with NumPyro development.
Integration of recent developments in Stein variational inference within the EinStein library. This includes support for non-linear optimization (Wang and Liu 2019b), wealth of ker- nels (Liu and Wang 2016] 2018 Gorham and Mackey| (017). matrix-valuue kernels (Wang) et al., 2019) supporting higher-order optimization, and factorization based on conditional independence between elements in the model (graphical kernels).
We support the application of transforms on the parameter space, including triangular (Parno and Marzouk 2018) and neural transforms (Hoffman et al. 2018) to improve the gradient geometry of the inference problem.
A series of examples demonstrate the power of an integrated library such as EinStein VI and the synergy between different Stein VI techniques. The examples include novel Stein Mixture version of Deep Markov Models (SM-DMM), Stein Mixture Latent Dirichlet Allocation (SM-LDA), and several examples using neural transforms and higher-order optimization.
The paper proceeds as follows. We first present a primer on the theory of Stein VI in Section P relating our integrated implementation in EinStein VI. We discuss the general details of the implementation of EinStein VI in NumPyro in Section 3 We present the various examples using EinStein VI in Section 4 and finally summarize our results and future work in Section
1.1 RELATED WORK
There has been proliferation of deep probabilistic programming languages based on tensor frame- works with automatic differentiation, supporting various inference techniques.
Pyro (Bingham et al. 2019) is universal PPL. based on PyTorch (Paszke et al. 2019). The main mode of inference in Pyro black-box Stochastic Variational Inference (Ranganath et al.] 2014) with guides, which are flexible programs that approximate the posterior distribution by repeating the (non-observed) sample statements as the probabilistic model and can contain deep neural networks
allowing amortization of inference (Kingma and Welling 2014; Gershman and Goodman 2014).
Pyro also supports various sampling algorithms like Hamiltonian Monte Carlo/NUTS (Neal 2011
Hoffman and Gelman 2014) and
MCMC (Zhu 2019) which provide more accurate posterior approximations but lack the scalability of VI techniques. NumPyro is version of Pyro that runs on the JAX framework (Frostig et al.] 2018), which allows it exploit the powerful program optimization and parallclizability of the JAX compiler. Our library, EinStein VI, extends the Stochastic VI mode of inference by adding Stein Variational Inference on top of NumPyro, allowing the optimizable parameters in the guide to be approximated by set of particles instead of a point estimate. Other languages with similar feature include PyMC3 (Salvatier et al., (016).. Edward (Tran et al., 2016 2019) and HackPPL (Ai et al. 2019).
2 PRIMER ON STEIN VI
The core idea of Stein VI Liu and Wang 2016) is to perform inference by approximating the target posterior distribution p(z(x) x) by an approximate distribution qz (z) = N - Σ, &2. (z) based on set of particles Z N Here, (y)) represents the Dirac delta measure, which is equal to 1 if x y and 0 otherwise. One could thus see the approximating distribution (z(z) as mixture of point estimates, each represented by a partiele E 2. The idea then that we minimize the Kullback-Leibler divergence DKL (q(z) || p(z[x)) between the approximation and the true posterior by iteratively updating the particles using the Stein forces:
2: - 2, + (Sz(z;)
where represents the learning rate and Sz the Stein forces.
3

---

Under review as a conference paper at ICLR 2020
The Two Forces of Stein VI Stein VI consists of two forces which work additively under the form S2 = 22 + Sz. where the attractive force
(z.) = (x) [k(z), V., pog
and the repulsive force
(z(z:) = (x))))(() k(z,, z:)l
Here k: 1d x Rd - IR is kernel. The attractive force can be seen as pushing the particles towards the direction that maximizes the true posterior distribution, smoothed by some kernel. For an example of kemnel, consider the standard RBF kernel k kzz,, 2z) exp (-) z; 2j Il)) with bandwidth parameter h, usually chosen as log
1 (z). The normalization constant becomes additive in the log-pose pog-posterior log p(z)(x) = log Z + log logp(x(z) + log p(z) and so does not need to be computed for the gradient.
The repulsive force can be seen as moving particles away from each other, ensuring that they do not all collapse to the same mode. For the RBF kemel, the repulsive force becomes (2))) [k= z.) (z 2,z]] and so particles that are close together (thus having a high kernel value) will be pushed away from each other.
Non-linear Stein In non-linear Stein Waang and Liu 2019a). the repulsive force can be scaled by a factor X, S2 2 XSZ which is often useful when dealing with multi-modal distributions. It is also useful in our framework since the repulsive force often vanishes compared to the likelihood for large datasets X' and therefore scaling the repulsive force by constant À X((2() proportional to the size of the dataset [ appears logical, e.g. 0.1 or 0.01.
Matrix-valuee kernels The choice of kernels can be extended matrix-valued ones (Wang et al. 2019 Rd x Rd TR ddd in which case the Stein forces become
(2 Egg (z) [K (2j, V,) log logp(2;|x))
and
(Z,)) Ezj (x) [K(z,, 2,) Vz.1
2 and so where the standalone del Vz, in the repulsive force represents the vector 02..1 Fa..
(K (z 2;) Vz,)z= Σx Ke 2.). The advantage of matrix-valued kernels that they allow preconditioning using the Hessian Fisher Information matrix (second-order derivatives), which can capture local curvature and thus achieve better optima than standard Stein VI. Furthermore, is easy to capture graphical kernels (Wang al. 2018b)
diag((((()) where the set of variables are partitioned with each their own local kernel K7(2)
The standard form of Stein VI guides the particles toward maximizing the log-pu reg-posterio log p(z(x) but in principle we could replace this with another objective (negative loss) Lo which we want to maximize. Assume our guide (z(z) comes from parametric distribution with parameters 0, as in traditional VI. We can then for example maximize the ELBO (Kingma and Welling 2014) Ce = EEz (x) [log p(z, x) - log qe (z) as in traditional VI but using set of particles for the parameters {0:}s guided by the Stein forces. One can see this as generalization of traditional VI using a Stein-based mixture-moder (Nalismick 2019), giving the parameters the flexibility to better capture correlations between distributions and to cover multiple modes (Wang and Liu 2019a). Another perspective that ELBO- S---wiihi---eeee adds uncertainty the approximating distribution ge (z), allowing work with distributions over latent variables that are richer than simply point masses.
ELBO-within-Stein
The objective to be maximized in our framework is compositional, which means that ELBO can be re- placed by other objectives, e.g., Rényi ELBO (Li and Turner 2016). Tail-adaptive f-divergence (Wang] et al.) 2018a) Wasserstein secudo-divergence (Ambrogioni et al., 2018) when these imple- mented in NumPyro. Einstein VI makes it possible to integrate all these developments within Stein VI in coherent fashion.
4

---

Under review as a conference paper at ICLR 2020
Parameter Transforms Deterministic transport maps have been suggested to be integrated in MCMC methods by Parno and Marzouk (2018) and Hoffman et al. (2018), to allow constructing a better geometry for the posterior approximation and thus improve inference. We provide novel adaptation of the idea to work with Stein-based methods, in the EinStein VI framework.
Consider the general objective optimize in EinStein VI, which has the form:
(e(O,) = -o [k(0,. Ve, Ce. Ve, ((0,, 0.)]
where Lo our negative loss objective (e.g., ELBO) parameterised by 0 which we update using set of Stein particles @ {.:)
The core idea is then that we would like to reparametrize the calculation so each Stein particle 0, expressed as an (invertible) deterministic transformation of another particle Φ;, so 0; = ( (thus = (T) l E })))
We do this reparametrization by first declaring the set of particles to be optimized + {o:); and the transformation be used T, and then define the reparametrized Stein force as follows:
(+)))) (o.))))) T(b:))
where Vo: Vo: T(b:): Raxe is the reparametrization Jacobian, that adjusts for the change variables (from 0, RR to Φ; R2).
The power of parameter transforms is amplified in EinStein VI framework, in that allow for its parameters to be Icarnable. We can thus use e.g., sparse triangular maps, invertible neural networks or normalizing flows, which have been shown to be particularly effective for Bayesian inference.
COMPOSITIONAL IMPLEMENTATION USING NUMPYRO
The system integrates well with the existing NumPyro API and programs using the standard Stochastic Variational Inference (SVI) interface can easily be converted to use EinStein VI instead. Fig. 1 shows the core usage of the EinStein VI interface; like the standard SVI interface it accepts a model, guide, optimizer and loss function. Additionally, EinStein VI accepts a kernel interface (e.g. RBFKernel) which is used to guide the Stein particles and to specify the number of Stein particles to use for inference.
We will discuss the major features of the implementation of Einstein VI, including allowing for re- initializable guides, the core algorithm that integrates Stein VI into NumPyro, and the compositional kernel interface.
3.1 RE-INITIALIABLE GUIDES
The Stein VI interface requires that inference programs can be re-initialized with different values for each parameter. The reason that we would like different Stein particles to be initialized differently in order for optimization to work correctly and avoid all particles collapsing into the posterior mode.
To support re-initializable guides we provide the Reini tGuide interface which requires implement- ing a function find_part Find_params that accepts list of random number generator (RNG) keys in addition to the arguments for the guide and returns a set of freshly initialized parameters for each RNG key.
Most guides in NumPyro written functions as in Fig. 4bg. and support those, we provide the WrappedGuide class which makes a callable guide re- re-initializable. works by running the provided guide multiple times, and reinitializing the parameters using NumPyro's interface follows:
WrappedGuide runs the guide transforming each parameter unconstrained space.
It replaces the values of the parameters with values provided by a NumPyro initialization strategy, e.g., ini t_to_tttttti form(r) which initializes each parameter with uniform random value in the range [-r;]].
It saves the parameter values for each particle and the required inverse transformations to constrained space, in order allow running the model correctly.
5

---

Under review as a conference paper at ICLR 2020
We also allow parameters without reinitialization, e.g., for stax neural network which have their own initializers.
3.2 STEIN VI IN NUMPYRO
The integration of Stein VI into the deep PPL NumPyro requires handling transformations between the parameter representation of NumPyro- NumPyro-wwhihh whhic is dictionary mapping parameters to their values, which can be arbitrary Python types-to vectorized Stein particles that can be used in EinStein VI. For this, we rely on NumPyro's integration with JAX PyTree& which can convert back and forth between Python collections and flattened vectorized representation. For example, the parameters {We, W ,, br,i, b. from Fig. I will be encoded into particles {a:} where each particle is 10-dimensional vector E R
Alg. shows the core algorithm of Einstein VI. Classical variational parameters & and ø' are updated by averaging the loss over Stein particles. For Stein parameters, the process is more elaborate. First, we convert set of individual parameters to set of vector-encoded particles using JAX PyTrees. Then we compute a kernel based on the vector-encoded Stein particles; the computation is kernel-depender but an example is the RBF kernel which requires a bandwidth parameter / that is dependent on the particle values (Liu and Wang 2016). We apply JAX's vmap operator (Frostig et al. 2018 Phan et al. 2019) to parallely compute the Stein forces for each particle a vectorized manner. Alg. presents how EinStein VI works with regular- valued kernels but the actual implementation is more general, since it takes into account the different features presented in Section 2 non-linear repulsion force, potential use of matrix-valued kernels, and parameter transforms. Finally, we convert the monolithic Stein particle back to their non-vectorized dictionary-based form and return the expected changes for both classical and Stein parameters.
Algorithm 1 EinStein VI
Input: Classical parameters & and d', Stein parameters }0,).. model Pe (z x), guide de (z).
loss L, kernel interface KI.
Output: Parameter changes based classical VI (AØ, AΦ') and Stein VI forces (120,).).
AΦ + E. V L(pa, 10 (0.0r)
Ad' <- Eo V d' L(Po, "0.00))
{a:h: + (ETREEFLATTEN(()))
kI((a:}+)
[Aath < VMAP(fa: a + Ea, k(a,, a, V., C(ps, (1 + Va, k(a,,a)) a.))
(A0:): <- PYTREERESTORE( (Aa;)i
return ΔΦ. AΦ', (A0:):
3.3 KERNEL INTERFACE
The Kernel interface SteinKernel is reasonably simple. Users must implement the compute function, which accepts as input the current set of particles, the mapping between model parameters and particles, and the loss function L and returns a differentiable kernel k. Currently, all kernels are disallowed from changing state, but one could imagine stateful interface in the future which relies on the history of particles, e.g., for computing conjugate gradients or quasi-Newton optimization. The list of implemented kernels in EinStein VI are described further in App.
EXAMPLES
We demonstrate the flexibility of EinStein by its application to two standard pathological examples (the Double Moon and Neal's Funnel), Stein Mixture Latent Dirichlet Allocation (SM-LDA).. and Stein Mixture Deep Markov Model (SM-DMM). All experimental code, and the EinStein package We plan to integrate our
available at hattps://gg thub. com/ work upstream to the main NumPyro repository. NoupPooooooooooooooooo

---

Under review as a conference paper at ICLR 2020
Neal's funnel We first consider toy model where MCMC methods often struggle sample the true posterior. The model was first described Neal (2003). and commonly known as Neal's funnel. The joint posterior is
V (0,3)
NO
which becomes difficult to sample as y grows increasingly negative. Figure 2 shows that EinStein VI captures the funnel well, and the particles are concentrated around the funnel area expected.
Figure 2: 100 particles fitted to Neal's Funnel using Stein in EinStein VI [left]. EinStein VI with neural transformation applied to the two moons model.
Dual Moon The Dual Moon example is classical model with multiple modes, which requires rich family of variational distributions (guides) to provide good approximation the posterior. thus serves good example of how integration of parameter transforms in EinStein VI.
In EinStein VI, each particle represents a variational parameter. Transformation of variables of the approximating distribution causes all particles to collapse modes of the distribution. This is because the repulsive force is applied in the preimage of the transformation (c.g. Inverse Autorcgressive flow, Planar flow, Radial flow). To remedy this we apply Stein in the transformed space instead of using it to transform our approximation distribution. Figure 4 shows how to use neural flows for inference.
Stein Mixture Latent Dirichlet Allocation Latent Dirichlet Allocation (Bei et al. 2003)
model that seeks to infer distribution over topics 0, that represents set of size N of related words, from a collection documents, each represented as a bag of words w. Each document is generated according to the following model:
Dir(o)
7 lid
Cat(@)
Wn Lid Cat(y,.)
n. E (1..N)
where the number of topics T 10) is fixed, is scalar that represents prior knowledge on the symmetric Dirichlet distribution, typically a = 1/T and {sp)]== is a distribution of words for each topic. Especially interesting is that the model contains set of discrete intermediate latent variables { which makes it non-differentiable per default. However, NumPyro supports automatic enumeration, which transforms the above generative model to the following equivalent model:
0 ~ Dir(a)
W 22
Cat((x))Cat(110)
E (1..N)
Figure 3
7

---

Under review as a conference paper at ICLR 2020
def fnows(n_ flows, in in_dim, h_dim): fs = [] for i in (n
def guide fl_wsls h didimm(2,22]
flows = flows (n_ 2, h_dim)
particle param('p's array([o... 0. 0.1).
arn = AutoRegNN (in__ii,, h_dim, fun-Efn=EEl)
ComposeTrans (flows)
numpyro.. numpyro.sample sample('x (elta(particle))
fs. a (arn))
fs. wpp.pedd arn
AutoregNN (in_diin h h_dim,
out_fun-Elu)
fs. apppedd (ppernd((Atttoeeggraassaaa)) return fs
(a) Inverse AutoRegressive Transformation Flow.
(b) Guide with neural transformation
Figure 4: Neural transformed guide using Inverse Autoregressive Flows in Numpyro for EinStein VI (function names are shortened for formatting purposes).
Here, all the potential topic choices t {1..T) are considered explicitly and marginalized (summed) out in the intermediate step. This way, we longer have discrete sampling step and we can differentiate the model again.
EinStein VI integrates with this automatic enumeration support in NumPyro, and thus provides more modelling power than traditional Stein methods that are not integrated into a corresponding powerful framework. In our example, the guide uses amortized inference (Gershman and Goodman 2014) with a multi-layer perceptron (MLP) with 100 nodes and 20 potential topics, mapping each document the distribution of topics, thus approximating the desired posterior p(0) [/]w). For each Stein particle, we get a separate MLP, and thus the resulting distribution would be a Stein mixture of neural networks.
Stein Mixture Deep Markoy Model Music generation requires model to learn complex temporal dependencies to achieve local consistency between notes. The Stein Mixture Deep Markov Model (SM-DMM) is a Deep Markov Model that uses a mixture of Stein particles to estimate distributions over model parameters. We consider vectorized version of DMM (Jankowiak and Karaletsos 2019) for the generation polyphonic music using the JSB chorales dataset.
The SM-DMM model consists of two feed-forward neural networks. The Transition network transi- tions between latent states in the Markov chain and the Emitter produces an observation at each time step from the current latent state. The variational distribution is of the form I[ (2)).r,, [f where the parametrized feature function firr. GRU (Chung al., 2014). with hidden dimension 400. In the JSB chorales dataset, each time step is an eighth of note.
We run the SM-DMM using the Adam optimizer (Kingma and Ba 2014) with learning rate of 10-5 5, using an RBF kernel and five Stein particles for one thousand epochs. The trained SM-DMM achieves negative log likelihood of 6.67 on the test set, comparable to the 6.85 originally reported in Krishnan et al. (2016) and 6.82 reported in the improved version in Dankowiak and Karaletsos (2019).
5 SUMMARY
We presented EinStein VI, novel library for Stein Variational Inference, built on top of NumPyro. It supports all recent techniques associated with Stein VI, including STLBO-within-Steer optimization, non-linear scaling of repulsive force and matrix-valued kernels, well adding novel features like parameter transforms. We have provided several examples that illustrate how EinStein VI has an easy-to-use high-level interface and performs well pathological examples as well as Stein Mixture versions of realistic models such as Latent Dirichlet Allocation and Deep Markov Model.
In the future, plan supporting updates inspired by Stein Points (Chen al. 2018 2019), to handle cases where initialization of particles sub-optimal, e.g. random restart of particles and MCMC updates to get better mode hopping properties.
8

---

Under review as a conference paper at ICLR 2020
REFERENCES
J. Ai, N. S. Arora, N. Dong. B. Gokkaya, T. Jiang, A. Kubendran, A. Kumar, M. Tingley, and N. Torabi. Hackppl: universal probabilistic programming language. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, MAPL 2019, page 20-28, New York, NY, USA, 2019. Association for Computing Machinery. doi: 10.1 10.1145/3315508.32299944 URL
ISBN 9781450367196. 9315508.322997
htt.//o...rrg///// 11457
L. Ambrogioni, U. Güçlüü Y. Güglütürk, M. Hinne, M. A. J. van Gerven, and E. Maris. Wasser- stein variational inference. CoRR, abs/1805. 11284, 2018. URL/ttpp:/arxiv.oo www.vvorrg//bs. 1805 11284
E. Bingham, J. P. Chen, M. Jankowiak, Obermeyer, N. Pradhan, T. Karaletsos, Singh, A.
Szerlip, P. Horsfall, and N. D. Goodman. Pyro: Deep universal probabilistic programming. J.
Mach, Learn. Res., 20:28: 20:28:1-28:66 1-28:6, 2019. URL.ttpp http://rrrrgppppps////////o// 1102200018144.110 1818844..htt html
D. M. Blei, A. Y. Ng. and M. I. Jordan. Latent dirichlet allocation. J. Mach. Learn. Res., 3:993-1022, 2003. URL httppp//mmaa..orr//ppess///////p
D. M. Blei, A. Kucukelbir, and J. D. McAuliffe. Variational inference: A review for statisticians. Journal of the American Statistical Association, 112(518):859-8777 2017. doi: 10. 1080/016214599 2017.1285773 URL https://:///oi. otttp:::/ddoioor/1000 org/10. 1080/01621400 2017. 1285773
W. Y. Chen, L. W. Mackey, J. Gorham, F. Briol, and C. J. Oates. Stein points. In J. G. Dy and
A. Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmüssann Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pages 843-852. PMLR, 2018. URL http:::///prrocedinng mlr press/v80/chss htt.////he118..mmmmm.m
W. Y. Chen, A. Barp, F. Briol, J. Gorham, M. A. Girolami, L. W. Mackey, and J. Oates. Stein Point Markov Chain Monte Carlo. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning. ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 1011-1021. .tml
PMLR, 2019. URL http:ttt////cceeeiiiiii httpeiiiinnnmlr..pees//////
Chung, C. Gulcehre, K. Cho, and Y. Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling, arXiv preprint arXiv: 14Xiv:1412.3555, 2014.
G. Detommaso, T. Cui, Y. M. Marzouk, A. Spantini, and R. Scheichl. A Stein variational New- ton method. In S. Bengio, H. M. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, 3-8 December 2018, Montréal, Canada, pages 9187-9197, 2018. URL http::/ppppers.niiiiie/p nips httpppsccccpappe//11131 8110- a- stein variational-newton method
R. Frostig, M. Johnson, and C. Leary. Compiling machine learning programs via high-level tracing. In SysML 2018, 2018. URL http:://www wwtpsslllll...//eeeee cctpcccccccccc/cc//eecc 146...df pdf
S. Gershman and N. D. Goodman. Amortized inference in probabilistic reasoning. In P. Bello, M. Guarini, M. McShane, and B. Scassellati, editors, Proceedings of the 36th Annual Meeting of the Cognitive Science Society, CogSci 2014, Quebec City, Canada, July 23-26, 2014. cognitive- sciencesociety.org 2014. http::ts:://minee
J. Gorham and L. W. Mackey. Measuring sample quality with kernels. In D. Precup and Y. W. Teh, ed- itors, Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-1 August 2017, volume 70 of P Proceedings of Machine Learning Research, pages 1292-1301. PMLR, 2017. URL Utt.hhttpp//rrcceedinng
mlr. www.ttpressssooroeeeee
7 html
M. D. Hoffman and A. Gelman. The No-U-tum sampler: adaptively setting path lengths in Hamilto- nian Monte Carlo. J. Mach. Learn. Res., 15(1) 1623, 2014. URL http: : ://ll. acm, org/ citation. c

---

Under review as a conference paper at ICLR 2020
M. D. Hoffman, D. M. Blei, C. Wang, and J. W. Paisley, Stochastic variational inference. J. Mach. Learn. Res., (4(1):1303-13444 2013. URL http://///dd.com acm. org/ wwt.tiin..ciiiiic. cfm.:m/2222mm222222222222
M. D. Hoffman, P. Sountsov, J. Dillon, I. Langmore, D. Tran, and S. Vasudevan. Neutra-lizing bad geometry in hamiltonian monte carlo using neural transport. 2018.
M. Jankowiak and T. Karaletsos. Pathwise derivatives for multivariate distributions, In K. Chaudhuri and M. Sugiyama, editors, The 22nd International Conference on Artificial Intelligence and Statistics, AISTATS 2019, 16-18 April 2019, Naha, Okinawa, Japan, volume 89 of Proceedings of Machine Learning Research, pages 333-342. PMLR, 2019. URL htt.httt://proceedinng .ml.
press/v89/js jawp://jjaawwiaaaaaaaaaa html
D. P. Kingma and J. Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1142.6600 1412.6980, 2014.
D. P. Kingma and M. Welling, Auto-encoding Variational Bayes. In Bengio and Y. LeCun, editors, 2nd International Conference on Learning Representations, ICLR 2014, Banff. AB, Canada, April 14-16, 16, 2014, Conference Track Proceedings, 2014. URI.htttp//arxiv org abs/ 1312.6114
R. G. Krishnan, U. Shalit, and D. Sontag. Structured inference networks for nonlinear state space models. arXiv preprint arXiv: 09869, 2016.
Y. Li and R. E. Turner. Rényi divergence variational inference. In D. D. Lee, M. Sugiyama, U. von Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pages 1073-1081, 2016. UtL/bttp://pppprs.nni httppp/e//pppppssssssccc//peeepp
inall www.oonaallnieeenenn inference
Q. Liu. Stein variational gradient descent as gradient flow. In I. Guyon, U. von Luxburg, Bengio, H. M. Wallach, R. Fergus, S. V. N. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, 4-9 December 2017, Long Beach, CA, USA, pages 31 3115-3123, 2017. URL httthh//ppppeerr wwtpccc/pappreeeee stein- www.vvarraiinaaggaaaaa www.daaeeeddeeeeeeeeeeeeeeeeeeee as gradi flow
Q. Liu and D. Wang. Stein variational gradient descent: A general purpose Bayesian inference algorithm. In D. D. Lee, M. Sugiyama, U. von Luxburg, I. Guyon, and R. Garnett, edi- tors, Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pages 2370- 2378, 2016. URL http://tp:/p aaicent-a-generaa -e--geeral-puuposeeeeee wwwppppppppppsssaaaaaeappa bayesian inference algori thm
cctpsssccpppppeeeeeaeeeee stein- iiii--variaiional wwwoooaall-rriiieeee ent
Q. Liu and D. Wang. Stein variational gradient descent as moment matching. In S. Bengio, H. M. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Pro- cessing Systems 2018, NeurIPS 2018, 3-8 December 2018, Montréal, Canada, pages 8868- 8877, 2018. httLhtttp//aaa wwwpppprrr.spp...oommmmpp www..nissssspppppeee wwt.cc/pppppppeeeeee (aaper/8101-steine www.sttiiiiiiiiinnaalllll -aaenallgrrienttt descent- as moment - matching
E. Nalisnick. Variational inference with Stein mixtures. In Advances in Approximate Bayesian Inference, NIPS 2017 Workshop, 2019.
R. M. Neal. Slice sampling. Annals of Statistics, pages 705-741, 2003.
R. M. Neal. MCMC using Hamiltonian dynamics. In S. Brooks, A. Gelman, G. Jones, and X.-L. Meng, editors, Handbook of Markov Chain Monte Carlo, chapter 5, Chapman & Hall / CRC Press, 2011.
Obermeyer, E. Bingham, M. Jankowiak, D. Phan, and J. P. Chen. Functional tensors for probabilistic programming. CoRR, abs/1910.10775, 2019. URL http:t:://ariiv..rraabb/////
M. D. Parno and Y. M. Marzouk. Transport map accelerated Markoy Chain Monte Carlo. SIAM/ASA Journal on Uncertainty Quantificatinn, 6(2):645-68) 2018.
10

---

Under review as a conference paper at ICLR 2020
A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, Gimelshein, L. Antiga, A. Desmaison, A. Köpf, E. Yang. Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. M. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché-Bue Buch E. B. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, 8-14 December 2019, Vancouver, BC, Camada, pages 8024-8035, 2019, http: httppp/s/pppppppppps../aaeeeee/p 901-- -0yytyy an- superaaivv-stylleee httpa//eeeeeteeeooooooooooooeree
per formance-ddee-
library.
D. Phan, N. Pradhan, and M. Jankowiak. Composable effects for flexible and accelerated probabilistic programming in numpyro. In Program Transformations for Machine Learning, NeurlPS 2019 Workshop, 2019.
R. Ranganath, S. Gerrish, and D. M. Blei. Black box variational inference. In Proceedings of the Seventeenth International Conference on Artificial Intelligence and Statisties, AISTATS 2014, Reykjavik, Iceland, April 22-25, 2014, volume 33 of JMLR Workshop and Conference Proceedings, pages 812. JMLR.org, 2014. html
lar.s.mml....... www. prrpes//v/3/raggaaatt..... http://ooeccnnggg.rrer/////aaaaa
Salvatier, T. V. Wiecki, and C. Fonnesbeck, Probabilistic programming in python using pymc3. PeerJ Comput, Sci., 2:e55, 2016. doi: 10.7717/peerr-cs.555 URL. https:///oi...... http://oi...r////// wwt.//..///......... 7717/
p
D. Tran, A. Kucukelbir, A. B. Dieng, M. R. Rudolph, D. Liang, and D. M. Blei. Edward: A library for probabilistic modeling, inference, and criticism. CoRR, abs// 2016. URL
09787
D. Tran, M. Dusenberry, M. van der Wilk, and D. Hafner. Bayesian layers: module for neural network uncertainty. In H. M. Wallach, H. Larochelle, A. Beygelzimer, F. d' d''Alché-De Buc, E. B. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference Neural Information Processing Systems 2019, NewrlPS 2019, 8- December 2019. Vancouver, BC, Canada, pages 14633-14643 2019. URL http:://pp/pppes..... wwpppprrr..ipp..p..mmp cctpppprrc//pppr 9607-bayesian-layee layers a- a-mouule e- www--nuuuallleeeee untcerrkkuuneetaintyy
J. van de Meent, B. Paige, H. Yang, and F. Wood. An introduction to probabilistic programming. CoRR, 1as/1809.10766. 2018.
10756
D. Wang and Q Liu. Nonlinear Stein variational gradient descent for learning diversified mixture models. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 6576-6585, PMLR, 2019a. URL mlr wwwpppeess/v/.. wang1 9h. html
D. Wang and Q. Liu. Nonlinear stein variational gradient descent for learning diversified mixture models. In K. Chaudhuri and R. Salakhutdinov, editors. Proceedings of the 36th International Conference on Machine Learning. volume 97 of Proceedings of Machine Learning Research. pages 6585, Long Beach, California, USA, 09-15 Jun 2019b. PMLR. URL /ttp:// proceedings mlr. pepps/sssvvv /97 1997///ggg9hhhmm
D. Wang, H. Liu, and Q. Liu. Variational inference with tail- tail-adaptive f-divergence. In S. Ben- gio, H. M. Wallach, H. Larochelle, K. Grauman, N. Resa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Informa- tion Processing Systems 2018, NeurIPS 2018, 3-8 December 2018, Montréal, Canada, pages 5752-5752, 2018a. URL http:http:pppppp www..iips.cco.... /et.scccappee//1--aaaaata tich-tail-aae wewall-daptivee---eee divergence
D. Wang, Z. Zeng, and Q. Liu. Stein variational message passing for continuous graphical models. In G. Dy and A. Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pages 5206-5214. PMLR, 2018b. http! //proceedings. mlr press//80/wwngg81. hmtmll

---

Under review as a conference paper at ICLR 2020
D. Wang, Z. Tang, C. Bajaj. and Q. Liu. Stein variational gradient descent with matrix- valued kernels. H. M. Wallach, H. Larochelle, A. Beygelzimer, F. Alché- Buc, E. B. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NewrlPS 20/9, 8-14 December 2019, Vancouver, BC, Canada, pages 7834-7844, 2019. URL http:/t:tpppppre.nini... earlational-gradienn
ent descent waw.reettttttt-ttttii--eeaaaeeee veattrix-vllued-keenelll hwtps/cccppppppeppeeee stein
M. Zhu. Sample adaptive MCMC. In H. Wallach, H. Larochelle, A. Beygelzimer, F. dÁlché- Buc, E. Fox, and R. Garnett, editors, Advances Neural Information Processing Systems 32, pages 9066-9077. Curran Associates, Inc., 2019. htttp::/papers..ips.cccc./ sample-adaptinn wwwptive-paaeeecceo mcme. pd.f
Zhuo, C. Liu, J. Shi, Zhu, N. Chen, and B. Zhang. Message passing Stein variational gradient descent. In J. G. Dy and A. Krause, editors, Proceedings the 35th International Conference on Machine Learning, ICML. 2018, Stockholmsmüsssan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pages 6013-6022. PMLR, 2018. URL http://proceedings mlr. prwss/v80/ zhuol8a. html
12

---

Under review as a conference paper at ICLR 2020
ACKNOWLEDGEMENTS
We thank the Pyro development team for answering our questions regarding the Pyro implementation. We especially thank Martin Jankowiak for providing us guidance with and access to the Deep Markov Model example.
A KERNELS IN EINSTEIN VI
Kerel
Definition
Detail
Type
Reference
Liu and Wang
Radial Basis Function (RBF)
exp( x y 113)
scalar
(2016)
exp( H/2 (x y))
vector
Pyrc
Inverse Multi-
Gorham and
Quadratic
(c2+ || x 12)0
8 E (-1,0) and > 0
scalar
Mackey (2017)
(IMQ)
Φ(x, w)
Random
V2 cos( à w x + Wo) where Wo Unif(0, 2T)
Liu and Wang (2018)
Feature
Ewlo(x, w)oly, [b]
scalar
Expansion
and W1 (0,
Liu and Wang (2018)
Linear
x y +1
scalar
scalar, vector,
{ individual kernels and we are weights
Liu and Wang
Mixture
E. (x,,
(2018)
matrix
Scalar-based
Wang et
I(x,y)I
I scalar-valued kernel
matrix
Matrix
(2019)
Vector-based
Wang et
diag(k(x,y))
A: vector-valued kernel
matrix
Matrix
(2019)
(K'()), ) matrix-valued kernels each for unique partition of latent variables K is inner matrix- matrix-valued kernel and Q is preconditioning matrix like the Hessian - log p(z|x)
Graphical
(lag((A("((xx))))
matrix
Wang et
(2019)
Constant Pre- conditioned
Q =
matrix
Wang et
(0019)
or Fischer information
E - (x) [V2 log p(z x)]
matrices
{a m
a set of anchor
points, Q₂ Q(a)) is preconditioning matrix for each anchor point, Kq. is an inner kernel conditioned using Qi, and wx(x) softmax(((AV ar, 1p)))e)
Anchor Point
Wang et
Precondi- tioned
Σ (x)
matrix
(0019)
13


Reviews:
[{"review_id": "yw1cGGf9hkU", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary\n========\nThe paper shows how a particle-based nonparameteric Variational Inference methodology known as Stein Variational Inference is integrated in a full-featured Probabilistic Programming Language, NumPyro. The paper goes into a fair amount detail describing a number of enhancements that have been made into numpyro using the general technique of particle-based representation of non-parameteric approximating distributions. They describe how geometric transforms of the parameter space can fit into their scheme, how matrix-valued kernels can be integrated. Also, they describe a new variant of Stein VI which they call ELBO-within-Stein. This introduces a new line of research for Stein VI. They also describe a Stein Mixture extension to Deep Markov Models (SM-DMM) and demonstrate on a very large dataset for the latter method.\n\nStrengths\n=========\n- Integrating a more powerful variational approximation has clear benefits for probabilistic inference. And integrating this into a full-featured PPL allows users of Bayesian modeling to get access to a cutting edge technique with minimal programmatic effort.\n- The integration of Stein VI into numpyro seems to have been very well designed given the very large number of ideas that have become easy to add including some innovative approaches.\n- Showing state-of-the-art results on the high dimensional JSB-Chorales-dataset is a very impressive achievement for any PPL, and it certainly lends credence to this work.\n\nWeaknesses\n==========\n- The only claim in the paper that is well supported is that the authors have extended NumPyro with SVI.\n- The presentation style in the paper sometimes fails to draw a clear distinction between implementations of prior work in NumPyro versus new innovations. It is somewhat unclear whether the authors are making claims about the following points in their paper:\n  * Non-linear Stein\n  * Matrix-valued kernels\n  * Parameter Transforms\n  * Enumeration-based integration of discrete random variables\n- The objective function of ELBO-within-Stein is not well motivated (see discussion below) and there is no direct comparison to the previous Stein Variational Gradient Descent which this method seeks to improve.\n- There is no way to objectively evaluate the results on the first three experiments.\n\nRecommendation\n===============\nReject\n\nRationale\n========\nExperiments don't directly validate the main innovations of the paper.\n\nSupporting Arguments\n====================\n- The main innovation in this paper appears to be the ELBO-within-Stein method. This appears to be different than SVGD (Stein Variational Gradient Descent). The difference appears to be that in the current paper both the entropy term and the Stein repulsion term are in the general objective (page 5 first equation) unlike in SVGD where the entropy term is not there. Philosophically, it doesn't look right to include both of these terms that are serving the same purpose (prevent the collapse of the variational approximation on the mode) . I could be mis-reading these equations, but if there are other difference the authors should clearly state and motivate these differences. Most importantly, the authors should show an experiment directly comparing to SVGD.\n[SVGD reference: Qiang Liu and Dilin Wang.  Stein variational gradient descent: A general purpose Bayesianinference algorithm.Neural Information Processing Systems (NIPS), 2016.]\n- The Neal's funnel should show the posterior marginal which is well known so that the reader can judge whether the samples are of good quality.\n- Not clear how to interpret the dual moons plot. What are we looking at in this plot (right plot figure 2)? The posterior density or the true density? How do we know if this is a good posterior?\n- For the LDA example there don't seem to be any results.\n\nQuestions for authors\n===================\n- Please provide motivation for the modification to SVGD objective.\n- Please clearly state which of the many enhancements to NumPyro are being claimed as novel extensions.\n- Any results worth sharing for the LDA?\n\nAdditional Suggestions Not Part of the Review Rating\n=============================================\n- The abstract mentions that this work is better than Stochastic VI but this claim is not actually supported explicitly. I had to read many of the referenced papers to realize that Jankowiak and Karaletsos (2019) had implemented a version of SVI. I'm assuming that this is what the abstract was referencing. Please do make such connections explicit!\n\n- In Figure 1b, variables X and y are not actually used in the guide. Space permitting you could make a note as to why these are there in the guide.\n\n- The first paragraph of the introduction mentions nuisance variables in Fig 4a. Not clear which variables in 4a were nuisance variables.", "extra_scores": {}, "extra_text": {}}, {"review_id": "xFK60pknP5S", "reviewer": "AnonReviewer1", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Unfortunately the authors link directly to the code, and the code is not anonymous. This might be a desk-reject as this is not a double blind review.\n\nThis work is a description of a library for developing variational inference algorithms using the ELBO-within-Stein framework developed in Nalisnick et al. (2017). The library is evaluated on on Neal's funnel and two moons, and on a polyphonic music dataset.\n\nComments\n\n- Nalisnick et al was published in 2017. I assume this was a typo on the authors' part.\n\n- Table A in the Appendix, describing different kernels, should include a column with computational and memory requirements for each kernel if they differ. This can affect the scalability.\n\n- The work describes LDA but does not evaluate it. It would be helpful to include held-out log likelihood numbers on a standard topic modeling dataset such as 20 newsgroups. This would help people compare to prior work. \n\n- Similarly, the library is evaluated by fitting to a standard polyphonic music dataset. Please report these numbers in a table, alongside a reasonable approach using standard variational inference and Stein VI (using the library) side-by-side. For example, the numbers here are much better, and use standard variational inference with the KL divergence: https://papers.nips.cc/paper/6039-sequential-neural-models-with-stochastic-layers.pdf (Stein Variational Inference can be difficult to understand, as can be Pyro, which is built on jax/pytorch, and the library developed here is built on top of all of these moving parts. Before embarking on using the library, a machine learning researcher should be very convinced that all this additional effort is worth it. Benchmarking this new library against existing work is important and will go a long way toward justifying its existence.)\n\n- The references are very poorly formatted. Please clean up.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "YM69dKGw_Xv", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\nThis paper introduces EinStein VI: a lightweight composable library for Stein Variational Inference (Stein VI). The library is built on top of NumPyro and can take advantage of many of NumPyro's capabilities. It supports recent techniques associated with Stein VI, as well as novel features. The paper provides examples of using the EinStein VI library on different probabilistic models. \n\n### Strengths\nI'm not aware of other PPLs that support Stein Variational Inference. EinStein VI can provide an easier way to compare different Stein VI algorithms, and make research in the area easily reproducible.  \n\n### Concerns \nThe paper states that it provides examples that demonstrate that EinStein VI's interface is easy-to-use and performs well on pathological examples and realistic models. While it is true that there are several examples described, in my opinion there are not enough details to support the claims that EinStein VI is easy to use and performs well. A concrete comparison between EinStein VI and other methods is missing. It would have been helpful to have, for example, some concrete numbers (e.g. time taken to do inference, posterior predictive checks, posterior mean convergence plots, etc) that showcase why it is useful to use Stein VI for those examples, as opposed to other, already existing methods. \n\nAnother concern is that it is difficult to judge from the paper what the difference to (standard) NumPyro is. There is only a high-level explanation of the examples in the paper, so it's hard to imagine what the actual code looks like. Most importantly, I would have liked to see a comparison between EinStein VI code and what the code would have looked like without EinStein VI. \n\n### Reasons for score\nUnfortunately, there is not enough to go on in this paper, which is why I recommend reject. There is no strong evidence to support either the usability of the system (through elaborate examples and contrasting EinStein VI to other systems) or its performance (through experiments). This paper will be much stronger, and will have a better chance of reaching more people, if it includes either 1) more elaborate code examples that demonstrate that using EinStein is indeed better and easier than vanilla NumPyro, or 2) experiments comparing different Stein VI techniques to other inference algorithms, as evidence that a dedicated Stein VI library is indeed empowering our inference toolkit.\n\nHowever, I do appreciate that writing a paper about tools / libraries is difficult, as the contribution of tools is typically a longer-term improvement in the workflow of developing new methods and techniques. I am open to increasing my score during rebuttal, depending on the answers of the questions listed below.\n\n### Questions for the authors\nWhy has Stein VI not been implemented in PPL systems previously? Is it a matter of timing, or is there something particularly challenging about integrating Stein VI into a PPL? \n\nThe paper mentions \"compositionality\" several times. I was a little confused about what you mean by that: can you explain, perhaps with an example? \n\nThe paper mentions novel features (second to last paragraph page 8): can you elaborate?\n\nThe paper shows an example of using NeuTra in combination with Stein VI. Can you elaborate on the kind of problems that NeuTra won't be able to handle on its own? What about more lightweight approaches that can be applied in the context of probabilistic programming, such as \"Automatic Reparameterisation of Probabilistic Programs\" (Gorinova, Maria I., Dave Moore, and Matthew D. Hoffman. ICML 2020)? When will we see benefits of *both* applying a reparameterization that improves the posterior geometry, *and* using a more sophisticated inference algorithm like Stein VI?\n\n### Suggestions for improvement and typos that have not affected the score I gave to the paper\nPerhaps the most important change that would improve the paper is adding more concrete examples that would showcase the importance of using EinStein VI as opposed to simply NumPyro / other libraries. It would be nice to see a model where Stein VI gives us better inference results than a range of other algorithms / techniques and compare the code to what the user would have to write otherwise to achieve the same results. The examples of composing Stein VI with reparameterization / marginalization in NumPyro can be improved by comparing the results to Stein VI without reparameterization / marginalization and to other inference algorithms with reparameterization / marginalization. \n\nTypos:\n* last line of the abstract should be 500 000 as opposed to 500.000.\n* URL in footnote 3 does not lead to the correct page\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "9Z5s9QlxIBq", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "In this paper, the authors developed a probabilistic programming framework for stein variational gradient descent and its variants using difference kinds of kernels, i.e. nonlinear kernels or matrix kernels. Simple experiments are included that the repository is effective and scalable for various problems. \n\nFollowings are a few of my questions and comments:\n\n1. How is the new implementation compared with other frameworks using black box variational inference? For example, What is the speed of the training comparing with previous frameworks such as edward in large scale dataset tasks?  And the report does not give us a more thorough guide of the performance of each kernels for difference tasks. \n\n2. The authors mentioned that the framework can be extended to use other objective function such as  Rényi ELB, , Tail-adaptive f-divergence, or Wasserstein pseudo-divergence. I am extremely confused about this part, since actually there is no objective function for svgd based methods (unless you design a new loss based on KSD or related things), how is this possible to combine other objective function using svgd? It would be great if the authors write down the derivations and have a detailed discussion.\n\n3.  Does the current framework implement amortized svgd and other related stein's paper that can be utilized to train neural networks based applications such as stein-vae, stein-gan or kernel stein generative modeling [1, 2, 3]? This implementation can be important since it can be quite helpful for many other applications such as meta learning. \n\nAlso, the authors give the public code link of their implementation in the paper, which may expose their identity, but I am not sure if this violates anonymous requirement of ICLR submissions. \n\n[1] Feng, Yihao, Dilin Wang, and Qiang Liu. \"Learning to draw samples with amortized stein variational gradient descent.\" arXiv preprint arXiv:1707.06626 (2017).\n\n[2] Wang, Dilin, and Qiang Liu. \"Learning to draw samples: With application to amortized mle for generative adversarial learning.\" arXiv preprint arXiv:1611.01722 (2016).\n\n[3] Chang, Wei-Cheng, et al. \"Kernel Stein Generative Modeling.\" arXiv preprint arXiv:2007.03074 (2020).", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers have carefully reviewed and discussed this paper. They are in consensus that this manuscript merits a strong revision. I encourage the authors to take these experts' thoughts into consideration in revising their manuscript.
Decision: Reject