Message passing

Message passing is how ReactiveMP.jl performs inference on a factor graph. The engine never forms the full joint distribution. Instead, factor nodes and variables exchange small local summaries with their neighbours, the messages. The posterior beliefs, the marginals, are formed from these messages.

Belief propagation

Belief propagation, the sum-product algorithm, computes exact marginals on tree-shaped graphs. A message from a factor node f towards a variable x summarises what f knows about x from the rest of the graph:

\[\mu_{f \to x}(x) = \int f(x, y, z) \; \mu_{y \to f}(y) \; \mu_{z \to f}(z) \; \mathrm{d}y \; \mathrm{d}z\]

The message from x back towards f is the product of the messages that arrive at x from every other factor. The marginal q(x) is the product of all the messages that arrive at x, normalised.

On a graph with cycles, the engine iterates the same procedure. This is loopy belief propagation, which typically converges to a good approximation.

A belief propagation message

Variational message passing

Variational message passing performs approximate inference. It minimises the Bethe free energy, a variational objective, under a factorisation of the posterior into clusters. ReactiveMP.jl implements this general form for three reasons:

1. It includes exact belief propagation, as the case with no factorisation constraints.
2. It handles non-conjugate and complex models, with mean-field or structured factorisations.
3. It has a local, message-level form that fits the reactive computation model.

Under the mean-field factorisation, q(x, y, z) = q(x) q(y) q(z), the message from f towards x is

\[\mu_{f \to x}(x) = \exp \int q(y) \, q(z) \log f(x, y, z) \; \mathrm{d}y \; \mathrm{d}z\]

It uses the marginals q(y) and q(z), not the messages μ(y) and μ(z).

A variational message under the structured factorisation q(x, y)q(z)

Which rule computes a message

Each message is computed by an update rule, an ordinary Julia function that a rule package defines for a node, a target interface and an algorithm. When you activate a node, the engine decides what each rule reads from the node's algorithm and factorisation. By default, a rule reads the messages on the other interfaces of its target's cluster, and the marginals of the other clusters. So one cluster over every interface gives belief propagation, and one interface per cluster gives mean-field variational message passing. A deterministic node's rules read the messages on every other interface.

At each update, the engine finds the rule with find_message_rule. It looks at the node type, the target, the node's algorithm and the types of the inputs. The algorithm is DefaultAlgorithm for most nodes. Messages and marginals of different types select different rules, and where no rule matches, the error lists the near misses. An algorithm may also declare its own dependencies with @define_dependencies, and you choose a node's algorithm at activation (see Activation options).

The same graph, both ways

The graph of Getting started has a latent x with the prior $\mathcal{N}(x \mid 0, 10)$ and an observation y with the likelihood $\mathcal{N}(y \mid x, 1)$. Both factors are the Gaussian node of the example node. The function below builds the graph with a factorisation of your choice, runs it on the observation y = 2, and prints every rule the engine runs through a callback:

function run_graph(factorisation)
    x, y = randomvar(label = :x), datavar(label = :y)
    constants = (constvar(0.0), constvar(10.0), constvar(1.0))
    prior = factornode(Gaussian, [(:out, x), (:μ, constants[1]), (:v, constants[2])], factorisation)
    likelihood = factornode(Gaussian, [(:out, y), (:μ, x), (:v, constants[3])], factorisation)

    io = IOContext(stdout, :compact => true, :module => @__MODULE__)
    callbacks = (after_message_rule_call = (event) -> println(io, event.mapping, " → ", event.result),)
    activate!(x, RandomVariableActivationOptions())
    activate!(y, DataVariableActivationOptions())
    for node in (prior, likelihood)
        activate!(node, FactorNodeActivationOptions(; callbacks))
    end

    posterior, energy = Ref{Any}(), Ref{Float64}()
    subscriptions = [
        subscribe!(get_stream_of_marginals(x), (q) -> posterior[] = q),
        subscribe!(bethe_free_energy(Float64, (prior, likelihood), (x, y, constants...)), (f) -> energy[] = f),
    ]
    new_observation!(y, 2.0)
    foreach(unsubscribe!, subscriptions)
    return getdata(posterior[]), energy[]
