# Probabilistic Programming with Programmable Variational Inference

arXiv: https://arxiv.org/abs/2406.15742 · HTML: https://mlsanity.com/paper/2406.15742/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 48 of its 79 references that are in the map.
- First posted on arXiv 22 Jun 2024.

## Builds on (most cited first, 30 of 48)

- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [Variational Diffusion Models](https://mlsanity.com/paper/2107.00630/) (2021)
- [NVAE: A Deep Hierarchical Variational Autoencoder](https://mlsanity.com/paper/2007.03898/) (2020)
- [Semi-Supervised Learning with Deep Generative Models](https://mlsanity.com/paper/1406.5298/) (2014)
- [Importance Weighted Autoencoders](https://mlsanity.com/paper/1509.00519/) (2015)
- [Pyro: Deep Universal Probabilistic Programming](https://mlsanity.com/paper/1810.09538/) (2018)
- [Monte Carlo Gradient Estimation in Machine Learning](https://mlsanity.com/paper/1906.10652/) (2019)
- [Attend, Infer, Repeat: Fast Scene Understanding with Generative Models](https://mlsanity.com/paper/1603.08575/) (2016)
- [Automatic Differentiation Variational Inference](https://mlsanity.com/paper/1603.00788/) (2016)
- [Variational Autoencoder for Deep Learning of Images, Labels and Captions](https://mlsanity.com/paper/1609.08976/) (2016)
- [Gradient Estimation Using Stochastic Computation Graphs](https://mlsanity.com/paper/1506.05254/) (2015)
- [GFlowNets and variational inference](https://mlsanity.com/paper/2210.00580/) (2022)
- [Variational Sequential Monte Carlo](https://mlsanity.com/paper/1705.11140/) (2017)
- [Hierarchical Variational Models](https://mlsanity.com/paper/1511.02386/) (2015)
- [Tighter Variational Bounds are Not Necessarily Better](https://mlsanity.com/paper/1802.04537/) (2018)
- [Auxiliary Deep Generative Models](https://mlsanity.com/paper/1602.05473/) (2016)
- [Reweighted Wake-Sleep](https://mlsanity.com/paper/1406.2751/) (2014)
- [Elements of Sequential Monte Carlo](https://mlsanity.com/paper/1903.04797/) (2019)
- [DiCE: The Infinitely Differentiable Monte-Carlo Estimator](https://mlsanity.com/paper/1802.05098/) (2018)
- [Markovian Score Climbing: Variational Inference with KL(p||q)](https://mlsanity.com/paper/2003.10374/) (2020)
- [Neural Adaptive Sequential Monte Carlo](https://mlsanity.com/paper/1506.03338/) (2015)
- [Automatic Differentiation of Programs with Discrete Randomness](https://mlsanity.com/paper/2210.08572/) (2022)
- [Deep Probabilistic Programming](https://mlsanity.com/paper/1701.03757/) (2017)
- [Venture: a higher-order probabilistic programming platform with programmable inference](https://mlsanity.com/paper/1404.0099/) (2014)
- [Importance Weighted Hierarchical Variational Inference](https://mlsanity.com/paper/1905.03290/) (2019)
- [Simple, Distributed, and Accelerated Probabilistic Programming](https://mlsanity.com/paper/1811.02091/) (2018)
- [A Convenient Category for Higher-Order Probability Theory](https://mlsanity.com/paper/1701.02547/) (2017)
- [Credit Assignment Techniques in Stochastic Computation Graphs](https://mlsanity.com/paper/1901.01761/) (2019)
- [Nested Variational Inference](https://mlsanity.com/paper/2106.11302/) (2021)
- [A Lambda-Calculus Foundation for Universal Probabilistic Programming](https://mlsanity.com/paper/1512.08990/) (2015)

## Built on by (most cited first, 4 of 4)

- [Building Machines that Learn and Think with People](https://mlsanity.com/paper/2408.03943/) (2024)
- [World Modeling with Probabilistic Structure Integration](https://mlsanity.com/paper/2509.09737/) (2025)
- [GenMatter: Perceiving Physical Objects with Generative Matter Models](https://mlsanity.com/paper/2604.22160/) (2026)
- [GradInf: Gradient Estimation as Probabilistic Inference](https://mlsanity.com/paper/2607.07840/) (2026)

## Abstract (arXiv)

Compared to the wide array of advanced Monte Carlo methods supported by modern probabilistic programming languages (PPLs), PPL support for variational inference (VI) is less developed: users are typically limited to a predefined selection of variational objectives and gradient estimators, which are implemented monolithically (and without formal correctness arguments) in PPL backends. In this paper, we propose a more modular approach to supporting variational inference in PPLs, based on compositional program transformation. In our approach, variational objectives are expressed as programs, that may employ first-class constructs for computing densities of and expected values under user-defined models and variational families. We then transform these programs systematically into unbiased gradient estimators for optimizing the objectives they define. Our design enables modular reasoning about many interacting concerns, including automatic differentiation, density accumulation, tracing, and the application of unbiased gradient estimation strategies. Additionally, relative to existing support for VI in PPLs, our design increases expressiveness along three axes: (1) it supports an open-ended set of user-defined variational objectives, rather than a fixed menu of options; (2) it supports a combinatorial space of gradient estimation strategies, many not automated by today's PPLs; and (3) it supports a broader class of models and variational families, because it supports constructs for approximate marginalization and normalization (previously introduced only for Monte Carlo inference). We implement our approach in an extension to the Gen probabilistic programming system (genjax.vi, implemented in JAX), and evaluate on several deep generative modeling tasks, showing minimal performance overhead vs. hand-coded implementations and performance competitive with well-established open-source PPLs.