A New Approach to Probabilistic Programming Inference

Introduction

Probabilistic programming differs substantially from traditional programming. In particular, probabilistic programs are written with parts not fixed in advance that instead take values generated at runtime by random sampling procedures. Inference in probabilistic programming characterizes the conditional distribution of such variables given observed data assumed to have been generated by executing the probabilistic program. Exploring the joint distribution over program execution traces that could have generated the observed data using Markov chain sampling techniques is one way to produce such a characterization.

We propose a novel combination of ideas from probabilistic programming and particle Markov chain Monte Carlo (PMCMC) that yields a new scheme for exploring and characterizing the set of probable execution traces. As our approach is based on repeated simulation of the probabilistic program, it is easy to implement and parallelize. We show that our approach supports accurate inference in models that make use of complex control flow, including stochastic recursion, as well as primitives from nonparametric Bayesian statistics. Our experiments also show that this approach can be more efficient than previously introduced single-site Metropolis-Hastings (MH) samplers .

Language

The probabilistic programming system Anglicanhttp://www.robots.ox.ac.uk/~fwood/anglican/ exists in two versions. The results in this paper were obtained using what is now called interpreted Anglicanhttp://bitbucket.org/probprog/interpreted-anglican which employed an interpreted execution model and a language syntax derived from the Venturehttp://probcomp.csail.mit.edu/venture/ modeling language. Since the time of original publication, both the syntax and execution model have been changed. The Venture style syntax is now deprecated, with the new version using a syntax much closer to that of Anglican’s host language Clojure. This new Anglican versionhttp://bitbucket.org/probprog/anglican, simply called Anglican, is a compiled language; a continuation-passing style (CPS) transformation compiles Anglican programs into Clojure programs that are then subsequently compiled to Java Virtual Machine (JVM) bytecode by the Clojure just-in-time compiler.

The change in language execution model and syntax affect neither the substance or the claims of this paper, however, readers wishing to experiment with the language starting from the code examples appearing in Section 4 would do well to note this evolution. What does change is the absolute time required to perform inference in the given models. In general, the latest compiled version of Anglican is ten to one hundred times faster than its interpreted ancestor.

The original, deprecated Anglican syntax is a Scheme/Lisp-dialect that is extended with three top-level special forms that we refer to as directives

Here, each is a Scheme/Lisp-syntax expression, symbol is a unique symbol, and is a constant-valued deterministic expression.

Semantically assume’s are (random) variable (generative) declarations, observe’s condition the distribution of assume’d variables by (noisily) constraining the output values of (random) functions of assume’d variables to match observed data, and predict’s are “watches” which report on (via print out) the values of variables in program traces as they are explored. In Anglican, probabilistic program interpretation is taken to be a forever-continuing exploration of the space of execution traces that obey (where hard) or reflect (where soft) the observe’d constraints in order to report functions (via predict’s) of the conditional distribution of subsets of the assume’d variables.

Like Scheme/Lisp, Anglican eagerly and exchangeably recursively “evaluates” subexpressions, for instance of = ( $\ldots$ ), before “applying” the procedure (which may be a random procedure or special form) resulting from evaluating the first expression to the value of its arguments. Anglican counts applications, reported as a computational cost proxy in Sec 4. Anglican supports several special forms, notably (lambda ( $\ldots$ ) )) which allows creation of new procedures and (if ) which supplies branching control flow; also begin, let, define, quote, and cond. Anglican exposes eval and apply. Built-in deterministic procedures include list, car, cdr, cons, mem, etc., and arithmetic procedures +, -, $\backslash$, etc.

All randomness in the language originates from built-in “random primitives” of which there are two types. Elementary random primitives, such as poisson, gamma, flip, discrete, categorical, and normal, generate independent and identically distributed (i.i.d.) samples when called repeatedly with the same arguments. Exchangeable random primitives, such as crp or beta-bernoulli, return a random procedure with internal state that generates an exchangeably distributed samples when called repeatedly.

In the interpreted version of Anglican the outer in all observe ’s must be a built-in random primitive, which guarantees that likelihood of output given arguments can be computed exactly.

2 New Clojure-Style Syntax

The compiled version of Anglican now supports a syntax that more closely integrates with the host language Clojure. The macro defquery defines an Anglican program from within Clojure

