Year: 2021

Abstract:
We present a meta-algorithm for learning a posterior-inference algorithm for restricted probabilistic programs. Our meta-algorithm takes a training set of probabilistic programs that describe models with observations, and attempts to learn an efficient method for inferring the posterior of a similar program. A key feature of our approach is the use of what we call a white-box inference algorithm that extracts information directly from model descriptions themselves, given as programs. Concretely, our white-box inference algorithm is equipped with multiple neural networks, one for each type of atomic command, and computes an approximate posterior of a given probabilistic program by analysing individual atomic commands in the program using these networks. The parameters of the networks are learnt from a training set by our meta-algorithm. We empirically demonstrate that the learnt inference algorithm generalises well to programs that are new in terms of both parameters and model structures, and report cases where our approach achieves greater test-time efficiency than alternative approaches such as HMC. The overall results show the promise as well as remaining challenges of our approach.

heading: 1 INTRODUCTION
text: 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.

heading: 2 SETUP
text: 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.

heading: 3 WHITE-BOX INFERENCE ALGORITHM
text: Given a program C = (A1; . . . ;Ak), our white-box inference algorithm views C as a sequence of its constituent atomic commands (A1, A2, . . . , Ak), and computes an approximate posterior and a marginal likelihood estimate for C by sequentially processing the Ai’s. Concretely, the algorithm starts by initialising its internal state to h0 = ~0 ∈ Rs and the current marginal-likelihood estimate to Z0 = 1. Then, it updates these two components based on the first atomic command A1 of C. It picks a neural network appropriate for the type of A1, applies it to h0 and gets a new state h1 ∈ Rs. Also, it updates the marginal likelihood estimate to Z1 by analysing the semantics of A1. This process is repeated for the remaining atomic commands A2, A3, . . . , Ak of C, and eventually produces the last state hk and estimate Zk. Finally, the state hk gets decoded to a probability density on the latent variables of C by a neural network, which together with Zk becomes the result of the algorithm.
Formally, our inference algorithm is built on top of three kinds of neural networks: the ones for transforming the internal state h ∈ Rs of the algorithm; a neural network for decoding the internal states h to probability densities; and the last neural network for approximately solving integration questions that arise from the marginal likelihood computation in observe statements. We present these neural networks for the programs that sample n-many latent variables z1, . . . , zn, and use at most m-many variables (so m ≥ n). Let V be [0, 1]m, the space of the one-hot encodings of those m variables, and P the set of procedure names. Our algorithm uses the following neural networks:
nnsa,φ1 :V3×Rs→Rs, nnob,φ2 :V2×R×Rs→Rs, nn if,φ3 :V5×Rs→Rs, nnc:=,φ4 :V×R×R s→Rs, nnv:=,φ5 :V 2×Rs→Rs, nnp,φp :V3×Rs→Rs for p ∈ P, nnde,φ6 :Rs→ (R×R)n, nn intg,φ7 :V2×R×Rs→R,
where φ1:7 and φp for p ∈ P are network parameters. The top six networks are for the six types of atomic commands in our language. For instance, when an atomic command to analyse next is a sample statement z ∼ N (v1, v2), the algorithm runs the first network nnsa on the current internal state h, and obtains a new state h′ = nnsa,φ1(z, v1:2, h), where z and v1:2 mean the one-hot encoded variables z, v1 and v2. The next nnde,φ6 is a decoder of the states h to probability densities over the latent variables z1, . . . , zn, which are the product of n independent normal distributions. The network maps h to the means and variances of these distributions. The last nn intg,φ7 is used when our algorithm updates the marginal likelihood estimate based on an observe statement obs(N (v0, v1), r). When we write the meaning of this observe statement as the likelihoodN (r; v0, v1), and the filtering distribution for v0 and v1 under (the decoded density of) the current state h as ph(v0, v1),1 the last neural network computes the following approximation: nn intg,φ7(v0:1, r, h) ≈ ∫ N (r; v0, v1)ph(v0, v1)dv0dv1. See Appendix D for the full derivation of the marginal likelihood.
Given a program C = (A1; . . . ;Ak) that draws n samples (and so uses latent variables z1, . . . , zn), the algorithm approximates the posterior and marginal likelihood of C as follows:
INFER(C) = let (h0, Z0) = (~0, 1) and (hk, Zk) = (INFER(Ak) ◦ . . . ◦ INFER(A1))(h0, Z0) in let ((µ1, σ21), . . . , (µn, σ2n)) = nnde,φ6(hk) in return (∏n i=1N (zi | µi, σ2i ), Zk ) ,
where INFER(Ai) : Rs × R→ Rs × R picks an appropriate neural network based on the type of Ai, and uses it to transform h and Z:
INFER(obs(N (v0, v1), r))(h, Z) = (nnob(v0:1, r, h), Z ×nn intg(v0:1, r, h)), 1The ph(v0, v1) is a filtering distribution, not prior.
INFER(v0 := if (v1>v2) v3 else v4)(h, Z) = (nn if(v0:4, h), Z), INFER(v0 := r)(h, Z) = (nn c :=(v0, r, h), Z), INFER(z∼N (v1, v2))(h, Z) = (nnsa(z, v1:2, h), Z), INFER(v0 := v1)(h, Z) = (nn v :=(v0:1, h), Z), INFER(v0 := p(v1, v2))(h, Z) = (nnp(v0:2, h), Z).
We remind the reader that v0:k refers to the sequence of the one-hot encodings of variables v0, . . . , vk. For the update of the state h, the subroutine INFER(A) relies on neural networks. But for the computation of the marginal likelihood estimate, it exploits prior knowledge that non-observe commands do not change the marginal likelihood (except only indirectly by changing the filtering distribution), and keeps the input Z for those atomic commands.

