ReactiveMP.jl: Reactive Message Passing-based Bayesian Inference
2021-07-28, 13:30–13:40 (UTC), Green

ReactiveMP.jl is a native Julia implementation of reactive message passing-based Bayesian inference in probabilistic graphical models. The package supports a large range of standard probabilistic models and can be extended to custom novel nodes and message update rules. In contrast to non-reactive (imperatively coded) Bayesian inference packages, ReactiveMP.jl scales easily to support inference on a standard laptop for large models with tens of thousands of variables and millions of nodes.


GitHub: https://github.com/biaslab/ReactiveMP.jl

Demos: https://github.com/biaslab/ReactiveMP.jl/tree/master/demo

YouTube: https://www.youtube.com/watch?v=twhTsKsXa_8

Experiments from the talk: https://github.com/biaslab/ReactiveMP_JuliaCon2021

Bayesian inference is one of the key computational mechanisms that underlies probabilistic model-based machine learning applications such as time series prediction, image and speech recognition, and robotics. Unfortunately, for many models of practical interest, Bayesian inference requires evaluating high-dimensional integrals that have no analytical solution. As a result, Probabilistic Programming (PP) tools for Automated Approximate Bayesian Inference (AABI) have become popular, e.g., Turing.jl, Soss.jl, ForneyLab.jl, Pyro, and others. These tools help researchers to define custom probabilistic models in a high-level domain-specific language and run AABI algorithms with minimal additional overhead.

An important issue in the development of PP frameworks is scalability of AABI algorithms for large models and large data sets. One solution approach concerns message passing-based inference in factor graphs. In this framework, relationships between model variables are represented by a graph of sparsely connected nodes, and inference proceeds efficiently by a sequence of nodes sending probabilistic messages to neighboring nodes. While the optimal message passing schedule is data-dependent, all existing factor graph frameworks (e.g., Infer.Net, ForneyLab.jl) use preset message sequence schedules. The potential benefits of massively parallel and asynchronous reactive message passing in a factor graph include scaling to large inference tasks, much smaller processing latency and processing of data samples that arrive at irregular time intervals.

We have developed ReactiveMP.jl, which is a native Julia package for automated reactive message passing-based (both exact and approximate) Bayesian inference. ReactiveMP.jl is based on Rocket.jl, which is a native Julia package for a reactive programming. In ReactiveMP.jl, there are no pre-scheduled messages. Instead, nodes subscribe to messages from connected nodes and react autonomously and asynchronously whenever a new message has been received. As a result, ReactiveMP.jl scales comfortably to inference tasks on factor graphs with tens of thousands of variables and millions of nodes.

The ReactiveMP.jl package comes with a collection of standard probabilistic models, including linear Gaussian state-space models, hidden Markov models, auto-regressive models and mixture models. Moreover, ReactiveMP.jl API supports various processing modes such as offline learning, online filtering of infinite data streams and protocols for handling missing data.

ReactiveMP.jl is customizable and provides an easy way to add new models, node functions and analytical message update rules to the existing platform. As a result, a user can extend built-in functionality with custom nodes to run automated inference in novel probabilistic models. The resulting inference procedures are differentiable with the ForwardDiff.jl or ReverseDiff.jl packages. In addition, the inference engine supports different types of floating point numbers, e.g., the built-in BigFloat Julia type.

We achieved excellent performance by relying on Julia's great multiple dispatch capabilities and advanced compile-time optimization techniques. Message passing-based inference requires computation of many messages by node-specific update rules. Some of these updates can be evaluated and in-lined at compile time, which results in a fast and accurate automated Bayesian inference realization with almost zero overhead when compared to manually hard-coded inference procedures.

We compared ReactiveMP.jl to other message passing and sampling-based inference packages. In terms of computation time and memory usage, specifically for conjugate models, the ReactiveMP.jl engine outperforms Turing.jl, ForneyLab.jl and Infer.Net significantly by orders of magnitude. Comparative performance benchmarks are available at the GitHub repository: https://github.com/biaslab/ReactiveMP.jl.

Automating scalable Bayesian inference is a key factor in the quest to apply Bayesian machine learning to useful applications. We developed ReactiveMP.jl as a package that enables developers to build novel probabilistic models and automate scalable inference in those models by asynchronous, reactive message passing in a factor graph. We are looking forward to presenting the ReactiveMP.jl package and discuss the advantages and drawbacks of the reactive message passing approach.

I am a PhD candidate in the SPS group of Electrical Engineering department in Eindhoven University of Technology. My research interests lie in the fields of computer science, software developing, numerical modeling, machine learning, computational optimization and high-performant parallelized applications.