Here symbol is the name of the program, and the arguments may be used to pass values of parameters or observed variables to the program. The set of allowable body expressions is a subset of the Clojure language, in which all basic Clojure language forms and first order primitives are supported. Macros and higher-order functions are not inherited from Clojure, although a subset of macroswhen, cond and higher-order functionsmap, reduce, filter, some, repeatedly, comp, partial has been implemented.

In the CPS version of Anglican, random primitives such as normal and discrete return first class distribution objects, rather than sampled values, and, furthermore, are user programmable without requiring modifications to the Anglican compiler. The special forms sample and observe associate values with random variables drawn from a distribution

The language additionally provides a data type for sequences of random variables that are not i.i.d., which we refer to as a random process. A random process implements two operations

The produce primitive returns a distribution object for the next random variable in the sequence. The absorb primitive returns an updated random process instance, in which a value has been associated with the next random variable in the sequence. A random process is most commonly used to represent an exchangeably distributed sequence of random variables, but it may be used to represent any sequence of random variables for which it is possible to construct a distribution on the next variable given preceding variables. Random process constructors are customarily identified with uppercase names (e.g. CRP) and, in the same manner as before, are user programmable.

Finally, the predict form may be used at any point in the program to generate labeled output values

Given a previously defined query, inference may be performed using the doquery macro

The inference algorithm may be specified using an algorithm keyword, and any options to the inference algorithm can be supplied as arguments as well. The doquery macro constructs a lazy sequence of program execution states that contain predicted values and optionally an importance weight, which may then be consumed and analyzed by, for instance, an outer Clojure or Java program.

Inference

An execution trace is the sequence of memory states resulting from the sequence of function applications performed during the interpretation of a program. In probabilistic programming systems like Anglican any variable may be declared as being the output of a random procedure. Such variables can take different values in independent interpretations of the program. This leads to a “many-worlds” computational trace tree in which, at interpretation time, there is a branch at every random procedure application.

To define the probability of a single execution trace, first fix an ordering of the exchangeable lines of the program and index the observe lines by $n.$ Let $p(y_{n}|\theta_{t_{n}},\mathbf{x}_{n})$ be the likelihood of the observe’d output $y_{n}$ where $t_{n}$ is a random procedure type (i.e. gamma, poisson, etc.), $\theta_{t_{n}}$ is its argument (possibly multi-dimensional), and $\mathbf{x}_{n}$ is the set of all random procedure application results computed before the likelihood of observation $y_{n}$ is evaluated. Both the type $t_{n}$ and the parameter $\theta_{t_{n}}$ can be functions of any in-scope subset of $\mathbf{x}_{n}.$ We can then define the probability of an execution trace to be

where $\mathbf{y}$ is the set of all observe’d quantities, $\mathbf{x}$ is the set of all random procedure application results, and $\sim$ marks distributions which we can only sample.

The number and type of all random procedure applications performed before the $n$th observe may vary in one program trace to the next. We define the probability of the sequence of their outputs $\mathbf{x}_{n}$ to be

Here $\mathbf{x}_{0}=\emptyset$ is the empty set and $x_{n,j:k}$ are the $j$ to $k$th values generated by random procedure applications in the trace up to observation likelihood computation $n$. The cardinality of the set $\mathbf{x}_{n}\setminus\mathbf{x}_{n-1}$, notated $|\cdot|$, arises implicitly as the total number of random procedure applications in a given execution trace. As before, $\theta_{t_{n,k}}$ are the arguments to, and $t_{n,k}$ the type of, the $(n,k)$th random procedure – both of which may be functions of subsets of in-scope subsets of variables $x_{n:1,(k-1)}\cup\mathbf{x}_{n-1}$. Note that $\mathbf{x}_{n-1}\subseteq\mathbf{x}_{n}.$ Also note that variable referencing defines a directed conditional dependency structure for the probability model encoded by the program, i.e. $x_{n,k}$ need not (and often cannot due to variable scoping) depend on the outputs of all previous random procedure applications.

All members of the set of all directed probabilistic models with fixed-structure joint distributions can be expressed as probabilistic programs that “unroll” in all possible execution traces into an equivalent joint distribution. As Church-like probabilistic programming frameworks, Anglican included, support recursive procedures and branching on the values returned by random procedures, the corresponding set of models is a superset of the set of all directed graphical models. Other related efforts eschew Turing-completeness and operate on a restricted set of models where inference techniques other than sampling can more readily be employed.

