Particle Gibbs with Ancestor Sampling for Probabilistic Programs

Introduction

Probabilistic programming languages extend traditional languages with primitives for sampling and conditioning on random variables. Running a program $F$ generates random variates $\boldsymbol{x}$ for some subset of program expressions, which can be thought of as a sample from a prior $p(\boldsymbol{x}\,|\,F)$ implicitly defined by the program. In a conditioned program $F[\boldsymbol{y}]$, a subset of expressions is constrained to take on observed values $\boldsymbol{y}$. This defines a posterior distribution $p(\boldsymbol{x}\,|\,F[\boldsymbol{y}])$ on the random variates $\boldsymbol{x}$ that can be generated by the program $F[\boldsymbol{\boldsymbol{y}}]$.

Probabilistic programs are in essence procedural representations of generative models. These representations are often both succinct and extendable, making it easy to iterate over alternative model designs. Any program that always samples and constrains a fixed set of variables $\{\boldsymbol{x},\boldsymbol{y}\}$ admits an alternate representation as a graphical model. In general, a program $F[\boldsymbol{y}]$ may also have random variables that are only generated when certain conditions on previously sampled variates are met. The random variables $\boldsymbol{x}$ therefore need not have the same entries for all possible executions of $F[\boldsymbol{y}]$. In languages that have recursion and higher-order functions (i.e. functions that act on other functions), it is straightforward to define models that can instantiate an arbitrary number of random variables, such as certain Bayesian nonparametrics, or models that are specified in terms of a generative grammar. At the same time this greater expressivity makes it challenging to design methods for efficient posterior inference in arbitrary programs.

In this paper we show how a recently proposed technique known as particle Gibbs with ancestor sampling (PGAS) Lindsten et al. (2012) can be adapted to inference in higher-order probabilistic programming systems Mansinghka et al. (2014); Wood et al. (2014); Goodman et al. (2008). A PGAS implementation requires combining a partial execution history, or prefix, with the remainder of a previously completed execution history, which we call a suffix. We develop a formalism for performing this operation in a manner that guarantees a consistent program state, correctly updates the probabilities associated with each sampled and observed random variable, and avoids unnecessary recomputation where possible. An empirical evaluation demonstrates that the increased statistical efficiency of PGAS can easily outweigh its greater computational cost.

Current generation probabilistic programming systems fall into two broad categories. On the one hand, systems like Infer.NET Minka et al. (2010) and STAN Stan Development Team (2014) restrict the language syntax in order to omit recursion. This ensures that the set of variables is bounded and has a well-defined dependency hierarchy. On the other hand languages such as Church Goodman et al. (2008), Venture Mansinghka et al. (2014), Anglican Wood et al. (2014), and Probabilistic C Paige and Wood (2014) do not impose such restrictions, which makes the design of general-purpose inference methods more difficult.

Arguably the simplest inference methods for probabilistic programming languages rely on Sequential Monte Carlo (SMC) Del Moral et al. (2006). These “forward” techniques sample from the prior by running multiple copies of the program, calculating importance weights using the conditioning expressions as needed. The only non-trivial requirement for implementing SMC variants is that there exists an efficient mechanism for “forking” a program state into multiple copies that may continue execution independently.

Another well-known inference technique for probabilistic programs is lightweight Metropolis-Hasting (LMH) Wingate et al. (2011). LMH methods construct a database of sampled values at run time. A change to the sampled values is proposed, and the program is rerun in its entirety, substituting previously sampled values where possible. The newly constructed database of random variables is then either accepted or rejected. LMH is straightforward to implement, but costly, since the program must be re-run in its entirety to evaluate the acceptance ratio. A more computationally efficient strategy is offered by Venture Mansinghka et al. (2014), which represents the execution history of a program as a graph, in which each evaluated expression is a node. A graph walk can then determine the subset of expressions affected by a proposed change, allowing partial re-execution of the program conditioned on all unaffected nodes.

SMC and MH based algorithms each have trade-offs. Techniques that derive from SMC can be run very efficiently, but suffer from particle degeneracy, resulting in a deteriorating quality of posterior estimates calculated from values sampled early in the program. In MH methods subsequent samples are typically correlated, and many updates may be needed to obtain an independent sample. As we will discuss in Section 3, PGAS can be thought of as a hybrid technique, in the sense that SMC is used to generate independent updates to the previous sample, which mitigates the degeneracy issues associated with SMC whilst increasing mixing rates relative to MH.

