Meta-Learning an Inference Algorithm for Probabilistic Programs
Authors: Gwonsoo Che
Hongseok Yang
Venue: ICLR 2022 Submitted (2022)
Keywords: Probabilistic Programming
Approximate Posterior Inference
Meta Learning

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 these networks are then 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 structures, and report cases where our approach
has advantages over alternative approaches such as HMC in terms of test-time efficiency.
The overall results show the promise as well as remaining challenges of our approach.

Under review as a conference paper at ICLR 2022
MEAR-ING AN INFERENCE ALGORITHM FOR PROBABILISTIC PROGRAMS
Anonymous authors Paper under double-blind review
ABSTRACT
We present a meta-algorithm for learning a posterior-inference algorithm for re- stricted probabilliticc programs. Our meta-algorithm takes training set of proba- bilistic 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 ex- tracts 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 given probabilistic program by analysing individual atomic com- mands 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.
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 al.J 201 2017]55 [aad & Mansinghka 2016 Kulkarni et al. 2015:) Young et al.,
2019
Jäger et al. 2022)) As way to realize this goal, researchers have extensively worked on the de-
velopment 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 ct al. 2008 Wood et al. [2014] Mansinghka et al. 2014; Minka et al.) 2018,, Narayanan al. 62016]
Salvatier et al. (2166) Carpenter et [1,[201 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 PPL.s include constructs for specifying conditional independence in 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 Aesumano-Townee et al. 2019). all enabling users to help inference or learning algorithms.
In this paper, we explore different approach. We present meta-algorithm for learning a posterior- inference algorithm itself from a given set of restricted probabilistic programs, which specifies class of probabilistic models, such as hierarchical or clustering models. The meta-algorithm aims at constructing customised inference algorithm for the given set of models, while ensuring universality to the extent that the constructed algorithm can generalise: 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 programs in a PPL. Concretely, our white-box inference algorithm is equipped with multiple neural networks, one for each type of atomic command in PPL, and computes an approximate posterior for 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. which describes simple model the Milky Way

---

Under review as a conference paper at ICLR 2022
mass (5, 10); // log of the mass of Milky Way
91 N mass x 2, 5); obs(W(g) 1); 10); // observed velocity we1==10 of the first satellite galaxy
92 (mass + 2) (bs(W(g2. 1).3) // observed velocity velz=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 vely and vel2.
galaxy, the white-box inference algorithm analyses the program as if an RNN handles 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 E R with ho 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 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 make the inference algorithm compute accurate posterior approximations over a training set of probabilistic programs. One can also view our white-box inference algorithm a message-passinng sense where transforming the internal state h corresponds to passing a message, and understand our meta-algorithm as method for learning how to pass a message for each type of atomic commands.
approximate filtering distribution
algorithm in broad
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 given model class; (ii) we describe a meta-algoriin mega-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 (Hoff) man Gelman 2014) to given inference task were implemented in popular tools. Recently, methods for meta-learning these techniques themselves from collection of inference tasks have been developed (Wang et al. 20188 Gong et al. 2019). The meta-learning approach also features in the work stochastic variational inference where 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 Takovleva et al. 2020). For a passsag--pssinng style variational--nferennce algorithm, such as expectation propagation (Minka 2001] Wainwright & Jordan 20088. Jitkrittum et al. (0155) studied the problem of learning a mechanism to pass message for a given single inference task. A natural follow-up question is how to meta-learn such mechanism from 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- arressag-paassin-seee passing-style variational inference algorithm that can meta-learn the representation of messages and mechanism for passing them for given probabilistic programs.
Amortised inference and inference compilation (Gershman Goodman 2014 Le et al. 2017 Paige
& Wood 2016) Stuhlmüülerr et al.) 2013] Kingma& & Welling | & Gregor] (014]) Rezende jet al. 20014, Ritchie et al. 2016; Marino et al.. 2018) are closely related to our approach in that they also attempt to learn a form of 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 compiler (to be used for multiple programs).
2

---

