From high-level inference algorithms to efficient code

TL;DR

This paper presents a compilation approach that transforms high-level probabilistic programs into efficient low-level code, enabling automatic, optimized execution of complex inference algorithms traditionally implemented manually.

Contribution

It introduces new transformations and optimizations that compile high-level probabilistic models into efficient loop code, matching manual implementation performance.

Findings

Compiled probabilistic programs achieve performance comparable to handwritten code.

New transformations enable direct expression of inference algorithms in high-level languages.

Optimizations and JIT compilation improve execution efficiency significantly.

Abstract

Probabilistic programming languages are valuable because they allow domain experts to express probabilistic models and inference algorithms without worrying about irrelevant details. However, for decades there remained an important and popular class of probabilistic inference algorithms whose efficient implementation required manual low-level coding that is tedious and error-prone. They are algorithms whose idiomatic expression requires random array variables that are latent or whose likelihood is conjugate. Although that is how practitioners communicate and compose these algorithms on paper, executing such expressions requires eliminating the latent variables and recognizing the conjugacy by symbolic mathematics. Moreover, matching the performance of handwritten code requires speeding up loops by more than a constant factor. We show how probabilistic programs that directly and concisely express these desired inference algorithms can be compiled while maintaining efficiency. We introduce new transformations that turn high-level probabilistic programs with arrays into pure loop code. We then make great use of domain-specific invariants and norms to optimize the code, and to specialize and JIT-compile the code per execution. The resulting performance is competitive with manual implementations.