Probabilistic Functional Programs

For the purposes of exposition we will consider a simple Lisp dialect, extended with primitives for sampling and observing random variables

An expression e is either a constant literal c, a symbol s, an application (e &e) with operator e and zero or more arguments &e, a function literal (lambda (&s) e) with argument list (&s), an if-statement (if e e e), or a quoted expression (quote e). Each expression e evaluates to a value v upon execution. In addition to the self-explanatory bool, int, float, and string types, values can be primitive procedures (i.e. language built-ins such +, -, etc.), compound procedures (i.e., closures), and lists of zero or more values (&v).

The stochastic type represents stochastic processes, whose samples must either be i.i.d. or exchangeable. A stochastic process sp supports two operations

The sample primitive draws a value from sp, whereas the observe primitive conditions execution on a sample v that is returned as passed. Both operations associate a value v to sp as a side-effect, which changes the probability of the execution state in the inference procedure. For exchangeable processes such as (crp 1.0) this also affects the probability of future samples.

Following the convention employed by Venture and Anglican, we define programs as sequences of three types of top-level statements

Variables in the global environment are defined using the assume directive. The observe directive conditions execution in the same way as its non-toplevel equivalent. The predict directive returns the value of the expression e as inference output.

2 Importance Sampling Semantics

In many probabilistic languages the (observe e v) form is semantically equivalent to imposing a rejection-sampling criterion. For example, a program subject to (observe (> a 0) true) can be interpreted as a rejection sampler that repeatedly runs the program and only returns predict values when (> a 0) holds.

The semantics of observe that we have defined here imply an interpretation of a probabilistic program as an importance sampler, where sample draws from the prior and observe assigns an importance weight. Instead of constraining the value of an arbitrary expression e, observe conditions on the value of (sample e). This restricted form of conditioning guarantees that we can calculate the likelihood $p($v$\,|\,$(sample e)$)$ for every observe, as long as we implement a density function for all possible stochastic values in the language.

More formally, we use the notation $F[\boldsymbol{y}]$ to refer to a program conditioned on values $\boldsymbol{y}$ via a sequence of top-level [observe e v] statements. Execution of $F[\boldsymbol{y}]$ will require the evaluation of a number of (sample e) expressions, whose values we will denote with $\boldsymbol{x}$. We now informally define $F[\boldsymbol{y},\boldsymbol{x}]$ as the program in which all (sample e) expressions are replaced by conditioned equivalents (observe e v), resulting in a fully deterministic execution. Similarly we can define $F[\boldsymbol{x}]$ as the program obtained from $F[\boldsymbol{y},\boldsymbol{x}]$ by replacing [observe e v] statements with unconditioned forms [assume s (sample e)] with a unique symbol s in each statement.

Whereas top-level observe statements are fixed in number and order, $F$ may not evaluate the same combination of sample calls in every execution. To provide a more precise definition of $F[\boldsymbol{x}]$, we associate a unique address $\alpha$ with each sample call that can occur in the execution of $F$. We here use a scheme in which the run-time address of each evaluation is a concatenation $\alpha^{\prime}\!\!::\!\!(t,p)$ of the address $\alpha^{\prime}$ of the parent evaluation and a tuple $(t,p)$ in which $t$ identifies the expression type and $p$ is the index of the sub-expression within the form. This particular scheme labels every evaluation, not just the sample calls, allowing us to formally represent $F:\mathcal{A}\to\mathcal{E}$ as a mapping from addresses $\mathcal{A}$ to program expressions $\mathcal{E}$. We represent $\boldsymbol{x}:\mathcal{S}\to\mathcal{V}$ as a mapping from a subset of addresses $\mathcal{S}\subset\mathcal{A}$ associated with sample calls to values $\mathcal{V}$. Similarly, $\boldsymbol{y}:\mathcal{O}\to\mathcal{V}$ is a mapping from addresses $\mathcal{O}\subset\mathcal{A}$ associated with observe statements to values. With these definitions in place, a conditioned program $F[\boldsymbol{x}]:\mathcal{A}\to\mathcal{E}$ simply replaces $F(\alpha)=\>$(sample e) with

The importance sampling interpretation of a program $F[\boldsymbol{y}]\leadsto W,\boldsymbol{x}$ (read as $F[\boldsymbol{x}]$ yields $W,\boldsymbol{x}$) is defined in terms of random variables $\boldsymbol{x}$ and a weight $W$