Under review as a conference paper at ICLR 2022
= 0; v:= 5; w w:= ~, 0); 22 ~ (
obs (W w), -1.)); // 21 / w). 1.9 24 ~N(u,, w); #4 := if (2) > u) 21 else 22: (o (μl w) -22); // 02) 2.2
29 ~ NN(u, w); μs = if (23 u) 21 else 22;
25 NN, w); 15 = if (2s > u) 21 else 22 obs (N (Jus: w),2.4); // 23 23~N([15 23
20 N(u, w); P6 : if (ze > u) 21 else 22; obs (V (ue, w), 2.2) // x4~N(66 w), I4 = 2.2
Figure 2: Probabilistic program for simple clustering model on four data points.
The idea of running programs with learnt neural networks also appears in the work training neural networks to execute programs (Zaremba & Sutskever] [2014] Bieber et al.] 2020) RRed & de Freitas 2016). As far as we know, however, we are the first 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, and e 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 == 2 ((11, U2) | obs(WN 01), 6o if (v) > U2) by else 24 vo := p()).
0o := 00 == 11
Programs in the language are constructed by sequentially composing atomic commands. The language supports six types of atomic commands. The first type is z ~ (N1,,, U2), which draws a sample from the normal distribution with mean in and variance U2, and assigns the sampled value z. The second command, obs (N (vo, U1), r), states that a random variable is drawn from ((vo, v1) and its value is observed to be r. The next is restricted form of conditional statement that selects one of v3 and vs depending on the result of the comparison i1 U2. The following two commands are different kinds of assignments, one for assigning constant and the other for copying a value from one variable to another. The last atomic command e == p(u, 02) is 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, 01 + 122 instead of + +(11, 02).
We permit only the programs where a variable does not appear more than once 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 compilers. This restriction lets us regard variables updated by ~ as latent random variables. We denote those variables by 21,..., 2n.
We this simple language for two reasons. First, the restriction imposed on 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 normal distribution or to procedure p should be variables, not general expression forms such as addition of two variables. As we will show soon, this restriction makes 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 compiler for 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. (0018) 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 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 21 and 22 store the centers of the two clusters, and 23, ........ ... 26 hold the random draws that decide cluster assignments for the data points. See Appendix Bor the Milky Way example in Fig. | compiled to a program in our language.
Probabilistic programs in the language denote unnormalised probability densities over R" for some n. Specifically, for a program C, if 21,-.. -1.-... En are all the variables assigned by the sampling statements
3

---

Under review as a conference paper at ICLR 2022
2, ......) ) in C' in that order and C' contains observe statements with observations 71 71..... mm then C denotes an unnormalised density Pe over the real-valued random variables ..,..., Zn. (2(2)m)) PC (21.m =11mm (2::a) I Pc(2+)213-1), where II...., 2m are variables not ap- pearing in C/ and are used to denote observed variables. This density is defined inductively over the structure of C. See Appendix for details. The goal of our white-hox inference algorithm is to compute efficiently accurate approximate posterior and marginal likelihood estimate for given C (that is, for the normalised version of PC and the normalising constant of pc), when pe has a finite non-zero marginal likelihood and, as result, well-defined posterior density. We next describe how the algorithm attempts to achieve this goal.
WHITE-BOX INFERENCE ALGORITHM
Given a program C = (Aii...; A Ax), our white-box inference algorithm views C as a sequence of its constituent atomic commands (A1, A Ax), and computes an approximate posterior and marginal likelihood estimate for C' by sequentially processing the A, 's. Concretely, the algorithm starts by initialising its internal state to lgg = Ô E R and the current marginal-
estimate to Zo 1. Then, it updates these two components based on the first atomic command A of C. It picks a neural network appropriate for the type of A1, applies it to hg and gets a new state h1 E R Also, updates the marginal likelihood estimate to Zi by analysing the semantics of A,. This process is repeated for the remaining atomic commands A2, Ag,..., Ag of C, and eventually produces the last probability density the latent
state hg and estimate Zy.- Finally, the state h& gets decoded variables of C' by neural network, which together with Zμ 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 E IR* of the algorithm; a neural network for decoding the internal states 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 21,...: Zns and use at most m-many variables (so m 2 n). Let V be [0. 1jmh 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:
o0 ./3 x R R nn V R R RR. dede.
Wab.. : V2 X R R
+ R2
MiFFO : V5 x R +R.
nn = V2
R
mpp V3 x R R for p P.P.
R
R R)"
Wnigg. V2
R
R +
where 11.7 and op for IP 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 2 ~ (N11,, 22). the algorithm runs the first network ma on the current internal state h, and obtains a new state h = nn (z, v1:2, h). where 2 and 71:2 mean the one-hot encoded variables 2, U1 and 02. The next midee is a decoder of the states h to probability densities variables 21..... Zn. which are the product of n independent normal distributions. The network maps h the means and variances of these distributions. The last Wiintg, is used when our algorithm updates the marginal likelihood estimate based on an observe statement obs (Ns(NN(up 21). r). When we write the meaning of this observe statement as the likelihood (r; Uo. 01), and the filtering distribution for Ug and 2 under (the decoded density of) the current state h as Ph ("o, en > the last neural network computes the following approximation: niiingg (1q:1 r. h) R f [(;; Up, 21 )P (Vo, in )dvgdu). See Appendix D for the full derivation of the marginal likelihood.
the latent
Given a program C = (Aii...; A&) that draws n samples (and so uses latent variables 2I,...,,,,) the algorithm approximates the posterior and marginal likelihood of C' as follows:
IC) let (ho, Zo) (ō, 1) and (lu, Zx) = (INFER(AR) c ... INFER(A:)) Zo) in let ((H1, of) ..... (Mm, o.)) de (hv) in return (IT= uis oi,,) Zx)
where INFER(A,) R R - R R picks appropriate neural network based on the type of A, and uses it to transform h and Z:
(NFER(obs(V(eo, v1), (h, (,, 2) = (nnob (Vo.I r, h),, Z x nminggg (Vo.I r, h)),
'The (vo, 01) is filtering distribution, not prior.
4

---

Under review as a conference paper at ICLR 2022
INFER (1) = if (t) (01 02) by else va) 2) (nng 0o:4: k): 2),
INFER( Up : (h, (nm v6 r, h), 2), INFER ~ NV (2),) 2)- (NM U1.2 h1), 2), (..
INFER(000 : [(,,)),,,))) 2) -
:= h), 2), INFER(vo:= (2)0,, b2)) Z) (nnp (T0.2, (on2,h), h), 2).
We remind the reader that vo:: refers to the sequence of the one-hot encodings of variables vo, .o..... Uk- 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.
4 META-LEARNING PARAMETERS
The parameters of our white-box inference algorithm are learnt from a collection of probabilistic programs our language. Assume that we are given a training set of programs DD = (Ci,..... Cv} such that each C₁ samples latent variables 21, 21...... Zn and uses at most m variables. Let = (Ó1:7: (op) be the parameters of all the neural networks used in the algorithm. We learn these parameters by solving the following optimisation problem
X Σ KL|Tc((zm)lllc(2n. + 2 (Nc -c)² CED
arg min Φ
where > 0 is a hyper-parameter Ne is the marginal likelihood (or the normalising constant) ( for pc, the next (= is the normalised posterior pc(#1..)/Nc for C, and the last qe and Zo are the approximate posterior and marginal likelihood estimate computed by the inference algorithm (that is. (2c() Zc) INFER(C")). Note that qa and Ze both depend on since INFER uses the ş- d-parameterised neural networks.
We optimise the objective by stochastic gradient descent.
The key component of the optimisa- tion is a gradient estimator derived as follows: (V ECED KL|Tc + (Nc Zc)²) =
SCCEDE V log
Zc) (c))) ~ ECED
Here Lc,s and Ne are sample estimates of Ex V log qc (=1m)) and the marginal likelihood, respectively. Both estimates can be computed using standard Monte-Carlo algorithms. For instance, we can run the self normalionn importance sampler with prior as proposal, and generate weighted sam- ples (w(()) 2(h) IsjsM for the unnormalised posterior Pc- Then, we can use these samples to com- pute the required estimates: Nc = + Σド= w(t) and LC = 1 M
Ei (ubi) V & log (2iii))) 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 (Martnno et al. 2017). Then, we compute Icc using posterior samples, and No using weighted importance samples. Note that neither тe in Exa.. ~nc I-V & log (c1 nor Ne depends on the parameters Φb Thus, for each C E DD No needs to be estimated only once throughout the entire optimisation process, and the posterior samples from ac need to be generated only once as well. We use this fact to speed up the computation of each gradient-update step.
M
EMPIRICAL EVALUATION
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 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 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. (017). 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 for the full list of our model classes and their details, and Appendix F for the detailed experimental setup.
Strietly speaking, assume that the marginal likelihood of any C E D is non-zero and finite.
5

---

Under review as a conference paper at ICLR 2022
Seed,
Seed.
Tr
To
3.
3. Te
(a) gauss
(b) hierl
(c) milky
Figure 3: Average training and test losses under three random seeds. The y-axes log-scaled. The increases in later epochs of Fig. 3c were due to only one or a few test programs out 50.
255
Before training
After IK epochs
(c) After 2K epochs
Figure 4: Comparisons of predicted and reference marginal posteriors recorded at different training steps: at the initial step, after 1K epochs, and after 2K epochs.
5.1 GENERALISATION NEW MODEL PARAMETERS AND OBSERVATIONS
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 hicrarchically (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-observation extension (milkyo) where five observations are made for each satellite galaxy; and (6) models (rb) using the Rosenbrock function which is expressed 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-of- the-art inference algorithm automatically, and also to identify the challenges the approach. These models are chosen for this purpose. For instance, 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 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, 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 ft (Gelman et al. 1992). All training and test programs were automatically generated by random program generator. This generator takes a program class and byperparameters (c.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 2 samples from the analytic (for gauss) or approximate (for the rest, by HMC) posterior distribution for the program Similarly, our meta-
'The function is often used to evaluate learning and inference algorithms Goodman & Weare 2010
Wang
& Li
2018 Pagani et al. (2019)
Except for rb; the discussion Rosenbrock models Appendix

---

Under review as a conference paper at ICLR 2022
Send,
Sond,
Seed.
Tr
To
(a) dep. graph.
(b) To 2nd dep. graph.
(c) To 3rd dep. graph.
Figure 5: Average losses for generalisation to dependency graphs in extl. The y-axes are log-scaled.
Seed, Tr or
Sead, Tr
Send, Tr To
Tb
(a) To 1st dep. graph.
(b) To 2nd dep. graph.
(c) To 3rd dep. graph.
Figure 6: Average losses for generalisation to dependency graphs in ext2. The y-axes are log-scaled.
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 HMC chain. We performed mini-batch training; single gradient update was done with training program and a mini batch of size 21 (out of 2 samples for the program). We used Adam (Kingma & Ba 2015) with its hyperparameters {31 0.9, B2 0.999, weight decaatdeeeyyeeeeeee 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. Fig.3 Fig.3sshow 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. compares, for 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. 4cp. The results show that our meta mlga-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 provide analysis for them in Appendix H
5.2 GENERALISATION TO NEW MODEL STRUCTURES
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 ul(x) 50/T x arctan(z/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 ul variable, which are grouped into five types based on their dependency graphs. The evaluation was done for ext1 and ext2 independently in $5.1 but here each of ext1 and ext2
7

---

Under review as a conference paper at ICLR 2022
Table 1: ESS per sec for the 60 test programs by HMC vs. IS IS-pred vs. IS IS-prior.
HMC
IS-pred
IS-prior
GM
Q1
Q3
GM
Q1
Q3
GM
Q1
Q3
ESS
204.8K
4.1K
4.6M
4.2K
2.2K
13.8K
2.8K
1.1K
9.5K
Time
48
27.4s
82.3s
22.7ms
21.4ms
23.0ms
23. 1ms
22.1ms
24.0ms
ESS / sec
4.3K
124
127.7K
196.5K
102.4K
646.5K
123.8K
52.6K
436.1K
15 pred 15-riior
100
2000
IS-ref HMC
0:00
ISeret
IS-pred
300
IS-prior
245
(a) Moments by HMC.
(b) Moments by IS-pred and IS-prior. ESS by IS-pred and IS- IS-prior.
Figure 7
itself has programs of multiple (12 for extl and 5 for ext2) model types. See Fig. 10 and in the appendix for visualisation of the different model types in extl 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 T;,; 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 ext 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, 3),, we used programs from T-,,+ , for training and those from Ti for testing. For generalisation to the j-th dependency graph (j = 1, 2, 3, 4), five
used programs from T for training and those from T..j for testing. For ext2, 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 x 3) and 15 (= 3) for extl 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 depen- dency graphs in extl. Fig. [6shows 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) and 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 the training losses decreased. In 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 extl. 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.
[
5.3 TEST-TIME EFFICIENCY COMPARISON WITH ALTERNATIVES
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 x x²/(100 +f)+ 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 meta- algorithm using 600 programs from all three types using importance samples (not HMC samples).
8

---

Under review as a conference paper at ICLR 2022
Marginal posteriors
(b) Chain
Chain 2
(d) Chain
Figure 8: Marginal posteriors for the comparison, and contours of three HMC chains for pgm_19 where the x-axis for z0 and the y-axis z1.
Then for 60 test programs from the last model type, 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.
shows 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 see" 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 both mean (GM) and the quartiles (QI and Q3).
Table
We provide further analysis for test program (pgm19; see Appendix L. Fig. 7 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 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- IS-prior (blue) tended to produce better estimates as the elapsed time increased, but each time, IS- -S-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 anland] 7b] IS-pred produced over 16K effective samples in 18ms, while HMC generated only 80 effective samples even after 84s. Similarly, IS IS-prior generated fewer than 1.4K effective samples in the approximately same elapsed time as in IS IS-pred. In fact, Fig. shows that as the time increases, the gap between the ESSes of IS- IS-pred and IS- -S-prior gets widen, because the former increases at a rate significantly higher than the latter. Note that IS- IS-pred has to scan program twice at test time, once for computing the proposal and another for IS with the predicted proposal. See Appendix for discussion.
Our manual inspection revealed that the programs in mulmod often have multimodal posteriors. Fig. Sal 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 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 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 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) the first column of Fig.
Conclusion In this paper, we presented a white-box inference algorithm that computes an approxi- mate posterior and marginal likelihood estimate by analysing the given program sequentially using neural networks, and a mega-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. 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.
$ 'Our inference algorithm in a multimodal multimodallposterior case leads to a good approximation in the following sense: the approximating distribution covers the regions of the modes well, and also approximates the mean and variance of the target distribution accurately. Note that such q useful when it is used as the proposal of importance sampler.

---

Under review as a conference paper at ICLR 2022
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 where 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 evaluation.
REFERENCES
Luca Ambrogioni, Gianluigi Silvestri, and Marcel van Gerven. Automatic variational inference with cascading flows. arXiv preprint arXiv:2102.04801, 2021.
Christophe Andrieu and Johannes Thoms. A tutorial on adaptive MCMC. Stat. Comput.. 18(4): 343-373, 2008.
Atilim Gunes Baydin, Lei Shao, Wahid Bhimji, Lukas Heinrich, Saeid Naderiparizi, Andreas Munk, Jialin Liu, Bradley Gram-Hansen, Gilles Louppe. Lawrence Meadows, Philip Torr, Victor Lee, Kyle Cranmer, Mr. Prabhat, and Frank Wood. Efficient probabilistic inference in the quest for physies beyond the standard model. In Advances in Neural Information Processing Systems, volume 32, pp. 5459-5477 5472. Curran Associates, Inc., 2019.
David Bieber, Charles Sutton, Hugo Larochelle, and Daniel Tarlow. Learning to execute programs with instruction pointer attention graph neural networks. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
Eli Bingham, Jonathan P. Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul Szerlip, Paul Horsfall, and Noah D. Goodman. Pyro: Deep universal probabilistic programming. Journal of Machine Learning Research, 2018.
Bob Carpenter, Andrew Gelman, Matthew D Hoffman, Daniel Lee, Ben Goodrich, Michael Be- tancourt, Marcus Brubaker, Jiqiang Guo, Peter Li, and Allen Riddell. Stan: A probabilistic programming language. Journal of statistical software, 76(1), 2017.
Marco Cusumano- Towner, Alexey Radul, David Wingate, and Vikash K Mansinghka. Probabilistic programs for inferring the goals of autonomous agents. arXiv preprint arXi:: 2017.
Marco F. Cusumano- -ooumann-TTowerr, Feras A. Saad, Alexander K. Lew, and Vikash K. Mansinghka. Gen: a general-purpose probabilistic programming system with programmable inference. Proceedings of the 40th ACM SIGPLAN Conference on Programming Language Design and Implementation, PLDI 2019, Phoenix, AZ, USA, June 22-26, 2019, pp. 221-236. ACM, 2019.
Simon Duane, Anthony D Kennedy, Brian J Pendleton, and Duncan Roweth. Hybrid Monte Carlo. Physies letters B, 195(2):216 222, 1987.
Hong Ge. Kai Xu, and Zoubin Ghahramani. Turing: Composable inference for probabilistic pro- gramming. International Conference on Artificial Intelligence and Statistics, AISTATS 2018, 9-11 April 2018, Playa Blanca, Lanzarote, Canary Islands, Spain, volume 84 of Proceedings of Machine Learning Research, pp. 1682-1690. PMLR, 2018
Andrew Gelman, Donald B Rubin, et al. Inference from iterative simulation using multiple sequences. Statistical science, 7(4):457472, 1992.
Samuel Gershman and Noah Goodman. Amortized inference in probabilistic reasoning. In Proceed- ings of the annual meeting of the cognitive science society, volume 36, 2014.
Wenbo Gong, Yingzhen Li, and José Miguel Hernández-Lobato. Meta-learning for stochastic gradient MCMC. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
Jonathan Goodman and Jonathan Weare. Ensemble samplers with affine invariance. Communications in applied mathematics and computational science, 5(1 80, 2010.
10

---

Under review as a conference paper at ICLR 2022
Noah D Goodman, Vikash K Mansinghka, Daniel Roy, Keith Bonawitz, and Joshua B Tenenbaum. Church: a language for generative models. In Proceedings of the Twenty-Fore Fourth Conference on Uncertainty in Artificial Intelligence. pp. 220-220-22022 229, 2008
Jonathan Gordon, John Bronskill, Matthias Bauer. Sebastian Nowozin, and Richard E. Turner. Meta- learning probabilistic inference for prediction. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenRevvew.net, 2019.
Matthew D. Hoffman and Andrew Gelman. The No-U-Tum sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. J. Mach. Learn. Res., 15(1):1593-123, 2014.
Ekaterina Iakovleva, Jakob Verbeek, and Karteek Alahari. Meta-learning with shared amortized variational inference. In Proceedings of the 37th International Conference on Machine Learning. ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning PMLR, 2020.
Research, pp.
Wittawat Jitkrittum,, Arthur Gretton, Nicolas Heess, S. M. Ali Eslami, Balaji Lakshminarayanan, Dino Sejdinovic, and Zoltán Szabó. Kernel-based just-in-time learning for passing expectation propagation messages. In Proceedings the Thirty-First Conference on Uncertainty in Artificial Intelligence, UAI 2015, July 12-16, 2015, Amsterdam, The Netherlands, pp. 405-414. AUAI Press, 2015.
Lena A. Jäger, Daniela Mertzen, Julie A. Van Dyke, and Shravan Vasishth. Interference patterns in subject-verb agreement and reflexives revisited: large-sample study. Journal of Memory and Language, 111:104063, 2020. ISSN 0749-596X.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
Diederik P Kingma and Max Welling, Auto-encoding variational Bayes. arXiv preprint arXiv:1312.6114, 2013.
Tejas D Kulkami, Pushmeet Kohli, Joshua B Tenenbaum, and Vikash Mansinghka, Picture: A probabilistic programming language for scene perception. In Proceedings the ieee conference on computer vision and pattern recognition, pp. 4390-4399, 2015.
Tuan Anh Le, Atilim Gunes Baydin, and Frank Wood. Inference compilation and universal proba- bilistic programming. In Artificial Intelligence and Statistics, pp. 1338-1348. PMLR, 2017.
Vikash Mansinghka, Daniel Selsam, and Yura Perov. Venture: a higher-order probabilistic program- ming platform with programmable inference. arXiv preprint arXiv:1404.000 0099, 2014.
Joe Marino, Yisong Yue, and Stephan Mandt. Iterative amortized inference. In International Conference on Machine Learning, pp. 3403-3412, 2018
Luca Martino, Victor Elvira, David Luengo, and Jukka Corander. Layered adaptive importance sampling. Statistics and Computing, 27(3):599-633, 2017.
Laurent Mazare, ocaml-torch: OCaml bindings for pytorch, 2018. URL httl:ttts///github.oom com/ Latzrentaazare/ccaml
T. Minka, J.M. Winn, J.P. Guiver, Y. Zaykov, D. Fabian, and Bronskill, /Infer.NET 0.3, 2018. Microsoft Research Cambridge. http://dotnet.gituu.ioooioi
Thomas P. Minka. Expectation propagation for approximate Bayesian inference. In UAI '01: Pro- ceedings of the 17th Conference Uncertainty Artificial Intelligence, University of Washington, Seattle, Washinglon, USA, August 2-5, 2001, pp. 362-369. Morgan Kaufmann, 2001.
Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. In International Conference Machine Learning, pp. 1791-1799, 2014.

---

Under review as a conference paper at ICLR 2022
Praveen Narayanan, Jacques Carette, Wren Romano, Chung-chieh Shan, and Robert Zinkov. Prob- abilistic inference by program transformation in hakaru (system description). In International Symposium Functional and Logic Programming 13th International Symposium, FLOPS 2016,
Kochi, Japan, March 4-6, 2016, Proceedings, pp. 62-79. Springer, 2016.
Filippo Pagani, Martin Wiegand, and Saralees Nadarajah. An n-dimensional rosenbrock distribution for meme testing. arXiv preprint arXiv: (et.:9903.099 09556, 2019.
Brooks Paige and Frank Wood. Inference networks for sequential Monte Carlo in graphical models. In International Conference on Machine Learning, pp. 3040-3049, 2016.
Yura Perov and Frank Wood. Automatic sampler discovery via probabilistic programming and approximate Bayesian computation. In Artificial General Intelligence, pp. 262-273, Cham, 2016. Springer International Publishing. ISBN 978-3-319-4164944
Scott E. Reed and Nando de Freitas. Neural mogrammerr-nnteppeeterr.. In 4th International Conference on Learning Representations, ICIR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning, pp. 1278-12866 2014.
Daniel Ritchie, Ben Mildenhall, Noah D. Goodman, and Pat Hanrahan. Controlling procedural modeling programs with soccastically-ooderee sequential Monte Carlo. ACM Trans. Graph., 34 (4), July 2015. ISSN 0730-0301.
Daniel Ritchie, Paul Horsfall, and Noah D Goodman. Deep amortized inference for probabilistic programs. arXiv preprint arXiv:1610.057:55 2016.
Feras Saad and Vikash K Mansinghka. A probabilistic programming approach to probabilistic data analysis. In Advances in Neural Information Processing Systems, pp. 201 1 1-2019, 2016.
John Salvatier, Thomas V Wiecki, and Christopher Fonnesbeck. Probabilistic programming in python
using pymc3. PeerJ Computer Science, 2:e55, 2016.
Ulrich Schaechtle, Feras Saad, Alexey Radul, and Vikash Mansinghka. Time series structure discovery via probabilistic program synthesis, arXiv preprint arXiv:1611.0700 2016.
N. Siddharth, Brooks Paige, Jan- Jan-Willem van de Meent, Alban Desmaison, Noah D. Goodman, Pushmeet Kohli, Frank Wood, and Philip Torr. Learning disentangled representations with semi- supervised deep generative models. In Advances in Neural Information Processing Systems 30, pp. 5927-5937. Curran Associates, Inc., 2017.
Andreas Stuhlmüller, Jacob Taylor, and Noah Goodman. Learning stochastic inverses. In Advances in neural information processing systems, pp. 3048-3056, 2013.
Dustin Tran, Alp Kucukelbir, Adji B. Dieng, Maja Rudolph, Dawen Liang, and David M. Blei. Edward: A library for probabilistic modeling, inference, and criticism. arXiv preprint arXiv:1610.09787, 2016.
Dustin Tran, Matthew D. Hoffman, Dave Moore, Christopher Suter, Srinivas Vasudevan, and Alexey Radul. Simple, distributed, and accelerated probabilistic programming. In Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pp. 7609-7620, 2018.
Man-Willem van de Meent, Brooks Paige, Hongseok Yang, and Frank Wood. An introduction to probabilistic programming. arXiv preprint arXiv: 10756, 2018.
Martin J. Wainwright and Michael I. Jordan. Graphical models, exponential families, and variational inference. Found. Trends Mach. Learn., 1(1-2):1-305, 2008
Hongqiao Wang and Jinglai Li. Adaptive gaussian process approximation for Bayesian inference with expensive likelihood functions. Neural computation, 50(11):37723004 2018.
12

---

Under review as a conference paper at ICLR 2022
Tongzhou Wang, Yi Wu, Dave Moore, and Stuart J. Russell. Meta-learning MCMC proposals. In Advances Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018. Montréal, Canada, pp. 4150-4160, 2018.
Frank Wood, Jan Willem van de Meent, and Vikash Mansinghka. A new approach to probabilistic pro- gramming inference. Proceedings of the 17th International conference on Artificial Intelligence and Statistics, pp. 1024-1032, 2014.
Mike Wu, Kristy Choi, Noah D Goodman, and Stefano Ermon. Meta-amortized variational inference and learning. In AAAI, pp. 6404-6412, 2020.
Yujun Yan, Kevin Swersky, Danai Koutra, Parthasarathy Ranganathan, and Milad Hashemi. Neural execution engines: Learning execute subroutines. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020,
December 6-12, 2020, virtual, 2020.
Jean-Gabriel Young. Fernanda Valdovinos, and Mark EJ Newman. Reconstruction of plant- pollinator networks from observational data. bioRxiv, pp. 754077, 2019.
Wojeiech Zaremba and Ilya Sutskever. Learning execute. CoRR, abs/1410.4615,, 2014. URL fttt:://arivy wwtp://///aabb/1....... 4615
13

---

Under review as a conference paper at ICLR 2022
A FURTHER DISCUSSION ABOUT THE TRANSLATION OF AN EXPRESSIVE PPL INTO OUR INTERMEDIATE LANGUAGE
Programs with recursion while loops cannot generally be translated into 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 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 future work.
in theory be translated into less expressive
one == 1; t== 2; f= 5; ten == 10;
21 ~N[ ten); // log of the mass of Milky Way
mass1 == 21 t;
22 >Mammss,, f); // for the first satellite galaxy
obs(N(22, one), ten); // x1 = 10 for 21 ~ (22,, one)
mass2 == 21 f:
23 N(mass2 t); // for the second satellite galaxy
obs(N(22) one).3) x2 3 for 22 N22,, one)
Figure 9: Milky Way example compiled to the probabilistic programming language used in the paper.
B MILKY WAY EXAMPLE IN THE PROBABILISTIC PROGRAMMING LANGUAGE
Fig. D shows the compiled version of the Milky way example to the probabilistic programming language of the paper.
C FORMAL SEMANTICS OF THE PROBABILISTIC PROGRAMMING LANGUAGE
In we stated that program C' in our language denotes an unnormalised density Pc that is factorised as follows:
pc(2mm) (0 Fi:m (1m))
II pC po
Here ......, are all the variables assigned by the sampling statements 24 that order, the program C contains m observe statements with observations ri...., T's and these observed random variables are denoted by æ1,... -m. The goal of this section is to provide the details of our statement. That is, we describe the formal semanties of our probabilistic programming language, and from it, derive map from programs C' to unnormalised densities pc.
.(...) in C' in
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) F C: 33), (S, V.) 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 a and 3 are sequences reals. The first judgement says that before running the program C, the latent variables in S are sampled in that order, the program variables in V are updated by non nam-sample statements, and the real values in the sequence a are observed in that order, then running C' changes these three data to T, W. and B. The second judgement means
14