heading: 4 META-LEARNING PARAMETERS
text: The parameters of our white-box inference algorithm are learnt from a collection of probabilistic programs in our language. Assume that we are given a training set of programs D = {C1, . . . , CN} such that each Ci samples n latent variables z1, . . . , zn and uses at most m variables. Let φ = (φ1:7, (φp)p∈P) be the parameters of all the neural networks used in the algorithm. We learn these parameters by solving the following optimisation problem:2
argmin φ ∑ C∈D KL[πC(z1:n)||qC(z1:n)] + λ 2 (NC − ZC)2
where λ > 0 is a hyper-parameter, NC is the marginal likelihood (or the normalising constant)∫ pC(z1:n)dz1:n for pC , the next πC(z1:n) is the normalised posterior pC(z1:n)/NC for C, and the last qC and ZC are the approximate posterior and marginal likelihood estimate computed by the inference algorithm (that is, (qC(z1:n), ZC) = INFER(C)). Note that qC and ZC both depend on φ, since INFER uses the φ-parameterised neural networks.
We optimise the objective by stochastic gradient descent. The key component of the optimisation is a gradient estimator derived as follows: (∇φ ∑ C∈D KL[πC ||qC ] + λ 2 (NC − ZC) 2) =
( ∑ C∈D Ez1:n∼πC [−∇φ log qC(z1:n)]−λ(NC−ZC)∇φZC) ≈ ∑ C∈D −L̂C,φ−λ(N̂C−ZC)∇φZC . Here L̂C,φ and N̂C are sample estimates of Ez1:n∼πC [∇φ log qC(z1:n)] and the marginal likelihood, respectively. Both estimates can be computed using standard Monte-Carlo algorithms. For instance, we can run the self-normalising importance sampler with prior as proposal, and generate weighted samples {(w(j), z(j)1:n)}1≤j≤M for the unnormalised posterior pC . Then, we can use these samples to compute the required estimates: N̂C = 1M ∑M j=1 w (j) and L̂C,φ = 1M ∑M j=1(w
(j)∇φ log qC(z(j)1:n))/N̂C . Alternatively, we may run Hamiltonian Monte Carlo (HMC) (Duane et al., 1987) to generate posterior samples, and use those samples to draw weighted importance samples using, for instance, the layered adaptive importance sampler (Martino et al., 2017). Then, we compute L̂C,φ using posterior samples, and N̂C using weighted importance samples. Note that neither πC in Ez1:n∼πC [−∇φ log qC(z1:n)] nor NC depends on the parameters φ. Thus, for each C ∈ D, NC needs to be estimated only once throughout the entire optimisation process, and the posterior samples from πC need to be generated only once as well. We use this fact to speed up the computation of each gradient-update step.

heading: 5 EMPIRICAL EVALUATION
text: An effective meta-algorithm should generalise well: the learnt inference algorithm should accurately predict the posteriors of programs unseen during training which have different parameters (§5.1) and model structures (§5.2), as long as the programs are similar to those in the training set. We empirically show that our meta-algorithm learns such an inference algorithm, and that in some cases using the learnt inference algorithm achieves higher test-time efficiency than alternative approaches such as HMC (Duane et al., 1987) (§5.3). We implemented our inference algorithm and meta-algorithm using ocaml-torch (Mazare, 2018), an OCaml binding for PyTorch. For HMC, we used the Python interface for Stan (Carpenter et al., 2017). We used a Ubuntu server with Intel(R) Xeon(R) Gold 6234 CPU @ 3.30GHz with 16 cores, 32 threads, and 263G memory. See Appendix E for the full list of our model classes and their details, and Appendix F for the detailed experimental setup.
2Strictly speaking, we assume that the marginal likelihood of any C ∈ D is non-zero and finite.

heading: 5.1 GENERALISATION TO NEW MODEL PARAMETERS AND OBSERVATIONS
text: We evaluated our approach with six model classes: (1) Gaussian models (gauss) with a latent variable and an observation where the mean of the Gaussian likelihood is an affine transformation of the latent; (2) hierarchical models with three hierarchically structured latent variables (hierl); (3) hierarchical or multi-level models with both latent variables and data structured hierarchically (hierd) where data are modelled as a regression of latent variables at different levels; (4) clustering models (cluster) where five observations are clustered into two groups; (5) Milky Way models (milky), and their multiple-observations extension (milkyo) where five observations are made for each satellite galaxy; and (6) models (rb) using the Rosenbrock function,3 which is expressed as an external procedure, to show that our approach can in principle handle models with non-trivial computation blocks.
The purpose of our evaluation is to show the feasibility of our approach, not to develop the state-ofthe-art inference algorithm automatically, and also to identify the challenges of the approach. These models are chosen for this purpose. For instance, an inference algorithm should be able to reason about affine transformations and Gaussian distributions (for gauss), and dependency relationships among variables (for hierl and hierd) to compute a posterior accurately. Successful outcomes in the classes indicate that our approach learns an inference algorithm with such capacity in some cases.
Setup For each model class, we used 400 programs to meta-learn an inference algorithm, and then applied the learnt algorithm to 50 unseen test programs. We measured the average test loss over the 50 test programs, and checked if the loss also decreases when the training loss decreases. We also compared the marginal posteriors predicted by our learnt inference algorithm with the reference marginal posteriors that were computed analytically, or approximately by HMC. When we relied on HMC, we computed the marginal sample means and standard deviations using one of the 10 Markov chains generated by independent HMC runs. Each chain consisted of 500K samples after 50K warmups. We ensured the convergence of the chains using diagnostics such as R̂ (Gelman et al., 1992). All training and test programs were automatically generated by a random program generator. This generator takes a program class and hyperparameters (e.g., boundaries of the quantities that are used to specify the models), and returns programs from the class randomly (see Appendix E).
For each training program, our meta-algorithm used 215 samples from the analytic (for gauss) or approximate (for the rest, by HMC) posterior distribution for the program.4 Similarly, our meta-
3The function is often used to evaluate learning and inference algorithms (Goodman & Weare, 2010; Wang & Li, 2018; Pagani et al., 2019)
4Except for rb; see the discussion on Rosenbrock models in Appendix H.
algorithm computed the marginal likelihood analytically (for gauss) or approximately (for the rest) using layered adaptive importance sampling (Martino et al., 2017) where the proposals were defined by an HMC chain. We performed mini-batch training; a single gradient update was done with a training program and a mini batch of size 212 (out of 215 samples for the program). We used Adam (Kingma & Ba, 2015) with its hyperparameters {β1 = 0.9, β2 = 0.999, weight_decay = 0}, and the initial learning rate was set to 0.001. When the average training loss converged enough, the training stopped. We repeated the same experiments three times using different random seeds.
Results Fig. 3 shows the training and test losses for gauss, hierl, and milky under three random seeds. The losses for the other model classes are in Appendix G. The training loss was averaged over the training set and 8 batch updates, and the test loss over the test set. The plots for training and test losses are drawn as solid and dotted lines, respectively, and the results with different random seeds are coloured differently. The training losses in all three experiments decreased rapidly, and more importantly, these decreases were accompanied by the downturns of the test losses, which shows that the learnt parameters generalised to the test programs well. The later part of Fig. 3c shows cases where the test loss increases. This was because the loss of only a few programs in the test set (of 50 programs) became large. Even in this situation, the losses of the rest remained small.
Fig. 4 compares, for 10 test programs in hierl, the reference marginal posteriors (blue) and their predicted counterparts (red) by the learnt inference algorithm instantiated at three different training epochs. The predicted marginals were initially around zero (Fig. 4a), evolved to cover the reference marginals (Fig. 4b), and finally captured them precisely in terms of both mean and standard deviation for most of the variables (Fig. 4c). The results show that our meta-algorithm improves the parameters of our inference algorithm, and eventually finds optimal ones that generalise well. We observed similar patterns for the other model classes and random seeds, except for cluster and rb; programs from these classes often have multimodal posteriors, and we provide an analysis for them in Appendix H.

