Bernoulli/Dirichelet bayesian model combination in RxInfer?

[meditans]

I'm trying to understand how to implement Bayesian Model Combination (BMC) in RxInfer, for a mixture of Bernoulli or Categorical variables (not too interested in the gaussian counterpart for now).

I found the description in the paper Automating Model Comparison in Factor Graphs The paper describes three possibilities, I think all of them implementable with RxInfer's Mixture node:

• Bayesian model averaging: a categorical prior over model index m, inference returns a soft posterior over m
• Bayesian model selection: same as BMA but q(m) is constrained to a Kronecker delta (hard selection)
• Bayesian model combination: uses a mixing weight pi with a Dirichlet prior, shared across observations. Each observation n gets its own selector m[n] ~ Bernoulli(pi), and pi accumulates evidence from all observations.

It seems that BMC would allow me to find the best weighted combination of models rather than converging to a single best one.

I tried starting from the John/Jane coin toss example in the Universal Mixtures docs, which looks like this:

@model function beta_model_mary(y)
    theta_jane ~ Beta(2.0, 7.0)
    theta_john ~ Beta(7.0, 2.0)

    john_is_right ~ Bernoulli(0.7)

    theta ~ Mixture(switch = john_is_right, inputs = [theta_jane, theta_john])

    y .~ Bernoulli(theta)
end

But as far as I can tell, this is BMS/BMA, not BMC: there is a single shared switch variable with a fixed prior probability.
If I understand correctly an BMC model should look more like:

@model function bmc_model(y)
    pi ~ Beta(1.0, 1.0)

    theta_jane ~ Beta(2.0, 7.0)
    theta_john ~ Beta(7.0, 2.0)

    for n in eachindex(y)
        m[n] ~ Bernoulli(pi)
        theta[n] ~ Mixture(switch = m[n], inputs = [theta_jane, theta_john])
        y[n] ~ Bernoulli(theta[n])
    end
end

init = @initialization begin
    q(pi) = vague(Beta)
    q(m)  = vague(Bernoulli)
end

rng      = MersenneTwister(0)
n_throws = 30
dataset  = Int.(rand(rng, Bernoulli(0.25), n_throws))

result = infer(
    model          = bmc_model(),
    data           = (y = dataset,),
    initialization = init,
    constraints    = MeanField(),
    addons         = AddonLogScale(),
    iterations     = 10,
)

but even adding MeanField() and AddonLogScale(), which I think were implied by the errors, I cannot get this to run.

So, my questions would be:

• Is the John/Jane beta_model_mary example actually BMA/BMS, or am I misreading it?
• Is BMC with a shared latent pi currently supported through the Mixture node, and if so, what is the correct way to set up the model, constraints, and initialization?

Thank you!

[bvdmitri]

Hey, as far as I remember the AddonLogScale and the Mixture node work specifically without the MeanField constraint. I think the problem here is that some messages are not defined (not sure if they can be). For example, if you look at the backward message for the Bernoulli node that computes the logscale, it explicitly wants the message to be a PointMass. In the bmc_model the message towards Bernoulli is Bernoulli. So there should a rule defined

@rule Bernoulli(:p, Marginalisation) (m_out::Bernoulli,) = begin
    @logscale ...
    return Beta(..., ...)
end

[bvdmitri]

Here I also removed the MeanField constraint and initialized the messages instead

julia> init = @initialization begin
           μ(pi) = vague(Beta)
           μ(m)  = vague(Bernoulli)
       end

julia> result = infer(
           model          = bmc_model(),
           data           = (y = dataset,),
           initialization = init,
           addons         = AddonLogScale(),
           iterations     = 10,
       )

[meditans]

Thank you @bvdmitri for guiding me towards the solution. In the end I initialized the message for m to μ(m) = vague(Categorical, 2) as that's what the Mixture node sends back on the :switch variable anyway. I calculated the expected value and variance of what the correct message should be in a Bernoulli node with an incoming Categorical{2} message:

from sympy import *

p, q = symbols("p q", positive=True)

unnorm_message = p * q + (1 - p) * (1 - q)
norm_constant = integrate(unnorm_message, (p, 0, 1))

message = unnorm_message / norm_constant

E_p = integrate(p * message, (p, 0, 1))
E_p2 = integrate(p**2 * message, (p, 0, 1))
V_p = simplify(E_p2 - E_p**2)

print(E_p) # q/3 + 1/3
print(V_p) # -q**2/9 + q/9 + 1/18

and then moment-matched a Beta. So, the working code:

using RxInfer, Distributions, Random, Plots, PrettyPrinting
import Pipe: @pipe as @p

@model function bmc_model(y)
    pi ~ Beta(100.0, 100.0)

    theta_jane ~ Beta(2.0, 7.0)
    theta_john ~ Beta(7.0, 2.0)

    for n in eachindex(y)
        m[n] ~ Bernoulli(pi)
        theta[n] ~ Mixture(switch = m[n], inputs = [theta_jane, theta_john])
        y[n] ~ Bernoulli(theta[n])
    end
end

init = @initialization begin
    μ(pi) = vague(Beta)
    μ(m)  = vague(Categorical, 2)
end

@rule Bernoulli(:p, Marginalisation) (m_out::Categorical,) = begin
    @logscale -log(2)

    q = probs(m_out)[2]

    e = q/3 + 1/3
    v = -q^2/9 + q/9 + 1/18

    c     = e * (1 - e) / v - 1
    alpha = c * e
    beta  = c * (1 - e)

    return Beta(alpha, beta)
end

result = infer(
           model          = bmc_model(),
           data           = (y = [0,0,0,1,1,1,1,1,1],),
           initialization = init,
           addons         = AddonLogScale(),
           iterations     = 10,
       )

marginal_content(x) = x.data.p
mixture_content(x) = x.data.components

@p result.posteriors[:m]     |> last |> marginal_content.(_) |> stack
@p result.posteriors[:theta] |> last |> mixture_content.(_)
@p result.posteriors[:theta_john] |> last
@p result.posteriors[:theta_jane] |> last
@p result.posteriors[:pi] |> last

That said, the posteriors seem a bit weird and I don't know how to interpret this from a modeling perspective. Is this on the right track for finding the best weights for the models?

[meditans]

On further debugging, I doubt the :p rule I wrote for the Bernoulli node is the right one, it seems I learned the hard way I cannot just do moment matching on the messages. I don't know the right mathematical framework I should apply then, as the other intuitive solution I can come up:

@rule Bernoulli(:p, Marginalisation) (m_out::Categorical,) = begin
    @logscale -log(2)
    q = probs(m_out)[2]
    return Beta(1 + q, 2 - q)
end

seems to have problems of its own. Could you point me towards the correct math?

[bvdmitri]

The "log-scale" addition was originally introduced a part of this paper as far as I remember https://www.mdpi.com/1099-4300/25/8/1138?utm_source=researchgate.net&utm_medium=article

I'm entirely sure how it should be computed, this is also largely an undocumented feature, so I'm inviting @bartvanerp @ThijsvdLaar to add something here as original authors of the paper