Inference execution

The RxInfer inference API supports different types of message-passing algorithms (including hybrid algorithms combining several different types). While RxInfer implements several algorithms to cater to different computational needs and scenarios, the core message-passing algorithms that form the foundation of our inference capabilities are:

• Belief Propagation
• Variational Message Passing

Whereas belief propagation computes exact inference for the random variables of interest, the variational message passing (VMP) is an approximation method that can be applied to a larger range of models.

The inference engine itself isn't aware of different algorithm types and simply does message passing between nodes. However, during the model specification stage user may specify different factorisation constraints around factor nodes with the help of the @constraints macro. Different factorisation constraints lead to different message passing update rules. See more documentation about constraints specification in the corresponding section.

Automatic inference specification

RxInfer exports the infer function to quickly run and test your model with both static and asynchronous (real-time) datasets. See more information about the infer function on the separate documentation section:

• Static Inference
• Streamlined Inference

infer(
    model; 
    data = nothing,
    datastream = nothing,
    autoupdates = nothing,
    initialization = nothing,
    constraints = nothing,
    meta = nothing,
    options = nothing,
    returnvars = nothing, 
    predictvars = nothing, 
    historyvars = nothing,
    keephistory = nothing,
    iterations = nothing,
    free_energy = false,
    free_energy_diagnostics = DefaultObjectiveDiagnosticChecks,
    showprogress = false,
    callbacks = nothing,
    annotations = nothing,
    logscales = nothing,
    postprocess = nothing,
    warn = true,
    events = nothing,
    uselock = false,
    autostart = true,
    catch_exception = false,
    benchmark = false,
    trace = false,
    session = RxInfer.default_session()
)

Note on NamedTuples

When passing NamedTuple as a value for some argument, make sure you use a trailing comma for NamedTuples with a single entry. The reason is that Julia treats returnvars = (x = KeepLast()) and returnvars = (x = KeepLast(), ) expressions differently. This first expression creates (or overwrites!) new local/global variable named x with contents KeepLast(). The second expression (note trailing comma) creates NamedTuple with x as a key and KeepLast() as a value assigned for this key.

(x = KeepLast()) # defines a variable `x` with the value `KeepLast()`

KeepLast()

(x = KeepLast(), ) # defines a NamedTuple with `x` as one of the keys and value `KeepLast()`

(x = KeepLast(),)

• model

Also read the Model Specification section.

The model argument accepts a model specification as its input. The easiest way to create the model is to use the @model macro. For example:

@model function beta_bernoulli(y, a, b)
    x  ~ Beta(a, b)
    y .~ Bernoulli(x)
end

result = infer(
    model = beta_bernoulli(a = 1, b = 1),
    data  = (y = [ true, false, false ], )
)

result.posteriors[:x]

Beta{Float64}(α=2.0, β=3.0)

• data

Either data or datastream keyword argument are required. Specifying both data and datastream is not supported and will result in an error.

The data keyword argument must be a NamedTuple (or Dict) where keys (of Symbol type) correspond to some arguments defined in the model specification. For example, if a model defines y in its argument list

@model function beta_bernoulli(y, a, b)
    x  ~ Beta(a, b)
    y .~ Bernoulli(x)
end

and you want to condition on this argument, then the data field must have an :y key (of Symbol type) which holds the data. The values in the data must have the exact same shape as its corresponding variable container. E.g. in the example above y is being used in the broadcasting operation, thus it must be a collection of values. The a and b arguments, however, could be just single numbers:

result = infer(
    model = beta_bernoulli(),
    data  = (y = [ true, false, false ], a = 1, b = 1)
)

result.posteriors[:x]

Beta{Float64}(α=2.0, β=3.0)

• datastream

Also read the Streamlined Inference section.

The datastream keyword argument must be an observable that supports subscribe! and unsubscribe! functions (e.g., streams from the Rocket.jl package). The elements of the observable must be of type NamedTuple where keys (of Symbol type) correspond to input arguments defined in the model specification, except for those which are listed in the @autoupdates specification. For example, if a model defines y as its argument (which is not part of the @autoupdates specification) the named tuple from the observable must have an :y key (of Symbol type). The values in the named tuple must have the exact same shape as the corresponding variable container.

