Title: MODULAR DEEP PROBABILISTIC PROGRAMMING

Abstract: Modularity is a key feature of deep learning libraries but has not been fully exploited for probabilistic programming. We propose to improve modularity of probabilistic programming language by offering not only plain probabilistic distributions but also sophisticated probabilistic model such as Bayesian nonparametric models as fundamental building blocks. We demonstrate this idea by presenting a modular probabilistic programming language MXFusion, which includes a new type of re-usable building blocks, called probabilistic modules. A probabilistic module consists of a set of random variables with associated probabilistic distributions and dedicated inference methods. Under the framework of variational inference, the pre-specified inference methods of individual probabilistic modules can be transparently used for inference of the whole probabilistic model. We demonstrate the power and convenience of probabilistic modules in MXFusion with various examples of Gaussian process models, which are evaluated with experiments on real data.

1. INTRODUCTION:

Deep neural networks (DNNs) have seen great successful in various areas such as computer vision, speech recognition, and artificial intelligence. A major driver of such success is the rise of deep learning libraries, starting from Theano and following by many popular libraries such as Ten-sorFlow, PyTorch. The unique benefits that these libraries bring are: (1) completely transparent training and prediction, (2) modular DNN construction with interchangeable components. A DNN component can vary dramatically in terms of size and complexity, ranging from a single matrix operator, to a highly complex deep neural network such as convolutional neural networks (CNN) or long short-term memory (LSTM). All the components share the same interface, which enables a user to replace one component by another with almost no changes on other parts of the code. In particular, modularity in the programming languages used by deep learning libraries revolutionizes the development, communication and deployment of DNNs. The ability to package a state-of-the-art implementation into a standard module and distribute enables rapid development with much lower chance of having bugs and performance issues comparing with re-implementing from scratch, which makes research results more reproducible. Communication about ongoing research method and results becomes much easier, as running a different experiment or modifying a network only requires to change a few lines of code. The standardized interface of components enables better testing and maintainability. It eases the transition of a research implementation into an industrial software.

Probabilistic programming language (PPL) is a similar domain specific programming language that aims at describing probabilistic models and automating inference in those models. Despite the different of DNNs and probabilistic models in mathematical formulations1 , a similar modularity structure exist in both probabilistic models and DNNs. However, the benefits of modularity has not been fully exploited in currently popular PPLs such as Stan (Carpenter et al., 2017) , Edward (Tran et al., 2016) , PyMC (Salvatier et al., 2016) , which is mostly about re-using plain probabilistic distributions and (configurable) generic inference algorithms such as Markov Chain Monte Carlo (MCMC) or Stochastic Variational Inference (SVI).

In this paper, we propose to further exploit the modularity in probabilistic models by offering not only plain probabilistic distributions but also sophisticated probabilistic model such as Bayesian non-parametric models as fundamental building blocks in a PPL. This would bring the a lot of benefits of modularities that we enjoy with deep learning libraries and improve the usabilities of PPLs on real world problems. We demonstrate this idea by presenting a modular probabilistic programming language MXFusion, which includes a new type of re-usable building blocks for probabilistic models called probabilistic modules. A probabilistic module consists of a set of random variables with associated probabilistic distributions and dedicated inference methods designed specifically for efficiency and accuracy on that set of random variables. This allows to package a sophisticated probabilistic model with their proposed inference methods. The reason for packaging inference methods within a probabilistic module is due to the fact that the exact inference of most of probabilistic models are intractable. Given a specific probabilistic model, a generic approximate inference method that is applicable to a wide range of probabilistic models often performs worse than a dedicated inference method. Most of current PPLs only rely on generic inference methods such as MCMC or black-box variational inference. We propose a framework of probabilistic modules based on variational inference. With the proposed framework, one could seamlessly combine generic inference methods with dedicated inference methods, which bridges the performance gap.

The rest of the paper is organized as follows: Section 2 describes relevant ideas in the literature of PPL; Section 3 provides details of probabilistic modules; Section 4 and ?? presents probabilistic modules in MXFusion; Section 5 demonstrates experiments of presented models with real data.