heading: 5.2 GENERALISATION TO NEW MODEL STRUCTURES
text: We let two kinds of model structure vary across programs: the dependency (or data-flow) graph for the variables of a program and the position of a nonlinear function in the program. Specifically, we considered two model classes: (1) models (ext1) with three Gaussian variables and one deterministic variable storing the value of the function nl(x) = 50/π × arctan(x/10), where the models have 12 different types — four different dependency graphs of the variables, and three different positions of the deterministic nl variable for each of these graphs; and (2) models (ext2) with six Gaussian variables and one nl variable, which are grouped into five types based on their dependency graphs. The evaluation was done for ext1 and ext2 independently as in §5.1, but here each of ext1 and ext2
itself has programs of multiple (12 for ext1 and 5 for ext2) model types. See Fig. 10 and 11 in the appendix for visualisation of the different model types in ext1 and ext2, respectively.
Setup For ext1, we ran seven different experiments. Three of them evaluated generalisation to unseen positions of the nl variable, and the other four to unseen dependency graphs. Let Ti,j be the type in ext1 that corresponds to the i-th position of nl and the j-th dependency graph, and T−i,∗ be all the types in ext1 that correspond to any nl positions except the i-th and any of four dependency graphs. For generalisation to the i-th position of nl (i = 1, 2, 3), we used programs from T−i,∗ for training and those from Ti,∗ for testing. For generalisation to the j-th dependency graph (j = 1, 2, 3, 4), we used programs from T∗,−j for training and those from T∗,j for testing. For ext2, we ran five different experiments where each of them tested generalisation to an unseen dependency graph after training with the other four dependency graphs. All these experiments were repeated three times under different random seeds. So, the total numbers of experiment runs were 21 (= 7 × 3) and 15 (= 5× 3) for ext1 and ext2, respectively. In each experiment run for ext1, we used 720 programs for training, and 90 (when generalising to new graphs) or 100 (when generalising to new positions of the nl variable) unseen programs for testing. In each run for ext2, we used 600 programs for training and tested the learnt inference algorithm on 50 unseen programs. We ran HMC to estimate posteriors and marginal likelihoods, and used 200K samples after 10K warmups to compute reference posteriors. We stopped training after giving enough time for convergence within a limit of computational resources. The rest was the same as in §5.1.
Results Fig. 5 shows the average training and test losses for generalisation to the first three dependency graphs in ext1. Fig. 6 shows the losses for generalisation to the first three dependency graphs in ext2. The losses for generalisation to the last dependency graph and to all positions of the nl variable in ext1, and those for generalisation to the 4th and 5th dependency graphs in ext2 are in Appendices I and J. In 17 runs (out of 21) for ext1, the decrease in the training losses eventually stabilised or reduced the test losses, even when the test losses were high and fluctuated in earlier training epochs. In 8 runs (out of 15) for ext2, the test losses were stabilised as the training losses decreased. In 4 runs out of the other 7, the test losses increased only slightly. In terms of predicted posteriors, we observed highly accurate predictions in 8 runs for ext1. For ext2, the predicted posteriors were accurate in 7 runs. For quantified accuracy, we refer the reader to Appendix K. Overall, the learnt algorithms generalised to unseen types of models well or fairly well in many cases.

heading: 5.3 TEST-TIME EFFICIENCY IN COMPARISON WITH ALTERNATIVES
text: We demonstrate the test-time efficiency of our approach using three-variable models (mulmod) where two latent variables follow normal distributions and the other stores the value of the function mm(x) = 100×x3/(10+x4). The models are grouped into three types defined by their dependency graphs and the positions of mm in the programs (see Fig.12 in the appendix). We ran our metaalgorithm using 600 programs from all three types using importance samples (not HMC samples).
Then for 60 test programs from the last model type, we measured ESS and the sum of second moments along the wall-clock time using three approaches: importance sampling (IS-pred; ours) with the predicted posteriors as proposal using 70K samples, importance sampling (IS-prior) with prior as proposal using 100K samples, and HMC with 1M samples after 500 warmups. As the reference sampler, we used importance sampling (IS-ref) with prior as proposal using 5M samples. All the approaches were repeated 10 times.
Table 1 shows the average ESS per unit time over the 60 test programs, by the three approaches. For HMC, “ESS” is the ESS computed using 10 Markov chains averaged over the 60 programs, and “ESS / sec” is the ESS per unit time, averaged over the programs. For IS-{pred, prior}, “ESS” and “ESS / sec” are the average ESS and ESS per unit time, respectively, both over the 10 trials and the 60 programs. GM is the geometric mean, and Q1 and Q3 are the first and third quartiles, respectively. We used the geometric mean, since the ESSes had outliers. The results show that IS-pred achieved the highest ESS per unit time in terms of both mean (GM) and the quartiles (Q1 and Q3).
We provide further analysis for a test program (pgm19; see Appendix L). Fig. 7a and 7b show the moments estimated by HMC and IS-{pred, prior}, respectively, in comparison with the same (across the two figures) reference moments by IS-ref. The estimates by IS-pred (red) quickly converged to the reference (green) within 18ms, while those by HMC (orange) did not converge even after 84s. IS-pred (red) and IS-prior (blue) tended to produce better estimates as the elapsed time increased, but each time, IS-pred estimated the moments more precisely with a smaller variance than IS-prior. In the same runs of the three approaches as in the last columns of Fig. 7a and 7b, IS-pred produced over 16K effective samples in 18ms, while HMC generated only 80 effective samples even after 84s. Similarly, IS-prior generated fewer than 1.4K effective samples in the approximately same elapsed time as in IS-pred. In fact, Fig. 7c shows that as the time increases, the gap between the ESSes of IS-pred and IS-prior gets widen, because the former increases at a rate significantly higher than the latter. Note that IS-pred has to scan a program twice at test time, once for computing the proposal and another for IS with the predicted proposal. See Appendix M for discussion.
Our manual inspection revealed that the programs in mulmod often have multimodal posteriors. Fig. 8a shows the posteriors for {pgm19, pgm30, pgm38} in the test set, computed by our learnt inference algorithm (without IS), HMC (200K samples after 10K warmups) , and IS-ref. The variable z0 in the three programs had multimodal posteriors. For pgm19, the learnt inference algorithm took only 0.6ms to compute the posteriors, while HMC took 120s on average to generate a chain. The predictions (red) from the learnt inference algorithm for z0 describe the reference posteriors (green) better than those (blue) by HMC in terms of mean, variance, and mode covering.5 The contour plots in Fig. 8b to 8d visualise three HMC chains for pgm19. Here, HMC failed to converge, and Fig. 8b explains the poor estimate (blue) in the first column of Fig. 8a.
Conclusion In this paper, we presented a white-box inference algorithm that computes an approximate posterior and a marginal likelihood estimate by analysing the given program sequentially using neural networks, and a meta-algorithm that learns the network parameters over a training set of probabilistic programs. In our experiments, the meta-algorithm learnt an inference algorithm that generalises well to similar but unseen programs, and the learnt inference algorithm sometimes had test-time advantages over alternatives. A moral of this work is that the description of a probabilistic model itself has useful information, and learning to extract and exploit the information may lead to an efficient inference. We hope that our work encourages further exploration of this research direction.
5Our inference algorithm in a multimodal-posterior case leads to a good approximation in the following sense: the approximating distribution q covers the regions of the modes well, and also approximates the mean and variance of the target distribution accurately. Note that such a q is useful when it is used as the proposal of an importance sampler.
Reproducibility statement Our paper provides detailed information that is needed for reproducing the results. In each of §5.1, §5.2, and §5.3 of the main text, we explain the key experimental design and setup clearly. More detailed experimental setup is in Appendix F, where we specify the hyperparameter and the design of the neural networks. In Appendix E, we provide full details of the model classes and how programs from the classes were automatically generated in our evaluation.

