Nothing Special   »   [go: up one dir, main page]

skip to main content
10.1145/3519941.3535066acmconferencesArticle/Chapter ViewAbstractPublication PagescpsweekConference Proceedingsconference-collections
research-article

JAX based parallel inference for reactive probabilistic programming

Published: 14 June 2022 Publication History

Abstract

ProbZelus is a synchronous probabilistic language for the design of reactive probabilistic models in interaction with an environment. Reactive inference methods continuously learn distributions over the unobserved parameters of the model from statistical observations. Unfortunately, this inference problem is in general intractable. Monte Carlo inference techniques thus rely on many independent executions to compute accurate approximations. These methods are expensive but can be parallelized.
We propose to use JAX to parallelize ProbZelus reactive inference engine. JAX is a recent library to compile Python code which can then be executed on massively parallel architectures such as GPUs or TPUs.
In this paper, we describe a new reactive inference engine implemented in JAX and the new associated JAX backend for ProbZelus. We show on existing benchmarks that our new parallel implementation outperforms the original sequential implementation for a high number of particles.

References

[1]
Guillaume Baudart, Louis Mandel, Eric Atkinson, Benjamin Sherman, Marc Pouzet, and Michael Carbin. 2020. Programming Reactive Probabilistic Applications. In PROBPROG.
[2]
Guillaume Baudart, Louis Mandel, Eric Atkinson, Benjamin Sherman, Marc Pouzet, and Michael Carbin. 2020. Reactive Probabilistic Programming. In PLDI. https://doi.org/10.1145/3385412.3386009
[3]
Guillaume Baudart, Louis Mandel, and Reyyan Tekin. 2022. Reproduction package for article JAX Based Parallel Inference for Reactive Probabilistic Programming. https://doi.org/10.1145/3462319
[4]
Albert Benveniste, Timothy Bourke, Benoît Caillaud, Bruno Pagano, and Marc Pouzet. 2014. A type-based analysis of causality loops in hybrid systems modelers. In HSCC. https://doi.org/10.1145/2562059.2562125
[5]
Albert Benveniste, Paul Caspi, Stephen A. Edwards, Nicolas Halbwachs, Paul Le Guernic, and Robert de Simone. 2003. The synchronous languages 12 years later. Proc. IEEE, 91, 1 (2003), 64–83. https://doi.org/10.1109/JPROC.2002.805826
[6]
Dariusz Biernacki, Jean-Louis Colaço, Grégoire Hamon, and Marc Pouzet. 2008. Clock-directed modular code generation for synchronous data-flow languages. In LCTES. https://doi.org/10.1145/1375657.1375674
[7]
Eli Bingham, Jonathan P. Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul A. Szerlip, Paul Horsfall, and Noah D. Goodman. 2019. Pyro: Deep Universal Probabilistic Programming. Journal of Machine Learning Research, 20 (2019), 28:1–28:6.
[8]
Timothy Bourke, Lélio Brun, and Marc Pouzet. 2020. Mechanized semantics and verified compilation for a dataflow synchronous language with reset. In POPL. https://doi.org/10.1145/3371112
[9]
Timothy Bourke and Marc Pouzet. 2013. Zelus, a Hybrid Synchronous Language. https://zelus.di.ens.fr
[10]
Timothy Bourke and Marc Pouzet. 2013. Zélus: a synchronous language with ODEs. In HSCC. https://doi.org/10.1145/2461328.2461348
[11]
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. 2018. JAX: composable transformations of Python+NumPy programs. https://github.com/google/jax
[12]
Paul Caspi and Marc Pouzet. 1998. A Co-iterative Characterization of Synchronous Stream Functions. In CMCS. https://doi.org/10.1016/S1571-0661(04)00050-7
[13]
Albert Cohen, Léonard Gérard, and Marc Pouzet. 2012. Programming parallelism with futures in Lustre. In EMSOFT. https://doi.org/10.1145/2380356.2380394
[14]
Jean-Louis Colaço, Bruno Pagano, and Marc Pouzet. 2017. SCADE 6: A formal language for embedded critical software development. In TASE. https://doi.org/10.1109/TASE.2017.8285623
[15]
Marco F. Cusumano-Towner, Feras A. Saad, Alexander K. Lew, and Vikash K. Mansinghka. 2019. Gen: a general-purpose probabilistic programming system with programmable inference. In PLDI. https://doi.org/10.1145/3314221.3314642
[16]
Pierre Del Moral, Arnaud Doucet, and Ajay Jasra. 2006. Sequential Monte Carlo samplers. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 68, 3 (2006), 411–436.
[17]
Alain Girault. 2005. A Survey of Automatic Distribution Method for Synchronous Programs. In SLAP.
[18]
Noah D. Goodman and Andreas Stuhlmüller. 2014. The Design and Implementation of Probabilistic Programming Languages. http://dippl.org
[19]
Léonard Gérard, Adrien Guatto, Cédric Pasteur, and Marc Pouzet. 2012. A modular memory optimization for synchronous data-flow languages: application to arrays in a lustre compiler. In LCTES. https://doi.org/10.1145/2345141.2248426
[20]
Nicolas Halbwachs, Paul Caspi, Pascal Raymond, and Daniel Pilaud. 1991. The Synchronous Dataflow Programming Language Lustre. Proc. IEEE, 79, 9 (1991), September, 1305–1320.
[21]
Thomas P. Minka. 2001. Expectation Propagation for approximate Bayesian inference. In UAI.
[22]
Lawrence M. Murray and Thomas B. Schön. 2018. Automated learning with a probabilistic programming language: Birch. Annual Reviews in Control, 46 (2018), 29–43.
[23]
Claire Pagetti, Julien Forget, Frédéric Boniol, Mikel Cordovilla, and David Lesens. 2011. Multi-task Implementation of Multi-periodic Synchronous Programs. Discrete Event Dynamic Systems, 21, 3 (2011), 307–338. https://doi.org/10.1007/s10626-011-0107-x
[24]
Du Phan, Neeraj Pradhan, and Martin Jankowiak. 2019. Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro. arXiv:1912.11554, https://doi.org/10.48550/arXiv.1912.11554
[25]
David Tolpin, Jan-Willem van de Meent, Hongseok Yang, and Frank D. Wood. 2016. Design and Implementation of Probabilistic Programming Language Anglican. In IFL. https://doi.org/10.1145/3064899.3064910
[26]
Dustin Tran, Matthew D. Hoffman, Rif A. Saurous, Eugene Brevdo, Kevin Murphy, and David M. Blei. 2017. Deep Probabilistic Programming. In ICLR.