• initialization

Also read the Initialization section.

For specific types of inference algorithms, such as variational message passing, it might be required to initialize (some of) the marginals before running the inference procedure in order to break the dependency loop. If this is not done, the inference algorithm will not be executed due to the lack of information and message and/or marginals will not be updated. In order to specify these initial marginals and messages, you can use the initialization argument in combination with the @initialization macro, such as

init = @initialization begin
    # initialize the marginal distribution of x as a vague Normal distribution
    # if x is a vector, then it simply uses the same value for all elements
    # However, it is also possible to provide a vector of distributions to set each element individually
    q(x) = vague(NormalMeanPrecision)
end

Initial state: 
  q(x) = NormalMeanPrecision{Float64}(μ=0.0, w=1.0e-12)

• algorithm

Also read the Algorithm specification section.

Every rule of a node belongs to an algorithm, which selects the rules that run and carries their parameters. Most nodes run under the default algorithm, and need nothing. Some need a choice, such as the approximation method of a deterministic node or the order of an autoregressive node. The algorithm argument accepts an @algorithm specification, which gives algorithms to the nodes of the model by their function and variables:

@model function nonlinear_model(y)
    x ~ Normal(mean = 0.0, variance = 1.0)
    z := exp(x)
    y ~ Normal(mean = z, variance = 0.1)
end

result = infer(
    model     = nonlinear_model(),
    data      = (y = 2.0,),
    algorithm = @algorithm(exp() -> DeltaApproximation(method = Linearization())),
)

result.posteriors[:x]

NormalWeightedMeanPrecision{Float64}(xi=9.999999999999996, w=10.999999999999996)

A single node can also take its algorithm in the model, as z := exp(x) where { algorithm = ... }. The meta argument is the old name of algorithm, and it still works with a deprecation warning.

• returnvars

returnvars specifies latent variables of interest and their posterior updates. Its behavior depends on the inference type: streamline or batch.

Batch inference:

• Accepts a NamedTuple or Dict of return variable specifications.
• Two specifications available: KeepLast (saves the last update) and KeepEach (saves all updates).
• When iterations is set, returns every update for each iteration (equivalent to KeepEach()); if nothing, saves the last update (equivalent to KeepLast()).
• Use iterations = 1 to force KeepEach() for a single iteration or set returnvars = KeepEach() manually.

@model function normal_gamma(y)
    x ~ Normal(mean = 0.0, variance = 10.0)
    τ ~ Gamma(shape = 1.0, rate = 1.0)
    y .~ Normal(mean = x, precision = τ)
end

result = infer(
    model          = normal_gamma(),
    data           = (y = [0.9, 1.1, 1.3, 0.8, 1.0],),
    constraints    = MeanField(),
    initialization = @initialization(q(τ) = Gamma(1.0, 1.0)),
    iterations     = 5,
    returnvars     = (
        x = KeepLast(),
        τ = KeepEach()
    )
)

result.posteriors[:τ]

5-element Vector{GammaShapeRate{Float64}}:
 GammaShapeRate{Float64}(a=3.5, b=1.5651960784313728)
 GammaShapeRate{Float64}(a=3.5, b=1.2958216933834892)
 GammaShapeRate{Float64}(a=3.5, b=1.2578972450186254)
 GammaShapeRate{Float64}(a=3.5, b=1.2525496216627718)
 GammaShapeRate{Float64}(a=3.5, b=1.2517954017984472)

Shortcut for setting the same option for all variables:

result = infer(
    model          = normal_gamma(),
    data           = (y = [0.9, 1.1, 1.3, 0.8, 1.0],),
    constraints    = MeanField(),
    initialization = @initialization(q(τ) = Gamma(1.0, 1.0)),
    iterations     = 5,
    returnvars     = KeepLast()  # or KeepEach()
)

result.posteriors[:τ]

GammaShapeRate{Float64}(a=3.5, b=1.2517954017984472)

Streamline inference:

• For each symbol in returnvars, infer creates an observable stream of posterior updates.
• Agents can subscribe to these updates using the Rocket.jl package.

