Meta-Learning an Inference Algorithm for Probabilistic Programs
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 efﬁcient 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 analyses the given program sequentially using multiple neural networks to compute an approximate posterior.
The parameters of these 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 efﬁciency than alternatives such as HMC.
Introduction
The development of performant probabilistic programming systems[Goodman et al., 2008, Wood et al., 2014, Mansinghka et al., 2014, Minka et al., 2018, Narayanan et al., 2016, Salvatier et al., 2016, Carpenter et al., 2017, Tran et al., 2016, Ge et al., 2018, Bingham et al., 2018] also revealed the difﬁculty of achieving efﬁciency and universality simultaneously, and the need for equipping probabilistic programming languages (PPLs) with mechanisms for customising inference or learning algorithms to a given domain.
In fact, recent PPLs include constructs for specifying conditional independence in a model [Bingham et al., 2018] or deﬁning proposals [Ritchie et al., 2015, Siddharth et al., 2017, Bingham et al., 2018, Tran et al., 2018, Cusumano-Towner et al., 2019], all enabling users to help the algorithms.
In this paper, we present a meta-algorithm for learning a posterior-inference algorithm itself from a given set of restricted probabilistic programs.
The meta-algorithm aims at constructing a customised inference algorithm for the given set of models, while ensuring universality to the extent that the constructed algorithm generalises to similar but unseen programs.
The distinguished feature of our approach is the use of what we call a white-box inference algorithm, which is equipped with multiple neural networks, one for each type of atomic command in a PPL, and computes an approximate posterior for a given program by analysing individual atomic commands in it sequentially using these networks.
For instance, our white-box inference algorithm regards the program in Fig.
1 as a sequence of the ﬁve atomic commands, initialises its internal state h ∈ R m with h 0 , and transforms the state over the sequence.
The internal state h is the encoding of an approximate posterior at the current program point, which corresponds to an approximate ﬁltering distribution of a state-space model.
The update of this state for each atomic command is directed by neural networks.
Our meta-algorithm trains the parameters of these networks by trying to make the inference algorithm compute accurate posterior approximations over a training set of probabilistic programs.
We discuss related work
[Andrieu and Thoms, 2008, Hoffman and Gelman, 2014, Wang et al., 2018, Gong et al., 2019, Wu et al., 2020, Gordon et al., 2019, Iakovleva et al., 2020, Minka, 2001, mass ∼ N (5 , 10); // log of the mass of Milky Way g 1 ∼ N ( mass × 2 , 5); obs ( N ( g 1 , 1) , 10); // observed velocity vel 1 =10 of the ﬁrst satellite galaxy g 2 ∼ N ( mass + 5 , 2); obs ( N ( g 2 , 1) , 3) // observed velocity vel 2 =3 of the second satellite galaxy
Wainwright and Jordan, 2008, Jitkrittum et al., 2015, Gershman and Goodman, 2014, Le et al., 2017,
Paige and Wood, 2016, Stuhlmüller et al., 2013, Kingma and Welling, 2013, Mnih and Gregor, 2014,
Rezende et al., 2014, Ritchie et al., 2016, Marino et al., 2018, Zaremba and Sutskever, 2014, Bieber et al., 2020, Reed and de Freitas, 2016] in Appendix A.
Our contributions are as follows: (i) we present a white-box posterior-inference algorithm, which works directly on model description and can be customised to a given model class; (ii) we describe a meta-algorithm for learning the parameters of the inference algorithm; (iii) we empirically analyse our approach with different model classes, and show the promise as well as the remaining challenges.
Our Approach
Setup Our results assume a simple PPL without loop and with a limited form of conditional statement.
The syntax of the language is given by the following grammar, where r represents a real number, z and v i variables storing a real, and p the name of a procedure taking two real-valued parameters and returning a real number:
Programs C ::=
A | C 1 ; C
Atomic Commands A ::=
z ∼ N ( v 1 , v 2 ) | obs ( N ( v 0 , v 1 ) , r ) | v 0 :
= if ( v 1 > v 2 ) v 3 else v| v 0 :=
r | v 0 :
= v 1 | v 0 : = p ( v 1 , v 2 )
Programs in the language are constructed by sequentially composing atomic commands.
The last atomic command v 0 :
= p ( v 1 , v 2 ) is a call to one of the known deterministic procedures, which may be standard binary operations such as addition and multiplication, or complex non-trivial functions that are used to build advanced, non-conventional models.
When p is a standard binary operation, we use the usual inﬁx notation and write, for example, v 1 + v 2 , instead of +( v 1 , v 2 ) .
We provide full discussion of (our choice of) the PPL in
Appendix B.
The formal semantics of the PPL is in
Appendix C. Our white-box inference algorithm aims at computing the approximate posterior and marginal likelihood estimate for a given program C , when the unnormalised density p C (denoted by
C ) has a ﬁnite non-zero marginal likelihood and, as a result, a well-deﬁned posterior density.
White-Box Inference Algorithm Consider, for presentation, programs that sample n -many latent variables z 1 , . . .
, z n and use at most m -many variables ( m ≥ n ).
Let V be [0 , 1] m , the space of the one-hot encodings of those m variables, and P the set of procedure names.
Given a program
C =
( A 1 ; . . .
; A k ) , our white-box inference algorithm computes an approximate posterior and a marginal likelihood estimate for C by sequentially processing the A i ’s. The computation for each A i is directed by one of the neural networks deﬁned below, according to the type of A i :
[Equation]
× R s → R , where φ 1:7 and φ p for p ∈ P are network parameters.
The top six networks are for the six types of atomic commands, nn de ,φ 6 is a decoder that maps the states h to the probability densities (means and variances in our setup) over latent variables, and nn intg ,φ 7 is used when the algorithm updates the marginal likelihood estimate for an observe statement obs ( N ( v 0 , v 1 ) , r ) .
The full details for nn intg ,φ 7 and the derivation of the marginal likelihood are provided in Appendix D.
Given a program C =
( A 1 ; . . .
; A k ) that draws n samples (and so uses latent variables z 1 , . .
. , z n ), the algorithm approximates the posterior and marginal likelihood of C as follows:
INFER ( C ) =
let ( h 0 , Z 0 ) =
( (cid:126) 0 , 1) and ( h k , Z k ) =
( INFER ( A k ) ◦ . . .
◦
INFER ( A 1 ))
( h 0 , Z 0 ) in let (( µ 1 , σ 21 ) , . . .
, ( µ n , σ 2 n ))
=
nn de ,φ 6 ( h k ) in return (cid:16) (cid:81) ni =1 N ( z i | µ i , σ 2 i ) , Z k (cid:17) , where INFER ( A i ) :
R s × R → R s ×
R picks an appropriate neural network based on the type of A i , and uses it to transform h and Z :
[Equation]
Seed, Tr or Te 1, Tr 1, Te 2, Tr 2, Te 3, Tr 3, Te
Epoch
L o ss (a) gauss
Seed, Tr or Te 1, Tr 1, Te 2, Tr 2, Te 3, Tr 3, Te
Epoch
L o ss (b) hierl   4000  10k
Seed, Tr or Te 1, Tr 1, Te 2, Tr 2, Te 3, Tr 3, Te
Epoch
L o ss (c) milky
[Figure] [Figure]
[Figure] [Figure]
[Equation]
h ) , Z ) .
where v 0: k refers to the sequence of the one-hot encodings of variables v 0 , . .
, v k .
Meta-Learning Parameters Given a training set of programs D = { C 1 , . . .
, C N } , we learn the parameters φ = ( φ 1:7 , ( φ p ) p ∈ P ) by solving the following optimisation problem:[-
-]
[Equation]
where λ > 0 is a hyper-parameter, N C is the marginal likelihood (cid:82) p C ( z 1: n ) dz 1: n for p C , the next π C ( z 1: n ) is the normalised posterior p C ( z 1: n ) /N
C for C , and the last q C and Z C are the approximate posterior and marginal likelihood estimate computed by the inference algorithm (that is, ( q C ( z 1: n ) , Z C ) =
INFER ( C ) ).
We optimise the objective by stochastic gradient descent, where the gradients can be approximated by sampling techniques.
We describe the full algorithm in Appendix E.
Empirical Results
This section describes our empirical results.
See Appendix F for the full list of our model classes,
Appendix G for the detailed experimental setup, and Appendix N for limitations and future work.
Generalisation to New Model Parameters and Observations We evaluated our approach on six model classes: Gaussian ( gauss ), two hierarchical ( hierl and hierd ), clustering ( cluster ), Milky Way ( milky and milkyo ), and Rosenbrock models ( rb ) (see Appendix F.1).
For each model class, we used 400 programs to meta-learn an inference algorithm, and applied the learnt algorithm to 50 unseen test programs.
We measured the average test loss over the 50 test programs, and also compared the marginal posteriors predicted by our learnt inference algorithm with the reference marginal posteriors that were computed analytically, or approximately by HMC using 500 K samples after 50 K warmups.
In case of HMC, we computed the marginal sample means and standard deviations using one of the 10 independent, converged Markov chains.
All training and test programs were automatically generated by a random program generator.
The training setup is described in Appendix G.1.
[Figure]
[Figure]
S u m o f s e c o n d m o m e n t s IS-
[Figure]
HMC (a) Moments by HMC.
1.86ms 2.33ms 4.9ms 6.65ms [Figure] Elapsed
[Figure]
S u m o f s e c o n d m o m e n t s
IS [Figure] IS
[Figure]
pred IS
[Figure]
prior
[Figure]
Fig.
2 shows the training and test losses for gauss , hierl , and milky under three random seeds.
The losses for the other model class are in Appendix H.
The training loss was averaged over the training set and over the 8 batch updates.
The test loss was averaged over the test set.
The decrease of training losses led to the downturns of the test losses.
The later part of Fig.
2c 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; the losses of the rest remained small.
Fig.
3 compares, for 10 test programs in hierl , the reference marginal posteriors (blue) and their predicted counterparts (red) by the learnt inference algorithm instantiated before training and after 2 K epochs.
It shows that the learnt inference algorithm predicts the reference posteriors precisely for most of the variables.
We observed similar patterns for the other cases, except for cluster and rb (see Appendix I).
Generalisation to New Model Structures We let two kinds of model structure vary across programs: the dependency (or data-ﬂow) graph for the variables of a program and the position of a nonlinear function nl( x )
= 50 /π × arctan( x/ 10) in the program.
We speciﬁcally considered two classes of models, ext1 (four-variable models that are grouped into 12 types according to the dependency graph and the position of nl ) and ext2 (seven-variable models that are grouped into ﬁve types according to the dependency graph).
See Appendix F.2 for full description.
We ran seven experiments for ext1 ; three evaluated generalisation to unseen dependency graphs, and four evaluated generalisation to unseen positions of nl .
We ran ﬁve experiments for ext2 , each of which evaluated generalisation to an unseen dependency graph.
All the experiments were repeated three times under different random seeds, and so the total numbers of experiment runs were 21 and 15 for ext1 and ext2 , respectively.
The rest of the setup was similar to the generalisation-to-new-parameters case, and we detail the full setup in Appendix G.2.
In 17 runs (out of 21 ) for ext1 , the test losses stabilised or reduced as the training losses decreased, even when the test losses were high and ﬂuctuated in earlier training epochs.
In 8 runs (out of 15 ) for ext2 , the test losses were stabilised as the training losses decreased; in 4 runs out of the other 7 , the test losses increased only slightly.
Appendix J and K show all the losses.
In terms of predicted posteriors, we observed highly accurate predictions in 8 runs of ext1 .
For ext2 , the predicted posteriors were precise in 7 runs.
Overall, the learnt algorithms generalised to unseen types of models well or fairly well in many cases.
Test-Time Efﬁciency in Comparison with Alternatives We considered three-variable models ( mulmod ) where two latent variables follow normal distributions and the other stores the value of the function mm( x ) = 100 × x 3 / (10 + x 4 ) .
The models are grouped into three types deﬁned by their dependency graphs and the positions of mm in the programs (see Appendix F.3).
We ran our meta-algorithm using 600 programs from all three types using importance samples (not HMC samples).
Then for 60 test programs from the last model type, we measured ESS and the sum of second moments along the wall-clock time using three approaches: importance sampling (IS-pred; ours) with the predicted posteriors as proposal, importance sampling (IS-prior) with prior as proposal, and HMC.
As the reference sampler, we used importance sampling (IS-ref) with prior as proposal using 5 M samples.
All the approaches were repeated 10 times.
We report the results for a program (see Appendix L) in the test set.
Fig. 4a and 4b show the moments estimated by HMC and IS-{pred, prior}, respectively, in comparison with the same (across the two ﬁgures) reference moments by IS-ref.
The estimates by IS-pred (red) converged to the reference (green) more quickly ( 18 ms) than those by HMC (orange; over 84 s) and IS-prior (blue; higher variance).
We obtained similar results for ESS, and the results are in Appendix M.
A Related Work
The difﬁculty of developing an effective posterior-inference algorithm has motivated active research on adapting key components of an inference algorithm.
Techniques for adjusting an MCMC proposal [Andrieu and Thoms, 2008] or an HMC integrator [Hoffman and Gelman, 2014] to a given inference task were implemented in popular tools.
Recently, methods for meta-learning these techniques themselves from a collection of inference tasks have been developed [Wang et al., 2018, Gong et al., 2019].
The meta-learning approach also features in the work on stochastic variational inference (VI) where a variational distribution receives observations for each inference task and is trained with a collection of datasets of observations [Wu et al., 2020, Gordon et al., 2019, Iakovleva et al., 2020].
For a message-passing-style VI, such as expectation propagation [Minka, 2001, Wainwright and
Jordan, 2008], Jitkrittum et al.
[2015] studied the problem of learning a mechanism to pass a message for a given single inference task.
A natural follow-up is to meta-learn such a mechanism from a dataset of multiple inference tasks that can generalise to unseen models.
Our white-box inference algorithm can be viewed as a message-passing-style VI that can meta-learn the representation of messages and a mechanism for passing them for given probabilistic programs.
Amortised inference and inference compilation [Gershman and Goodman, 2014, Le et al., 2017,
Paige and Wood, 2016, Stuhlmüller et al., 2013, Kingma and Welling, 2013, Mnih and Gregor, 2014,
Rezende et al., 2014, Ritchie et al., 2016, Marino et al., 2018] are closely related to our approach in that they also attempt to learn a form of a posterior-inference algorithm.
However, the learnt algorithm by them and that by ours have different scopes.
The former is designed to work for unseen inputs or observations of a single model, while the latter for multiple models with different structures.
The idea of running programs with learnt neural networks also appears in the work on training neural networks to execute programs [Zaremba and Sutskever, 2014, Bieber et al., 2020, Reed and de Freitas, 2016].
As far as we know, however, we are the ﬁrst to frame the problem of learning a posterior-inference algorithm as the one of learning to execute.
B Full Discussion about the Choice of the PPL
Programs C ::=
A | C 1 ; C
Atomic Commands A ::=
z ∼ N ( v 1 , v 2 ) | obs ( N ( v 0 , v 1 ) , r ) | v 0 :
= if ( v 1 > v 2 ) v 3 else v| v 0 := r | v 0 : = v 1 | v 0 : = p ( v 1 , v 2 )
Programs in the language are constructed by sequentially composing atomic commands.
The language supports six types of atomic commands .
The ﬁrst type is z ∼ N ( v 1 , v 2 ) , which draws a sample from the normal distribution with mean v 1 and variance v 2 , and assigns the sampled value to z .
The second command, obs ( N ( v 0 , v 1 ) , r ) , states that a random variable is drawn from N ( v 0 , v 1 ) and its value is observed to be r .
The next is a restricted form of a conditional statement that selects one of v 3 and vdepending on the result of the comparison v 1 > v 2 .
The following two commands are different kinds of assignments, one for assigning a constant and the other for copying a value from one variable to another.
We permit only the programs where a variable does not appear more than once on the left-hand side of the := and ∼ symbols.
This means that no variable is updated twice or more, and it corresponds to the so-called static single assignment assumption in the work on compilers.
This restriction lets us regard variables updated by ∼ as latent random variables.
We denote those variables by z 1 , . . .
, z n .
We use this simple language for two reasons.
First, the restriction imposed on our language enables the simple deﬁnition of our white-box inference algorithm.
For instance, the language supports only a limited form of conditional statements and restricts the syntactic forms of atomic commands; the arguments to a normal distribution or to a procedure p should be variables, not more general expression forms such as addition of two variables.
As we will show soon, this restriction makes it easy to exploit information about the type of each atomic command in our inference algorithm; we use different neural networks for different types of atomic commands in the algorithm.
Second, the language is intended to serve as an intermediate language of a compiler for a high-level PPL, not the one to be used directly by the end user.
The compilation scheme in, for instance, [van de Meent et al., 2018] from high-level probabilistic programs with general conditional statements and for loops to graphical models can be adopted to compile such programs into our language.
Programs with recursion or while loops cannot generally be translated into our intermediate language, since such programs may go into inﬁnite loops while the programs in our language always terminate.
Programs with for loops and general branches can in theory be translated into a less expressive language such as ours.
For example, van de Meent et al.
[2018] explain a language called FOPPL (Section 2), which has for loops and branches, and the translation of FOPPL into graphical models (Section 3).
We think that these graphical models can be translated into programs in our language.
Of course, this does not mean that the learnt inference algorithm would interact well with the compilation; the interaction between compilation and inference in the context of meta-learning is something to be explored in future work.
C Formal Semantics of the PPL
Probabilistic programs in the language denote unnormalised probability densities over R n for some n .
Speciﬁcally, for a program C , if z 1 , . . .
, z n are all the variables assigned by the sampling statements z i ∼ N ( . . . )
in C in that order and C contains m observe statements with observations r 1 , . . .
, r m , then C denotes an unnormalised density p C over the real-valued random variables z 1 , . . .
, z n :
[Equation]
where x 1 , . .
, x m are variables not appearing in C and are used to denote observed variables.
This density is deﬁned inductively over the structure of C .
Here z 1 , . .
z n are all the variables assigned by the sampling statements z i ∼ N ( . . . )
in C in that order, the program C contains m observe statements with observations r 1 , . . .
, r m , and these observed random variables are denoted by x 1 , . . .
, x m .
The goal of this section is to provide the details of our statement.
That is, we describe the formal semantics of our probabilistic programming language, and from it, we derive a map from programs C to unnormalised densities p C .
To deﬁne 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 the main text.
The type system lets us derive the following judgements for programs C and atomic commands A :
[Equation]
where S and T are sequences of distinct variables, V and W are sets of variables that do not appear in S and T , respectively, and α and β are sequences of reals.
The ﬁrst judgement says that if before running the program C , the latent variables in S are sampled in that order, the program variables in V are updated by non-sample statements, and the real values in the sequence α are observed in that order, then running C changes these three data to T , W , and β .
The second judgement means the same thing except that we consider the execution of A , instead of C .
The triples ( S, V, α ) and ( T, W, β ) serve as types in this type system.
The rules for deriving the judgements for C and A follow from the intended meaning just explained.
We show these rules below, using the notation @
for the concatenation operator for two sequences and also set( S ) for the set of elements in the sequence S :
[Equation]
v 2 )) : ( S @
[ z ] , V, α ) v 0 , v 1 ∈ set( S ) ∪ V ( S, V, α ) (cid:96) 2 obs ( N ( v 0 , v 1 ) , r ) : ( S, V, α @
[ r ]) v 0 (cid:54)∈ set( S ) ∪ V v 1 , v 2 , v 3 , v 4 ∈ set( S ) ∪ V ( S, V, α ) (cid:96) 2 ( v 0 :
= if ( v 1 > v 2 ) v 3 else v 4 ) : ( S, V ∪ { v 0 } , α ) v 0 (cid:54)∈ set( S ) ∪ V ( S, V, α ) (cid:96) 2 ( v 0 :
= r ) : ( S, V ∪ { v 0 } , α ) v 0 (cid:54)∈ set( S ) ∪ V v 1 ∈ set( S ) ∪ V ( S, V, α ) (cid:96) 2 ( v 0 :
= v 1 ) : ( S, V ∪ { v 0 } , α ) v 0 (cid:54)∈ set( S ) ∪ V v 1 , v 2 ∈ set( S ) ∪ V ( S, V, α [Equation] (cid:96) 2 ( v 0 :
= p ( v [Equation]
, v 2 )) : ( S, V ∪ { v 0 } , α )
We now deﬁne our semantics, which speciﬁes mappings from judgements for C and A to mathematical entities.
First, we interpret each type ( S, V, α ) as a set, and it is denoted by (cid:74) ( S, V, α ) (cid:75) :
(cid:74) ( S, V, α ) (cid:75) =
{ ( p, f, l ) | p is a (normalised) density on R | S | , f = ( f v ) v ∈ set( S ) ∪ V , each f v is a measurable map from R | S | to R , l is a measurable function from R | S | × R | α | to R + } , where | S | and | α | are the lengths of the sequences S and α , and R + means the set of positive reals.
Next, we deﬁne the semantics of the judgements ( S, V, α ) (cid:96) 1 C : ( T, W, β ) and ( S, V, α ) (cid:96) [-
-]2 A : ( T, W, β ) that can be derived by the rules from above.
The formal semantics of these judgements, denoted by the (cid:74) − (cid:75) notation, are maps of the following type:
[Equation]
The semantics is given by induction on the size of the derivation of each judgement, under the assumption that for each procedure name p ∈ P , we have its interpretation as a measurable map from Rto R : [Equation]
We spell out the semantics below, ﬁrst the one for programs and next that for atomic commands.
[Equation]
Let N ( a ; b, c ) be the density of the normal distribution with mean b and variance c when c > 0 andwhen c ≤ 0 .
For a family of functions f = ( f v ) v ∈ V , a variable w (cid:54)∈ V , and a function f (cid:48) w , we write f ⊕ f (cid:48)
w for the extension of f with a new w -indexed member f (cid:48) w .
[Equation]
( where p (cid:48) ( a 1: | S | +1 )
= p ( a 1: | S | ) × N ( a | S | +1 ; f v 1 ( a 1: | S | ) , f v 2 ( a 1: | S | )) , f (cid:48) v ( a 1: | S | +1 ) =
f v ( a 1: | S | ) for all v ∈ V , f (cid:48)
z ( a 1: | S | +1 )
= a | S | +1 , l (cid:48)
( a 1: | S | +1 , b 1: | α | )
=
l ( a 1: | S | , b 1: | α | )) , (cid:74) ( S, V, α ) (cid:96) 2 obs
N ( v 0 , v 1 ) , r ) : ( S, V, α @
[ r ]) (cid:75) ( p, f, l ) =
( p, f, l (cid:48) )
( where l (cid:48) ( a 1: | S | , b 1: | α | +1 )
=
l ( a 1: | S | , b 1: | α | ) × N ( b | α | +1 ; f v 1 ( a 1: | S | ) , f v 2 ( a 1: | S | )) , (cid:74) ( S, V, α ) (cid:96) 2 ( v 0 :
= if ( v 1 > v 2 ) v 3 else v 4 ) : ( S, V ∪ { v 0 } , α ) (cid:75) ( p, f, l ) =
( p, f ⊕ f (cid:48) v 0 , l ) ( where f (cid:48) v 0 ( a 1: | S | )
= if ( f v 1 ( a 1: | S | ) >
f v 2 ( a 1: | S | ))
then f v 3 ( a 1: | S | ) else f v 4 ( a 1: | S | )) , (cid:74) ( S, V, α ) (cid:96) 2 ( v 0 := r ) : ( S, V ∪ { v 0 } , α ) (cid:75) ( p, f, l ) =
( p, f ⊕ f (cid:48) v 0 , l ) ( where f (cid:48) v 0 ( a 1: | S | ) =
r ) , (cid:74) ( S, V, α ) (cid:96) 2 ( v 0 :
= v 1 ) : ( S, V ∪ { v 0 } , α ) (cid:75) ( p, f, l ) =
( p, f ⊕ f (cid:48) v 0 , l ) ( where f (cid:48) v 0 ( a 1: | S | ) =
f v 1 ( a 1: | S | )) , (cid:74) ( S, V, α ) (cid:96) 2 ( v 0 :=
p (cid:48) ( v 0 , v 1 )) :
( S, V ∪ { v 0 } , α ) (cid:75) ( p, f, l ) =
( p, f ⊕ f (cid:48) v 0 , l ) ( where f (cid:48) v 0 ( a 1: | S | ) =
(cid:74) p (cid:48) (cid:75)
( f v 0 ( a 1: | S | ) , f v 1 ( a 1: | S | ))) .
Finally, we deﬁne p C for the well-initialised well-typed programs C , i.e., programs C for which we can derived
[Equation]
For such a C , the deﬁnition of p C is given below:
[Equation]
where ( p, _ , l ) =
(cid:74) ([] , ∅ , []) (cid:96) 1 C : ( S, V, α ) (cid:75) ( p 0 , f 0 , l 0 ) for the constant- 1 functions p 0 and l 0 of appropriate types and the empty family f 0 of functions.
D Marginal Likelihood Computation: Derivation and Correctness
What nn intg ,φ Aims at Computing When we write the meaning of this observe statement as the likelihood N ( r ; v 0 , v 1 ) , and the ﬁltering distribution for v 0 and v 1 under (the decoded density of) the current state h as p h ( v 0 , v 1 ) , 2 the neural network nn intg ,φ computes the following approximation:
[Equation]
where v 0:1 mean the one-hot encoded variables v 0:1 .
Derivation and Correctness of the Marginal Likelihood Let x n be the random variable (RV) that is observed by the command obs ( N ( v 0 , v 1 ) , r ) and x 1:( n − 1) be the ( n − 1) RVs that are observed before the command.
When our algorithm is about to analyse this observe command, we have (an estimate of) p ( x 1:( n − 1) ) by induction.
Then, the marginal likelihood of x 1: n can be computed as follows:
[Equation]
// The ﬁltering distribution p ( v 0 , v 1 | x 1:( n − 1) ) is approximated by p h .
[Equation]
p h ( v 0 , v 1 ) p ( x n | v 0 , v 1 ) dv 0 dv// The RV x n is conditionally independent of x 1:( n − 1) given v 0 , v 1 .
= p ( x 1:( n − 1) ) (cid:90)(cid:90) p h ( v 0 , v 1 ) N ( r ; v 0 , v 1 ) dv 0 dv// p ( x 1:( n − 1) ) is Z in the description of INFER
( A i ) in the main text, and the neural network // nn intg ,φ 7 aims at approximating the integral term accurately.
This derivation leads to the ﬁrst equation in this appendix section.
In a setting of probabilistic programming where observations are allowed to be different in true and false branches, the marginal likelihood may fail to be deﬁned, and such a setting is beyond the scope of our language.
Using variables multiple times or having observe commands spread out in the program does not make differences in the derivation above.
E Description of Our Optimisation Algorithm
We optimise the objective by stochastic gradient descent.
The key component of the optimisation is a gradient estimator derived as follows: ( ∇ φ (cid:80) C ∈D KL[ π C || q C ] + λ 2 ( N C − Z C ) 2 ) = ( (cid:80) C ∈D E z 1: n ∼ π C
[ −∇ φ log q C ( z 1: n )] − λ ( N C − Z C ) ∇ φ Z C )
≈ (cid:80) C ∈D − (cid:91) L C,φ − λ
(cid:100) N C − Z C ) ∇ φ Z C .
Here (cid:91) L C,φ and (cid:100) N C are sample estimates of E z 1: n ∼ π C [ ∇ φ log q C ( z 1: n )] and the marginal likelihood, respectively.
Both estimates can be computed using standard Monte-Carlo algorithms.
For instance, we can run the self-normalising importance sampler with prior as proposal, and generate weighted samples { ( w ( j ) , z ( j ) 1: n ) } 1 ≤ j ≤ M for the unnormalised posterior p C .
Then, we can use these samples to compute the required estimates: (cid:100) N C = 1 M (cid:80) M j =1 w ( j ) and (cid:91) L C,φ = 1 M (cid:80) M j =1 ( w ( j ) ∇ φ log q C ( z ( j ) 1: n )) / (cid:100) N C .
Alternatively, we may run Hamiltonian Monte Carlo [Duane et al., 1987] to generate posterior samples, and use those samples to draw weighted importance samples using, for instance, the layered adaptive importance sampler [Martino et al., 2017].
Then, we compute (cid:91) L C,φ using posterior samples,
[Table]
Generalisationtonewmodelparametersandobservations
[Table]
Gaussian models with a latent variable and an observation where the mean of the Gaussian likelihood is an afﬁne transformation of the latent.
Appendix F.
[Table]
Appendix F.
[Table]
Hierarchical or
[Table]
models with both latent variables and data structured hierarchically where data are modelled as a regression of latent variables of different levels.
Appendix F.1.
[Table]
Clustering models where ﬁve observations are clustered into two groups.
Appendix F.1.
milky and milkyo
Milky Way models, and their multipleobservations extension where ﬁve observations are made for each satellite galaxy.
Appendix F.1.
[Table]
Models with the Rosenbrock function, which is expressed as an external procedure.
Appendix F.1.
Generalisationtonewmodelstructures
[Table]
Models with three Gaussian variables and one deterministic variable storing the value of the function nl( x ) = 50 /π × arctan( x/ 10) , where the models have 12 different types — four different dependency graphs of the variables, and three different positions of the deterministic nl variable for each of these graphs.
Appendix F.2.
(and Fig. 5)
[Table]
Models with six Gaussian variables and one nl variable, which are grouped into ﬁve model types based on their dependency graphs.
Appendix F.2.
(and Fig. 6)
[Table]
comparisonwithalternatives mulmod
Three-variable models where two latent variables follow normal distributions and the other stores the value of the function mm( x ) =
× x 3 / (10 + x 4 ) .
The models in this class are grouped into three types deﬁned by their dependency graphs and the positions of mm in the programs.
Appendix F.
(and Fig. 7) and (cid:100) N C using weighted importance samples.
Note that neither π C in E z 1: n ∼ π C [ −∇ φ log q C ( z 1: n )] nor N
C depends on the parameters φ .
Thus, for each C ∈ D , N C needs to be estimated only once throughout the entire optimisation process, and the posterior samples from π C need to be generated only once as well.
We use this fact to optimise the computation of each gradient-update step.
F Detailed Descriptions for Probabilistic Models Used in the Empirical Evaluation
Table 1 shows the full list of the model classes that we considered in our empirical evaluation (§3).
We detail the program speciﬁcations for the classes using our probabilistic programming language, and then describe how our program generator generated programs from those classes randomly.
In the program speciﬁcations to follow, randomly-generated constants are written in the Greek alphabets ( θ ), and latent and other program variables in the English alphabets.
Also, we often use more intuitive variable names instead of using z i for latent variables and v i for the other program variables, to improve readability.
When describing random generation of the parameter values, we let
U( a, b ) denote the uniform distribution whose domain is ( a, b )
⊂ R ; we use this only for describing the random program generation process itself, not the generated programs (only normal distributions are used in our programs, with the notation N ).
F.1 Generalisation to New Model Parameters and Observations F.1.1 gauss
The model class is described as follows:
[Equation]
For each program of the class, our random program generator generated the parameter values as follows:
[Equation]
and then generated the observation o by running the program forward where the value for z 1 was sampled from z 1 ∼ U( m z − 2 × √ v z , m z + 2 × √ v z ) .
F.1.2 hierl
The model class is described as follows:
[Equation]
For each program of the class [Equation] our generator generated the parameter values as follows:
[Equation]
and then generated the observations o
and o 2 by running the program
(i
[Equation]
e., simulating the model) forward.
F.1.3 hierd
The model class is described as follows:
[Equation]
For each program of the class
[Equation]
our generator generated the parameter values as follows:
[Equation]
and then generated the observations o
and o 2 by running the program forward where the values for a 0 , a 1 , a [Equation] and b in this speciﬁc simulation were sampled as follows:
[Equation]
F.1.4 cluster
The model class is described as follows: m g 1 :
= θ 1 ; v g 1 :
= θ (cid:48) 2 ; m g 2 :
= θ 3 ; v g 2 :
= θ (cid:48) 4 ; v x := θ (cid:48) 5 ; g 1 ∼ N ( m g 1 , v g 1 ); g 2 ∼ N ( m g 2 , v g 2 ); zero := 0; hund := 100; t 1 ∼ N ( zero , hund ); m 1 := if ( t 1 > zero ) g 1 else g 2 ; obs ( N ( m 1 , v x ) , o 1 ); t 2 ∼ N ( zero , hund ); m 2 :
= if ( t 2 > zero ) g 1 else g 2 ; obs ( N ( m 2 , v x ) , o 2 ); t 3 ∼ N ( zero , hund ); m 3 :
= if ( t 3 > zero ) g 1 else g 2 ; obs ( N ( m 3 , v x ) , o 3 ); t 4 ∼ N ( zero , hund ); m 4 := if ( t 4 > zero ) g
1 else g 2 ; obs ( N ( m 4 , v x ) , o 4 ); t 5 ∼ N ( zero , hund ); m 5 :
= if ( t 5 > zero ) g 1 else g 2 ; obs ( N ( m 5 , v x ) , o 5 )
For each program of the class, our generator generated the parameter values as follows:
[Equation]
and then generated the observations o 1:5 by running the program forward.
F.1.5 milky and milkyo
The model class milky is described as follows:
[Equation]
[Equation] [Equation]
g 1 ∼ N ( mass
[Equation]
v g 1 )
[Equation] [Equation]
[Equation]
g 2 ∼ N
[Equation]
mass
[Equation]
v g
[Equation]
For each program of milky , our generator generated the parameter values as follows:
[Equation]
and then generated the observations o
and o 2 by running the program forward.
Everything remained the same for the milkyo class, except that the two obs commands were extended to obs ( N ( g 1 , v x 1 ) , [ o 1 , o 2 , o 3 , o 4 , o 5 ]) and obs ( N ( g 2 , v x 2 ) , [ o 6 , o 7 , o 8 , o 9 , o 10 ]) , respectively, and all the observations were generated similarly by running the extended model forward. F.1.6 rb
The model class rb is described as follows:
[Equation]
where Rosenbrock( z 1 , z 2 ) = 0 .
05 × ( z 1 − 1) 2 + 0 .
005 × ( z 2 − z 12 ) 2 .
The function is often used to evaluate learning and inference algorithms [Goodman and Weare, 2010, Wang and Li, 2018, Pagani et al., 2019].
For each program of the class, our generator generated the parameter values as follows:
[Equation]
and then generated the observation o by running the program forward where the values for z 1 and zin this speciﬁc simulation were sampled as follows:
[Equation]
F.2 Generalisation to New Model Structures
For readability, we present canonicalised dependency graphs where variables are named in the breadth-ﬁrst order.
In the experiments reported in the corresponding part of the main text, we used a minor extension of our probabilistic programming language with procedures taking one parameter
. F.2.1 ext
Fig.
5 shows the dependency graphs for all model types in ext1 .
The variables z 0 , z 1 , . . .
and x 1 , x 2 , . .
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.
5, while pointing out that the similar process is applied to the other eleven types.
To generate programs of the model type (1,1), we use the following program template:
[Equation]
where nl( z ) = 50 /π × arctan( z/ 10) .
The generation involves randomly sampling the parameters of this template, converting the template into a program in our language [Equation] and creating synthetic observations.
Speciﬁcally, our generator generates the parameter values as follows:
[Equation]
and generates the observation o 1 by running the program forward where the values for z 0:3 in this speciﬁc simulation were sampled (and ﬁxed to speciﬁc values) as follows:
[Equation]
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.
F.2.2 ext
Fig.
6 shows the dependency graphs for all ﬁve model types in ext2 .
Programs of these ﬁve types are randomly generated by our program generator.
As in the ext1 case, we explain the generator only for one model type, which corresponds to the ﬁrst dependency graph in Fig.
To generate programs of this type, we use the following program template:
[Equation] [Equation]
In order to generate a program of this model type and observations
[Equation]
our generator instantiates the parameters of the template as follows:
[Equation]
Then [Equation] it generates the observations o 1:4 by running the program forward where the values for z 0:6 in this speciﬁc simulation were sampled (and ﬁxed to speciﬁc values) as follows:
[Equation]
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.
F.3 Test-Time Efﬁciency in Comparison with Alternatives
This section details the mulmod class, which has three different model types.
Fig.
7 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 [Equation] and compared the results with those returned by HMC.
We similarly explain the generator only using the model type corresponding to the ﬁrst dependency graph in Fig.
To generate programs of this type [Equation] we use the following program template:
[Equation]
where mm( x ) = 100 × x 3 / (10+ x 4 ) .
For each program in this model type
[Equation]
our generator instantiates the parameter values as follows:
[Equation]
and synthesises the observation o 1 by running the program forward where the values for z 0:2 in this speciﬁc simulation were sampled (and ﬁxed to speciﬁc values) as follows:
[Equation]
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.
G Detailed Evaluation Setup
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.,
We used a Ubuntu server with Intel(R)
Xeon(R)
Gold 6234 CPU @ 3.30GHz with 16 cores, 32 threads, and 263G memory.
In our evaluation, the dimension s of the internal state h was 10 (i.e., h ∈ R 10 ).
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 nn de where the hidden dimensions were 50 .
The hyper-parameter in our optimisation objective was set to λ = 2 in the evaluation.
For HMC, we used the NUTS sampler [Hoffman and Gelman, 2014].
G.1 Detailed Training Setup for Generalisation to New Model Parameters and Observations in §
For each training program, our meta-algorithm used 2 15 samples from the analytic (for gauss ) or approximate (for the rest, by HMC) posterior distribution for the program.
3 Similarly, our metaalgorithm 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 deﬁned by an HMC chain.
We performed mini-batch training with the batch size 2 12 .
We used Adam [Kingma and Ba, 2015] with its hyperparameters { β 1 = 0 .
9 , β 2 = 0 .
999 , weight_decay = 0 } and 0 .
001 as the initial learning rate.
We repeated the same experiments three times using different random seeds, and the training stopped when the average training loss converged enough.
G.2 Detailed Setup for Generalisation to New Model Structures in §
We let two structural aspects vary across programs: the dependency (or data-ﬂow) graph for the variables of a program and the position of a nonlinear function in the program.
Speciﬁcally, we considered two kinds of models: (1) models ( ext1 ) with three Gaussian variables and one deterministic variable storing the value of the function nl( x )
= 50 /π × arctan( x/ 10) , where the models havedifferent types — four different dependency graphs of the variables, and three different positions of the deterministic nl variable for each of these graphs; and (2) models ( ext2 ) with six Gaussian variables and one nl variable, which are grouped into ﬁve types based on their dependency graphs.
See Fig.
5 and 6 in the appendix for visualisation of the different types in ext1 and ext2 , respectively.
Before running our inference algorithm, we canonicalise the names of variables in a given program based on its dependency (i.e., data-ﬂow) graph.
Although not perfect, this preprocessing removes a superﬁcial difference between programs caused by different variable names, and enables us to avoid the unnecessary complexity caused by variable-renaming symmetries at training and inference times.
For ext1 , we ran seven different experiments.
Four of them evaluated generalisation to unseen dependency graphs.
Among all the four possible dependency graphs of programs we considered, these experiments used only three of them during training, and the other for testing: the dependency graph of each training program is one of these three (with the nl variable positioned in all the three possible places in the graph), and that of each test program is the other dependency graph.
The remaining three experiments evaluated generalisation to unseen positions of the nl variable where we ﬁxed the position of the nl variable in the four dependency graphs, and used programs with the variable in one of those ﬁxed positions for testing, while using programs with the variables in all the other positions for training.
For ext2 , we ran ﬁve different experiments where each of them tested generalisation to an unseen dependency graph after training with the other four types of dependency graphs.
All of these experiments were repeated three times under different random seeds.
So, the total numbers of experiment runs were 21 (= 7 × 3) and 15 (= 5 × 3) for ext1 and ext2 , respectively.
In each experiment run for ext1 , we used 720 programs for training, and 90 (when generalising to new graphs) or 100 (when generalising to new positions of the nl variable) unseen programs for testing.
In each run for ext2 , we used 600 programs for training and tested the learnt inference algorithm on 50 unseen programs.
We ran HMC to estimate posteriors and marginal likelihoods, and used 200 K samples after 10 K 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 Appendix G.1.
H Losses for hierd , cluster , milkyo , and rb
Fig.
8 shows the average training and test losses under three random seeds for hierd , cluster , milkyo , and rb .
The later part of Fig. 8a, 8c and 8d shows cases where the test loss surges.
This was when the loss of only a few programs in the test set (of 50 programs) became large.
Even in this situation, the losses of the rest remained small.
We give analyses for cluster and rb separately in §I.
I Multimodal Posteriors: cluster and rb
The cluster and rb classes posed another challenge: the models often had multimodal posteriors, and it was signiﬁcantly 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 an inference algorithm.
Second, our random program generator placed some restriction on the programs it generated (e.g., by using tight boundaries on some model parameters), guided by the analysis of the geometry of the Rosenbrock function [Pagani et al., 2019].
Consequently, HMC (with 500K samples after 50K warmups) failed to converge for only one ﬁfth of the test programs.
Fig.
9 shows the similar comparison plots between reference and predicted marginal posteriors fortest 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 signiﬁcant multimodality.
The latent variable pgm75 _ z0 had at least two modes at around − 10 (visible in the ﬁgure) and around 10 (hidden in the ﬁgure) 4 .
Our inference algorithm showed a mode-seeking behavior for this latent variable.
Similarly, the variable pgm79 _ z0 had at least two modes in the similar domain region (one shown and one hidden), but this time our inference algorithm showed a mode-covering behavior.
The multimodality issue raises two questions.
First, how can our meta-algorithm generate samples from the posterior more effectively so that it can optimise the inference algorithm for classes of models with multimodal posteriors?
For example, our current results for cluster suffer from the fact that the samples used in the training are often biased (i.e., only from a single mode of the posterior).
One possible direction would be to use multiple Markov chains simultaneously and apply ideas from the mixing-time research.
Second, how can our white-box inference algorithm catch more information from the program description and ﬁnd non-trivial properties that may be useful for computing the posterior distributions having multiple modes?
We leave the answers for future work.
J Training and Test Losses for ext
Fig.
10 shows the average training and test losses in the ext1 experiment runs (under different random seeds).
K Training and Test Losses for ext
Fig.
11 shows the average training and test losses in the ext2 experiment runs (under different random seeds).
L Program from mulmod That was Reported in the Test-Time Comparisons in §
Fig.
12 shows the program that is reported in the test-time comparisons in our evaluation, written in our probabilistic programming language.
[Table]
M ESS Results in the Comparisons with Alternatives in Terms of Test-Time Efﬁciency in §
Table 2 shows ESS by the same runs of the three approaches as in the last columns of Fig.
4a and 4b in §3.
IS-pred produced over 16 K effective samples in 18 ms, while HMC generated onlyeffective samples even after 84 s. Similarly, IS-pred produced effective samples 10 times more than
IS-prior in the approximately same elapsed time.
In fact, Fig. 13 shows that as the time increases, the gap between the ESSes of IS-pred and IS-prior gets widen, because the former increases at a rate signiﬁcantly higher than the latter.
N Limitation and Future Work
A learnt inference algorithm in this work does not generalise to programs with different sizes [Yan et al., 2020], e.g., from clustering models with two clusters to those with ten clusters.
Each model class assumes a ﬁxed number of variables, and the neural networks crucially exploit the assumption.
Second, our meta-algorithm does not scale in practice to large programs, e.g., state-space models with a few hundred time steps, and cannot learn an optimal inference algorithm within a reasonable amount of time.
Third, it is a good future work to relax the (meanﬁeld Gaussian) assumption on the approximating distribution in our inference algorithm, and to meta-learn the form of the approximation itself by, e.g., incorporating the ideas proposed by Ambrogioni et al.
This extension is closely related to automatic guide generation in the Pyro community [Bingham et al., 2018].
[Figure] [Figure]
[Figure] [Figure]
Seed
[Figure]
Tr or Te
[Figure]
Tr
[Figure]
Tr
[Figure]
Te
[Figure] [Figure]
[Figure] [Figure]
a := 3 . 93; b := 348 . 16; c := 57 .
5; d := 14 . 04; e := 40 .
z 1 ∼ N ( a, b ); z 2 ∼ N ( z 1 , c ); z 3 :=
mm( z 1 ); obs ( N ( z 2 , d ) , 53 . 97); obs ( N ( z 3 , e ) , 0 . 12)
[Figure]