Inference with CVI for a sparse discrete transition model

[ThijsvdLaar]

I'm trying to implement inference for a transition model with some sparse structure.

using RxInfer
using StableRNGs
using Optimisers


#-------------------------
# Extensions of ReactiveMP
#-------------------------

import Base: size
size(d::MatrixDirichlet) = size(d.a)

import Distributions: _logpdf
function _logpdf(d::MatrixDirichlet, X::AbstractMatrix{<:Real})
    c = 0.0
    for i=1:size(X)[2] # For each column
        c += _logpdf(Dirichlet(d.a[:,i]), X[:,i])
    end
    return c
end

using Base.Broadcast: BroadcastFunction
@rule Transition(:out, Marginalisation) (m_in::Categorical, q_a::Any) = begin
    a = clamp.(exp.(mean(BroadcastFunction(log), q_a)) * probvec(m_in), tiny, Inf)
    return Categorical(a ./ sum(a))
end


#-----------------
# Generative model
#-----------------

f(a, b, c) = [1.0-a-c b       0.1
              a       1.0-b-c 0.1
              c       c       0.8] # Some sparse structure

cvi = CVI(StableRNG(42), 100, 100, Optimisers.Descent(0.01))

@model function sparse()
    p_x_min = datavar(Vector{Float64})

    a ~ Beta(1.0, 1.0)
    b ~ Beta(1.0, 1.0)
    c ~ Beta(1.0, 1.0)

    T ~ f(a, b, c) where {meta = DeltaMeta(method=cvi)}

    x_min ~ Categorical(p_x_min)
    x ~ Transition(x_min, T)
    x ~ Categorical(ones(3)./3) # Uninformative
end


#--------------------------
# Constraints and inference
#--------------------------

constraints = @constraints begin
    q(x, x_min, T) = q(x, x_min)q(T)
end

initmarginals = (T=MatrixDirichlet(ones(3,3)),)

result = inference(
    model         = sparse(),
    data          = (p_x_min=[0.2, 0.2, 0.6],),
    constraints   = constraints,
    initmarginals = initmarginals,
    iterations    = 5
)

I'm getting warnings and an error

┌ Warning: CVI approximation has not updated the initial state. The method did not converge. Set `warn = false` to supress this warning.
└ @ ReactiveMP [C:\Users\tvdlaar\.julia\packages\ReactiveMP\qCT2A\src\approximations\cvi.jl:177](file:///C:/Users/tvdlaar/.julia/packages/ReactiveMP/qCT2A/src/approximations/cvi.jl:177)

DomainError with -0.5069832093228921:

log will only return a complex result if called with a complex argument. Try log(Complex(x)).

[11] (::ReactiveMP.var"#184#185"{BroadcastFunction{typeof(log)}})(::Tuple{Float64, StaticArraysCore.SMatrix{3, 3, Float64, 9}})

    @ ReactiveMP [C:\Users\tvdlaar\.julia\packages\ReactiveMP\qCT2A\src\distributions\sample_list.jl:210](file:///C:/Users/tvdlaar/.julia/packages/ReactiveMP/qCT2A/src/distributions/sample_list.jl:210)

@Nimrais Perhaps you could tell me more about what I'm missing here? In particular, why is the CVI state not updating? And how is it possible that I obtain samples from the MatrixDirichlet distribution with negative entries?

[wouterwln]

@Nimrais I think this is exactly the same problem as what we discussed today with the TrueSkill node, I recognize this error message from earlier experiments ;)

[Nimrais]

@ThijsvdLaar, unfortunately, I need to say that there is a bug in the current implementation of ReactiveMP.jl. Unfortunately, it's not easy to hotfix it there. I am writing an extension for ReactiveMP.jl right now, where this problem will be resolved.

I think you would be able to try it early next week, I already have a Direchlet so should be fine for you.

[Nimrais]

The state is not updating because all steps are improper distributions, and that's is because we can not compute ForwardDiff gradient trough parameters of distributions.

[Nimrais]

@ThijsvdLaar @albertpod I have experimented with this model a bit with new CVI. And in a closer look this model is ill-defined.

The main issue arises when we run CVI and need to assume a mean-field approximation over the variables $a$, $b$, and $c$.

The problem becomes evident when we attempt to normalize the message update over the mean-field marginal distributions of $a$, $b$, and $c$. Specifically, we encounter the following integral:

$$\int q(a) q(b) q(c) \cdot \mu_T(p_{MD}(a, b, c)) \mathrm{d}a \mathrm{d}b \mathrm{d}c = 0$$

Here, $q(a)$, $q(b)$, and $q(c)$ represent the approximate posterior distributions for variables $a$, $b$, and $c$, respectively, and $\mu_T(f(a, b, c))$ is a message from $\delta_{f}$ towards joint $q(a)q(b)q(c).$

The issue is that this integral evaluates to zero, which suggests that the model is not well-defined under the current assumptions.

To resolve this problem, further investigation is needed to determine the appropriate assumptions for the variables $a$, $b$, and $c$, or/and to change $f$.

[albertpod]

Thanks for checking. @ThijsvdLaar didn't you have a working example of this model + inference in ForneyLab.jl?

[ThijsvdLaar]

@albertpod ForneyLab didn't have this exact example, and I doubt that Semih's CVI implementation in FL would have been able to handle this.

To phrase the problem in another way, the domain for the random variables in the prior $p(a,b,c)$ is a cube, because the prior is defined as three independent beta distributions. However, because of the deterministic constraint $f(a,b,c)$, the domain for the posterior $q(a,b,c)$ is a lower-dimensional (non-Euclidian) manifold that transsects this cube. Therefore, samples from the prior have a virtually zero chance to land on the posterior manifold.

@Nimrais Another way to resolve this problem might be to soften the deterministic relationship $f$, so that the domain for $f(a,b,c)$ also becomes the cube (instead of a lower-dimensional manifold). This way samples from the prior always have a non-zero chance to land on the posterior manifold (which is now also the cube).