Cited By

View all
  • (2023)ViX: Analysis-driven Compiler for Efficient Low-Precision Variational Inference2023 Design, Automation & Test in Europe Conference & Exhibition (DATE)10.23919/DATE56975.2023.10137324(1-6)Online publication date: Apr-2023

Recommendations

Comments

Please enable JavaScript to view thecomments powered by Disqus.

Information & Contributors

Information

Published In

cover image ACM Conferences
LCTES 2022: Proceedings of the 23rd ACM SIGPLAN/SIGBED International Conference on Languages, Compilers, and Tools for Embedded Systems
June 2022
161 pages
ISBN:9781450392662
DOI:10.1145/3519941
Publication rights licensed to ACM. ACM acknowledges that this contribution was authored or co-authored by an employee, contractor or affiliate of a national government. As such, the Government retains a nonexclusive, royalty-free right to publish or reproduce this article, or to allow others to do so, for Government purposes only.

Sponsors

Publisher

Association for Computing Machinery

New York, NY, United States

Publication History

Published: 14 June 2022

Permissions

Request permissions for this article.

Check for updates

Badges

Author Tags

  1. Compilation
  2. Parallel Computing
  3. Probabilistic Programming
  4. Reactive Programming
  5. Streaming inference

Qualifiers

  • Research-article

Conference

LCTES '22

Acceptance Rates

Overall Acceptance Rate 116 of 438 submissions, 26%

Contributors

Other Metrics

Bibliometrics & Citations

Bibliometrics

Article Metrics

  • Downloads (Last 12 months)26
  • Downloads (Last 6 weeks)0
Reflects downloads up to 12 Nov 2024

Other Metrics

Citations

Cited By

View all
  • (2023)ViX: Analysis-driven Compiler for Efficient Low-Precision Variational Inference2023 Design, Automation & Test in Europe Conference & Exhibition (DATE)10.23919/DATE56975.2023.10137324(1-6)Online publication date: Apr-2023

View Options

Get Access

Login options

View options

PDF

View or Download as a PDF file.

PDF

eReader

View online with eReader.

eReader

Media

Figures

Other

Tables

Share

Share

Share this Publication link

Share on social media