PGMax: Factor Graphs for Discrete Probabilistic Graphical Models and Loopy Belief Propagation in JAX
Abstract
PGMax is an open-source Python/ JAX package for (a) easily specifying discrete Probabilistic Graphical Models (PGMs) as factor graphs; and (b) automatically running efficient and scalable differentiable Loopy Belief Propagation (LBP). PGMax supports general factor graphs with tractable factors, and leverages modern accelerators like GPUs for inference. Compared with alternative libraries, PGMax obtains higher-quality inference results with up to three orders-of-magnitude inference time speedups. PGMax interacts seamlessly with the growing JAX ecosystem, opening up new research possibilities. Our source code, examples and documentation are available at https://github.com/google-deepmind/PGMax
Cite
Text
Zhou et al. "PGMax: Factor Graphs for Discrete Probabilistic Graphical Models and Loopy Belief Propagation in JAX." Machine Learning Open Source Software, 2024.Markdown
[Zhou et al. "PGMax: Factor Graphs for Discrete Probabilistic Graphical Models and Loopy Belief Propagation in JAX." Machine Learning Open Source Software, 2024.](https://mlanthology.org/mloss/2024/zhou2024jmlr-pgmax/)BibTeX
@article{zhou2024jmlr-pgmax,
title = {{PGMax: Factor Graphs for Discrete Probabilistic Graphical Models and Loopy Belief Propagation in JAX}},
author = {Zhou, Guangyao and Dedieu, Antoine and Kumar, Nishanth and Lehrach, Wolfgang and Kushagra, Shrinu and George, Dileep and Lázaro-Gredilla, Miguel},
journal = {Machine Learning Open Source Software},
year = {2024},
pages = {1-25},
volume = {25},
url = {https://mlanthology.org/mloss/2024/zhou2024jmlr-pgmax/}
}