The SMC procedure described is, to first approximation, the inner loop of PMCMC. It corresponds to a procedure whereby the probabilistic program is interpreted in parallel (possibly each particle in its own thread or process) between observation likelihood calculations. Unfortunately, SMC with a finite set of particles is not itself directly viable for probabilistic programming inference for all the familiar reasons: particle degeneracy, inefficiency in models with global, continuous parameters, etc.

PMCMC, on the other hand, is directly viable. PMCMC for probabilistic programming inference is a MH algorithm for exploring the space of execution traces that uses SMC proposals internally. This, unlike prior art, allows sampling of execution traces with changes to potentially many more than one variable at a time. The particular variant of PMCMC we discuss in this paper is Particle-Gibbs (PG) although we have developed engines based on other PMCMC variants including particle independent Metropolis Hastings and conditional sequential Monte Carlo too. PG works by iteratively re-running SMC, with, on all but the first sweep, reinsertion of a “retained” particle trace into the set of particles at every stage of SMC. PG is theoretically justified as an MH transition operator that, like the Gibbs operator, always accepts . In this paper we describe PMCMC for probabilistic programming inference algorithmically in Alg.1 and experimentally demonstrate its relative efficacy for probabilistic programming inference.

In Alg. 1 the function r$(N,\mathcal{S})$ stands for multi-nomial sample $N$ items from the set of pairs $\mathcal{S}=\{\{w_{1},\theta_{1}\},\ldots,\{w_{M},\theta_{M}\}\}$, where each element of $\mathcal{S}$ consists of an unnormalized weight $w_{m}$ and interpreter memory states $\theta_{m}$. For each sample value $\theta_{m}$ returned, the function r also returns the original, corresponding unnormalized weight $w_{m}$. This kind of weight bookkeeping retains, for all particles, the results of the outermost observe likelihood function applications so that the unnormalized weights are available in the retained particle $\{w^{*},\mathbf{x}^{*}\}$ in the next sweep.

In Alg. 1 “Run SMC” means running one sweep of the $s$ loop with $L$ particles and no retained particle, $d$ is a program line, and fork$(\cdot)$ means to copy the entire interpreter memory datastructure (efficient implementations have characteristics similar to POSIX fork() ). The command interpret$(d,\cdot)$ means execute line $d$ in the given interpreter. Only when interpreting an observe must the interpreter return a weight, that being result of the outermost apply of the observe statement. Bars indicate temporary data structures, not averages. All sets are ordered with unions implemented by append operations.

Note that there are more efficient PMCMC algorithms for probabilistic programming inference. In particular, there is no reason to fork unless an observe has just been interpreted. Alg. 1 is presented in this form for expositional purposes.

2 Random Database

We refer to the MH approach to sampling over the space of all traces proposed in as “random database” (RDB). A RDB sampler is a MH sampler where a single variable drawn in the course of a particular interpretation of a probabilistic program is modified via a standard MH proposal, and this modification is accepted by comparing the value of the joint distribution of old and new program traces. For completeness we review RDB here, noting a subtle correction to the acceptance ratio proposed in the original reference which is proper for a larger family of models.

The RDB sampler employs a data structure that holds all random variables $\mathbf{x}$ associated with an execution trace, along with the parameters and log probability of each draw. Note that interpretation of a program is deterministic conditioned on $\mathbf{x}$. A new proposal trace is initialized by picking a single variable $x_{m,j}$ from the $|\mathbf{x}|$ random draws, and resampling its value using a reversible kernel $\kappa(x^{\prime}_{m,j}|x_{m,j})$. Starting from this initialization, the program is rerun to generate a new set of variables $\mathbf{x}^{\prime}$ that correspond to a new valid execution trace. In each instance where the random procedure type remains the same, we reuse the existing value from the set $\mathbf{x}$, rescoring its log probability conditioned on the preceding variables where necessary. When the random procedure type has changed, or a new random variable is encountered, its value is sampled in the usual manner. Finally, we compute the probability $p(\mathbf{y}|\mathbf{x^{\prime}})$ by rescoring each observe as needed, and accept with probability

