Reactive Message Passing for Scalable Bayesian Inference

This section briefly describes the essential concepts of the Reactive Programming (RP) paradigm such as observables, subscriptions, actors, subjects, and operators. Generally speaking, RP provides a set of guidelines and ideas to simplify working with reactive systems and asynchronous streams of data and/or events. Many modern programming languages have some support for RP through extra libraries and/or packages¹¹1ReactiveX programming languages support http://reactivex.io/languages.html. For example, the JavaScript community features a popular RxJS²²2JavaScript reactive extensions library https://github.com/ReactiveX/rxjs library (more than 26 thousands stars on GitHub) and the Python language community has its own RxPY³³3Python reactive extensions library https://github.com/ReactiveX/RxPY package.

In the next section we show how to formulate message passing-based inference in the reactive programming paradigm. First, we cast all variables of interest to observables. These include messages $\mu_{f_{j}\rightarrow x_{i}}(x_{i})$ and $\mu_{x_{i}\rightarrow f_{j}}(x_{i})$, marginals $q(s_{i})$ over all model variables, and local variational posteriors $q_{a}(s_{a})$ for all factors in the corresponding TFFG of a probabilistic generative model. Next, we show that nodes and edges are, in fact, just special types of subjects that multicast messages with the combineLatest operator and compute corresponding integrals with the map operator. We use the subscription mechanism to start listening for new posterior updates in a model and perform actions based on new observations. The resulting reactive system resolves message (passing) updates locally, automatically reacts on flowing messages and updates itself accordingly when new data arrives.

ReactiveMP.jl uses a @model macro for joint model and variational family constraints specification. The @model macro returns a function that generates the model with specified constraints. It has been designed to resemble existing probabilistic model specification libraries like Turing.jl closely, but also to be flexible enough to support reactive message passing-based inference as well as factorisation and form constraints for variational families $\mathcal{Q}_{B}$ from (7). As an example of the model specification syntax, consider the simple example in Listing 5.1:

⬇# # In this example, the model accepts an additional argument `n` that denotes # the number of observations in the dataset @model function coin_model(n) # Reactive data inputs for Beta prior over $\theta$ random variable # so we don't need actual data at model creation time a = datavar(Float64) b = datavar(Float64) # We use the tilde operator to define a probabilistic relationship between # random variables and data inputs. It automatically creates a new random # variable in the current model and the corresponding factor nodes $\theta$ ~ Beta(a, b) # creates both `$\theta$` random variable and Beta factor node # A sequence of observations with length `n` y = datavar(Float64, n) # Each observation is modelled by a `Bernoulli` distribution # that is governed by a `$\theta$` parameter for i in 1:n y[i] ~ Bernoulli($\theta$) # Reuses `y[i]` and `$\theta$` and creates `Bernoulli` node end # `@model` function must have a `return` statement. # Later on, the returned references might be useful to obtain the marginal # updates over variables in the model and to pass new observations to data # inputs so the model can react to changes in data return y, a, b, $\theta$ end # The `@model` macro modifies the original function to return a reference # to the actual graphical probabilistic model as a factor graph # and the exactly same output as in the original `return` statement model, (y, a, b, $\theta$) = coin_model(100) \lst@TestEOLChar

The tilde operator ($\sim$) creates a new random variable with a given functional relationship with other variables in a model and can be read as "is sampled from ..." or "is modelled by ...". Sometimes it is useful to create a random variable in advance or to create a sequence of random variables in a probabilistic model for improved efficiency. The ReactiveMP.jl package exports a randomvar() function to manually create random variables in the model.

⬇x x = randomvar(n) # A collection of n random variables \lst@TestEOLChar

To create observations in a model, the ReactiveMP.jl package exports the datavar(T) function, where T refers to a type of data such as Float64 or Vector{Float64}. Later on it is important to continuously update the model with new observations to perform reactive message passing.

⬇y y = datavar(Vector{Float64}) # A single data input of type Vector{Float64} y = datavar(Float64, n) # A sequence of n data inputs of type Float64 \lst@TestEOLChar

The tilde operator supports an optional where statement to define local node- or variable-context objects. As discussed in Section 4.2, local context objects specify local factorisation constraints for the variational family of distributions $\mathcal{Q}_{B}$, or may include extra pipeline stages to the default computational pipeline of messages and marginals updates, see Listing 5.1.

⬇# # constraints for the local variational distribution $\theta$ ~ Beta(a, b) where { q = q($\theta$)q(a)q(b) } # Structured factorisation assumption x_next ~ NormalMeanVariance(x_previous, tau) where { q = q(x_next, x_previous)q(tau) } # We can also use `where` clause to modify the default # computational pipeline for all outbound messages from a node # In this example `LoggerPipelineStage` modifies pipeline such # that it starts to print all outbound messages into standard output y ~ Gamma(a, b) where { pipeline = LoggerPipelineStage() } \lst@TestEOLChar

The ReactiveMP.jl framework API exports the getmarginal() function, which returns a reference for an observable of marginal updates. We use the subscribe!() function to subscribe to future marginal updates and perform some actions.

⬇m # `subscribe!` accepts an observable as its first argument and # a callback function as its second argument $\theta$_subscription = subscribe!(getmarginal($\theta$), (update) -> println(update)) \lst@TestEOLChar

For variational Bayesian inference and loopy belief propagation algorithms, it might be important to initialise marginal distributions over the model’s variables and to initialise messages on edges. The framework exports the setmarginal!() and setmessage!() functions to set initial values to the marginal update observables and the message update observables, respectively.

⬇s setmessage!($\theta$, Bernoulli(0.5)) \lst@TestEOLChar

The ReactiveMP.jl framework API exports the score() function to get a stream of Bethe Free Energy values. By default, the Bethe Free Energy observable emits a new update only after all variational posterior distributions $q_{a}(s_{a})$ and $q_{i}(s_{i})$ have been updated.

