Add Continuous Transition Node for Latent State Transformation

[albertpod]

This draft pull request (PR) addresses an issue related to variational message passing (VMP) in the context of updating a full-rank rectangular matrix within a "latent state transformation" model.

In this model, an m-dimensional latent state (x) is transformed into an n-dimensional observation (y) (or vice versa) through a linear transformation (y = K(a) x), where K(.) is a function that transforms a vector of arbitrary length into a matrix A (i.e., K(a) = A). The objective is to infer the variational posterior distribution of a, x, y, and process noise precision W.

To achieve this, we introduced a new node called ContinuousTransition. The update rules are implemented using structured variational inference. The complete derivation stack is based on unpublished material, which is available for reviewers.

The use of the ContinuousTransition node is primarily in two regimes:

1. When no structure on A is specified:

transformation = a -> reshape(a, 2, 2)
...
a ~ MvNormalMeanCovariance(zeros(2), Diagonal(ones(2)))
y ~ ContinuousTransition(x, a, W) where {meta = CTMeta(transformation, a)}
...

2. When structure is known:

transformation = a -> [cos(a[1]) -sin(a[1]); sin(a[1]) cos(a[1])]
...
a ~ MvNormalMeanCovariance(zeros(1), Diagonal(ones(1)))
y ~ ContinuousTransition(x, a, W) where {meta = CTMeta(transformation, a)}
...

Milestones:

• Finalize derivations for the ContinuousTransition node.
• Add documentation for the ContinuousTransition node.
• Write unit tests for the implemented update rules.
• [ ] Provide examples and benchmarks showcasing the ContinuousTransition node's capabilities.
• Refine the name and API of the Transfominator node based on feedback.

[ThijsvdLaar]

Very nice work, I think this node may prove very useful. Two comments from my side:

1. In the CTMeta it's not clear to me what the \hat{a} represents. Why does it need to be passed explicitly? And why in your example of the post above is a a Gaussian while the CTMeta constructor requires a Vector{<:Real}? A more explicit naming could already help here.
2. It may be good to mention the two use cases of the post above in the docs as well.

[bartvanerp]

Very cool new feature! I think we can start a lot of nice projects with this addition. I left some comments, but mostly found a lot of places where we can optimize the code. I did not annotate all of them, but optimizing the code is definitely required. I leave it in the middle whether it should be part of this PR.

[albertpod]

@ThijsvdLaar, thanks for the suggestion. Both points are addressed in the latest commit.

[albertpod]

@bartvanerp, thanks for checking. I've optimized where evident. Perhaps further optimization can be addressed with new issues.

[albertpod]

@bartvanerp or @bvdmitri only remaining issue with this PR now is that in case of linear transformation (e.g. reshape) we compute Jacobians, which is unnecessary). Should I improve on it within this PR or we a corresponding issue? In case of a former, what do you think is best way of doing it?

My idea is to make meta smth like:

struct ContinuousTransitionMeta
    f::Union{Vector{<:AbstractMatrix}, <:Function} 

    function ContinuousTransitionMeta(transformation::F) where {F}
        return new{F}(transformation)
    end
   
   function ContinuousTransitionMeta(transformation::F, a0::Vector{<:Real}) where {F}
        dy = size(transformation(a0), 1)
        Js = [ForwardDiff.jacobian(a -> transformation(a)[i, :], 1:length(a0)) for i in 1:dy]
        return new(Js)
    end

end

[bartvanerp]

I am fine with creating an issue for this and assigning someone. For later reference: an example solution could be to create an anonymous function and pre-specify its jacobian, such as

foo(x) = reshape(x, d1, d2)
jacobian_sliced(f::foo, i) = .... # this should be specified by the user in their experiment
jacobian_sliced(f::Function, i) = ForwardDiff.jacobian(a -> transformation(a)[i, :] # generic case (fallback in case specialization not specified)

I would say that this kind of optimization has to be performed by the user himself as an advanced feature and only needs to be documented.

Also note that are a couple of unresolved comments.

[albertpod]

@bartvanerp I've resolved the rest of the comments.
There's an issue that must be opened regarding AE and Jacobian computations.
I will open it after the merge.

[bvdmitri]

What do you mean unnecessary? In either case you can have a boolean flag, which controls whether to compute the jacobians or not (true by default for example)

[albertpod]

@bvdmitri with unnecessary computations, I meant that in the case of a linear transformation, we don't need to compute Jacobians each time when ctcompanion_matrix is called. As for now, it happens regardless of transformation type.

[bvdmitri]

But your code cannot assume anything about the type of the transformation, how can you decide if it is necessary or not?

[albertpod]

@bvdmitri, you are correct. That's why I disregarded this option of precomputed Jacobians and kept it generic.
However, the semi-experienced user can specify a flag that precomutes the Jacobians. It's a burden of the user that should be appropriately documented.