In order to calculate the ratio of the proposal probabilities $q(\mathbf{x^{\prime}}|\mathbf{x})$ and $q(\mathbf{x}|\mathbf{x^{\prime}})$, we need to account for the variables that were resampled in the course of constructing the proposal, as well as the fact that the sets $\mathbf{x^{\prime}}$ and $\mathbf{x}$ may have different cardinalities $|\mathbf{x^{\prime}}|$ and $|\mathbf{x}|$. We will use the (slightly imprecise) notation $\mathbf{x^{\prime}}\backslash\mathbf{x}$ to refer to the set of variables that were resampled, and let $\mathbf{x^{\prime}}\cap\mathbf{x}$ represent the set of variables common to both execution traces. The proposal probability is now given by

In our implementation, the initialization $x^{\prime}_{m,j}$ is simply resampled conditioned on the preceding variables, such that $\kappa(x^{\prime}_{m,j}|x_{m,j})=p(x^{\prime}_{m,j}|\mathbf{x^{\prime}}\cap\mathbf{x})$. The reverse proposal density can now be expressed in a similar fashion in terms of $\mathbf{x}\backslash\mathbf{x^{\prime}}$ and $\mathbf{x}\cap\mathbf{x^{\prime}}=\mathbf{x^{\prime}}\cap\mathbf{x}$, allowing the full acceptance probability to be written as

Testing

Programming probabilistic program interpreters is a non-trivial software development effort, involving both the correct implementation of an interpreter and the correct implementation of a general purpose sampler. The methodology we employ to ensure correctness of both involves three levels of testing; 1) unit tests, 2) measure tests, and 3) conditional measure tests.

In the context of probabilistic programming, unit testing includes verifying that the interpreter correctly interprets a comprehensive set of small deterministic programs. Measure testing involves interpreting short revealing programs consisting of assume and predict statements (producing a sequence of ancestral, unconditioned samples, i.e. no observe’s). Interpreter output is tested relative to ground truth, where ground truth is computed via exhaustive enumeration, analytic derivation, or some combination, and always in a different, well-tested independent computational system like Matlab. Various comparisons of the empirical distribution constructed from the accumulating stream of output predicts’s and ground truth are computed; Kulback-Leibler (KL) divergences for discrete sample spaces and Kolmogorov Smirnov (KS) test statistics for continuous sample spaces. While it is possible to construct distribution equality hypothesis tests for some combinations of test statistic and program we generally are content to accept interpreters for which there is clear evidence of convergence towards zero of all test statistics for all measure tests. Anglican passed all unit and measure tests.

2 Conditional Measure Tests

Measure tests involving conditioning provide additional information beyond that provided by measure and unit tests. Conditioning involves endowing programs with observe statements which constrain or weight the set of possible execution traces. Interpreting observe statements engages the full inference machinery. Conditional measure test performance is measured in the same way as measure test performance. They are also how we compare different probabilistic programming inference engines.

Inference Engine Comparison

We compare PMCMC to RDB measuring convergence rates for an illustrative set of conditional measure test programs. Results from four such tests are shown in Figure 1 where the same program is interpreted using both inference engines. PMCMC is found to converge faster for conditional measure test programs that correspond to expressive probabilistic graphical models with rich conditional dependencies.

The four test programs are: 1) a program that corresponds to state estimation in a hidden Markov model (HMM) with continuous observations (HMM: Program 5.1, Figure 1(a)), 2) a program that corresponds to learning an uncollapsed Dirichlet process (DP) mixture of Gaussians with fixed hyperparameters (DP Mixture: Program 5.2, Figure 1(b)), 3) a multimodal branching with deterministic recursion program that cannot be represented as a graphical model in which all possible execution paths can be enumerated (Branching: Program 5.3, Figure 1(c)), and 4) a program that corresponds to inferring the mean of a univariate normal generated via an Anglican-coded Marsaglia rejection-sampling algorithm that halts with probability one and generates an unknown number of internal random variables (Marsaglia: Program 5.4, Figure 1(d)). We refer to 1) and 2) as “expressive” models because they have complex conditional dependency structures and 3) and 4) as simple models because the programs encode models with very few free parameters. 1) and 2) illustrate our claims; 3) and 4) are included to document the correctness and completeness of the Anglican implementation while also demonstrating that the gains illustrated in 1) and 2) do not come at too great a cost even for simple programs for which, a priori, PMCMC might be reasonably be expected to underperform.

