Probabilistic Programming with Programmable Variational Inference

Explore a groundbreaking approach to variational inference in probabilistic programming languages through this 21-minute conference talk from PLDI 2024. Delve into a modular method based on compositional program transformation that enhances expressiveness and flexibility in variational inference. Learn how variational objectives can be expressed as programs, utilizing first-class constructs for computing densities and expected values under user-defined models and variational families. Discover how these programs are systematically transformed into unbiased gradient estimators for objective optimization. Understand the benefits of this design, including the ability to prove unbiasedness through modular reasoning and increased expressiveness in defining variational objectives, gradient estimation strategies, and supported models and variational families. Examine the implementation of this approach in the Gen probabilistic programming system extension, genjax.vi, and its evaluation on deep generative modeling tasks. Gain insights into the performance comparisons with hand-coded implementations and established open-source probabilistic programming languages.