mlsanityData to 14 Oct 2026

PGMax: Factor Graphs for Discrete Probabilistic Graphical Models and Loopy Belief Propagation in JAX

arXiv 2202.04110 ↗First posted 8 Feb 2022cs.LG cs.AI stat.MLcited by 6

Not on the map. Why?

Abstract (arXiv)

PGMax is an open-source Python package for (a) easily specifying discrete Probabilistic Graphical Models (PGMs) as factor graphs; and (b) automatically running efficient and scalable loopy belief propagation (LBP) in JAX. PGMax supports general factor graphs with tractable factors, and leverages modern accelerators like GPUs for inference. Compared with existing alternatives, PGMax obtains higher-quality inference results with up to three orders-of-magnitude inference time speedups. PGMax additionally interacts seamlessly with the rapidly growing JAX ecosystem, opening up new research possibilities. Our source code, examples and documentation are available at https://github.com/deepmind/PGMax.

Key facts. Cited by 6 papers in the map (data through 14 Oct 2026). Builds on 12 of its 38 references that are in the map. First posted on arXiv 8 Feb 2022.

Builds on

12 of 38 references.

References · 12

  1. 2018Pyro: Deep Universal Probabilistic Programmingcited by 380
  2. 2019Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyrocited by 191
  3. 2019BERT has a Mouth, and It Must Speak: BERT as a Markov Random Field Language Modelcited by 185
  4. 2013Loopy Belief Propagation for Approximate Inference: An Empirical Studycited by 76
  5. 2018An Introduction to Probabilistic Programmingcited by 63
  6. 2017Pomegranate: fast and flexible probabilistic modeling in pythoncited by 17
  7. 2014A Comparative Study of Modern Inference Techniques for Structured Discrete Energy Minimization Problemscited by 16
  8. 2015Boolean Matrix Factorization and Noisy Completion via Message Passingcited by 12
  9. 2018A Factor Graph Approach to Automated Design of Bayesian Signal Processing Algorithmscited by 11
  10. 2021Perturb-and-max-product: Sampling and learning in discrete energy-based modelscited by 9
  11. 2016Hierarchical compositional feature learningcited by 4
  12. 2021Graphical Models with Attention for Context-Specific Independence and an Application to Perceptual Groupingcited by 3

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2022: 1; 2023: 3; 2024: 1; 2025: 0; 2026: 1.

6 papers.

  1. 2023Learning from Label Proportions: Bootstrapping Supervised Learners via Belief Propagationcited by 7
  2. 2023SynJax: Structured Probability Distributions for JAXcited by 2
  3. 2022Approximate Inference for Stochastic Planning in Factored Spacescited by 1
  4. 2023Learning noisy-OR Bayesian Networks with Max-Product Belief Propagationcited by 0
  5. 2024Scalable Data Assimilation with Message Passingcited by 0
  6. 2026A Hardware-oriented Approach for Efficient Bayesian Inference Computation and Deploymentcited by 0