In Figures 1(a)-1(d) there are three panels that report similar style findings across test programs and a fourth that is specific to the individual test program. In all, PMCMC results are reported for a single-threaded interpreter with 100 particles. The choice of 100 particles is largely arbitrary; our results are stable for a large range of values. PMCMC is dark blue while RDB is light orange. We report the 25% (lower dashed) median (solid) and 75% (upper dashed) percentiles over 25 runs with differing random number seeds. We refer to KL/KL∗/KS as distances and compute each via a running average of the empirical distribution of predict statement outputs to ground truth starting at the first predict output. Note that lower is better. We define the number of simulations to be the number of times the program is interpreted in its entirety. For RDB this means that the number of simulations is exactly the number of sampler sweeps; for PMCMC it is the number of particles multiplied by the number of sampler sweeps. The time horizontal axes report wall clock time; the apply axes report the number of function applications performed by the interpreter. In the distance vs. time plots, observed single-threaded PMCMC wall-clock times are reported via filled circles; the left-ward dotted lines illustrate hypothetically what should be achievable via parallelism. Carefully note that PMCMC requires completing a number of simulations equal to the number of particles (here 100) before emitting batched predict outputs. This means that single-threaded implementations of PMCMC suffer from latency that RDB does not. Still, for some programs both the quality of PMCMC’s predict outputs and PMCMC’s convergence rate is faster even in direct wall clock time comparison to RDB. PMCMC appears to converge faster for some programs than RDB even relative to the number of function applications. Equivalent results were obtained relative to eval counts.

The HMM program corresponds to a latent state inference problem in an HMM with three states, one-dimensional Gaussian observations (.9, .8, .7, 0, -.025, 5, 2, 0.1, 0, .13, .45, 6, .2, .3, -1, -1), with known means and variances, transition matrix, and initial state distribution. The lines of the program were organized with the observe’s in “time” sequence.

The $KL^{*}$ axis reports the sum of the Kulback-Leibler divergences between the running sample average state occupancy across all states of a HMM including the initial state and one trailing predictive state $KL^{*}_{S}=\sum_{i}D_{KL}(\frac{1}{S}\sum_{s=1}^{S}\delta_{z_{i}^{(s)}}||\gamma_{i})$. Here $\delta_{z_{i}^{(s)}}(k)$ returns one if simulation $s$ has latent state $z$ at time step $i$ equal to $k$ and $\gamma_{i}(k)$ is the true marginal probability of the latent state indicator $z_{i}$ taking value $k$ at time step $i$. The vertical red line in the apply plot indicates the time and number of applies it takes to run forward-backward in the Anglican interpreter.

The fourth plot shows the learned posterior distribution over the latent state value for all time steps including both the initial state and a trailing predictive time step. While RDB produces a reasonable approximation to the true posterior, it does so more slowly and with greater residual error.

2 DP Mixture

The DP mixture program corresponds to a clustering with unknown mean and variance problem modelled via a Dirichlet process mixture of one-dimensional Gaussians with unknown mean and variance (normal-gamma priors). The KL divergence reported is between the running sample estimate of the distribution over the number of clusters in the data and the ground truth distribution over the same. The ground truth distribution over the number of clusters was computed for this model and data by exhaustively enumerating all partitions of the data (1.0, 1.1, 1.2, -10, -15, -20, .01, .1, .05, 0), analytically computing evidence terms by exploiting conjugacy, and conditioning on partition cardinality. The fourth plot shows the posterior distribution over the number of classes in the data computed by both methods relative to the ground truth.

This program was written in a way that was intentionally antagonistic to PMCMC. The continuous likelihood parameters were not marginalized out and the observe statements were not organized in an optimal ordering. Despite this, PMCMC outperforms RDB per simulation, wall clock time, and apply count.

3 Branching

The branching program has no corresponding graphical model. It was designed to test for correctness of inference in programs with control logic and execution paths that can vary in the number of sampled values. It also illustrates mixing in a model where, as shown in the fourth plot, there is a large mismatch between the prior and the posterior, so rejection and importance sampling are likely to be ineffective. Because there is only one observation and just a single named random variable PMCMC and RDB should and does achieve essentially indistinguishable performance normalized to simulation, time and apply count.

4 Marsaglia