heading: A FURTHER DISCUSSION ABOUT THE TRANSLATION OF AN EXPRESSIVE PPL INTO OUR INTERMEDIATE LANGUAGE
text: Programs with recursion or while loops cannot generally be translated into our intermediate language, since such programs may go into infinite loops while the programs in our language always terminate. Programs with for loops and general branches can in theory be translated into a less expressive language such as ours. For example, van de Meent et al. (2018) explain a language called FOPPL (Section 2), which has for loops and branches, and the translation of FOPPL into graphical models (Section 3). We think that these graphical models can be translated into programs in our language. Of course, this does not mean that the learnt inference algorithm would interact well with the compilation; the interaction between compilation and inference in the context of meta-learning is something to be explored in future work.

heading: B MILKY WAY EXAMPLE IN THE PROBABILISTIC PROGRAMMING LANGUAGE
text: Fig. 9 shows the compiled version of the Milky way example to the probabilistic programming language of the paper.

heading: C FORMAL SEMANTICS OF THE PROBABILISTIC PROGRAMMING LANGUAGE
text: In §2, we stated that a program C in our language denotes an unnormalised density pC that is factorised as follows:
pC(z1:n) = pC(x1:m = r1:m|z1:n)× n∏ i=1 pC(zi|z1:i−1).
Here z1, . . . , zn are all the variables assigned by the sampling statements zi ∼ N (. . .) in C in that order, the program C contains m observe statements with observations r1, . . . , rm, and these observed random variables are denoted by x1, . . . , xm. The goal of this section is to provide the details of our statement. That is, we describe the formal semantics of our probabilistic programming language, and from it, we derive a map from programs C to unnormalised densities pC .
To define the formal semantics of programs in our language, we need a type system that tracks information about updated variables and observations, and also formalises the syntactic conditions that we imposed informally in §2. The type system lets us derive the following judgements for programs C and atomic commands A:
(S, V, α) `1 C : (T,W, β), (S, V, α) `2 A : (T,W, β), where S and T are sequences of distinct variables, V and W are sets of variables that do not appear in S and T , respectively, and α and β are sequences of reals. The first judgement says that if before running the program C, the latent variables in S are sampled in that order, the program variables in V are updated by non-sample statements, and the real values in the sequence α are observed in that order, then running C changes these three data to T , W , and β. The second judgement means
the same thing except that we consider the execution of A, instead of C. The triples (S, V, α) and (T,W, β) serve as types in this type system.
The rules for deriving the judgements for C and A follow from the intended meaning just explained. We show these rules below, using the notation @ for the concatenation operator for two sequences and also set(S) for the set of elements in the sequence S:
(R,U, α) `1 C1 : (S, V, β) (S, V, β) `1 C2 : (T,W, γ) (R,U, α) `1 (C1;C2) : (T,W, γ)
(S, V, α) `2 A : (T,W, β) (S, V, α) `1 A : (T,W, β)
z 6∈ set(S) ∪ V v1, v2 ∈ set(S) ∪ V (S, V, α) `2 (z ∼ N (v1, v2)) : (S@[z], V, α)
v0, v1 ∈ set(S) ∪ V (S, V, α) `2 obs(N (v0, v1), r) : (S, V, α@[r])
v0 6∈ set(S) ∪ V v1, v2, v3, v4 ∈ set(S) ∪ V (S, V, α) `2 (v0 := if (v1 > v2) v3 else v4) : (S, V ∪ {v0}, α)
v0 6∈ set(S) ∪ V (S, V, α) `2 (v0 := r) : (S, V ∪ {v0}, α)
v0 6∈ set(S) ∪ V v1 ∈ set(S) ∪ V (S, V, α) `2 (v0 := v1) : (S, V ∪ {v0}, α)
v0 6∈ set(S) ∪ V v1, v2 ∈ set(S) ∪ V (S, V, α) `2 (v0 := p(v1, v2)) : (S, V ∪ {v0}, α)
We now define our semantics, which specifies mappings from judgements forC andA to mathematical entities. First, we interpret each type (S, V, α) as a set, and it is denoted by J(S, V, α)K:
J(S, V, α)K = {(p, f, l) | p is a (normalised) density on R|S|, f = (fv)v∈set(S)∪V , each fv is a measurable map from R|S| to R, l is a measurable function from R|S| × R|α| to R+},
where |S| and |α| are the lengths of the sequences S and α, and R+ means the set of positive reals. Next, we define the semantics of the judgements (S, V, α) `1 C : (T,W, β) and (S, V, α) `2 A : (T,W, β) that can be derived by the rules from above. The formal semantics of these judgements, denoted by the J−K notation, are maps of the following type:
J(S, V, α) `1 C : (T,W, β)K : J(S, V, α)K→ J(T,W, β)K, J(S, V, α) `2 A : (T,W, β)K : J(S, V, α)K→ J(T,W, β)K.
The semantics is given by induction on the size of the derivation of each judgement, under the assumption that for each procedure name p ∈ P, we have its interpretation as a measurable map from R2 to R: JpK : R2 → R. We spell out the semantics below, first the one for programs and next that for atomic commands.
J(S, V, α) `1 A : (T,W, β)K(p, f, l) = J(S, V, α) `2 A : (T,W, β)K(p, f, l), J(R,U, α) `1 (C1;C2) : (T,W, γ)K(p, f, l) = (J(S, V, β) `2 C2 : (T,W, γ)K
◦ J(R,U, α) `2 C1 : (S, V, β)K)(p, f, l). Let N (a; b, c) be the density of the normal distribution with mean b and variance c when c > 0 and 1 when c ≤ 0. For a family of functions f = (fv)v∈V , a variable w 6∈ V , and a function f ′w, we write f ⊕ f ′w for the extension of f with a new w-indexed member f ′w.
J(S, V, α) `2 z ∼ N (v1, v2) : (S@[z], V, α)K(p, f, l) = (p′, f ′, l′) (where p′(a1:|S|+1) = p(a1:|S|)×N (a|S|+1; fv1(a1:|S|), fv2(a1:|S|)),
f ′v(a1:|S|+1) = fv(a1:|S|) for all v ∈ V , f ′z(a1:|S|+1) = a|S|+1, l′(a1:|S|+1, b1:|α|) = l(a1:|S|, b1:|α|)),
J(S, V, α) `2 obs(N (v0, v1), r) : (S, V, α@[r])K(p, f, l) = (p, f, l′)
(where l′(a1:|S|, b1:|α|+1) = l(a1:|S|, b1:|α|)×N (b|α|+1; fv1(a1:|S|), fv2(a1:|S|)),
J(S, V, α) `2 (v0 := if (v1 > v2) v3 else v4) : (S, V ∪ {v0}, α)K(p, f, l) = (p, f ⊕ f ′v0 , l) (where f ′v0(a1:|S|) = if (fv1(a1:|S|) > fv2(a1:|S|)) then fv3(a1:|S|) else fv4(a1:|S|)),
J(S, V, α) `2 (v0 := r) : (S, V ∪ {v0}, α)K(p, f, l) = (p, f ⊕ f ′v0 , l) (where f ′v0(a1:|S|) = r),
J(S, V, α) `2 (v0 := v1) : (S, V ∪ {v0}, α)K(p, f, l) = (p, f ⊕ f ′v0 , l) (where f ′v0(a1:|S|) = fv1(a1:|S|)),
J(S, V, α) `2 (v0 := p′(v0, v1)) : (S, V ∪ {v0}, α)K(p, f, l) = (p, f ⊕ f ′v0 , l) (where f ′v0(a1:|S|) = Jp ′K(fv0(a1:|S|), fv1(a1:|S|))).
Finally, we define pC for the well-initialised well-typed programs C, i.e., programs C for which we can derived
([], ∅, []) `1 C : (S, V, α). For such a C, the definition of pC is given below:
pC(z1:|S|) = p(z1:|S|)× l(z1:|S|, α)
where (p, _, l) = J([], ∅, []) `1 C : (S, V, α)K(p0, f0, l0) for the constant-1 functions p0 and l0 of appropriate types and the empty family f0 of functions.

heading: D MARGINAL LIKELIHOOD COMPUTATION: DERIVATION AND CORRECTNESS
text: Let xn be the random variable (RV) that is observed by the command obs(N (v0, v1), r) and x1:(n−1) be the (n−1) RVs that are observed before the command. When our algorithm is about to analyse this observe command, we have (an estimate of) p(x1:(n−1)) by induction. Then, the marginal likelihood of x1:n can be computed as follows:
p(x1:(n−1), xn)
= ∫∫ p(x1:(n−1), xn, v0, v1) dv0 dv1
= ∫∫ p(x1:(n−1)) p(v0, v1|x1:(n−1)) p(xn|x1:(n−1), v0, v1) dv0 dv1
≈ p(x1:(n−1)) ∫∫ ph(v0, v1) p(xn|x1:(n−1), v0, v1) dv0 dv1
// The filtering distribution p(v0, v1|x1:(n−1)) is approximated by ph.
= p(x1:(n−1)) ∫∫ ph(v0, v1) p(xn|v0, v1) dv0 dv1
// The RV xn is conditionally independent of x1:(n−1) given v0, v1.
= p(x1:(n−1)) ∫∫ ph(v0, v1)N (r; v0, v1) dv0 dv1
// p(x1:(n−1)) is Z in the description of INFER(Ai) in Section 4, and the neural network
// nn intg,φ7 aims at approximating the integral term accurately.
This derivation leads to the equation in the main text.
In a setting of probabilistic programming where observations are allowed to be different in true and false branches, the marginal likelihood may fail to be defined, and such a setting is beyond the scope of our language. Using variables multiple times or having observe commands spread out in the program does not make differences in the derivation above.

heading: E DETAILED DESCRIPTIONS FOR PROBABILISTIC MODELS USED IN THE EMPIRICAL EVALUATION
text: Table 2 shows the full list of the model classes that we considered in our empirical evaluation (§5). We detail the program specifications for the classes using the probabilistic programming language in §2, and then describe how our program generator generated programs from those classes randomly.
In the program specifications to follow, randomly-generated constants are written in the Greek alphabets (θ), and latent and other program variables in the English alphabets. Also, we often use more intuitive variable names instead of using zi for latent variables and vi for the other program variables, to improve readability. When describing random generation of the parameter values, we let U(a, b) denote the uniform distribution whose domain is (a, b) ⊂ R; we use this only for describing the random program generation process itself, not the generated programs (only normal distributions are used in our programs, with the notation N ).
E.1 GENERALISATION TO NEW MODEL PARAMETERS AND OBSERVATIONS
This section details the model classes in §5.1.
E.1.1 gauss
The model class is described as follows:
mz := θ1; vz := θ ′ 2; c1 := θ3; c2 := θ4; vx := θ ′ 5; z1 ∼ N (mz, vz); z2 := z1 × c1; z3 := z2 + c2; obs(N (z3, vx), o)
For each program of the class, our random program generator generated the parameter values as follows:
θ1 ∼ U(−5, 5), θ2 ∼ U(0, 20), θ′2 = (θ2)2, θ3 ∼ U(−3, 3) θ4 ∼ U(−10, 10), θ5 ∼ U(0.5, 10), θ′5 = (θ5)2
and then generated the observation o by running the program forward where the value for z1 was sampled from z1 ∼ U(mz − 2× √ vz,mz + 2× √ vz).
E.1.2 hierl
The model class is described as follows:
mg := θ1; vg := θ ′ 2; vt1 := θ ′ 3; vt2 := θ ′ 4; vx1 := θ ′ 5; vx2 := θ ′ 6; g ∼ N (mg, vg); t1 ∼ N (g, vt1); t2 ∼ N (g, vt2); obs(N (t1, vx1), o1); obs(N (t2, vx2), o2)
For each program of the class, our generator generated the parameter values as follows:
θ1 ∼ U(−5, 5), θ2 ∼ U(0, 50), θ′2 = (θ2)2, θ3 ∼ U(0, 10) θ′3 = (θ3)
2, θ4 ∼ U(0, 10), θ′4 = (θ4)2, θ5 ∼ U(0.5, 10) θ′5 = (θ5) 2, θ6 ∼ U(0.5, 10), θ′6 = (θ6)2
and then generated the observations o1 and o2 by running the program (i.e., simulating the model) forward.
E.1.3 hierd
The model class is described as follows:
ma0 := θ1; va0 := θ ′ 2; va1 := θ ′ 3; va2 := θ ′ 4; mb := θ5; vb := θ ′ 6; d1 = θ7; d2 = θ8; vx1 := θ ′ 9; vx2 := θ ′ 10;
a0 ∼ N (ma0 , va0); a1 ∼ N (a0, va1); a2 ∼ N (a0, va2); b ∼ N (mb, vb); t1 := b× d1; t2 := a1 + t1; obs(N (t2, vx1), o1); t3 := b× d2; t4 := a2 + t3; obs(N (t4, vx2), o2)
For each program of the class, our generator generated the parameter values as follows:
θ1 ∼ U(−10, 10), θ2 ∼ U(0, 100), θ′2 = (θ2)2, θ3 ∼ U(0, 10) θ′3 = (θ3)
2, θ4 ∼ U(0, 10), θ′4 = (θ4)2, θ5 ∼ U(−5, 5) θ6 ∼ U(0, 10), θ′6 = (θ6)2, θ7 ∼ U(−5, 5), θ8 ∼ U(−5, 5) θ9 ∼ U(0.5, 10), θ′9 = (θ9)2, θ10 ∼ U(0.5, 10), θ′10 = (θ10)2
and then generated the observations o1 and o2 by running the program forward where the values for a0, a1, a2, and b in this specific simulation were sampled as follows:
a0 ∼ U(ma0 − 2× √ va0 , ma0 + 2× √ va0) a1 ∼ U(a0 − 2× √ va1 , a0 + 2× √ va1) a2 ∼ U(a0 − 2× √ va2 , a0 + 2× √ va2) b ∼ U(mb − 2× √ vb, mb + 2× √ vb)
E.1.4 cluster
The model class is described as follows:
mg1 := θ1; vg1 := θ ′ 2; mg2 := θ3; vg2 := θ ′ 4; vx := θ ′ 5; g1 ∼ N (mg1 , vg1); g2 ∼ N (mg2 , vg2); zero := 0; hund := 100;
t1 ∼ N (zero, hund); m1 := if (t1 > zero) g1 else g2; obs(N (m1, vx), o1); t2 ∼ N (zero, hund); m2 := if (t2 > zero) g1 else g2; obs(N (m2, vx), o2); t3 ∼ N (zero, hund); m3 := if (t3 > zero) g1 else g2; obs(N (m3, vx), o3); t4 ∼ N (zero, hund); m4 := if (t4 > zero) g1 else g2; obs(N (m4, vx), o4); t5 ∼ N (zero, hund); m5 := if (t5 > zero) g1 else g2; obs(N (m5, vx), o5)
For each program of the class, our generator generated the parameter values as follows:
θ1 ∼ U(−15, 15), θ2 ∼ U(0.5, 50), θ′2 = (θ2)2 θ3 ∼ U(−15, 15), θ4 ∼ U(0.5, 50), θ′4 = (θ4)2 θ5 ∼ U(0.5, 10), θ′5 = (θ5)2
and then generated the observations o1:5 by running the program forward.
E.1.5 milky AND milkyo
The model class milky is described as follows:
mmass := θ1; vmass := θ ′ 2; c1 := θ3; vg1 := θ ′ 4; c2 := θ5; vg2 := θ ′ 6; vx1 := θ ′ 7; vx2 := θ ′ 8; mass ∼ N (mmass , vmass); mass1 := mass × c1; g1 ∼ N (mass1, vg1); mass2 := mass + c2; g2 ∼ N (mass2, vg2);
obs(N (g1, vx1), o1); obs(N (g2, vx2), o2) For each program of milky, our generator generated the parameter values as follows:
θ1 ∼ U(−10, 10), θ2 ∼ U(0, 30), θ′2 = (θ2)2, θ3 ∼ U(−2, 2) θ4 ∼ U(0, 10), θ′4 = (θ4)2, θ5 ∼ U(−5, 5), θ6 ∼ U(0, 10) θ′6 = (θ6)
2, θ7 ∼ U(0.5, 10), θ′7 = (θ7)2, θ8 ∼ U(0.5, 10) θ′8 = (θ8) 2
and then generated the observations o1 and o2 by running the program forward.
Everything remained the same for the milkyo class, except that the two obs commands were extended to obs(N (g1, vx1), [o1, o2, o3, o4, o5]) and obs(N (g2, vx2), [o6, o7, o8, o9, o10]), respectively, and all the observations were generated similarly by running the extended model forward.
E.1.6 rb
The model class rb is described as follows:
mz1 := θ1; vz1 := θ ′ 2; mz2 := θ3; vz2 := θ ′ 4; vx := θ ′ 5; z1 ∼ N (mz1 , vz1); z2 ∼ N (mz2 , vz2); r := Rosenbrock(z1, z2); obs(N (r, vx), o)
where Rosenbrock(z1, z2) = 0.05 × (z1 − 1)2 + 0.005 × (z2 − z12)2. For each program of the class, our generator generated the parameter values as follows:
θ1 ∼ U(−8, 8), θ2 ∼ U(0, 5), θ′2 = (θ2)2, θ3 ∼ U(−8, 8) θ4 ∼ U(0, 5), θ′4 = (θ4)2, θ5 ∼ U(0.5, 10), θ′5 = (θ5)2
and then generated the observation o by running the program forward where the values for z1 and z2 in this specific simulation were sampled as follows:
z1 ∼ U(mz1 − 1.5× √ vz1 , mz1 + 1.5× √ vz1) z2 ∼ U(mz2 − 1.5× √ vz2 , mz2 + 1.5× √ vz2)
E.2 GENERALISATION TO NEW MODEL STRUCTURES
This section details the model classes, and different types in each model class in §5.2. For readability, we present canonicalised dependency graphs where variables are named in the breadth-first order. In the experiments reported in this section, we used a minor extension of our probabilistic programming language with procedures taking one parameter.
E.2.1 ext1
Fig. 10 shows the dependency graphs for all model types in ext1. The variables z0, z1, . . . and x1, x2, . . . represent latent and observed variables, respectively, and observed variables are colored in gray. The red node in each graph represents the position of the nl variable.
Our program generator in this case generates programs from the whole model class ext1; it generates programs of all twelve different types in ext1. We explain this generation process for the model type (1,1) in Fig. 10, while pointing out that the similar process is applied to the other eleven types. To generate programs of the model type (1,1), we use the following program template:
mz0 := θ1; vz0 := θ ′ 2; vz2 := θ ′ 3; vz3 := θ ′ 4; vx1 := θ ′ 5; z0 ∼ N (mz0 , vz0); z1 := nl(z0); z2 ∼ N (z1, vz2); z3 ∼ N (z2, vz3); obs(N (z3, vx1), o1)
where nl(z) = 50/π × arctan(z/10). The generation involves randomly sampling the parameters of this template, converting the template into a program in our language, and creating synthetic observations. Specifically, our generator generates the parameter values as follows:
θ1 ∼ U(−5, 5), θ2 ∼ U(0, 20), θ′2 = (θ2)2, θ3 ∼ U(0, 20), θ′3 = (θ3)2
θ4 ∼ U(0, 20), θ′4 = (θ4)2, θ5 ∼ U(0.5, 10), θ′5 = (θ5)2
and generates the observation o1 by running the program forward where the values for z0:3 in this specific simulation were sampled (and fixed to specific values) as follows:
z0 ∼ U(mz0 − 2× √ vz0 , mz0 + 2× √ vz0)
z1 = nl(z0) z2 ∼ U(z1 − 2× √ vz2 , z1 + 2× √ vz2) z3 ∼ U(z2 − 2× √ vz3 , z2 + 2× √ vz3).
The generator uses different templates for the other eleven model types in ext1, while sharing the similar process for generation of the parameters and observations.
E.2.2 ext2
In order to generate a program of this model type and observations, our generator instantiates the parameters of the template as follows:
θ1 ∼ U(−5, 5), θ2 ∼ U(0, 10), θ′2 = (θ2)2, θ3 ∼ U(0, 10), θ′3 = (θ3)2, θ4 ∼ U(0, 10), θ′4 = (θ4)2 θ5 ∼ U(0, 10), θ′5 = (θ5)2, θ6 ∼ U(0, 10), θ′6 = (θ6)2, θ7 ∼ U(0, 10), θ′7 = (θ7)2 θ8 ∼ U(0, 10), θ′8 = (θ8)2, θ9 ∼ U(0, 10), θ′9 = (θ9)2, θ10 ∼ U(0, 10), θ′10 = (θ10)2 θ11 ∼ U(0, 10), θ′11 = (θ11)2.
Then, it generates the observations o1:4 by running the program forward where the values for z0:6 in this specific simulation were sampled (and fixed to specific values) as follows:
z0 ∼ U(mz0 − 2× √ vz0 , mz0 + 2× √ vz0) z1 ∼ U(z0 − 2× √ vz1 , z0 + 2× √ vz1)
z2 = nl(z0) z3 ∼ U(z0 − 2× √ vz3 , z0 + 2× √ vz3) z4 ∼ U(z1 − 2× √ vz4 , z1 + 2× √ vz4) z5 ∼ U(z1 − 2× √ vz5 , z1 + 2× √ vz5) z6 ∼ U(z2 − 2× √ vz6 , z2 + 2× √ vz6).
The generator uses different templates for the other four model types in ext2, while sharing the similar process for generation of the parameters and observations.
E.3 TEST-TIME EFFICIENCY IN COMPARISON WITH ALTERNATIVES
This section details the mulmod class in §5.3, which has three different model types. Fig. 12 shows the dependency graphs for all the model types. The red node in each graph represents the position of the mm variable. We used all the three types in training, applied the learnt inference algorithm to programs in the third model type, and compared the results with those returned by HMC.
We similarly explain the generator only using the model type corresponding to the first dependency graph in Fig. 12. To generate programs of this type, we use the following program template:
mz0 := θ1; vz0 := θ ′ 2; vz1 := θ ′ 3; vx1 := θ ′ 4;
z0 ∼ N (mz0 , vz0); z1 ∼ N (z0, vz1); z2 := mm(z1); obs(N (z2, vx1), o1)
where mm(x) = 100×x3/(10+x4). For each program in this model type, our generator instantiates the parameter values as follows:
θ1 ∼ U(−5, 5), θ2 ∼ U(0, 20), θ′2 = (θ2)2, θ3 ∼ U(0, 20), θ′3 = (θ3)2 θ4 ∼ U(0.5, 10), θ′4 = (θ4)2
and synthesises the observation o1 by running the program forward where the values for z0:2 in this specific simulation were sampled (and fixed to specific values) as follows:
z0 ∼ U(mz0 − 2× √ vz0 , mz0 + 2× √ vz0) z1 ∼ U(z0 − 2× √ vz1 , z0 + 2× √ vz1)
z2 = mm(z1).
The generator uses different templates for the other two model types in mulmod, while sharing the similar process for instantiation of the parameters and observations.

heading: F DETAILED EVALUATION SETUP
text: In our evaluation, the dimension s of the internal state h was 10 (i.e., h ∈ R10). We used the same neural network architecture for all the neural network components of our inference algorithm INFER. Each neural network had three linear layers and used the tanh activation. The hidden dimension was 10 for each layer in all the neural networks except for nnde where the hidden dimensions were 50. The hyper-parameter in our optimisation objective (§4) was set to λ = 2 in the evaluation. For HMC, we used the NUTS sampler (Hoffman & Gelman, 2014). We did not use GPUs.
Before running our inference algorithm, we canonicalise the names of variables in a given program based on its dependency (i.e., data-flow) graph. Although not perfect, this preprocessing removes a superficial difference between programs caused by different variable names, and enables us to avoid unnecessary complexity caused by variable-renaming symmetries at training and inference times.
G LOSSES FOR hierd, cluster, milkyo, AND rb
Fig. 13 shows the average training and test losses under three random seeds for hierd, cluster, milkyo, and rb. The later part of Fig. 13a, 13c and 13d shows cases where the test loss surges. This was when the loss of only a few programs in the test set (of 50 programs) became large. Even in this situation, the losses of the rest remained small. We give analyses for cluster and rb separately in §H.
H MULTIMODAL POSTERIORS: cluster AND rb
The cluster and rb classes in §5.1 posed another challenge: the models often had multimodal posteriors, and it was significantly harder for our meta-algorithm to learn an optimal inference algorithm. To make the evaluation partially feasible for rb, we changed two parts of our metaalgorithm slightly, as well as increasing the size of the test set from 50 to 100. First, we used importance samples instead of samples by HMC, which often failed to converge, to learn an inference algorithm. Second, our random program generator placed some restriction on the programs it generated (e.g., by using tight boundaries on some model parameters), guided by the analysis of the geometry of the Rosenbrock function (Pagani et al., 2019). Consequently, HMC (with 500K samples after 50K warmups) failed to converge for only one fifth of the test programs.
variable. Similarly, the variable pgm79_z0 had at least two modes in the similar domain region (one shown and one hidden), but this time our inference algorithm showed a mode-covering behavior.
The multimodality issue raises two questions. First, how can our meta-algorithm generate samples from the posterior more effectively so that it can optimise the inference algorithm for classes of models with multimodal posteriors? For example, our current results for cluster suffer from the fact that the samples used in the training are often biased (i.e., only from a single mode of the posterior). One possible direction would be to use multiple Markov chains simultaneously and apply ideas from the mixing-time research. Second, how can our white-box inference algorithm catch more information from the program description and find non-trivial properties that may be useful for computing the posterior distributions having multiple modes? We leave the answers for future work.
I TRAINING AND TEST LOSSES FOR THE OTHER CASES IN ext1
Fig. 15 shows the average training and test losses in the ext1 experiment runs (under three different random seeds) for generalisation to the last (4th) dependency graph and to all three positions of the nl variable.
J TRAINING AND TEST LOSSES FOR THE OTHER CASES IN ext2
Fig. 16 shows the average training and test losses in the ext2 experiment runs (under three different random seeds) for generalisation to the 4th and 5th dependency graphs.

heading: K QUANTIFIED ACCURACY OF PREDICTED POSTERIORS
text: For accuracy, it would be ideal to report KL[p||q], where p is the fully joint target posterior and q is the predicted distribution. It is, however, hard to compute this quantity since often we cannot compute the density of p. One (less convincing) alternative is to compute KL[p′(z)||q(z)] for a latent variable z where p′(z) is the best Gaussian approximation (i.e., the best approximation using the mean and standard deviation) for the true marginal posterior p(z), and average the results over all the latent variables of interest. We computed KL[p′(z)||q(z)] for the test programs from ext1 and ext2 in §5.2, and for the three from mulmod that are reported in §5.3.
For ext1 and ext2, we measured the average KL[p′(z)||q(z)] over all the latent variables z in the test programs. For instance, if there were 90 test programs and each program had three latent variables, we averaged 90×3 = 270 KL measurements. In an experiment run for ext1 (which tested generalisation to an unseen dependency graph), the average KL was around 1.32. In an experiment run for ext2, the estimation was around 0.95. When we replaced q with a normal distribution that is highly flat (with mean 0 and standard deviation 10K), the estimation was 7.11 and 7.57, respectively. The results were similar in all the other experiment runs that were reported in §5.2.
For the three programs from mulmod in §5.3, p′ was the best Gaussian approximation whose mean and standard deviation were estimated by the reference importance sampler (IS-ref), and q was either the predicted marginal posterior by the learnt inference algorithm or the best Gaussian approximation whose mean and standard deviation were estimated by HMC. The average KL was around 1.19 when q was the predicted posterior, while the estimation was 40.9 when q was the best Gaussian approximation by HMC. The results demonstrate that the predicted posteriors were more accurate on average than HMC at least in terms of p′.

heading: L PROGRAM IN §5.3
text: Fig. 17 shows the program that is reported in §5.3, written in our probabilistic programming language.

heading: M DISCUSSION OF THE COST OF IS-PRED VS. IS-PRIOR
text: Our approach (IS-pred) must scan the given program “twice” at test time, once for computing the proposal using the learnt neural networks and another for running the importance sampler with the
predicted proposal, while IS-prior only needs to scan the program once. Although it may seem that IS-prior has a huge advantage in terms of saving the wall-clock time, our observation is that the effect easily disappears as the sample size increases. In fact, going through the neural networks in our approach (i.e., the first scanning of the program) does not depend on the sample size, and so its time cost remains constant given the program; the time cost was 0.6ms for the reported test program (pgm19) in §5.3.

heading: N LIMITATIONS AND FUTURE WORK
text: Currently, a learnt inference algorithm in our work does not generalise to programs with different sizes (Yan et al., 2020), e.g., from clustering models with two clusters to those with ten clusters. Each model class assumes a fixed number of variables, and the neural networks crucially exploit the assumption. Also, our meta-algorithm does not scale in practice. When applied to large programs, e.g., state-space models with a few hundred time steps, it cannot learn an optimal inference algorithm within a reasonable amount of time. Overcoming these limitations is a future work. Another direction that we are considering is to remove the strong independence assumption (via mean field Gaussian) on the approximating distribution in our inference algorithm, and to equip the algorithm with the capability of generating an appropriate form of the approximation distribution with rich dependency structure, by, e.g., incorporating the ideas from Ambrogioni et al. (2021). This direction is closely related to automatic guide generation in Pyro (Bingham et al., 2018).

Review summaries:
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.

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. 	

The paper proposes a new restrictive class of probabilistic programs with fixed number of random variables and without loops. The authors then propose an inference technique that learns the parameters of a neural network for sampling from the programs posterior distribution by composing it from individual neural networks for each atomic command in the language. This technique is shown to perform well during inference once the neural network has been trained on training programs. At the very least the inference speed is shown to be very high.