Automatic Alignment in Higher-Order Probabilistic Programming Languages

Abstract

Probabilistic Programming Languages (PPLs) allow users to encode statistical inference problems and automatically apply an inference algorithm to solve them. Popular inference algorithms for PPLs, such as sequential Monte Carlo (SMC) and Markov chain Monte Carlo (MCMC), are built around checkpoints—relevant events for the inference algorithm during the execution of a probabilistic program. Deciding the location of checkpoints is, in current PPLs, not done optimally. To solve this problem, we present a static analysis technique that automatically determines checkpoints in programs, relieving PPL users of this task. The analysis identifies a set of checkpoints that execute in the same order in every program run—they are aligned. We formalize alignment, prove the correctness of the analysis, and implement the analysis as part of the higher-order functional PPL Miking CorePPL. By utilizing the alignment analysis, we design two novel inference algorithm variants: aligned SMC and aligned lightweight MCMC. We show, through real-world experiments, that they significantly improve inference execution time and accuracy compared to standard PPL versions of SMC and MCMC.

TreePPL: A Universal Probabilistic Programming Language for Phylogenetics

“…In the left side of Fig. CorePPL then adds code for the selected inference strategy: one of importance sampling, bootstrap particle filter, alive particle filter [30], aligned lightweight MCMC [29,43], trace MCMC, naive MCMC, or particle MCMC-particle independent Metropolis-Hastings (PMCMC-PIMH, [44]). The code obtained in this way is converted either into C++ (upcoming feature) or pure Miking code, which is subsequently processed by nvcc or mi to obtain respectively a CUDA-enabled executable (upcoming feature) or a regular executable.…”

“…Factorizing the likelihood. Since observe-type statements indicate a potential likelihood update (a checkpoint-see [29] for an in-depth discussion), another way of thinking about TreePPL (and PPLs in general) is that it provides a way of specifying the likelihood function L(θ; y) in a programmatic way, instead of using a closed-form expression. In fact, we do not have to explicitly construct the data distribution and then observe from it-we can instead directly factor in the likelihood via weight or logWeight as we have shown on lines 4 (ensuring that the side-branch dies) and line 19 (correcting for the rotation factor of the tree).…”

“…If SMC is used as the inference algorithm, one key optimization that can be performed by the compiler is to automatically decide where to do resampling in the program. This alignment analysis will be performed by the TreePPL compiler (see [29,20] for details, also Section 3.1).…”

“…Diversification models and similar phylogenetic problems require so-called alignment of checkpoints for inference to be efficient [18]. Strategies for automating checkpoint alignment and exploiting it in SMC and MCMC inference have only been described very recently, and are only available in the Miking PPL framework currently, as far as we are aware [29]. Another automatic inference technique, which is not yet widely available but can lead to considerable efficiency improvements for some models is the alive particle filter (APF) [30].…”

“…Thus, it is easily extensible with new optimizations and inference strategies. It already supports importance sampling, the bootstrap particle filter (BPF) and the alive particle filter (APF) [30], aligned lightweight MCMC [29,43], trace MCMC, naive MCMC, and particle MCMC-particle independent Metropolis-Hastings (PMCMC-PIMH, [44]). However, as an intermediate language, Miking CorePPL is not designed to be used by end-users unfamiliar with functional programming and without a deep interest in compiler optimizations.…”