---

measurable map from Under review as a conference paper at ICLR 2022
the same thing except that consider the execution of A, instead of C. The triples (S, V, and (T, W, (3 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 (Ct( for the set of elements in the sequence S:
(R, a) Fr C1 (S. 3,9) (S.
C2 W.),
(S.V.a) F2 A: (3)
(R,U,a) Fi (C1:0)) (T,
(S. F1 A: (T,
2 sct(S)) UV 01 U2 E (S)
Up 21 E (ct UV (S.V obs(N 21 21),r) (S, a@(r))
(S, F2 (2 ((11 (S@[2).V,o
Uo E set (S) 01, U2, U3, 04 E set UV (S, (o0 if (v1 > 22) by else vn) (S, V (vo).a)
Uo E (S)
Up sel UV in E (e (S. (co == u)) U
(5.V.a) = (S, (voo).a))
up E (e 01, E set(S) UV (S, (No p(u, e2)) (S (w)),a
We now define semanties, which specifies mappings from judgements for C and mathematical entities. First, we interpret each type (S, a) as a set, and it is denoted by [(S, a,
[[S, a) = ((P. p is a (normalised) density on |R|S| f = (Fo))ese((S)UV
each fo is measurable map from R to R,
is a measurable function from || x Rllo to R+,,
where |S| and |a| are the lengths of the sequences S and a and 1 means the set of positive reals. Next, define the semantics of the judgements (S, V.a) F C': (T, W. 3) and (S, V, F2 A (T, W, B) that can be derived by the rules from above. The formal semantics of these judgements, denoted by the [-] notation, are maps of the following type:
[(S, V. a)) -1 C: (T, W. 3) [(S, V.a)] > [[T, 3).
[(S, v.) t2 A: (T, W, 3)] [(S, V,a)] - [[T, W. B)].
The semantics is given by induction on the size of the derivation of each judgement, under the assumption that for each procedure name p E P. we have its interpretation
R to IR:
R +RR. R.
We spell out the semantics below, first the one for programs and next that for atomic commands.
[[S, a)) Fi A: (T, W, (p, f.l) = [[S, V, a) Ha A: (T, W. (,, f. (RR, (G:C) (T, W. Wl f,1) ([(S, V, B) T2 C2 (T, w) o [[R, U, 22 C1
Let Na;; b, c) be the density of the normal distribution with mean and variance when > 0 and 1 when e I 0. For a family of functions f = (fo)vev, variable w & V, and a function f we write f @ fl, for the extension of f with new w-indexed member ft-
[[(S, V, 22 N((1, : (S@ 2)] V, (p')
(where p'(a1: (11)) = p(( ((qsstii (0r1s)), (o)(s:)))) f((11111411) fu(():19)) for all V, (((1118441) a ((01)))-1:b)11)) [(a1))))))))))))))))))))))) b14a1))
[S, V, obs (N m)) ( (1-
15

---

Under review as a conference paper at ICLR 2022
(where 1' ('(11:1s),b b1:jo|+1) [(11s) (1:ja)) V(b (f
(an.s))) (a1:ss));
[[S, V,, (vo = if (m) > 02) 03 else v4) (8, V (wo al) f. (0.f + 1 (where Fn (111))) = if (fe, (Q11))) > (4m4s)) then fuz (a11s)) elsc for (11:))))
= [[S,, V2 :=)::)): (S, : = (p. @ fao 4)
(where (21: = r))
[[S, V, (vD = v1) (S,V aln),,) = (p, @ !
(where Lag (41: st) fn (111)))),
[[S, v, (vD = p' m1) U (v)) (p, + 1
(where fny (a:: st) = Ir' (anss)), fu, (an,s)))))
Finally, we define PC for the well-initialised well-typed programs C,, i.c., programs C for which we can derived
(0. 0. D) C.
For such a C, the definition of pC is given below:
pC(214)]) p(z14s) (2145),
where (p._.) I[]. 0, ()) F, C: (S. V, (po. for lo) for the constant-1 functions po and lo of appropriate types and the empty family fo of functions.
D MARGINAL LIKELIHOOD COMPUTATION, DERIVATION AND CORRECTNESS
Let zn be the random variable (RV) that is observed by the command (Ns(N(oo. r) and (1:(n=1) I) 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:(nn-1) by induction. Then, the marginal likelihood of Fi:n can be computed as follows:
// p(Z1:(n-1), In, vo, 21) dog dui
// p(Z1:(n-1)) pluo. v =) p(z ((zn1 0o. 01) dug dug
po ~ p(11(1-1) /f ph (vo, vi) p(= (1: Up, vi duo dey
1/ The filtering distribution p(vo, U1 (1)(()1)) is approximated by Ph-
-)) If pa (vo, v1) p( lo, 01) dug dug
The RV on is conditionally independent of #1:(n-1) given uo V1-
= =((F1(n-1)) If pa (vo, vi) (r; uo, v1) duo dui
Il p(#1(n=1)) is Z in the description of INFER(A,) A:) in Section 4 and the neural network I/ niintggd7 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 different in true and false branches, the marginal likelihood may fail to be defined, and such setting 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.
16

---

Under review as a conference paper at ICLR 2022
Table Full list of the model classes the empirical evaluation.
Section
Model class
Description
Detail
gauss
Gaussian models with a latent variable and an observation where the mean of the Gaussian likelihood is an affine transformation of the latent.
Appendix 1.1
5.1
hierl
Hierarchical models with three hierarchically structured latent variables.
Appendix E.I. 1.2
hierd
Hierarchical multi-level models with both latent variables and data structured hierarchi- cally where data are modelled of latent variables of different levels.
Appendix E.1.3
regression
cluster
Clustering models where five observations are clustered into two groups
Appendix E.1.4
milky and milkyo
Milky Way models, and their multiple- observations extension where five observations are made for each satellite galaxy.
Appendix 1.5
rb
Models with the Rosenbrock function, which is expressed as an external procedure.
Appendix E. 1.
ext1
Models with three Gaussian variables and one deterministic variable storing the value of the function nl(x) 50/m / arctan(z/10), 100), where the models have 12 diffferent types - four dif- ferent dependency graphs of the variables, and three different positions of the deterministic nl variable for each these graphs.
Appendix
H.2.1
5.2
(and Fig. 10
ext2
Models with six Gaussian variables and one nl variable, which are grouped into five model types based on their dependency graphs.
Appendix E.2.2 (and Fig. 11)
5.3
mulmod
Three-variable models where two latent vari- ables follow normal distributions and the other stores the value of the function mmm 100 x² 122(10 + )²).. The models in this class are grouped into three types defined by their
Appendix
E.3
(and Fig. 12
dependency graphs and the positions of mm in the programs.
17

---

Under review as a conference paper at ICLR 2022
E DETAILED DESCRIPTIONS FOR PROBABILISTIC MODELS USED IN THE EMPIRICAL EVALUATION
Table 2 shows the full list of the model classes that considered in empirical evaluation ($5) We detail the program specifications for the classes using the probabilistie programming language in 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 (0), and latent and other program variables in the English alphabets. Also, we often use more intuitive variable names instead of using for latent variables and vú for the other program variables, improve readability. When describing random generation of the parameter values, we let U(a b) denote the uniform distribution whose domain is (a, b) C I; 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 S.I
E.1.1 gauss
The model class is described follows:
m: := 01; U2 :- 02j 03i 02 == 04i Ux 0'ii
21 ~ (m Uz)i ; 22 = 21 x c1; 23 == 22 + C2
obs (Ns( o)
For each program of the class, our random program generator generated the parameter values as follows:
U(-5,5), 02 2(0, 02- (02)², U(-3,3)
0. U(- U(-10,10), 0g U(0.5, 10), 6 (0₂)²
and then generated the observation o by running the program forward where the value for z was sampled from 21 ~ U(m, 2 x VVz, m, + 2 vv.).
E.1.2 hierl
The model class is described follows:
mg := 01; Ug = 823 et, = 03 ub, = 8/4: Un = 03: E. 0g ~ N N(m, "g); t ~ ~Ngg v. 3 12 \gg, ve2);
obs (N Un, ), o1); obs obs(N (N(t,, Uzy), 02)
For each program of the class, our generator generated the parameter values as follows:
0, ~ U(-5,5), 02 U(0, 50), 0g = (02)², 0g U(0, 10)
6 = (0z)², 0, ~ U(0, 10), 0, = =0,)2, 0g ~ U(0.5,10)
6 = (0s)2, 0g ~ U(0.5, 10), 0g = (0&)²
and then generated the observations on and og by running the program (i.e., simulating the model) forward.
1.1 hierd
The model class is described as follows:
Mao := 01; Voo : 022; Vos = =0/; Voa : 0/1 my := 0s:
2b = 0g: d1 07; d2 0s: Uxz = 0'g; Vzz 010i
18

---

Under review as a conference paper at ICLR 2022
ag Mmmo: Van): a1 Nass var): d2 ~N(as, Voz):
b ~ N(mb, vbb);
= x d1i tg == a1 + t1: obs(NV(t,, Un 01);
t3 = b x d2: t1 := 02 + t3: obs(NV(ta, vx 02)
For each program of the class, generator generated the parameter values follows:
0 ~ U(-10, 0g ~ UO, 100), 8/ = (02)², 0g ~ U(O,
0'3 = (0z)², 04 ~ U0), - 004(),)² 0s ~ U(-5,5)
0. ~ U(0, 10), 8/ = -66),, 07 ~ U(-5,5), 0s ~ U(-5,5)
0g U(0.5, 10), 0g = (0a)², 010 U(0.5, 10), 010 (01a)²
and then generated the observations 0₁ and o2 by running the program forward where the values for ao, a1, 02, and 6 in this specific simulation were sampled as follows:
ag ~ U(mas x Voag Mao + 2 x VUao)
a1 U(as 2 x Vaa ag + 2 VVs,)
02 ~ U(as - 2 x Vuaz ag + 2 X V UU22)
U(m- 2 x / Tb. my + 0b)
E.1.4 cluster
The model class is described follows:
my. := 01; Uog == =6; mg. == 03: Vo2 := 8/:; U2 == 03::
91 N 2g, ): 92 ~ ~N(m992 192);
zero == 0; hund == 100;
t >N(erro,, hund); mi if (t1 > zero) 91 else g2:
(Ns,N)(,, 01);
12 N((eroo hund); ma if (t2 > zero) 91 else 92i
obs(N((m2) 0z);
t3 ~ NN(zero, hund); ma if (ts > zero) 91 else 92:
(Nm(, os))
: if (t4 > zero) 91 else 92: N((erro hund); m=:=
(N 04):
ts N((eroo hund); ms if (ts > zero) 91 else 92i
obs(N (Mm(s,, 0s)
For each program of the class, our generator generated the parameter values follows:
0, ~ U(- 15), 0g ~ U(0.5,50) 022 = =02)²
0g U(-15, 15), 02 U(0.5,50) 04- = (0,)²
0g U(0.5, 10), 0' = (0s)²
and then generated the observations 01:5 by running the program forward.
E. milky AND milkyo
The model class milky is described follows:
Mmmas := 01; Umass == 022; ci :1:=000 03: Fog = 0/: 02 = 05:
Uga = 0gi Uxi == Oz; Vag == 0'/i
mass ~ (mmmass Umass); );
mass1 := mass x C1: 91 ~ N (massi, );
mas82 mass + 02 g2 N (massz, Ug2
19

---

Under review as a conference paper at ICLR 2022
obs( U21). 01); obs(N(g2 U22).
For each program of milky, our generator generated the parameter values as follows:
0 ~ U 10), 02 U(O, 30), 0/2 (02)², 03 ~ U(-
0. 10), 04 = (0:)², 0g -(-5,5)) 0s U(0, 10)
0 (0n)², 07 U(0.5, 10), 8, (0r)², 0g U(0.5, 10)
8'g = (0&)
and then generated the observations and 02 by running the program forward.
Everything remained the for the milkyo class, except that the obs commands were extended to (N U2, ), [01. 02, 03, 01, 0sl) and obs(NN (g2, Ves), [oo, 07, 0g, Og, 010]), 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:
my, == 01; Us, == 82: mæ, = 03: Usa == 8/; U2 := =8/;
21 (m U222i 22 NM U22); ); = Rosenbrock(2) 22);
(br, ox),
where Rosenbrock(z1, = 0.05 x (ž) 1)² + 0.005 x (#2 - 21 2) For each program of the class, our generator generated the parameter values as follows:
U)- 02 ~ U(0, 5), 02 = (02)², , 0g U(-8,8)
01 ~ 50, 01 = (0,)², 0g ~ U(0.5, 10), 03 = (0z)²
and then generated the observation by running the program forward where the values for ø and 22 in this specific simulation were sampled as follows:
21 U(ma, 1.5 x UZi Mai 1.5 x U22)
22 ~ U(m22 1.5 x Uear M22 + 1.5 x / U02)
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-firs. order. In the experiments reported in this section, used minor extension of our probabilistic programming language with procedures taking one parameter.
F.2.1
extl
Fig. 10 shows the dependency graphs for all model types in ext1. The variables 2p, 2 and I1, 22.... 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 the other eleven types. To generate programs of the model type (1,1), we use the following program template:
m = =,; U2o == 03z U2 Ozg Uzg = 0/i Va := 03;;
~ Nm )z == nl(z0); z2 ~ N(1,, U2y); za No22, U2y);
obs(A((s,, ve ),)))
where nl(2) = 50/T x arctan(z/10), The generation involves randomly sampling the parameters of this template, converting the template into program in our language, and creating synthetic observations. Specifically, our generator generates the parameter values as follows:
0, U(-5,5), O2 ((0,20), 62 (02)². Og U(0, 20), 0's (0g)²
20

---

Under review as a conference paper at ICLR 2022
0 ~ U(0, 82 (0,)², 0g U(0.5, 10), 83 (8s)²
and generates the observation 0₁ by running the program forward where the values for 20:3 in this specific simulation were sampled (and fixed to specific values) as follows:
2u ~ U(m20 - x Vuzo: Mao + 2 X /uzo)
21 nl(20)
22 U( 2 x U2g1 21 + 2 X U22)
23 U(22 -2 2 Ueat 22 + (Uaz).
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
Fig. II shows the dependency graphs for all five model types in ext2. Programs of these five types are randomly generated by our program generator. As the extl case, we explain the generator only for model type, which corresponds to the first dependency graph in Fig. 11 To generate programs of this type, we use the following program template:
12z0 01; U2n = 02j V= 03; U23 = 0/5 Uz4 0zi Usa = 06: Uzo = 04;
U2, f'': U22 == fg: Uag = 8io: Oxs = f
2a N(mz0 Uzo): 21 ~ N2o Uz.); ; 22 == nl(z)); 23 ~ N N Uzs):
N221, Uza); 25 M22 Uz0); 26 N22,, )za);
(W U,,). o1): obs(NV(2,. Uzz), ).002) 02); obs (bs(N(z)) Uny), ) 0g): obs(N (2g, U2,), 04)
In order to generate a program of this model type and observations, generator instantiates the parameters of the template as follows:
0, U(-5,5), 02 U(0, 10), 0/2 (02)², 03 U(0, 10), 02 = (0z)², 04 U(0, 03 - (04)²
8. U(0, 10). 8 (0g)², 0g U(0, 10). 8/ (0,)², 0, 10), 8, (07)²
0g U00, 10), 0g = (0s)², 0g ~ U(0, 10), 0g = (0a)², 010 ~ U0), 00 = (01o)
0n1 ~ U(O, 10), 81 = (0,1)2.
Then, it generates the observations 01:4 by running the program forward where the values for 20.6 in this specific simulation were sampled (and fixed to specific values) as follows:
20 U(m
U2or M00 /UZz)
U(zo - 2 x VUZi, 20 + 2 x (C',)
22 nl(20)
23 U(zo - 2 x Vusgi + 2+ x /Uzz)
24 U(z1 2 / Uuzi 21 + /Uzu)
25 U(21 - 2 x VUZi 21 + 2 x /Uz,)
26 U(22 2 Uzer 22
/Uza).
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 the first dependency To generate programs of this type, we use the following program template:
graph in Fig. 12
01; Uzo
022 Uz = 0'g: Vay
0/=
21

---

Under review as a conference paper at ICLR 2022
20 N( Uzo); 21 N U21); 22 mm(21); obs(N(22 (22,221),00 u21).01
where mm (z) 100 x 2²9 /100 + For each program in this model type, our generator instantiates
the parameter values as follows:
0 ~ (),, 02 ~ U(0, 20), 82 = (02)², 0z U(0, 20), 0g = (0a)²
01 ~ U005, 10), 0/ = (0,)²
and synthesises the observation 0n by running the program forward where the values for 20.2 in this specific simulation were sampled (and fixed to specific values) as follows:
20 U( - Vuzoi Ma0 + x /(zz)
U(zo - 2 x VULi 20 x vU2,)
22 mm(21).
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.
F DETAILED EVALUATION SETUP
In our evaluation, the dimension of the internal state h was 10 (i.e., h E R 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 nide where the hidden dimensions were 50. The hyper-parameter in our optimisation objective (54) was set to À = 2 in the evaluation. For HMC, we used the NUTS sampler (Hoffman & Gelman 2014 We did 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 superficial difference between programs caused by different variable names, and enables us to avoid unnecessary complexity caused by variahle-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. [3a 13c and 13d shows cases where the test loss surges. This was when the loss of only 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 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 meta- algorithm 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 inference algorithm. Second, our random program generator placed some restriction on the programs generated (e.g., by using tight boundaries 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 converge for only one fifth of the test programs.
Fig. 14 shows the similar comparison plots between reference and predicted marginal posteriors for 10 test programs of the rb type, after 52.4K epochs. Our inference algorithm computed the posteriors precisely for most of the programs except two (pgm75 and pgm79) with significant multimodality. The latent variable 7g z0 had least two modes around 10 (visible the figure) and around 10 (hidden in the figure Our inference algorithm showed mode-seeking behavior for this latent
TThe blue reference plots were drawn using an HMC chain, but the HMC chain got stuck in the mode around -10 for this variable.
22

---

Under review as a conference paper at ICLR 2022
variable. Similarly, the variable pgm79_ z0 had least 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 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, that the samples used in the training are often biased (i.e., only from 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.
current results for cluster suffer from the fact
I TRAINING AND TEST LOSSES FOR THE OTHER CASES IN ext1
Fig. I5 shows the average training and test losses in the extl experiment runs (under three different random sceds) 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. 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.
K QUANTIFIED ACCURACY OF PREDICTED POSTERIORS
For accuracy, would ideal to report [L[p|lg]. where p is the fully joint target posterior and is the predicted distribution. It is, however, hard to compute this quantity since often cannot compute the density of p. One (less convincing) alternative to compute (=) for a latent variable 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 [L[p' (=) (=)) 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 (=)ll())) over all the latent variables in the test programs. For instance, if there were 90 test programs and each program had three latent variables, we averaged 90 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
270 KL measurements, In an experiment run for ext1 (which tested generalisation
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 the predicted posterior, while the estimation 40.9 when 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
L PROGRAM $5.3
Fig.1
17 shows the program that is reported in 55.3 written in our probabilistic programming language.
M DISCUSSION OF THE COST OF IS-PRED vs. IS-PRIOR
Our approach (IS- (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
23

---

Under review as a conference paper at ICLR 2022
predicted proposal, while IS-prior only needs to scan the program once. Although it may that IS- 1S-prior has a huge advantage in terms of saving the wall-clock time, our observation 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
N LIMITATIONS AND FUTURE WORK
Currently, a learnt inference algorithm in our work does not generalise to programs with different sizes (Yan al, 2020). e.g., from clustering models with clusters those with ten clusters. Each model class assumes fixed number of variables, and the neural networks crucially exploit the assumption. Also, our meta-algorithm does not scale practice. When applied to large programs, e.g., state-space models with few hundred time steps, it cannot learn optimal inference algorithm within a reasonable amount of time. Overcoming these limitations 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 appropriate form of the approximation distribution with rich dependency structure, by, e.g., incorporating the ideas from Ambrogioni et al. (202T). This direction is closely related to automatic guide generation in Pyro (Bingham et al. 2018).
24

---

Under review as a conference paper at ICLR 2022
50
21
20
22
11
11
22
23
22
23
T2
23
T1
I1
22
22
T'
(a) Model type (1,1).
(b) Model type (1,2).
(c) Model type (1,3).
(d) Model type (1.4).
ZU
21
22
23
22
23
23
32
2/2
Z/
F1
F1
22
J1
3T
J3
(c) Model type (2,1).
(f) Model type (2,2).
(g) Model type (2,3).
(h) Model type (2.4).
20
21
22
1
21
22
ZO
23
22
Z3
23
T2
Z
Z2
Z3
F1
F1
F2
J1
F22
Σ3
Model type (3.1).
ii) Model type (3,2).
(k) Model type (3,3).
() Model type (3,4).
Figure 10: Canonicalised dependency graphs for all 12 model types in ext1. The rows are for different positions of the nl variable, and the columns are for different dependency graphs: the model type (i.j) means one of the 12 model types in ext1 that corresponds to the i-th position of the nl variable and j-th dependency graph.
25

---

Under review as a conference paper at ICLR 2022
20
20
1
22
23
21
2/2
Z3
24
25
Z6
F4
24
25
26
11
T2
F3
22
T'ss
(a) Model type for the 1st dependency graph.
(b) Model type for the 2nd dependency graph.
Z0
20
1
7/2
21
22
23
A4
23
24
26
Z6
25
26
I
TA
T2
J3
T:4
31
J/2
(c) Model type for the 3rd dependency graph.
(d) Model type for the 4th dependency graph
zz
21
22
23
ZA
25
26
Z
F2
J3
AA
(c) Model type for the 5th dependency graph.
Figure 11: Canonicalised dependency graphs for all five model types in ext2.
20
20
11
22
22
Z
Z/2
F1
F1
F1
F2
(a) 1st model type.
(b) 2nd model type.
(c) 3rd model type.
Figure 12: Canonicalised dependency graphs for all three model types in the mulmod class.
26

---

Under review as a conference paper at ICLR 2022
Seed, Tr Te
Seed, Tr
1.
1. Tr
1. Te
1. Te
2. Tr
2. Tr
2. Te
2. Te
3.
3. Tr
B.
300
3, Te
(a) hierd
(b) cluster
Seed. Tr or Té
Seed, Tr or Te
Tr
1.
Te
Tr
Tr
5000
Opech
Eooch
(c) milkyo
(d) rb
Figure 13: Losses for hierd, cluster, milkyo, and rb. The y-axes are log scalec. The surges in later epochs of Fig.
13a
13c
and 13d were due to only single or few test programs out of 50.
Figure 14: Comparisons reference and predicted marginal posteriors for 10 programs in the rb test set.
27

---

Under review as a conference paper at ICLR 2022
197
Seed, Tr or
Seed,
207
1.
1.
Te
2. Tr
2. Tr
130
2. Te
2. Te
3. Tr
3. Tr
3, Te
(a) To 4th dep. graph.
(b) nl position.
Seed, Tr Te
107
Seed. Tr or To
1. Tr
1
1. Te
2.
2.
1808
2000
1000
3.00
2000
Opech
(c) To 2nd ul position.
(d) To 3rd ul position.
Figure 15: Average training and test losses for generalisation to the last (4th) dependency graph and all three positions of the nl variable ext1. The y-axes are log-scaled.
Seed, Tr or Te
Seed, or Te
2.
100
Tr
0
2.
2. Te
3,
3, Te
Spoch
(a) To 4th dep. graph.
(b) To 5th dep. graph.
Figure 16: Average training and test losses for generalisation to the 4th and 5th dependency graphs in ext2. The y-axes are log-scaled.
= 3.93; b = 348.16; c c: 57.5; d = 14.04; e:= 40.34; 21 (a, 6); 22 NN(2, c); 23 mm(21):
obs (N(22, d), 53.97); abs (N),.
Figure 17: The program that is reported in 5.3 written in our probabilistic programming language.
28


Reviews:
[{"review_id": "T-HuJ0gO239", "reviewer": "Reviewer_mmqX", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- The paper claims that the learned inference algorithm works well for tasks which are \"similar\" to the training problems, but the notion of similarity is not fully defined, nor is there an example of how the system fails when applied to a dissimilar program. Are there diagnostics for checking the output when the model is applied to a program which is too dissimilar?\n- The experimental study tests performance within a family of programs, each using small neural networks to infer each program statement. Does training on all families of programs allow the system to make accurate inferences on any of those families?  Does it allow it to generalise across families? Without some notion of how the system generalises I'm not sure when I would choose to use this rather than running HMC on my program, given a single HMC run will be faster than training the neural networks on multiple HMC runs for different programs \"close\" to the program of interest.\n- The experiment in section 5.3 shows that the system is around 2x faster than importance sampling from the prior, but this doesn't take into account the time necessary to train the neural nets nor the time taken for the importance sampling runs used to generate the training data.\n- How are losses propagated through the program? If each neural net is 3 layers, then programs with 10 statements have at most 30 layers, which is usually past the point where some amount of regularization or normalization is necessary to stabilize training or prevent vanishing gradients. Could the authors comment on the stability of training?\n- What's the failure mode when the test loss diverges? Is it detectable without having HMC or other high quality samples?\n- How robust is this approach to differing choices of neural net architecture?  The paper uses a 10 dimensional state space when parsing the program, but it's not clear how this value should be modified as the number of latent variables or the program complexity changes.\n- Overall the paper is well written and explained, and the experimental study is detailed for the areas it covers.", "summary_of_the_paper": "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.", "summary_of_the_review": "The paper presents a learned inference algorithm, but it's not clear how it generalises either across program types (which is necessary to amortize the training cost wrt HMC), or across neural network architectures (e.g., changing the internal state space in response to increased program complexity).  Additionally it's noted that occasionally the test error diverges, but there's no discussion of how to detect this in practice if the system was used for inference.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "OmUIjdQuT-r", "reviewer": "Reviewer_mFVo", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This work is very interesting and novel. It's a unique attempt\nto learn a general inference algorithm. I think meta learning\nis something that is a good fit for many Bayesian approaches\nand I want to see more work like this.\n\nI'm curious about the expressivity of the language. The\ngrammar suggestions a modelling language that consists of\nsequence of commands. I don't see how this language would be\nable to express recursive programs. It seems one would need\nsomething like a label and jump commands to accomplish\nthat. This is admitted in the appendix but not really\nacknowledged in the main text. I think the main paper should\nreflect the present limitations and not over promise the\nexisting contribution.\n\nI have some concerns about the experiments. Some examples are\nfairly simple and the results for the more substantial ones\nare not shown (like hierd and rb). The test losses look fairly\nbad for the experiments that are shown, so I'm not fully sure\ngeneralisation has been demonstrated. If the issue is a few bad\n generated programs maybe show median loss? The results in Figure 4\nseem to be for during training, but feasibility presumably\nrequires similar results on unseen programs at test time?\n\n I also worry about correctness. As the neural networks are\nused as is and not as a proposal, how much can we trust\nthe posteriors that come out of this method? Is there anything\nthat can be said about the learnt posterior? It seems right now\nthat the learned distribution can recover the mean of the true\nparameters and maybe the variance?\n\nSome of the writing is slightly sloppy. For example the phrase\n\"so-called static single assignment assumption\" is used. I don't\nknow what that means, but I do know there is a static single\nassignment intermediate representation that exists within many\ncompilers. I think that's what the paper meant to refer to.\n\nRelated work should cover how this approach differs from Stites et al.\n(https://arxiv.org/abs/2103.00668). ", "summary_of_the_paper": "This introduces a meta-learning algorithm for learning\ninference algorithms applicable to any probabilistic\nprogram. This is accomplished by associating a neural network\nwith every grammar rule of a probabilistic programming\nlanguage and outputting posterior draws. This generalisation is\npossible because each neural network component is feed marginal\nlikelihood information for each PPL instruction.\n\t", "summary_of_the_review": "I recommend this paper for rejection. While I think the approach has the potential to work, right now it's very hard to get a sense of what was learned by the meta learning algorithm, how well any of this generalises when model structure or even observed data changes significantly, or even if the language is too restricted to make this is a significant enough contribution.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "MXt1h2BBnf6", "reviewer": "Reviewer_BQrG", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The main drawback in this work is a lack of novelty. The use of neural networks to train a proposer for a model is not new. While the authors attempt to cast their work as something different than IC this doesn’t quite come out. The claim that a neural network for one program can be used for a different program even though the neural network takes a one hot encoded representation of the variables in the model is hard to see. A clear technical statement of what kind of cross model generalization is possible is needed. \n\nThe paper shows results across model structures where the dependency graph and the position of a function changes but the number of variables is the same. It is not clear why IC can’t deal with this minor variation in the same probabilistic program. \n\nThese models are so simple, how hard would it be to train ic on these models and then do inference? I would like to have seen ic results in this paper to believe that this work is different. \n\nThe language chosen is not a universal ppl. I can’t follow how this can be used as an intermediate language by a compiler for a universal ppl as claimed on page 3. Please show an example of how a program with unbounded random variables can be compiled into this language. \n\nThe inference in this paper looks like mean field variational inference. Which makes me wonder whether hmc is really such a good comparison. Please show some comparisons to vi. In Stan this would be trivial to run since you are already running hmc. \n\nThe models are very simplistic with no discrete variables and no multimodal posteriors. It is not a meaningful claim to make (footnote on page 9) that the inference algorithm provides a good coverage of the posterior by covering all the modes. The posteriors shown for multimodal models should at least look multimodal.  Variational distance or symmetric kL divergence results would be needed to make claims about correctness of the posterior. \n\nRegarding ess per second results these can be misleading. The algorithm might have a cap of ess for example. It would be better to run your algorithm for the same duration as hmc and show higher ess numbers.\n", "summary_of_the_paper": "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. \n\n", "summary_of_the_review": "The claim that the paper provides generalization of compiled inference across models is not supported by the description or the simple examples. These appear to be covered by existing work on inference compilation. \n\nThe focus on a very restricted class of ppls makes this work very limited. I don’t believe I learned anything from this paper. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "9Wiab7HRVGc", "reviewer": "Reviewer_VGsf", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "There were quite a few terms that I was not familiar with - for example, what is the \"state\" of an algorithm?  I did not find a formal definition of this term in the paper.  In the definition of the INFER function, what does the keyword \"in\" mean?  It did not seem clear how well the white-box inference algorithm performs when the number of commands in the program is very large. ", "summary_of_the_paper": "This paper proposes an algorithm for computing an approximation of the posterior and marginal likelihood by analysing the sequence of programs using neural networks, as well as a meta-algorithm for learning the network parameters over a training set of probabilistic programs.  Experiments demonstrate the feasibility of the meta-algorithm for learning inference algorithms that generalise well to new but similar programs; these learnt algorithms were sometimes found to outperform alternatives in terms of time-efficiency. ", "summary_of_the_review": "I am not at all familiar with probabilistic programming; this does look like a serious piece of work, though I do not know how novel it is or whether the claims in the paper are correct.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents a meta-algorithm for learning a posterior-inference algorithm for restricted probabilistic programs. While the reviews agree that this is a very interesting research direction, they also reveal that there are several questions still open. One reviewer points out that there learning to infer should take both the time for learning+inference and the generalization to other programs into account, i.e., what happens if the program is too different from the training set? Is benefit than vanishing? Moreover, as pointed out by another review, recursion as well as while loops are not yet supported. Also, the relation to IC needs some further clarification. These issues show that the paper is not yet ready for publication at ICLR. However we would like to encourage the authors to improve the work and submit it to one of the next AI venues.
Decision: Reject