By definition, the generated samples $\boldsymbol{x}$ are drawn from the prior $p(\boldsymbol{x}\,|\,F)$. The weight $W=p(\boldsymbol{y}\,|\,F[\boldsymbol{x}])$ is the joint probability of all top-level observe statements in $F[\boldsymbol{y}]$. Repeated execution of $F[\boldsymbol{y}]$ yields a weighted sample set $\{W^{l},\boldsymbol{x}^{l}\}$ that may be used to approximate the posterior as

More generally the importance weight is defined as the joint probability of all observe calls (top-level and transformed) in the program

Particle MCMC methods

SMC methods are importance sampling techniques that target a posterior $p(\boldsymbol{x}\,|\,\boldsymbol{y})$ on as space $\mathcal{X}$ by performing importance sampling on unnormalized densities $\{\gamma_{n}(\boldsymbol{x}_{n})\}_{n=1}^{N}$ defined on spaces of expanding dimensionality $\{\mathcal{X}_{n}\}_{n=1}^{N}$, where each $\mathcal{X}_{n}\subseteq\mathcal{X}_{n+1}$ and $\mathcal{X}_{N}=\mathcal{X}$. This results in a series of intermediate particle sets $\{w_{n}^{l},\boldsymbol{x}_{n}^{l}\}_{n=1}^{N}$, which we refer to as generations. Each generation is sampled via two steps,

Here $R(a\,|\,w)$ is a resampling procedure that returns index $a=l$ with probability $w^{l}/\sum_{l^{\prime}}w^{l^{\prime}}$ and $\rho_{n}$ is a transition kernel. The samples $\boldsymbol{x}_{n}^{l}$ are assigned weights

In the context of probabilistic programs, we can define a series of partial programs $F_{n}[\boldsymbol{y}_{n}]$ that truncate at each top-level [observe e v] statement. We can then sequentially sample $F_{n}[\boldsymbol{y}_{n},\boldsymbol{x}_{n-1}]\leadsto W_{n},\boldsymbol{x}_{n}$ by partially conditioning on $\boldsymbol{x}_{n-1}$ at each generation. Since $\boldsymbol{x}_{n}$ is a sample from the prior, this results in an importance weight Wood et al. (2014)

Note that this is simply the likelihood of the $n$-th top-level observe. In practice we continue execution relative to $F_{n-1}[\boldsymbol{y}_{n-1},\boldsymbol{x}_{n-1}]$ to avoid rerunning $F_{n}[\boldsymbol{y}_{n},\boldsymbol{x}_{n-1}]$ in its entirety. The means that the inference backend must include a routine for forking multiple independent executions from a single state.

2 Iterative Conditional SMC

An advantage of SMC methods is that they provide a generic strategy for joint proposals in high dimensional spaces. An importance sampling scheme where $F[\boldsymbol{y}]\leadsto W,\boldsymbol{x}$ draws from the prior has a vanishingly small probability of generating a high-weight sample. By sampling the smallest possible set of variables $\boldsymbol{x}^{l}_{n}$ at each generation and selecting ancestors $a^{l}_{n+1}$ according to the likelihood of the next observed data point, we ensure that $\boldsymbol{x}^{l}_{n+1}$ is sampled conditioned on high-weight values of $\boldsymbol{x}_{n}$ from the previous generation. At the same time this strategy has a drawback: each time the particle set is resampled, the number of unique values at previous generations decreases, typically resulting in coalescence to a single common ancestor in $O(L\log L)$ generations Jacob et al. (2013).

In many applications it is not practically possible (due to memory requirements) to set $L$ to a value large enough to guarantee a sufficient number of independent samples at all generations. In such cases, particle variants of MCMC techniques Andrieu et al. (2010) can be used to combine samples from multiple SMC sweeps. An iterated conditional SMC (ICSMC) sampler repeatedly selects a retained particle $k$ with probability $w_{N}^{k}/\sum_{k^{\prime}}w_{N}^{k^{\prime}}$, and then performs a conditional SMC (CSMC) sweep, where the resampling step is conditioned on the inclusion of the retained particle at each generation. Formally, this procedure is a partially collapsed Gibbs sampler that targets a density $\phi(\boldsymbol{x},a,k)$ on an extended space

