Meta-Learning an Inference Algorithm for Probabilistic Programs

1 INTRODUCTION . One key objective of probabilistic programming is to automate reasoning about probabilistic models from diverse domains ( Ritchie et al. , 2015 ; Perov & Wood , 2016 ; Baydin et al. , 2019 ; Schaechtle et al. , 2016 ; Cusumano-Towner et al. , 2017 ; Saad & Mansinghka , 2016 ; Kulkarni et al. , 2015 ; Young et al. , 2019 ; Jäger et al. , 2020 ) . As a way to realize this goal , researchers have extensively worked on the development of posterior-inference or parameter-learning algorithms that are efficient and universal ; the algorithms can be applied to all or nearly all models written in probabilistic programming languages ( PPLs ) . This line of research has led to performant probabilistic programming systems ( Goodman et al. , 2008 ; Wood et al. , 2014 ; Mansinghka et al. , 2014 ; Minka et al. , 2018 ; Narayanan et al. , 2016 ; Salvatier et al. , 2016 ; Carpenter et al. , 2017 ; Tran et al. , 2016 ; Ge et al. , 2018 ; Bingham et al. , 2018 ) . Yet , it also revealed the difficulty of achieving efficiency and universality simultaneously , and the need for equipping PPLs with mechanisms for customising inference or learning algorithms to a given domain . In fact , recent PPLs include constructs for specifying conditional independence in a model ( Bingham et al. , 2018 ) or defining proposal or variational distributions ( Ritchie et al. , 2015 ; Siddharth et al. , 2017 ; Bingham et al. , 2018 ; Tran et al. , 2018 ; Cusumano-Towner et al. , 2019 ) , all enabling users to help inference or learning algorithms . In this paper , we explore a different approach . We present a meta-algorithm for learning a posteriorinference algorithm itself from a given set of restricted probabilistic programs , which specifies a class of probabilistic models , such as hierarchical or clustering models . The meta-algorithm aims at constructing a customised inference algorithm for the given set of models , while ensuring universality to the extent that the constructed algorithm can generalise : it works well for models not in the training set , as long as the models are similar to the ones in the set . The distinguished feature of our approach is the use of what we call a white-box inference algorithm , which extracts information directly from model descriptions themselves , given as programs in a PPL . Concretely , our white-box inference algorithm is equipped with multiple neural networks , one for each type of atomic command in a PPL , and computes an approximate posterior for a given program by analysing ( or executing in a sense ) individual atomic commands in it using these networks . For instance , given the probabilistic program in Fig . 1 , which describes a simple model on the Milky Way mass ∼ N ( 5 , 10 ) ; // log of the mass of Milky Way g1 ∼ N ( mass × 2 , 5 ) ; obs ( N ( g1 , 1 ) , 10 ) ; // observed velocity vel1=10 of the first satellite galaxy g2 ∼ N ( mass +5 , 2 ) ; obs ( N ( g2 , 1 ) , 3 ) // observed velocity vel2=3 of the second satellite galaxy Figure 1 : Probabilistic program for a model for Milky Way and its two satellite galaxies . The obs statements refer to the observations of ( unnamed ) random variables vel1 and vel2 . galaxy , the white-box inference algorithm analyses the program as if an RNN handles a sequence or an interpreter executes a program . Roughly , the algorithm regards the program as a sequence of the five atomic commands ( separated by the “ ; ” symbol ) , initialises its internal state h ∈ Rm with h0 , and transforms the state over the sequence . The internal state h is the encoding of an approximate posterior at the current program point , which corresponds to an approximate filtering distribution of a state-space model . How to update this state over each atomic command is directed by neural networks . Our meta-algorithm trains the parameters of these networks by trying to make the inference algorithm compute accurate posterior approximations over a training set of probabilistic programs . One can also view our white-box inference algorithm as a message-passing algorithm in a broad sense where transforming the internal state h corresponds to passing a message , and understand our meta-algorithm as a method for learning how to pass a message for each type of atomic commands . This way of exploiting model descriptions for posterior inference has two benefits . First , it ensures that even after customisation through the neural-network training , the inference algorithm does not lose its universality and can be applied to any probabilistic programs . Thus , at least in principle , the algorithm has a possibility to generalise beyond the training set ; its accuracy degrades gracefully as the input probabilistic program diverges from those in the training set . Second , our way of using model descriptions guarantees the efficiency of the inference algorithm ( although it does not guarantee the accuracy ) . The algorithm scans the input program only once , and uses neural networks whose input dimensions are linear in the size of the program . As a result , its time complexity is quadratic over the size of the input program . Of course , the guaranteed speed also indicates that the customisation of the algorithm for a given training set , whose main goal is to achieve good accuracy for probabilistic programs in the set , is a non-trivial process . Our contributions are as follows : ( i ) we present a white-box posterior-inference algorithm , which works directly on model description and can be customised to a given model class ; ( ii ) we describe a meta-algorithm for learning the parameters of the inference algorithm ; ( iii ) we empirically analyse our approach with different model classes , and show the promise as well as the remaining challenges . Related work The difficulty of developing an effective posterior-inference algorithm is well-known , and has motivated active research on learning or adapting key components of an inference algorithm . Techniques for adjusting an MCMC proposal ( Andrieu & Thoms , 2008 ) or an HMC integrator ( Hoffman & Gelman , 2014 ) to a given inference task were implemented in popular tools . Recently , methods for meta-learning these techniques themselves from a collection of inference tasks have been developed ( Wang et al. , 2018 ; Gong et al. , 2019 ) . The meta-learning approach also features in the work on stochastic variational inference where a variational distribution receives information about each inference task in the form of its dataset of observations and is trained with a collection of datasets ( Wu et al. , 2020 ; Gordon et al. , 2019 ; Iakovleva et al. , 2020 ) . For a message-passingstyle variational-inference algorithm , such as expectation propagation ( Minka , 2001 ; Wainwright & Jordan , 2008 ) , Jitkrittum et al . ( 2015 ) studied the problem of learning a mechanism to pass a message for a given single inference task . A natural follow-up question is how to meta-learn such a mechanism from a dataset of multiple inference tasks that can generalise to unseen models . Our approach provides a partial answer to the question ; our white-box inference algorithm can be viewed as a message-passing-style variational inference algorithm that can meta-learn the representation of messages and a mechanism for passing them for given probabilistic programs . Amortised inference and inference compilation ( Gershman & Goodman , 2014 ; Le et al. , 2017 ; Paige & Wood , 2016 ; Stuhlmüller et al. , 2013 ; Kingma & Welling , 2013 ; Mnih & Gregor , 2014 ; Rezende et al. , 2014 ; Ritchie et al. , 2016 ; Marino et al. , 2018 ) are closely related to our approach in that they also attempt to learn a form of a posterior-inference algorithm . However , the learnt algorithm by them and that by ours have different scopes . The former is designed to work for unseen inputs or observations of a single model , while the latter for multiple models with different structures . The relationship between these two algorithms is similar to the one between a compiled program ( to be applied to multiple inputs ) and a compiler ( to be used for multiple programs ) . u : = 0 ; v : = 5 ; w : = 1 ; z1 ∼N ( u , v ) ; z2 ∼N ( u , v ) ; z3 ∼N ( u , w ) ; µ3 : = if ( z3 > u ) z1 else z2 ; obs ( N ( µ3 , w ) , −1.9 ) ; // x1∼N ( µ3 , w ) , x1 =−1.9 z4 ∼N ( u , w ) ; µ4 : = if ( z4 > u ) z1 else z2 ; obs ( N ( µ4 , w ) , −2.2 ) ; // x2∼N ( µ4 , w ) , x2 =−2.2 z5 ∼N ( u , w ) ; µ5 : = if ( z5 > u ) z1 else z2 ; obs ( N ( µ5 , w ) , 2.4 ) ; // x3∼N ( µ5 , w ) , x3 =2.4 z6 ∼N ( u , w ) ; µ6 : = if ( z6 > u ) z1 else z2 ; obs ( N ( µ6 , w ) , 2.2 ) // x4∼N ( µ6 , w ) , x4 =2.2 Figure 2 : Probabilistic program for a simple clustering model on four data points . The idea of running programs with learnt neural networks also appears in the work on training neural networks to execute programs ( Zaremba & Sutskever , 2014 ; Bieber et al. , 2020 ; Reed & de Freitas , 2016 ) . As far as we know , however , we are the first to frame the problem of learning a posterior-inference algorithm as the one of learning to execute . 2 SETUP . Our results assume a simple probabilistic programming language without loop and with a limited form of conditional statement . The syntax of the language is given by the following grammar , where r represents a real number , z and vi variables storing a real , and p the name of a procedure taking two real-valued parameters and returning a real number : Programs C : := A | C1 ; C2 Atomic Commands A : := z ∼ N ( v1 , v2 ) | obs ( N ( v0 , v1 ) , r ) | v0 : = if ( v1 > v2 ) v3 else v4 | v0 : = r | v0 : = v1 | v0 : = p ( v1 , v2 ) Programs in the language are constructed by sequentially composing atomic commands . The language supports six types of atomic commands . The first type is z ∼ N ( v1 , v2 ) , which draws a sample from the normal distribution with mean v1 and variance v2 , and assigns the sampled value to z . The second command , obs ( N ( v0 , v1 ) , r ) , states that a random variable is drawn from N ( v0 , v1 ) and its value is observed to be r. The next is a restricted form of a conditional statement that selects one of v3 and v4 depending on the result of the comparison v1 > v2 . The following two commands are different kinds of assignments , one for assigning a constant and the other for copying a value from one variable to another . The last atomic command v0 : = p ( v1 , v2 ) is a call to one of the known deterministic procedures , which may be standard binary operations such as addition and multiplication , or complex non-trivial functions that are used to build advanced , non-conventional models . When p is a standard binary operation , we use the usual infix notation and write , for example , v1+ v2 , instead of + ( v1 , v2 ) . We permit only the programs where a variable does not appear more than once on the left-hand side of the : = and ∼ symbols . This means that no variable is updated twice or more , and it corresponds to the so-called static single assignment assumption in the work on compilers . This restriction lets us regard variables updated by ∼ as latent random variables . We denote those variables by z1 , . . . , zn . We use this simple language for two reasons . First , the restriction imposed on our language enables the simple definition of our white-box inference algorithm . The language supports only a limited form of conditional statements and restricts the syntactic forms of atomic commands ; the arguments to a normal distribution or to a procedure p should be variables , not general expression forms such as addition of two variables . As we will show soon , this restriction makes it easy to exploit information about the type of each atomic command in our inference algorithm ; we use different neural networks for different types of atomic commands in the algorithm . Second , the language is intended to serve as an intermediate language of a compiler for a high-level PPL , not the one to be used directly by the end user . The compilation scheme in , for instance , §3 of ( van de Meent et al. , 2018 ) from high-level probabilistic programs with general conditional statements and for loops to graphical models can be adopted to compile such programs into our language . See Appendix A for further discussion . Fig . 2 shows a simple model for clustering four data points { −1.9 , −2.2 , 2.4 , 2.2 } into two clusters , where the cluster assignment of each data point is decided by thresholding a sample from the standard normal distribution . The variables z1 and z2 store the centers of the two clusters , and z3 , . . . , z6 hold the random draws that decide cluster assignments for the data points . See Appendix B for the Milky Way example in Fig . 1 compiled to a program in our language . Probabilistic programs in the language denote unnormalised probability densities over Rn for some n. Specifically , for a program C , if z1 , . . . , zn are all the variables assigned by the sampling statements zi ∼ N ( . . . ) in C in that order and C contains m observe statements with observations r1 , . . . , rm , then C denotes an unnormalised density pC over the real-valued random variables z1 , . . . , zn : pC ( z1 : n ) = pC ( x1 : m= r1 : m|z1 : n ) × ∏n i=1 pC ( zi|z1 : i−1 ) , where x1 , . . . , xm are variables not appearing in C and are used to denote observed variables . This density is defined inductively over the structure of C. See Appendix C for details . The goal of our white-box inference algorithm is to compute efficiently accurate approximate posterior and marginal likelihood estimate for a given C ( that is , for the normalised version of pC and the normalising constant of pC ) , when pC has a finite non-zero marginal likelihood and , as a result , a well-defined posterior density . We next describe how the algorithm attempts to achieve this goal .

Summary: The paper presents an algorithm for composing inference algorithms out of simpler neural net building blocks, one per unique statement type in the probabilistic programming language presented. The language is simple without recursion or loops, reducing issues due to feedback problems from the approximation. The networks are trained using HMC or importance sampling samples from many programs in a similar space. There is an empirical study of several classes of small Gaussian probabilistic programs.