2. RELATED WORKS:

Probabilistic programming languages typically focus on offering an expressive programming language that describes probabilistic models and automating inference of probabilistic models through a "compilation" step. Due to the intractability of inference in probabilistic models, existing PPLs either focus on the expressiveness and employs a generic inference engine which suffers from scalability issues (Pfeffer, 2001; 2009; Goodman et al., 2012) , or focus on efficiency by restricting down to a specific class of models and inference algorithms (Spiegelhalter et al., 1995; Murphy, 2001; Salvatier et al., 2015; Carpenter et al., 2016) . By exploiting automatic-differentiation and hardware acceleration, Edward (Tran et al., 2017) and Pyro (Goodman, 2017 ) take a middle approach by offering the capability of customized inference methods, which is often referred to as an inference model. Both propose a clear separation between model and inference and typically rely on stochastic variational inference methods such as amortized inference.

The idea of encapsulating models and approximate inference programs in probabilistic modules has also been proposed by Cusumano-Towner & Mansinghka (2017) , in which they demonstrate with Markov Chain Monte Carlo (MCMC) and Sequential Monte Carlo (SMC) inference programs. Venture and Anglican define inference as a collection of program fragments corresponding to local inference problems, although they do not support customizable posterior approximation (Mansinghka et al., 2014; Wood et al., 2014) .

3. MODULARITY IN PROBABILISTIC PROGRAMMING:

DNN has a very nice modularity property that is exploited by deep learning libraries, which is a result of function composition. The formulation of a function can be broken into the composition of several sub-functions, e.g., f (x) = g 1 (• • • g k (x)). This enables a deep learning library to implement building blocks at various granularities as individual sub-functions. Then, a user can easily construct a DNN by putting these building blocks via function composition.

A similar modularity property also exists in probabilistic models, but in sightly different form. The joint probability distribution of a probabilistic model can often be decomposed into the product of the conditional distribution of individual variables with all the latent variables being marginalized out, e.g., p(y|z) = x p(y|x)p(x|z). This allows us to implement these individual conditional distributions as building blocks for a probabilistic model. So far, most of PPLs exploits this modularity and provides standard probabilistic distribution as building blocks. In this paper, we present an approach to provide more sophisticated probabilistic models as building blocks.

Inference for a DNN is mostly based on gradient optimization, in which a gradient can be computed transparently through automatic-differentiation. For a probabilistic model, inference is very chal-lenging, as the exact inference for most of probabilistic models is intractable. Research has been focused on developing approximate inference or sampling methods for resolving this intractability. The literature of approximate inference methods can be broadly categorized into two groups: (1) proposing a better approximate inference method for a specific probabilistic model (2) developing a generic inference algorithm that can be applied to a family of probabilistic models as big as possible, sometimes being referred to as black-box inference. Both types of inference methods have received great attention and been developed by a lot of researchers. Obviously, given a specific model, a specialized inference method will always give a better performance than a generic inference method, as it is able to take advantages of the specific mathematical properties of the given model. Even among generic inference methods, a method only applicable to a small family of models performs better than a more generic method. For example, stochastic variational inference with reparameterization trick that is only applicable to a location-scale family of distributions produces a lower variance estimate of gradient than the score function approach that is applicable to mostly of distribution including discrete distributions.

Currently existing PPLs only use generic inference methods. Although those generic inference methods may be configurable, e.g., providing a customized variational posterior, it often performs worse for the probabilistic models, of which a specialized inference method exists. In MXFusion, we bridge the performance gap between generic and specialized inference methods by further exploiting the modularity of probabilistic models and offering probabilistic modules as building blocks. A probabilistic module encapsulates a sophisticated probabilistic model, e.g., a Gaussian process, and a set of specialized inference algorithms. With the latent variables of a probabilistic module being invisible from outside, a probabilistic module defines a conditional distribution over its exposed variables, which has no difference from a plain probabilistic distribution. A probabilistic module offers the same interface as a plain probabilistic distribution, so that an external inference method can transparently treat a probabilistic module as a plain probabilistic distribution. Internally, the interface of a probabilistic module is realized by calling a pre-specified inference method.

