Suspension Analysis and Selective Continuation-Passing Style for Universal Probabilistic Programming Languages

TL;DR

The paper addresses the overhead of suspension in universal probabilistic programming languages by introducing a static suspension analysis and a selective CPS transformation. Grounded in a formal core calculus and $0$-CFA, the method identifies exactly which program parts may suspend under a given inference algorithm and applies CPS only there, yielding substantial runtime improvements over full CPS across multiple inference algorithms and real-world models. The authors implement these techniques in the Miking CorePPL compiler, provide an inference-problem extraction pass, and empirically validate gains on CRBD, ClaDS, LDA, and VBD, with improvements persisting across likelihood weighting, SMC variants, and MCMC methods. The work demonstrates practical impact for high-level PPLs, enabling performance-competitive suspensions without sacrificing generality or correctness, and offers a framework extensible to additional inference algorithms and models.

Abstract

Universal probabilistic programming languages (PPLs) make it relatively easy to encode and automatically solve statistical inference problems. To solve inference problems, PPL implementations often apply Monte Carlo inference algorithms that rely on execution suspension. State-of-the-art solutions enable execution suspension either through (i) continuation-passing style (CPS) transformations or (ii) efficient, but comparatively complex, low-level solutions that are often not available in high-level languages. CPS transformations introduce overhead due to unnecessary closure allocations -- a problem the PPL community has generally overlooked. To reduce overhead, we develop a new efficient selective CPS approach for PPLs. Specifically, we design a novel static suspension analysis technique that determines parts of programs that require suspension, given a particular inference algorithm. The analysis allows selectively CPS transforming the program only where necessary. We formally prove the correctness of the analysis and implement the analysis and transformation in the Miking CorePPL compiler. We evaluate the implementation for a large number of Monte Carlo inference algorithms on real-world models from phylogenetics, epidemiology, and topic modeling. The evaluation results demonstrate significant improvements across all models and inference algorithms.

Figures (14)

• Figure 1: A probabilistic program $\textbf{\upshapet}_\textrm{example}$ modeling the bias of a coin. Fig. (a) gives the program. The function $f_\textrm{Bernoulli}$ is the probability mass function of the Bernoulli distribution. Fig. (b) illustrates the distribution for $a_1$ at line \ref{['line:base:assume']} in (a). Fig. (c) shows the set of (weighted) samples resulting from conceptually running $\textbf{\upshapet}_\textrm{example}$ infinitely many times. Fig. (d) and Fig. (e) show the selective CPS transformations required for suspension at assume and weight, respectively. Fig. (f) gives $\textbf{\upshapet}_\textrm{example}$ in full CPS, with suspensions at assume and weight. The $_\textrm{CPS}$ subscript indicates CPS-versions of intrinsic functions such as $\mathit{head}$ and $\mathit{tail}$.
• Figure 2: A big-step operational semantics for $\textbf{\upshapet} \in T$. We omit the rule (If-False) for brevity; it is analogous to (If-True). The environment $\rho, x \mapsto \textbf{\upshapev}$ denotes $\rho$ extended with a binding $\textbf{\upshapev}$ for $x$. For each $d \in D$, the function $f_d$ is its probability density or probability mass function. E.g., $f_{\mathcal{N}(0,1)}(x) = e^{x^2/2}/\sqrt{2\pi}$, the density function of the standard normal distribution. We use the following notation: $\mathbin\Vert$ for sequence concatenation, $\cdot$ for multiplication, and $\lor$ for logical disjunction.
• Figure 3: The running example $\textbf{\upshapet}_\textrm{example}$ from Fig. \ref{['fig:running:base']} transformed to ANF.
• Figure 4: The running example from Fig. \ref{['fig:runninganf']} after selective CPS transformation. The program is semantically equivalent to Fig. \ref{['fig:running:weight']}.
• Figure 5: Overview of the Miking CorePPL compiler implementation. We divide the overall compiler into two parts, (i) suspension analysis and selective CPS (Section \ref{['sec:susimpl']}) and (ii) inference problem extraction (Section \ref{['sec:extractimpl']}). The figure depicts artifacts as gray rectangular boxes and transformation units and libraries as blue rounded boxes. Note how the inference extractors transformation separates the program into two different paths that are combined again after the inference-specific compilation. The white inheritance arrows (pointing to suspension analysis and selective CPS transformations) mean that these libraries are used within the inference-specific compiler transformation.

Theorems & Definitions (13)

• definition 1: Terms, values, and environments
• definition 2: Target language terms
• definition 3: Intrinsic arities and the $\delta$-function
• definition 4: Traces
• definition 5: Suspension requirement
• definition 6: A-normal form
• definition 7: Abstract values
• lemma 1: 0-CFA soundness
• theorem 1: Suspension analysis soundness
• proof