Probabilistic Programming with Programmable Variational Inference
Name
3656463.pdf
Size
887.21 KB
Format
Adobe PDF
Checksum (MD5)
db7c9e6df05466c165491e56b4f00af4
Author(s)
Becker, McCoy R.
Lew, Alexander K.
Wang, Xiaoyan
Ghavami, Matin
Huot, Mathieu
Rinard, Martin C.
Mansinghka, Vikash K.
Date Issued
June 20, 2024
Journal
Proceedings of the ACM on Programming Languages
Publisher
Association for Computing Machinery
Citation
Becker, McCoy R., Lew, Alexander K., Wang, Xiaoyan, Ghavami, Matin, Huot, Mathieu et al. 2024. "Probabilistic Programming with Programmable Variational Inference." Proceedings of the ACM on Programming Languages, 8 (PLDI).
Version
Final published version
Abstract
Compared to the wide array of advanced Monte Carlo methods supported by modern probabilistic programming languages (PPLs), PPL support for variational inference (VI) is underdeveloped: users are typically limited to a small selection of predefined variational objectives and gradient estimators, which are implemented monolithically (and without explicit correctness arguments) in PPL backends. In this paper, we propose a modular approach to supporting VI in PPLs, based on compositional program transformation. First, we present a probabilistic programming language for defining models, variational families, and compositional strategies for propagating gradients. Second, we present a differentiable programming language for defining variational objectives. Models and variational families from the first language are automatically compiled into new differentiable functions that can be called from the second language, for estimating densities and expectations. Finally, we present an automatic differentiation algorithm that differentiates these variational objectives, yielding provably unbiased gradient estimators for use during optimization. We also extend our source language with features not previously supported for VI in PPLs, including approximate marginalization and normalization. This makes it possible to concisely express many models, variational families, objectives, and gradient estimators from the machine learning literature, including importance-weighted autoencoders (IWAE), hierarchical variational inference (HVI), and reweighted wake-sleep (RWS). We implement our approach in an extension to the Gen probabilistic programming system (genjax.vi, implemented in JAX), and evaluate our automation on several deep generative modeling tasks, showing minimal performance overhead vs. hand-coded implementations and performance competitive to well-established open-source PPLs.
MIT Department
Massachusetts Institute of Technology. Department of Electrical Engineering and Computer Science
Massachusetts Institute of Technology. Computer Science and Artificial Intelligence Laboratory
Terms of Use
Creative Commons Attribution
Persistent DSpace Link
DOI of Published Version
https://doi.org/10.1145/3656463
