FlowMessagePassingRules

The Flow node of the ReactiveMP ecosystem: a deterministic, invertible transformation $y = f(x)$ of a multivariate variable, with $f$ a normalising-flow model built from layers. Use it to give a Gaussian latent variable a flexible, non-Gaussian image, or to learn a nonlinear map whose inverse and Jacobian are cheap. The flow models themselves, and their layers, live in this package too: see Flow models.

FlowMessagePassingRules

julia> model = compile(FlowModel(2, (AdditiveCouplingLayer(PlanarFlow(); permute = false),)), [1.0, 2.0, 3.0]);

julia> result = @call_message_update_rule(
           node = Flow, target = :out, algorithm = FlowApproximation(model),
           m = (in = MvNormalMeanCovariance([1.0, 2.0], [1.0 0.0; 0.0 1.0]),),
       );

julia> μ, Σ = mean_cov(getresult(result));

julia> μ ≈ [1.0, 3.0 + tanh(5.0)]   # y₂ = x₂ + x₁ + tanh(2x₁ + 3) at the mean
true

julia> J = [1.0 0.0; 1 + 2 * (1 - tanh(5.0)^2) 1.0];   # the Jacobian at the mean

julia> Σ ≈ J * J'
true

The site has three pages: this one, the node; Flow models, the models and their layers; and Internals, the helpers the rules are built on.

Overview

Flow is a deterministic node: it connects an input in and an output out of the same dimension through a compiled flow model. A message arriving on either side is pushed through the model, forwards through $f$ or backwards through $f^{-1}$. It leaves as a multivariate normal, by linearisation with the model's own Jacobians or by the unscented transform. Both directions are exact for an affine flow and approximate otherwise.

Definition

\[p(y \mid x) = \delta\big(y - f(x)\big), \qquad x, y \in \mathbb{R}^d,\]

where $f = f_L \circ \dots \circ f_1$ is the composition of the model's layers, each invertible with a known Jacobian. The Dirac delta puts all the probability on $y = f(x)$: the output is a function of the input, with no noise.

Interfaces

using FlowMessagePassingRules, MessagePassingRulesBase, MessagePassingRulesApproximations, BayesBase, ExponentialFamily
MessagePassingRulesBase.nodespec(Flow)

The node has two interfaces.

Under Linearization the rules take MvNormalMeanCovariance, MvNormalMeanPrecision and MvNormalWeightedMeanPrecision, and return an MvNormalMeanCovariance for the first and an MvNormalMeanPrecision for the others; under Unscented they take any multivariate normal and return an MvNormalMeanCovariance. No rule reads a marginal.

Algorithm

The node's algorithm, FlowApproximation(model; method = Linearization()), carries the CompiledFlowModel and the method, Linearization() or Unscented(). A model must name it for every Flow node: the node has no default algorithm, so under DefaultAlgorithm no rule exists. The node declares no dependencies, so each rule takes the inputs of the default scheme: the message on the other interface.

Supported rules

MessagePassingRulesBase.rule_coverage(Flow)

The two FlowApproximation columns are its two methods. The first is linearisation, with three rules per edge, one per normal form; the second is the unscented transform. Every rule is a belief propagation message. There is no marginal rule, since the marginal of in is the product of its two messages. There is no average energy, since a deterministic node needs none: its contribution to the Bethe free energy is minus the entropy of q(in), as good as the approximation behind it.

Example

A two-dimensional model of one coupling layer with a planar flow, given its three parameters, and the message towards out from a standard normal centred at [1, 2]:

model = compile(FlowModel(2, (AdditiveCouplingLayer(PlanarFlow(); permute = false),)), [1.0, 2.0, 3.0])

@call_message_update_rule(
    node = Flow, target = :out, algorithm = FlowApproximation(model),
    m = (in = MvNormalMeanCovariance([1.0, 2.0], [1.0 0.0; 0.0 1.0]),),
)

(algo, args)->begin
        (μ_in, Σ_in) = mean_cov(args.m[:in])
        (μ_out, J) = forward_jacobian(getmodel(algo), μ_in)
        MvNormalMeanCovariance(μ_out, J * Σ_in * J')
    end

@call_message_update_rule runs the rule and draws the message it read. Under linearisation the mean of the result is the image of the input's mean, FlowMessagePassingRules.forward of it. The unscented transform gives a mean close to it. getresult extracts both for the comparison:

julia> model = compile(FlowModel(2, (AdditiveCouplingLayer(PlanarFlow(); permute = false),)), [1.0, 2.0, 3.0]);

julia> m_in = MvNormalMeanCovariance([1.0, 2.0], [1.0 0.0; 0.0 1.0]);

julia> linearised = getresult(@call_message_update_rule(node = Flow, target = :out, m = (in = m_in,), algorithm = FlowApproximation(model)));

julia> mean(linearised) ≈ FlowMessagePassingRules.forward(model, [1.0, 2.0])
true

julia> unscented = getresult(@call_message_update_rule(node = Flow, target = :out, m = (in = m_in,), algorithm = FlowApproximation(model; method = Unscented())));

julia> isapprox(mean(unscented), mean(linearised); atol = 1e-2)
true

Limitations

• Multivariate normal messages only, in and out; no rule takes a marginal.
• No average energy, as for every deterministic node, and a free energy only as good as the linearisation or the unscented transform behind the messages.
• The model's parameters are fixed once compiled: the node does not learn them.
• Under Linearization, the rules for a message in precision form evaluate the Jacobian at the message's mean read as a point on the other side of the flow, where the covariance form evaluates it at the mean itself: the two forms agree for a two-dimensional model of one coupling layer, and differ in general for a nonlinear flow.

API

The node and its algorithm, and the algorithm's accessors, called qualified:

Flow

julia> model = compile(FlowModel(2, (AdditiveCouplingLayer(PlanarFlow(); permute = false),)), [1.0, 2.0, 3.0]);

julia> y = [1.0, 3.0 + tanh(5.0)];   # f([1.0, 2.0])

julia> result = @call_message_update_rule(
           node = Flow, target = :in, algorithm = FlowApproximation(model),
           m = (out = MvNormalMeanPrecision(y, [1.0 0.0; 0.0 1.0]),),
       );

julia> mean(getresult(result)) ≈ [1.0, 2.0]
true

FlowApproximation(model::CompiledFlowModel; method = Linearization())
FlowApproximation(model::CompiledFlowModel, method)

julia> model = compile(FlowModel(2, (AdditiveCouplingLayer(PlanarFlow()),)), [1.0, 2.0, 3.0]);

julia> FlowApproximation(model).method isa Linearization
true

julia> FlowApproximation(model; method = Unscented(2)).method isa Unscented
true

FlowMessagePassingRules.getmodel(algo::FlowApproximation) -> CompiledFlowModel

FlowMessagePassingRules.getmethod(algo::FlowApproximation)