Automatic Alignment in Higher-Order Probabilistic Programming Languages

TL;DR

This paper introduces a static analysis method for automatically aligning checkpoints in higher-order probabilistic programming languages, enhancing the efficiency and accuracy of inference algorithms like SMC and MCMC.

Contribution

It presents a formal alignment analysis, proves its correctness, and integrates it into Miking CorePPL, enabling improved inference algorithms.

Findings

Aligned SMC and MCMC outperform standard versions in speed

Alignment improves inference accuracy in real-world tests

Automatic checkpoint determination simplifies user experience

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.