We use the shorthand $\boldsymbol{x}^{k}=\boldsymbol{x}^{b_{1}:b_{N}}_{1:N}$ and $a^{k}=a^{b_{2}:b_{N}}_{2:N}$ to refer to the sampled values and ancestor indices of the retained particle, whose index $b_{n}$ at each generation can be recursively defined via $b_{N}=k$ and $b_{n-1}=a_{n}^{b_{n}}$. The notation $\boldsymbol{x}^{-k}$ and $a^{-k}$ refers to the complements where the retained particle is excluded. An ICSMC sampler iterates between two updates

$\{\boldsymbol{x}^{*,-k},a^{*,-k}\}\sim\phi(\boldsymbol{x}^{-k},a^{-k}\,|\,\boldsymbol{x}^{k},a^{k},k)$

$k^{*}\sim\phi(k\,|\,\boldsymbol{x}^{*,-k},a^{*,-k},\boldsymbol{x}^{k},a^{k})$

If we interpret $\boldsymbol{x}^{-k}$, $a$ and $k$ as auxilliary variables, then the marginal on $\boldsymbol{x}_{N}^{k}$ leaves the density $\gamma_{N}(\boldsymbol{x}_{N})$ invariant Andrieu et al. (2010).

The advantage of ICSMC samplers is that the target space is iteratively explored over subsequent CSMC sweeps. A disadvantage is that consecutive sweeps often yield partially degenerate particles, since many newly generated particles will coalesce to the retained particle with high probability. For this reason ICSMC samplers mix poorly when the number of particles is not large enough to generate at least two completely independent lineages in a single CSMC sweep.

3 Particle Gibbs with Ancestor Sampling

PGAS is a technique that augments the CSMC sweep with a resampling procedure for the index $a_{n}^{b_{n}}$ of the retained particle Lindsten et al. (2012). At a high level, this sampling scheme performs two updates

$\{\boldsymbol{x}^{*,-k},a^{*}\}\sim\phi(\boldsymbol{x}^{-k},a\,|\,\boldsymbol{x}^{k},k)$

$k^{*}\sim\phi(k\,|\,\boldsymbol{x}^{*,-k},a^{*,-k},\boldsymbol{x}^{k},a^{k})$

Here update 1 differs from the normal CSMC update in that it samples a complete set of ancestor indices $a^{*}$, not the complement to the retained indices $a^{*,-k}$. At each generation the ancestor $a_{n}^{b_{n}}$ of the retained particle is resampled according to a weight

Here $\boldsymbol{x}_{N}^{k}[\boldsymbol{x}_{n}^{l}]$ denotes the substitution of $\boldsymbol{x}_{n}^{l}$ into $\boldsymbol{x}_{N}^{k}$, which is defined as a mapping $\boldsymbol{x}_{N}^{k}[\boldsymbol{x}_{n}^{l}]:\mathcal{A}_{N}^{k}\cup\mathcal{A}_{n}^{l}\to\mathcal{V}$ where entries in $\boldsymbol{x}_{n}^{l}:\mathcal{A}_{n}^{l}\to\mathcal{V}$ augment or replace entries in $\boldsymbol{x}_{N}^{k}:\mathcal{A}_{N}^{k}\to\mathcal{V}$.

Intuitively, ancestor resampling can be thought of as proposing new program executions by complementing random variables $\boldsymbol{x}_{n}^{l}$ of a partial execution, or prefix, with retained values from $\boldsymbol{x}_{N}^{k}$ to specify the future of the execution, or suffix. This step is performed at each generation, allowing the retained lineage to potentially be resampled many times in the course of one sweep.

Incorporation of the ancestor resampling step results in the following updates at each $n=2,\ldots,N$

Update the particles for $l\in\{1,...,L\}\setminus b_{n}$

Rescoring Probabilistic Programs

PGAS for probabilistic programs requires calculation of $p(\boldsymbol{y}_{N},\boldsymbol{x}^{k}_{N}[\boldsymbol{x}_{n}^{l}]\,|\,F[\boldsymbol{y}_{n},\boldsymbol{x}^{l}_{n}])$. This probability is, by definition, the importance weight of $F_{N}[\boldsymbol{y}_{N},\boldsymbol{x}^{k}_{N}[\boldsymbol{x}^{l}_{n}]]$ and can therefore in principle be obtained by executing this re-conditioned form. The main drawback of this naive approach is that it requires $LN$ evaluations of the program in its entirety. This results in an $O(LN^{2})$ computational cost, which quickly becomes prohibitively expensive as the number of generations $N$ increases. A second complicating factor is that naively rewriting part of the execution history of the program may not yield a set of random values $\boldsymbol{x}_{N}^{k}[\boldsymbol{x}^{l}_{n}]$ that could be generated by running $F_{N}[\boldsymbol{y}_{N}]$.