⬇b bfe_subscription = subscribe!(bfe_updates, (update) -> println(update)) \lst@TestEOLChar

An important problem is how to select the most efficient update rule for an outbound message, given the types of inbound messages and/or marginals. In general, we want to use known analytical solutions if they are available and otherwise select an appropriate approximation method. In the context of software development, the choice of which method to execute when a function is applied is called dispatch.

Locality is a central design choice of the reactive message passing approach, but it makes it impossible to infer the correct inbound message types locally before the data have been seen. Fortunately, the Julia language supports dynamic multiple dispatch, which enables an elegant solution for this problem. Julia allows the dispatch process to choose which of a function’s methods to call based on the number of arguments given, and on the types of all of the function’s arguments at run-time. This feature enables automatic dispatch of the most suitable outbound message computation rule, given the functional form of a factor node and the functional forms of inbound messages and/or posterior marginals.

Julia’s built-in features also support the ReactiveMP.jl implementation to dynamically dispatch to the most efficient message update rule for both exact or approximate variational algorithms. If no closed-form analytical message update rule exists, Julia’s multiple dispatch facility provides several options to select an alternative (more computationally demanding) update as discussed in the Section 4.2.1.

The ReactiveMP.jl uses the following arguments to dispatch to the most efficient message update rule (see Listing. 7):

• •

  Functional form of a factor node: e.g. Gaussian or Gamma;

• •

  Name of the connected edge: e.g. :out;

• •

  Local variable constraints: e.g. Marginalisation or MomentMatching;

• •

  Input messages names and types with m_ prefix: e.g. m_mean::Gaussian;

• •

  Input marginals names and types with q_ prefix: e.g. q_precision::Any;

• •

  Optional context object, e.g. a strategy to correct nonpositive definite matrices, an order of autoregressive model or an optional approximation methods used to compute messages.

All of this together allows the ReactiveMP.jl framework to automatically select the most efficient message update rule. It should be possible to emulate this behavior in other programming languages, but Julia’s core support for efficient dynamic multiple dispatch is an important reason why we selected it as the implementation language. It is worth noting that the Julia programming language generates efficient and scalable code with similar run-time performance as for C and C++.

As our first example, we consider the linear Gaussian state space model (LG-SSM) that is widely used in the signal processing and control communities, see (Sarkka, 2013). The simple LG-SMM is specified by

where $y_{t}$ is the observation at time step $t$, $x_{t}$ represents the latent state at time step $t$, $A$ and $B$ are state transition and observation model matrices, respectively, and $P$ and $Q$ are Gaussian noise covariance matrices.

In this example, the goal is to perform Bayesian inference in a Hidden Markov model (HMM). An HMM can be viewed as a specific instance of a state space model in which the latent variables are discretely valued. HMMs are widely used in speech recognition (Jelinek, 1998; Rabiner and Juang, 1993), natural language modelling (Manning and Schütze, 1999) and in many other related fields.

We consider an HMM specified by

where $A$ and $B$ are state transition and observation model matrices respectively, $z_{t}$ is a discrete $M$-dimensional one-hot coded latent state, $y_{t}$ is the observation at time step $t$, $\text{MatrixDirichlet}(\;\cdot\;|\;P)$ denotes a matrix variate generalisation of Dirichlet distribution with concentration parameters matrix $P$ (Zhang et al., 2014) and $\text{Cat}(\;\cdot\;|\;p)$ denotes a Categorical distribution with concentration parameters vector $p$. One-hot coding of $z_{t}$ implies that $z_{t,i}\in\{0,1\}$ and $\sum_{i=1}^{M}z_{t,i}=1$. With this encoding scheme, (14c) is short-hand for

For our last example, we consider Bayesian inference in the Hierarchical Gaussian Filter (HGF) model. The HGF is popular in the computational neuroscience literature and is often used for the Bayesian modelling of volatile environments, such as uncertainty in perception or financial data (Mathys et al., 2014). The HGF is essentially a Gaussian random walk, where the time-varying variance of the random walk is determined by the state of a higher level process, e.g., a Gaussian random walk with fixed volatily or another more complex signal. Specifically, a simple HGF is defined as

where $y_{t}$ is the observation at time step $t$, $s_{t}^{(j)}$ is the latent state for the $j$-th layer at time step $t$, and $f$ is a link function, which maps states from the (next higher) $(j+1)-th$ layer to the non-negative variance of the random walk in the $j$-th layer. The HGF model typically defines the link function as $f(s^{(j+1)}_{t})=\exp(\kappa s^{(j+1)}_{t}+\omega)$, where $\kappa$ and $\omega$ are either hyper-parameters or random variables included in the model.

As in the previous example, the exact Bayesian inference for this type of model is not tractable. Moreover, the variational message update rules (9) are not tractable either since the HGF model contains non-conjugate relationships between variables in the form of the link function $f$. However, inference in this model is still possible with custom message update rules approximation (Senoz and De Vries, 2018). The ReactiveMP.jl package supports a straightforward API to add custom nodes with custom message passing update rules. In fact, a big part of the ReactiveMP.jl library is a collection of well-known factor nodes and message update rules implemented using the same API as for custom novel nodes and message update equations. The API also supports message update rules that involve extra nontrivial approximation steps. In this example, we added a custom Gaussian controlled variance (GCV) node to model the nonlinear time-varying variance mapping from different hierarchy levels (17a), with a set of approximate update rules based on the Gauss-Hermite cubature rule from (Kokkala et al., 2015). The main purpose of this example is to show that ReactiveMP.jl is capable of running inference in complex non-conjugate models, but requires creating custom factor nodes and choosing appropriate integral calculation approximation methods.