In the following section, we demonstrate the idea of a probabilistic module with variational inference. This idea can also be applied to other inference algorithms.

3.1. PROBABILISTIC MODULE WITH VARIATIONAL INFERENCE:

The concept of hiding the details of inference of a probabilistic module by specialized inference methods is nice. The challenge is that not all the approximate inference method are compatible with each other. In this section, we present an approach to implement probabilistic module with the framework of variational inference.

The main idea of variational inference is to approximate an intractable posterior distribution of latent variables with a parametric approximation, referred to as a variational posterior distribution. VI is often framed as a lower bound of the logarithm of the marginal distribution, e.g,

EQUATION

where p(y|x)p(x|z) forms a probabilistic model with x as a latent variable, q(x|y) is the variational posterior distribution, and the lower bound is denoted as L(y, z). By then taking a natural exponentiation of L(y, z), we get a lower bound of the marginal probability denoted as p(y|z) = e L(y,z) .

Assume we are interested in plugging p(y|z) into another probabilistic model p(l|y)p(y|z) where y is a latent variable. With variational inference, the lower bound of the overall model can be derived as log p(l|z) ≥ y q(y) log p(l|y)p(y|z) q(y) ≥ y q(y) log p(l|y)p(y|z) q(y) ,

where the usual variational lower bound is further lower bounded by replacing p(y|z) with its lower bound p(y|z). By substituting (1) into (2), we then derive the variational lower bound for the whole model,

log p(l|z) ≥ y,x

q(x|y, z)q(y|z) log p(l|y)p(y|x)p(x|z) q(x|y, z)q(y|z) = L(l, z).

This example shows that variational inference has a recursive property that enables inference modularity. A technical challenge with VI is that the integral of the lower bound of a probabilistic module with respect to external latent variables, such as (2), may not always be tractable. Stochastic variational inference (SVI) offers an approximated solution to this new intractability by applying Monte Carlo Integration. Monte Carlo Integration is applicable to generic probabilistic distributions and lower bounds as long as we are able to draw samples from variational posterior.

In this case, the lower bound is approximated as

EQUATION

where y i |z ∼ q(y|z), x j |y i , z ∼ q(x|y i , z) and N is the number of samples of y and M is the number of samples of x given y. Note that if there is a closed form solution of p(y i |z), the calculation of L(y i , z) can be replaced with the closed-form solution.

MXFusion offers modularity via probabilistic modules, which combine the definition of probabilistic distributions and specialized inference methods in a concise interface. If using VI as the primary inference algorithm, probabilistic modules used in a model automatically compute variational lower bounds, such as p(y|z) in the above example.

4. MODEL CONSTRUCTION WITH PROBABILISTIC MODULES:

In PPLs, a probabilistic model is often represented as a graph of random variables, where connections between random variables denotes corresponding probabilistic distributions. In such a representation, a probabilistic distribution is a fundamental re-usable component for constructing various probabilistic models. A probabilistic distribution typically have two reusable operations required by an inference method: drawing a sample of the random variable and computing the logarithm of its probabilistic density/mass function. Then, a user specifies a probabilistic model by defining individual variables and their corresponding probabilistic distributions if they are random variables. Each specified probabilistic distribution includes the form of the distribution and all the variables that the distribution is conditioned on, which corresponds to all the directed connections towards the random variable.

MXFusion follows this design of model specification. Figure 1 implements a Bayesian linear regression (BLR) model using a deep neural network as a feature extractor. There are N data points, each of which contains an input vector x n ∈ R Q and an output variable y n ∈ R. Each input vector is first applied to a DNN, h n = f φ (x n ), where φ denotes the parameters of the DNN. The output variable, y = (y 1 , . . . , y N ), follows a multi-variate normal distribution, y ∼ N (0, H H + σ 2 I), where H is the stack of the outcome of DNN, H = (h 1 , . . . , h N ) . As there are no latent variables in BLR, a MAP inference method is created to estimate the model parameters with maximum likelihood given the observed data, i.e., φ * , σ * = arg max φ,σ log p(y|f φ (X), σ). The maximum likelihood estimate is achieved through optimizing the log-likelihood via a gradient-based optimization algorithm. The log-likelihood of BLR is computed via the log pdf method of the probabilistic distribution of the variable y, which is a multi-variate normal distribution. The gradient of the log-likelihood is obtained through auto-differentiation using MXNet.