end

Under belief propagation, each node is one cluster, (out, μ, v):

run_graph(((:out, :μ, :v),))

(ExponentialFamily.NormalWeightedMeanPrecision{Float64}(xi=2.0, w=1.1), 2.2997043514220397)

Each printed line is a ReactiveMP.MessageMapping: the node, the target, and the inputs the rule reads. msgs lists messages. The prior's rule towards out reads the messages on μ and v, and the likelihood's rule towards μ reads the messages on out and v.

Under mean field, each interface is a cluster of its own:

run_graph(((:out,), (:μ,), (:v,)))

(ExponentialFamily.NormalWeightedMeanPrecision{Float64}(xi=2.0, w=1.1), 2.2997043514220397)

The same targets read marginals here, so the engine runs the node's variational rules. The two rules towards μ differ in their inputs, which which_message_update_rule shows without running them. With messages, it finds the belief propagation rule:

which_message_update_rule(Gaussian, :μ; m = (out = PointMass(2.0), v = PointMass(1.0)))

args->NormalMeanVariance(mean(args.m[:out]), var(args.m[:out]) + mean(args.m[:v]))

With marginals, it finds the variational rule:

which_message_update_rule(Gaussian, :μ; q = (out = PointMass(2.0), v = PointMass(1.0)))

args->NormalMeanVariance(mean(args.q[:out]), mean(args.q[:v]))

Both runs give the same posterior, $\mathcal{N}(x \mid 2/1.1, 1/1.1)$ in its natural form, and the same free energy, $-\log \mathcal{N}(2 \mid 0, 11)$. Mean field is exact here because every other edge of each node is a constant or an observation: a point mass, under which the expectation of the log-density and the integral agree. Where two latent variables share a node, the two factorisations give different messages. Mean field then needs initial marginals (ReactiveMP.set_initial_marginal!) and several iterations.

For the theory in depth, see the PhD dissertation that ReactiveMP.jl is based on.

Messages as streams

The word reactive refers to how messages are scheduled. Many message passing libraries build an explicit schedule, such as forward and backward passes, before inference starts. ReactiveMP.jl builds none:

• Every connection between a variable and a node carries a stream of messages, a ReactiveMP.MessageObservable. Every variable carries a stream of marginals, a ReactiveMP.MarginalObservable. Each emits a new value whenever its inputs change.
• When data arrives, with new_observation!, the change propagates through the graph, and only the rules that depend on it run.
• The order of the updates follows from the graph and the data, not from a plan. In variational message passing, the order in which a rule's inputs are declared is the update schedule.

A message is computed lazily. A node emits a DeferredMessage, which is computed when a variable first reads it, so a message nobody needs is never computed.

The streams are Rocket.jl observables. Nothing runs until something subscribes: a subscription to a marginal, or to the free energy, pulls the computation through the graph. This is why you first build a graph and then activate it, and why the same graph serves streaming data, each observation propagating as it arrives (see Inference lifecycle).

Messages and marginals

The engine wraps the two in their own types:

• A Message travels along one edge, from a node towards a variable or back.
• A Marginal is a belief about a variable, formed from the product of the messages that arrive at it. It can also be a belief about a cluster of a node's variables, computed by the node's marginal rule.

Both hold a distribution and forward its statistics, such as mean and var. Both record whether the value is clamped, computed from constants and observations only, and whether it is initial, set before inference. Both carry their log scale when log scales are tracked (see Log scales), and optional annotations.

Next steps

• Messages and Marginals: the types and their streams.
• Inference lifecycle: construction, activation and observation.
• MessagePassingRulesBase: how nodes, rules and algorithms are defined.