We here develop a formalism that allows regeneration of a self-consistent program execution starting from a partial execution with random variables $\boldsymbol{x}_{n}^{l}$, assuming all future random samples are inherited from retained values $\boldsymbol{x}_{N}^{k}$. To do so we introduce the notion of a trace, a data structure that annotates each evaluation with information necessary to re-execute the expression relative to a new program state. Given a trace for the suffix, i.e. the remaining top-level statements in a program, it becomes possible to re-execute conditioned on future random values, in a manner that avoid unnecessary recomputation where possible.

In higher order languages with recursion and memoization, regeneration is complicated by two factors:

The program $F_{N}[\boldsymbol{y}_{N},\boldsymbol{x}^{k}_{N}[\boldsymbol{x}^{l}_{n}]]$ may be underconditioned, in the sense that $\boldsymbol{x}^{k}_{N}[\boldsymbol{x}^{l}_{n}]$ does not contain values for some sample calls that can be evaluated in its execution. It can also be overconditioned, when $\boldsymbol{x}^{k}_{N}[\boldsymbol{x}^{l}_{n}]$ contains values for sample calls that will never be evaluated.

The expression for the stochastic argument to each observe in $F_{N}[\boldsymbol{y}_{N},\boldsymbol{x}^{k}_{N}[\boldsymbol{x}^{l}_{n}]]$ may need to be re-evaluated if it in some way depends on global variables defined in $F_{n}[\boldsymbol{y}_{n},\boldsymbol{x}^{l}_{n}]$. Because each variable may in turn reference other variables, we must be able to reconstruct the program environment recursively in order to rescore each observe.

The first non-triviality arises from the existence of if expressions. As an example, consider the program

The sample expression in line 1 is only evaluated when the sample expression in line 0 evaluates to true. More generally, any program that contains (sample e) inside an if expression will not be guaranteed to instantiate the same random variables in cases where the predicate of the if expression itself depends on previously sampled values. In programming languages that lack recursion we have the option of evaluating both branches and including or excluding the associated probabilities conditioned on the predicate value. This is essentially the strategy that is employed to handle if expressions in Infer.NET and BUGS variants. In languages that do permit recursion, this is in general not possible. For example, the following program would require evaluation of an infinite number of branches:

In other words, we cannot in general pre-evaluate the values associated with both branches in the suffix. When a predicate in the suffix no longer takes on the same value, we have a choice of either rejecting the regenerated suffix outright, or updating it using a regeneration procedure that evaluates the newly chosen branch and removes reference to any values sampled in the invalidated branch. We here consider the former strict form of regeneration, which guarantees that the regenerated suffix references precisely the same set of sample values as before.

A second aspect that complicates rescoring is the existence of memoized procedures. As an example, consider the following infinite mixture model

Here each observe makes a call to class, which samples an integer class label k from a Chinese restaurant Process (CRP). The call to class-dist either retrieves an existing stochastic value, or generates one when a new value k is encountered. This type of memoization pattern allows us to delay sampling of the parameters until they are in fact required in the evaluation of a top-level observe, and makes it straightforward to define open world models with unbounded numbers of parameters. At the same time it complicates analysis when performing rescoring. Memoized procedure calls are semantically equivalent to lazily defined variables. Programs that rely on memoization can therefore essentially define variables in a non-deterministic order. A regeneration operation must therefore dynamically determine the set of variables that need to be re-evaluated at run time.

2 Traced Evaluation

The operations that need to happen during a rescoring step are (1) the regeneration of a consistent set of global environment variables, which includes any bindings in the prefix, augmented with any bindings defined in the suffix (some of which may require re-evaluation as a result of changes to bindings in the prefix), (2) a verification that the flow control path in the suffix is consistent with the environment bindings in the prefix, and (3) the recomputation of the probabilities of any sampled and observed values whose density values depend on bindings in the prefix.

