Modular Deep Probabilistic Programming

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 TensorFlow , 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 1A DNN represents a deterministic function , while a probabilistic model describes a probabilistic distribution . 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 ) = g1 ( · · · gk ( 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 , log p ( y|z ) = log ∫ x p ( y|x ) p ( x|z ) ≥ ∫ x q ( x|y , z ) log p ( y|x ) p ( x|z ) q ( x|y , z ) = L ( y , z ) , ( 1 ) 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 ) = eL ( 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 ) , ( 2 ) 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 ) . ( 3 ) 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 L ( l , z ) ≈ 1 N ∑ i log p ( l|yi ) eL ( yi , z ) q ( yi|z ) , L ( yi , z ) ≈ 1 M ∑ j log p ( yi|xj ) p ( xj |z ) q ( xj |yi , z ) , ( 4 ) where yi|z ∼ q ( y|z ) , xj |yi , z ∼ q ( x|yi , z ) andN is the number of samples of y andM is the number of samples of x given y . Note that if there is a closed form solution of p̃ ( yi|z ) , the calculation of L ( yi , 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 .

Summary: The paper presents an extension of the MXFusion language that allows the use of probabilistic modules. These modules are defined as a set of random variables and a specific probabilistic distribution. The modules also contain dedicated inference methods. Using these modules, one can use probabilistic distributions with inference methods tailored to the distribution, which are usually more efficient than generic inference systems.