A noticeable difference from PyMC and Edward is that all the variables in a probabilistic model are defined as explicit members of a Model class, which offers a clean memory management by avoiding any global states. Like Edward, the information about which variables are observed is not given until inference time. This allows reuse a model specification across multiple inference methods. When applying a machine learning model to a problem, one almost always uses a model with at least two inference methods that have different observed variables.: one for estimating model parameters/posterior distribution from training data and the other for making prediction with new data. Taking BLR as an example, during training, the input variable X and output variable y are both observed and the aim of inference is to estimate the model parameters that maximizes the loglikelihood, while at prediction stage, only the input variable X is observed and the inference is to estimate the output variable with the optimal model parameters p(y|X, σ * , θ * ).

4.1. PROBABILISTIC MODULE AS MODEL COMPONENT:

A stronger modularity for PPL is to enable reusing probabilistic models as building blocks. Ideally, one should be able to construct a probabilistic model by putting together well-known sophisticated probabilistic models such as Gaussian process, just like building a deep neural work with convolutional layers or LSTM layers. MXFusion offers such modularity by encapsulating the specification of a probabilistic model and a set of specialized inference algorithms, which is called a probabilistic module. A probabilistic module implements the same interface that a probabilistic distribution has, i.e., drawing a sample of output random variables and computing the logarithm of the probabilistic density/mass function. By following the same interface, a probabilistic module can be transparently used by a probabilistic model just like an usual probabilistic distribution.

In Figure 2 , we demonstrate the convenience of a probabilistic module by modifying a BLR into deep kernel learning (DKL) (Wilson et al., 2016) . This only requires to change one line of code, which replaces a multi-variate normal distribution with a sparse Gaussian process regression (SGPR) module. The syntax of specifying a probabilistic module is similar to the one for a normal distribution. The input arguments include the input variable to the SGPR module and the choice of kernel function for GP, which is radius basis function (RBF) in this case, and the noise variance for a Gaussian likelihood distribution. The function returns a variable representing the output of the SGPR module, which is assigned as the variable m.y in this example. In (Wilson et al., 2016) , a particular GP approximation, known as KISS-GP, is used for scalability, which is replaced by a variational sparse GP approximation (Titsias, 2009) in this example.

Variational sparse GP speeds up the computation of log-likelihood by replacing it with a variational lower bound. It introduces a set of auxiliary variables u ∈ R M , known as inducing data, and a set of corresponding inducing inputs Z ∈ R M ×Q . The resulting variational lower bound is

