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

TL;DR

This paper introduces a static suspension analysis and selective CPS transformation technique for probabilistic programming languages, reducing overhead and improving efficiency in Monte Carlo inference algorithms.

Contribution

It presents a novel static analysis method that identifies suspension points, enabling selective CPS transformation to optimize probabilistic programming execution.

Findings

Significant performance improvements across various models and algorithms.

Formal proof of analysis correctness.

Effective reduction of CPS overhead in real-world applications.

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.