@model function gaussian_filter(y, x_mean, x_var)
    x ~ Normal(mean = x_mean, variance = x_var)
    y ~ Normal(mean = x, variance = 1.0)
end

my_autoupdates = @autoupdates begin
    x_mean, x_var = mean_var(q(x))
end

datastream   = RecentSubject(Float64)
observations = labeled(Val((:y,)), combineLatest(datastream))

engine = infer(
    model          = gaussian_filter(),
    datastream     = observations,
    autoupdates    = my_autoupdates,
    initialization = @initialization(q(x) = NormalMeanVariance(0.0, 10.0)),
    returnvars     = (:x,),
    autostart      = false
)

RxInferenceEngine:
  Posteriors stream    | enabled for (x)
  Free Energy stream   | disabled
  Posteriors history   | unavailable
  Free Energy history  | unavailable
  Enabled events       | [  ]

• predictvars

predictvars specifies the variables which should be predicted. Similar to returnvars, predictvars accepts a NamedTuple or Dict. There are two specifications:

• KeepLast: saves the last update for a variable, ignoring any intermediate results during iterations
• KeepEach: saves all updates for a variable for all iterations

result = infer(
    model          = normal_gamma(),
    data           = (y = [0.9, 1.1, 1.3, 0.8, missing],),
    constraints    = MeanField(),
    initialization = @initialization(q(τ) = Gamma(1.0, 1.0)),
    iterations     = 5,
    predictvars    = (y = KeepLast(),)
)

result.predictions[:y][end]

NormalMeanPrecision{Float64}(μ=1.0141213250047167, w=2.330525834820038)

• historyvars

Also read the Keeping the history of posteriors.

historyvars specifies the variables of interests and the amount of information to keep in history about the posterior updates when performing streamline inference. The specification is similar to the returnvars when applied in batch setting. The historyvars requires keephistory to be greater than zero.

historyvars accepts a NamedTuple or Dict or return var specification. There are two specifications:

• KeepLast: saves the last update for a variable, ignoring any intermediate results during iterations
• KeepEach: saves all updates for a variable for all iterations

engine = infer(
    model          = gaussian_filter(),
    datastream     = observations,
    autoupdates    = my_autoupdates,
    initialization = @initialization(q(x) = NormalMeanVariance(0.0, 10.0)),
    historyvars    = (
        x = KeepLast(),
    ),
    keephistory    = 10
)

for y in (0.9, 1.1, 1.3)
    next!(datastream, y)
end

engine.history[:x]

3-element DataStructures.CircularBuffer{Any}:
 NormalWeightedMeanPrecision{Float64}(xi=0.9, w=1.1)
 NormalWeightedMeanPrecision{Float64}(xi=2.0, w=2.1)
 NormalWeightedMeanPrecision{Float64}(xi=3.3, w=3.1)

It is also possible to set either historyvars = KeepLast() or historyvars = KeepEach() that acts as an alias and sets the given option for all random variables in the model.

engine = infer(
    model          = gaussian_filter(),
    datastream     = observations,
    autoupdates    = my_autoupdates,
    initialization = @initialization(q(x) = NormalMeanVariance(0.0, 10.0)),
    historyvars    = KeepLast(),
    keephistory    = 10
)

• keephistory

Specifies the buffer size for the updates history both for the historyvars and the free_energy buffers in streamline inference.

• iterations

Specifies the number of variational (or loopy belief propagation) iterations. By default set to nothing, which is equivalent of doing 1 iteration. However, if set explicitly to 1 the default setting for returnvars changes from KeepLast to KeepEach.

See Early stopping for an opt-in callback example that implements early stopping criterion.

• free_energy

Batch inference:

Specifies if the infer function should return Bethe Free Energy (BFE) values.

• Optionally accepts a floating-point type (e.g., Float64) for improved BFE computation performance, but restricts the use of automatic differentiation packages.

Streamline inference:

Specifies if the infer function should create an observable stream of Bethe Free Energy (BFE) values, computed at each VMP iteration.