log p(y|X, Z, θ) ≥ L SGP (y, X, Z, θ) = f ,u p(f |u, X, Z)q(u) log p(y|f ) ( ( ( ( ( ( p(f |u, X, Z)p(u|Z) ( ( ( ( ( ( p(f |u, X, Z)q(u)

where q(f , u) = p(f |u, X, Z)q(u) is the variational posterior. The above variational lower bound has a closed form solution with a computational complexity O(N M 2 ), which is significantly lower than the cubic complexity of GP, O(N 3 ). The log pdf method of the SGPR module compute the above variational lower bound. The MAP inference method of the DKL model finds the parameters of SGPR θ, Z and the parameters of DNN φ that maximize the lower bound of DKL, i.e., θ * , φ * , Z * = arg max θ,φ,Z L SGP (y, f φ (X), Z, θ). m = Model() m.x = Normal.define_variable(mean=0, variance=1, shape=(N, Q)) m.sigma2 = Variable(shape=(1,), transformation=PositiveTransformation()) m.y = SparseGaussianProcessRegression.define_variable( shape=(N, D), X=m.x, kernel=RBF(Q), noise_var=m.sigma2) q = Posterior(m) q.mu = Variable(shape=(N, Q)) q.S = Variable(shape=(N, Q), transformation=PositiveTransformation()) q.x.assign_factor(Normal(mean=q.mu, variance=q.S)) infr = GradientBasedInference(SVI(m, q, [m.x, m.y])) The variational inference method of SGPR transparently provides a specialized implementation that exploits all the knowledge about the SGPR module and delivers good accuracy and efficiency. To achieve the same task with a standard PPL, one needs to explicit the prior distribution of a sparse Gaussian process regression model and the corresponding variational posterior. Even by doing so, the same quality of inference is not achievable with a generic variational inference method. Stochastic variational inference with reparameterization tricks would be the best choice in this case. The inference result will be slow and has high variance. This is because: (1) The inference algorithm is not able to figure out the cancellation of p(f |u, X, Z) in both the nominator and denominator. A direct evaluation of p(f |u, X, Z) results into a cubic complexity O(N 3 ), which fails to deliver any speed-up comparing with GP. (2) The variational distribution p(f |u, X, Z) is high dimensional R N and highly correlated. The variance of the Monte Carlo integration via sampling will be high.

(3) The optimal value of q(u) in the variational lower bound has a closed form solution, which is exploited by variational sparse GP (Titsias, 2009) . With SVI, one needs to explicitly optimize it with gradient optimization. The speed and quality of the two inference approaches reflects the gap between a generic inference algorithm and a specialized inference algorithm. The probabilistic module aims at bridging this gap.

4.2. NESTED VARIATIONAL INFERENCE:

A probabilistic module can be treated transparently as a probabilistic distribution. Then, it is straightforward to construct a probabilistic model consists of multiple probabilistic modules such as Deep GPs (Dai et al., 2016) . In these models, some of the exposed variables of probabilistic modules is also latent variables. As shown in Section 3.1, as long as a variational inference method is used for the whole probabilistic model, the inference of individual probabilistic modules can be transparently handled.

Figure 3 implements Bayesian Gaussian process latent variable model (BGPLVM) (Titsias & Lawrence, 2010), which is an example of this kind of models. BGPLVM can be constructed by assign the input variable X a Gaussian distribution with zero-mean and unit-variance. As the input variable X is a latent variable, the marginal log-likelihood is not tractable anymore. A variational lower bound can be written as

log p(Y) ≥ X q(X) log p(Y|X)p(X) q(X) ≥ X q(X) log L SGP (y, X, Z, θ)p(X) q(X)

where p(Y|X) is a GP and q(X) = N (X|µ, diag(s)) is the variational posterior of X, which is assumed to be a Gaussian distribution with a diagonal covariance matrix. By applying a variational sparse GP (Titsias, 2009) approximation to p(Y|X), we further lower bound the original lower bound by replace p(Y|X) with the variational sparse GP lower bound mentioned above L SGP (y, X, Z, θ). For the expectation with respect to q(X), we can apply stochastic variational inference (SVI) by drawing samples from q(X). In this way, we result into a nested variational inference combining a generic inference method (SVI) with a specialized inference method (variational sparse GP). Following the same approach, it is also straight-forward to extend BGPLVM into a variationally auto-encoded GPLVM/deep GP (Dai et al., 2016) by parameterizing µ and s in q(X) as the outcome of a DNN. 

5. EXPERIMENTS:

In the previous examples, we demonstrate that various sophisticated probabilistic models can be easily constructed with a few lines of codes by using probabilistic modules. In this section, we evaluate the models that are constructed with probabilistic modules deliver good performance on real data.

5.1. GAUSSIAN PROCESS REGRESSION AND DEEP KERNEL LEARNING:

We evaluate the variational sparse GP (SGP) (Titsias, 2009) module implemented in MXFusion and a variant of deep kernel learning (DKL) (Wilson et al., 2016) shown in Figure 2 , comparing with a standard sparse GP implementation from GPy (2012). All the GP models use RBF kernels. By taking the advantage of GPU acceleration that is provided by MXNet, it allows to scale to a significantly higher number of inducing points. For DKL, we use one hidden layer with 20 hidden units and tanh activation. It outputs a 5-dimensional representation and, then, feeds into our SGP module.

In this experiment, we compare all the methods on four standard regression benchmarks. Naval1 and Naval2 are the same regression dataset with two different regression targets, they have 11934 data points and 16 dimensional inputs. Kin8nm has 8192 data points and 8 dimensional inputs. We used 90% data points for training and 10% for testing. All the datasets are normalized column-wise. We use 50 inducing points for GPy implementation which is a standard choice in the literature, and use 50 and 3200 inducing points for our SGP module. The Adam optimizer (Kingma & Ba, 2015) was used for training in MXFusion, while GPy models are trained with L-BFGS. Results with 10-fold averages measured in terms of root mean squared error (RMSE) and test set log likelihood (TLL) are provided in Table 1 . GPy and MXFusion with 50 inducing points give similar performance for both RMSE and TLL. With 3200 inducing points, it gives significantly better performance on all the datasets for both RMSE and TLL. For DKL, with a relative small number of inducing points (1000), it produces a comparable RMSE with the help of a neural network, but does not give as good TLL.

6. DISCUSSION:

Despite the expressiveness and flexibility provided by a PPL, a major limitation of PPL is that they rely on generic inference algorithms. Recent work such as Edward and Pyro support customizable inference algorithms, in particular, stochastic variational inference with a customized variational posterior. Once an inference algorithm is chosen, it remains the same across a probabilistic model. However, given a specific probabilistic model, e.g., a conjugate model, a specialized inference algorithm that exploits the mathematical properties of that particular model will always produce inference results that are as good or better than the generic inference in terms of both accuracy and efficiency. In practice, when doing probabilistic modeling there is a tradeoff between implementing a specialized inference algorithm for improved speed and performance at the cost of maintainability and flexibility by introducing specialized code for each model. MXFusion aims at closing the gap between having specialized, highly performant algorithms and generic, easily maintained generic algorithms by introducing probabilistic modules. A probabilistic module consists of both a model and inference definition, defined together and wrapped up in a modular, plug-and-play package. This allows specialized inference algorithms for corresponding probabilistic modules to be smoothly integrated into the inference algorithm of a larger probabilistic model. By bringing modularity into probabilistic programming, MXFusion offers a flexible maintainable framework for doing probabilistic modeling while keeping the accuracy and efficiency of specialized inference algorithms.

:

A DNN represents a deterministic function, while a probabilistic model describes a probabilistic distribution.

Figures and Tables:

Figure 1: Bayesian linear regression. The variable net references a deep neural network defined in the MXNet Gluon syntax.

Figure 2: Deep kernel learning(Wilson et al., 2016). The variable net references a deep neural network defined in the MXNet Gluon syntax, of the output dimensionality is 100.

Figure 3: Bayesian Gaussian process latent variable model (Titsias & Lawrence, 2010).

Performance comparison of variational sparse GP (SGP) and deep kernel learning (DKL) on four standard regression benchmarks. SGP † is the GPy implementation. SGP is the SGP module in MXFusion. The implementation of DKL is shown in Figure2. The numbers in the parentheses in the method column show the number of inducing points. Both rooted mean square error (RMSE) and test log-likelihood (TLL) are the mean of measure with 10-fold cross-validation. The number in the parentheses next to performance measure is its standard deviation.

Table content: <table><tr><td colspan="2">Metric Method</td><td>naval1</td><td>naval2</td><td>kin8nm</td><td>power</td></tr><tr><td/><td>SGP  † (50)</td><td colspan="2">3.5e-5 (1.0e-5) 3.1e-4 (9.0e-6)</td><td>0.087 (3.1e-3)</td><td>3.98 (0.19)</td></tr><tr><td>RMSE</td><td>SGP (50)</td><td colspan="2">3.7e-5 (1.1e-5) 3.1e-4 (8.0e-6)</td><td colspan="2">0.089 (2.96e-3) 3.98 (0.19)</td></tr><tr><td/><td>SGP (3200)</td><td colspan="4">0.6e-5 (0.6e-5) 2.97e-4 (5.0e-6) 0.068 (2.48e-3) 3.08 (0.28)</td></tr><tr><td/><td colspan="5">DKL (1000) 2.0e-5 (1.2e-5) 4.85e-4 (5.5e-4) 0.062 (1.51e-3) 3.39 (0.26)</td></tr><tr><td/><td>SGP  † (50)</td><td>8.58 (0.22)</td><td>6.66 (0.03)</td><td>0.98 (0.02)</td><td>-2.80 (0.05)</td></tr><tr><td>TLL</td><td>SGP (50)</td><td>8.69 (0.19)</td><td>6.67 (0.02)</td><td>0.98 (0.02)</td><td>-2.80 (0.05)</td></tr><tr><td/><td>SGP (3200)</td><td>10.70 (0.51)</td><td>6.70 (0.02)</td><td>1.28 (0.04)</td><td>-2.53 (0.10)</td></tr><tr><td/><td colspan="2">DKL (1000) 5.26 (4.35)</td><td>-83.8 (43.3)</td><td>-19.4 (1.73)</td><td>-12.7 (3.84)</td></tr></table>

Peer review (iclr2019):
 In this paper authors present a new Probabilistic Programming Language (PPL) MXFusion. Similarly to the languages for the deep learning (TensorFlow, PyTorch, etc.), this language introduce probabilistic modules that are used as building blocks for complex probabilistic models. Introducing modularity to the probabilistic programming, raises the problem of inference for probabilistic models. Since, we cannot obtain the exact solution on practice we have to resort to approximate inference methods. The approximate inference methods can be either generic, thus, being suitable for many probabilistic models but resulting in poor approximation, or specific, thus, having good approximation quality, but only for specific probabilistic models. Authors propose to address this problem by encapsulating specific inference methods in corresponding probabilistic modules. Doing so, one can perform approximate inference for every module with the best suitable inference technique. Authors demonstrate interface of MXFusion for three well known probabilistic models: Bayesian linear regression, deep kernel learning, Bayesian Gaussian process latent variable model. Approaching the problem of building complex probabilistic models by introducing modular PPL is an important direction of study. But, regarding this paper I have the following concerns. - In my opinion, the structure of the paper can be greatly improved. From general words about modularity and approximate inference authors dive to the very specific cases of probabilistic models. Following such structure, authors don’t give a clear answer to the following questions. Why the paradigm of encapsulating inference methods in probabilistic modules is legitimate for constructing complex probabilistic models? What inference methods and probabilistic models can we use as building blocks? Do we need to be aware of specific inference methods that are encapsulated or we can use any blocks in any order as we do in deep learning frameworks? - Novelty of that paper is the new design of PPL. That is an interesting and important question for the community, but maybe ICLR paper is not the best format to present such kind of novelty. - From the specific examples in the paper, legitimacy of such modular structure is clear only for variational inference (that seems to be a common knowledge) and variational approximation of gaussian processes. But the application area of MXFusion remains unclear. Verbatim examples of code for the specific examples doesn’t make the difference between MXFusion and other PPLs clear, because it can be treated as encapsulation of the code into some classes, that can be implemented in other languages as well. - Comparison with other frameworks can be improved. In experimental section authors provide comparison with GPy framework in terms of RMSE and log-likelihood for gaussian process with 50 inducing points. As I understood both frameworks use the same inference methods and achieve the same performance, so the experiment can be considered as sanity check for MXFusion. The paper could benefit from comparison between different inference methods and providing benchmarks for inference time. Overall, the paper proposes a new PPL that is an important direction of study, but have several drawbacks and conference paper format is not the best way to present such kind of novelty. Typos: - Page 1, “despite the different of DNNs…” -> “despite the difference of DNNs…”? - Page 2, missing reference of the section - Page 2, section 3, “... sightly different form.” -> “... slightly different form”?