Possibility to include MatrixNormal?

[grafaelw]

Is there any possible way to include MatrixNormal distribution? I am currently working on a project that requires the inference of matrices.

[albertpod]

Hi @grafaelw, the short answer is no. The inference for matrix normal isn't supported out of the box in RxInfer.jl. There are some developments on the continuous transition node, which enables work with MatrixNormal distributions when they play a role of linear transformation of a random variable x, i.e. y = Hx, where H ~ MatrixNormal(...). However, I am unsure when this branch will merge with the main.

There are ways to circumvent the absence of MatrixNormal within the current version of RxInfer.jl but for that, I'd need a minimum model specification from your side.

[grafaelw]

Thank you for the reply @albertpod. What I am trying to infer is a linear state-space model, where A and C are matrices with DD and MD dimension respectively. The model specification is as follows:

@model function lssm(N, D, M)
    
    x = randomvar(N)
    y = datavar(Vector{Float64}, N)
    
    x_prior ~ MvNormalMeanCovariance(zeros(D), 0.01*diageye(D))
    x_prev = x_prior

    α ~ InverseWishart(D^2, 1.0e-5*diageye(D))
    γ_M ~ InverseWishart(M*D, 1.0e-5*diageye(M))
    γ_D ~ InverseWishart(D*M, 1.0e-5*diageye(D))
    τ ~ InverseWishart(M^2, 1.0e-5*diageye(M))
   

    A ~ MatrixNormal(zeros(D, D), α, α)
    C ~ MatrixNormal(zeros(M, D), γ_M, γ_D)

    for i in 1:N
        x[i] ~ MvNormalMeanCovariance(A * x_prev, diageye(D))
        y[i] ~ MvNormalMeanCovariance(C * x[i], τ) 
        x_prev = x[i]
    end
end

[albertpod]

@grafaelw that's almost precisely what CTransition aims to achieve, except for the priors on the parameters of MatrixNormal. We do it by working out the derivations involving MatrixNormal in its vector representation. The feature will be added to RxInfer somewhere this quarter.

For now, if you want to continue faster with RxInfer.jl, you can do the following trick with the delta node (I simplified your model a little bit for the sake of brevity).

# simple example of a function that translates a vector into the matrix and multiplies it with state x
function f(x, a)
    A = [a[1:2] a[3:4]]
    A*x
end

@model function lssm(N, D, C)
    
    x = randomvar(N)
    x̂ = randomvar(N)
    y = datavar(Vector{Float64}, N)
    
    x_prior ~ MvNormal(μ=zeros(D), Σ=0.01*diageye(D))
    x_prev = x_prior

    A ~ MvNormal(μ=vec(zeros(D, D)), Σ=kron(diageye(D), 0.01*diageye(D))) # this is equivalent to MatrixNormal

    for i in 1:N
        x̂[i] ~ f(x_prev, A)
        x[i] ~ MvNormal(μ=x̂[i], Σ=diageye(D))
        y[i] ~ MvNormal(μ=C*x[i], Σ=diageye(D)) 
        x_prev = x[i]
    end
end

I am unsure how familiar you are with graphical models, but even with this simple tweak, we have introduced a loop in the model. So the inference call will require initmessages:

delta_meta = @meta begin 
    f() -> Linearization()
end

D = 2
N = 100 
dataset = (y = [randn(D) for i in 1:N],)

result = inference(model = lssm(N, D, diageye(D)), 
initmessages=(A = MvNormalMeanCovariance(vec(zeros(D, D)), kron(diageye(D), 0.01*diageye(D))),), 
meta=delta_meta, data = dataset, 
free_energy=true, showprogress=true, iterations=10, returnvars=KeepLast())

You can check free_energy values for convergence.

EDIT: Please check @meta macro in the documentation.

[grafaelw]

@albertpod Thank you for the amazing explanation. I'll try this out and see if things are working well.