In order to make it possible to perform the above operations, we begin by introducing a set of annotations for each value v that is returned upon evaluation of an expression e. In practical terms, each value in the language is boxed into a data structure which we call a trace. We represent a trace $\tau$ as tuple $({\tt v},\epsilon,l,\rho,\omega,\sigma,\phi)$. v is the value of the expression. $\epsilon$ is a partially evaluated expression, whose sub-expressions are themselves represented as traces. $l$ is the accumulated log-weight of observes evaluated within the expression. $\rho$ is a mapping $\{{\tt s}\to\tau\}$ from symbols to traces, containing the subset of the global environment variables that were referenced in the evaluation of e. $\omega$ is a mapping $\{\alpha\mapsto(\tau,{\tt v},l)\}$. It contains an entry at the address $\alpha$ of each observe that was evaluated in e and its sub-expressions. This entry is represented as a tuple $(\tau,{\tt v},l)$ containing a trace of the first argument to the observe (which must be of the stochastic type), the observed value v, and the associated log-weight $l$. Similarly $\sigma$ is a mapping $\{\alpha\mapsto(\tau,{\tt v})\}$ that contains an entry for each evaluated sample expression (which omits the associated log-weight). The last component $\phi$ is again a mapping $\{\alpha\mapsto\tau\}$ that records all traces that appear as conditions in if expressions and thereby influence the control flow of program execution.

We now describe the semantics of the traced evaluation $({\tt e},\alpha,R,\Lambda)\Downarrow\tau$. The evaluation operator $\Downarrow$ returns the trace $\tau$ of an expression ${\tt e}$ at address $\alpha$, relative to a global environment $R$ (i.e. variables defined via assume statements) and local bindings $\Lambda$ (i.e. variables bound in compound procedure calls), resulting in a trace $\tau=({\tt v},\epsilon,l,\rho,\omega,\sigma,\phi)$. We assume the implementation provides a standard evaluation function for primitive procedures $\mathit{eval}(\mathtt{prim}~{}\mathtt{v}_{1}\ldots\mathtt{v_{n}})=\mathtt{v}$. We reiterate that the notation $\alpha{::}(t,p)$ denotes an evaluation address composed of a parent address $\alpha$, a type identifier $t$ and a sub-expression index $p$, where $t$ is one of i for if, l for lambda, q for quote, a for applications, and b when evaluating compound procedure bodies.

We call this type of trace “transparent”, since it contains no references to other traces that may take on different values or probabilities in another execution.

Symbol lookups s in the global environment return the value of s stored in $R$:

Lookups from the local environment are inlined:

Calls to sample inherit annotations from the trace $\tau$ passed as an argument, and add an entry in $\sigma$:

Calls to observe inherit from $\tau$ and add an entry in $\omega$:

Here $\mathcal{L}({\tt v}_{1},{\tt v}_{2})$ is used to denote the log-density of value ${\tt v}_{2}$ relative to ${\tt v}_{1}$ (which must be of type stochastic).

Primitive procedure applications $({\tt primop}~{}{\tt e}_{1}~{}{\tt e}_{2})$ evaluate to $\mathit{eval}({\tt prim}~{}{\tt v}_{1}~{}{\tt v}_{2})$ where ${\tt v}_{i}$ is the result of evaluating ${\tt e}_{i}$:

if $({\tt e}_{i},\alpha{::}({\tt p},i),R,\Lambda)\Downarrow\tau_{i}$, and $\rho$, $\omega$, $\sigma$, and $\phi$ are obtained by merging the corresponding components of the $\tau_{i}$, and the log-density $l$ is the sum of the $l_{i}$.

Application of a single-argument compound procedure (i.e. closure) ${\tt e}_{1}$ leads to the evaluation of the body e of the procedure relative to the environments $R_{1},\Lambda_{1}$ in which the compound procedure was defined:

and $l,\rho,\omega,\sigma,\phi$ are obtained by combining the corresponding components of the $\tau_{i}$. The case with multiple arguments is defined similarly.

Quote expressions $({\tt quote}~{}\tau)$ simply return $\tau$:

Finally, an if expression $({\tt if}~{}{\tt e}~{}{\tt e}_{1}~{}{\tt e}_{2})$ returns either the result of evaluating ${\tt e}_{1}$ or ${\tt e}_{2}$. We show only the case where the true branch is taken:

and $l,\rho,\omega,\sigma,\phi$ are obtained by combining the corresponding components of $\tau$ and $\tau_{1}$. The other case is that the value of $\tau$ is false, and has the semantics similar to the one above.

3 Regeneration and Rescoring