• When free_energy = true and keephistory > 0, additional fields are exposed in the engine for accessing the history of BFE updates.
  • engine.free_energy_history: Averaged BFE history over VMP iterations.
  • engine.free_energy_final_only_history: BFE history of values computed in the last VMP iterations for each observation.
  • engine.free_energy_raw_history: Raw BFE history.

• free_energy_diagnostics

This settings specifies either a single or a tuple of diagnostic checks for Bethe Free Energy values stream. By default checks for NaNs and Infs. See also RxInfer.ObjectiveDiagnosticCheckNaNs and RxInfer.ObjectiveDiagnosticCheckInfs. Pass nothing to disable any checks.

• options

ReactiveMPInferenceOptions(; kwargs...)

• catch_exception

The catch_exception keyword argument specifies whether exceptions during the batch inference procedure should be caught in the error field of the result. By default, if exception occurs during the inference procedure the result will be lost. Set catch_exception = true to obtain partial result for the inference in case if an exception occurs. Use RxInfer.issuccess and RxInfer.iserror function to check if the inference completed successfully or failed. If an error occurs, the error field will store a tuple, where first element is the exception itself and the second element is the caught backtrace. Use the stacktrace function with the backtrace as an argument to recover the stacktrace of the error. Use Base.showerror function to display the error.

• callbacks

The inference function and underlying reactive message passing procedure both have their own lifecycle. The user is free to provide some (or none) of the callbacks to inject extra logic during the inference procedure. Callbacks can be a NamedTuple, Dict, or any custom structure that implements ReactiveMP.handle_event. For example:

@model function beta_bernoulli(y)
    x ~ Beta(1.0, 1.0)
    y .~ Bernoulli(x)
end

result = infer(
    model     = beta_bernoulli(),
    data      = (y = [true, false, true],),
    callbacks = (
        on_marginal_update = (event) -> println("$(event.variable_name) has been updated: $(event.update)"),
        after_inference    = (event) -> println("Inference has been completed")
    )
)

x has been updated: Marginal(Beta{Float64}(α=3.0, β=2.0))
Inference has been completed

For the full list of available events, supported callback types, model metadata, and built-in callback handlers, see the Callbacks section. See also Early stopping and Benchmark callbacks.

• annotations

Sets a tuple of annotation processors that attach extra information to messages and marginals during inference, such as InputArgumentsAnnotations(). When annotations are enabled, the inference results preserve the Marginal wrapper type so that annotation data remains accessible via ReactiveMP.getannotations. See ReactiveMP.jl documentation for available annotation types and how to implement custom annotation processors.

Automatically changes the default value of the postprocess argument to NoopPostprocess.

• logscales

logscales = true makes messages and marginals carry log scales: the log of the constant each rule's result leaves out, where the rule can compute it, combined through every product of messages. The Mixture node needs them, and in a model inferred exactly by belief propagation a posterior's log scale, getlogscale(result.posteriors[:x]), is the log evidence of the data, which is useful for Bayes factors and model comparison. Where a log scale cannot be computed, for instance after a variational rule, it is a ReactiveMP.UndefinedLogScale that says why. It can also be given as options = (logscales = true,).

Automatically changes the default value of the postprocess argument to NoopPostprocess, so the results keep the Marginal wrapper that carries the log scale.

• postprocess

Also read the Inference results postprocessing section.

The postprocess keyword argument controls whether the inference results must be modified in some way before exiting the inference function. By default, the inference function removes the Marginal wrapper type from the results. Change this setting to NoopPostprocess if you would like to keep the Marginal wrapper type. If the annotations or logscales argument has been used, automatically changes the default strategy value to NoopPostprocess.

• Error hints

By default, RxInfer provides helpful error hints when an error occurs during inference. This, for example, includes links to relevant documentation, common solutions and troubleshooting steps, information about where to get help, and suggestions for providing good bug reports.

Use RxInfer.disable_inference_error_hint! to disable error hints or RxInfer.enable_inference_error_hint! to enable them. Note that the change requires a Julia session restart to take effect.

disable_inference_error_hint!()

enable_inference_error_hint!()

Where to go next?

Read more explanation about the other keyword arguments in the Streamlined (online) inference section or check out the Static Inference section or check some more advanced examples.