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: This introduces a meta-learning algorithm for learning inference algorithms applicable to any probabilistic program. This is accomplished by associating a neural network with every grammar rule of a probabilistic programming language and outputting posterior draws. This generalisation is possible because each neural network component is feed marginal likelihood information for each PPL instruction.