We now define an operation $\mathcal{R}(\tau,R)=\tau^{\prime}$ that regenerates a traced value relative to an environment $R$. This operation performs the following steps:

Re-evaluate predicates: Compare $\tau=\phi(\alpha)$ to $\tau^{\prime}=\mathcal{R}(\tau,R)$ for all $\alpha$. Abort if $\tau^{\prime}$ and $\tau$ have different values v. Otherwise update $\phi[\alpha\mapsto\tau^{\prime}]$.

Re-score observe expressions and statements: Let $(\tau,{\tt v},l)=\omega(\alpha)$, and $\tau^{\prime}=\mathcal{R}(\tau,R)$. If $\tau^{\prime}$ and $\tau$ have different values ${\tt v}^{\prime}_{0}$ and ${\tt v}_{0}$, recalculate $l^{\prime}=\mathcal{L}({\tt v}^{\prime}_{0},{\tt v})$ and update $\omega[\alpha\to(\tau^{\prime},{\tt v},l^{\prime})]$. Otherwise, update $\omega[\alpha\to(\tau^{\prime},{\tt v},l)]$.

Re-score samples: Let $(\tau,{\tt v})=\sigma(\alpha)$, and $\tau^{\prime}=\mathcal{R}(\tau,R)$. Calculate $l^{\prime}=\mathcal{L}(\tau^{\prime},{\tt v})$ and update $\omega[\alpha\mapsto(\tau^{\prime},{\tt v},l^{\prime})]$.

Regenerate the environment bindings: For all symbols s that do not exist in $R$, let $\tau^{\prime}=\mathcal{R}(\rho({\tt s}),R)$ and update $R[{\tt s}\mapsto\tau^{\prime}]$ and $\rho[{\tt s}\mapsto\tau^{\prime}]$. For existing symbols update $\rho[{\tt s}\mapsto R({\tt s})]$. We for convenience assume that $R$ is updated in place, though this may be avoided by having $\mathcal{R}$ return a tuple $(R^{\prime},\tau^{\prime})$.

If any bindings were changed with new values in step 4, regenerate all sub-expressions $\tau_{i}$ in $\epsilon$ to reconstruct $\epsilon^{\prime}$.

If $\epsilon^{\prime}$ was reconstructed in step 5, evaluate $\epsilon^{\prime}$ and update v to the result of this evaluation.

Rescoring an individual trace may be performed as part of the regeneration sweep by calculating a difference in log-density $\Delta l$. This $\Delta l$ is the sum of all terms $l^{\prime}-l$ in step 2 and all terms $l^{\prime}$ in step 3, and any $\Delta l^{\prime}$ values return from recursive calls to $\mathcal{R}$.

We have omitted a few technical details in this high-level description. The first is that we build a map $\mathcal{C}=\{\tau\to\tau^{\prime},\ldots\}$ on a call to $\mathcal{R}$, which is passed as an additional argument in recursive calls to $\mathcal{R}$, effectively memoizing the computation relative to a given initial environment $R$. This reduces the computation on recursive calls, which potentially expand the same traces $\tau$ many times as sub-expressions of $\epsilon$.

4 Ancestor Resampling

Given an implementation of a traced evaluator and a regenerating/rescoring procedure $\mathcal{R}$, an implementation for PGAS in probabilistic programs becomes straightforward. We represent the programs $F_{n}[\boldsymbol{y}_{n},\boldsymbol{x}_{n}^{l}]$ evaluated up to the first $n$ top-level statements as pairs $(R_{n}^{l},\tau^{l}_{n})$. We now construct a concatenated suffix $\mathcal{T}_{n}$ by recursively re-evaluating $\mathcal{T}_{n}=({\tt cons}~{}\tau^{b_{n}}_{n}~{}\mathcal{T}_{n+1})$. For $n={1,\ldots,N}$, we then define $p(\boldsymbol{y}_{N},\boldsymbol{x}^{k}_{N}[\boldsymbol{x}_{n}^{*,l}]\,|\,F_{N}[\boldsymbol{y}_{n},\boldsymbol{x}_{n}^{l}])$ as the log-weight of the rescored trace $\mathcal{R}(\mathcal{T}^{l}_{n},R^{l}_{n-1})$. Note here that by construction, we may extract $\mathcal{T}_{n+1}$ from $\mathcal{T}_{n}$ without additional computation.

Experiments