Marsaglia is a test program included here for completeness. It is an example of a type of program for which PMCMC sometimes may not be more efficient. Marsaglia is the name given to the rejection form of the Box-Muller algorithm for sampling from a Gaussian . The Marsaglia test program corresponds to an inference problem in which observed quantities are drawn from a Gaussian with unknown mean and this unknown mean is generated by an Anglican implementation of the Marsaglia algorithm for sampling from a Gaussian. The KS axis is a Kolmogorov-Smirnov test statistic computed by finding the maximum deviation between the accumulating sample and analytically derived ground truth cumulative distribution functions (CDF). Equal-cost PMCMC, RDB, and ground truth CDFs are shown in the fourth plot.

Because Marsaglia is a recursive rejection sampler it may require many recursive calls to itself. We conjecture that RDB may be faster than PMCMC here because, while PMCMC pays no statistical cost, it does pay a computational cost for exploring program traces that include many random procedure calls that lead to rejections whereas RDB, due to the implicit geometric prior on program trace length, effectively avoids paying excess computational costs deriving from unnecessarily long traces.

5 Line Permutation

Syntactically and semantically observe and predict’s are mutually exchangeable (so too are assume’s up to syntactic constraints). Given this and the nature of PMCMC it is reasonable to expect that line permutations could effect the efficiency of inference. We explored this by randomly permuting the lines of the HMM and DP Mixture programs. The results are shown in Fig. 2 where blue lines correspond to median (out of twenty five) runs of PMCMC for each of twenty five program line permutations (including unmodified (dark) and reversed (light)) and orange are the same for RDB. For the HMM we found that the natural, time-sequence ordering of the lines of the program resulted in the best performance for PMCMC relative to RDB. This is because in this ordering the observe’s cause re-weighting to happen as soon as possible in each SMC phase of PMCMC. The effect of permuting code lines interpolates inference performance between optimal where PMCMC is best and adversarial orderings where RDB is instead. RDB performance is demonstrated here to be independent of the program line ordering.

The DP Mixture results show PMCMC outperforming RDB on all program reorderings. Further, it can be seen that the original program ordering was not optimal with respect to PMCMC inference.

While PMCMC presents the opportunity for significant gains in inference efficiency, it does not prevent programmers from seeking to further optimize performance manually. Programmers can influence inference performance by reordering program lines, in particular pushing observe statements as near to the front of the program as syntactically allowed, or restructuring programs to lazily rather than eagerly generate latent variables. Efficiency gains via automatic transformations or online adaptation of the ordering may be possible.

6 Number of Particles

Fig. 3 shows the number of particles in the PMCMC inference engine affects performance. Performance improves as a function of the number of particles. In this plot the red line indicates 100 particles. Increasing from dark to light the number of particles plotted is 2, 5, 10, 20, 50, (red, 100), 200, and 500.

Discussion

The PMCMC approach to probabilistic program interpretation appears to converge faster to the true conditiona distribution over program execution traces for programs that correspond to expressive models with dense conditional dependencies than single-site Metropolis-Hastings methods. That PMCMC converges faster in our tests even after normalizing for computational time is, to us, both surprising and striking. To reiterate, this includes all the computation done in all the particle executions.

Current versions of Anglican include a multi-threaded inference core but work remains to achieve optimal parallelism due, in particular, to memory organization sub-optimality leading to excessive locking overhead. Using just a single thread, PMCMC surprisingly sometimes outperforms RDB per simulation anyway in terms of wall clock time.

Our specific choice of syntactically forcing observe’s to be noisy requires language and interpreter level restrictions and checks. Hard constraint observe’s can be supported in Anglican by exposing a Dirac likelihood to programmers. Persisting in not doing so should help programmers avoid writing probabilistic programs where finding even a single satisfying execution trace is NP-hard or not-computable, but also could be perceived as requiring a non-intuitive programming style.

As explored in concurrent work on Venture, there may be opportunities to improve inference performance by partitioning program variables, either automatically or via syntax constructs, and treating them differently during inference. Then, for instance, some could be sampled via plain MH, some via conditional SMC. One simple way to do this would be to combine conditional SMC and RDB MH.

We thank Xerox and Google for their generous support. We also thank Arnaud Doucet, Brooks Paige, and Yura Perov for helpful discussions about PMCMC and probabilistic programming in general.

References