To evaluate the mixing properties of PGAS relative to ICSMC, we consider a linear dynamical system (i.e. a Kalman smoothing problem) with a 2-dimensional latent space and a D-dimensional observational space,

We impose additional structure by assuming that the transition matrix $\boldsymbol{A}$ is a simple rotation with angular velocity $\omega$, whereas the transition covariance $\boldsymbol{Q}$ is a diagonal matrix with constant coefficient $q$,

We simulate data with $D=36$ dimensions, $T=100$ time points, $\omega=4\pi/T$, $q=0.1$, $\alpha=0.1$, and $r=0.01$. We now consider an inference setting where $\boldsymbol{C}$ and $\boldsymbol{R}$ are assumed known and estimate the state trajectory $\boldsymbol{z}_{1:T}$, as well as the parameters of the transition model $\omega$ and $q$, which are given mildly informative priors, $\omega\sim{\rm Gamma}(10,2.5)$ and $q\sim{\rm Gamma}(10,100)$.

While this is a toy problem where an expectation maximization (EM) algorithm could likely be derived, it is illustrative of the manner in which probabilistic programs can extend models by imposing additional structure, in this case the dependency of $\boldsymbol{A}$ on $\omega$. This modified Kalman smoothing problem can be described in a small number of program lines

Here [...] refers to a vector or matrix literal.

We compare results for PGAS with 10 particles to ICSMC with $10,20,50,100,200,$ and $500$ particles. In each case we run 100 PMCMC sweeps and 25 restarts with different random seeds. To characterize mixing rates we calculate the effective sample size (ESS) of the aggregate sample set $\{w_{t}^{s,l},\boldsymbol{z}_{t}^{s,l}\}$ over all sweeps $s=\{1,\ldots,100\}$,

Here $V_{t}^{k}$ represents the total importance weight associated with each unique value $\boldsymbol{z}_{t}^{k}$ in $\{\boldsymbol{z}_{t}^{s,l}\}$.

Figure 1 shows the ESS as a function of $t$. ICSMC shows a decreasing ESS as $t$ approaches 0, indicating poor mixing for values sampled at early generations. PGAS, in contrast, exhibits an ESS that fluctuates but is otherwise independent of $t$. ESS estimates varied approximately 15% relative to the mean across independent restarts. This suggests that fluctuations in the ESS reflect variations in the prior probability of latent transitions $p(\boldsymbol{z}_{t}\,|\,\boldsymbol{z}_{t-1},\boldsymbol{A})$.

For this model, our PGAS implementation with 10 particles has a computational cost per sweep comparable to that of ICSMC with 300 particles. However, when we consider the ESS per computation time, the increases in mixing efficiency outweigh increases in computational cost for state estimates below $t\simeq 50$. This is further illustrated in Figure 2, which shows the standard deviation of sample estimates of the parameters $\omega$ and $q$. PGAS shows better convergence per sweep, particularly for estimates of $q$. ICSMC with 500 particles performs similarly to PGAS with 10 particles when estimating $\omega$, though ICSMC with 500 particles notably has a higher cost per sweep.

Discussion

Relative to other PMCMC methods such as ICSMC, PGAS methods have qualitatively different mixing characteristics, particularly for variables sampled early in a program execution. Implementing PGAS in the context of probabilistic programs poses technical challenges when programs can make use of recursion and memoization. The technique for traced evaluation developed here incurs an additional computational overhead, but avoids unnecessary recomputation during regeneration. When the cost of recomputation is large this will result in computational gains relative to a naive implementation that re-executes the suffix fully. Note that our approach tracks upstream, not downstream dependencies. In other words, we know what environment variables affect the value of a given expression, but not which expressions in a suffix depend on a given variable. All referenced symbol values must therefore be checked during regeneration, which can require a $O(LN^{2})$ computation in itself. Further gains could be obtained constructing a downstream dependency graph for the suffix, allowing more targeted regeneration via graph walk techniques analogous to those employed in Venture Mansinghka et al. (2014). At the same time, the empirical results presented here are indicative of the fact that, even without these additional optimizations, PGAS can easily yield better statistical results in cases where the parameter space is large and ICSMC sampling fails to mix.

Acknowledgements

We would like to thank our anonymous reviewers, as well as Brooks Paige and Dan Roy for their comments on this paper. JWM was supported by Google and Xerox. HY was supported by the EPSRC. FW and VKM were supported under DARPA PPAML. VKM was additionally supported by the ARL and ONR.

References