Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add MvNormalMeanScalePrecision distribution #206

Open
wants to merge 28 commits into
base: main
Choose a base branch
from

Conversation

albertpod
Copy link
Member

@albertpod albertpod commented Aug 21, 2024

This PR was initially aimed at addressing ReactiveBayes/ReactiveMP.jl#387, which it still does. The distribution in question is parametrized by the mean and scale parameter of the precision matrix.

Initially, I implemented it as part of the MultivariateNormalDistributionsFamily. However, the conversions between MvNormalMeanScalePrecision and other distributions in this "class" don't always hold.

During the process, @Nimrais suggested that this distribution could be particularly interesting for ExponentialFamilyProjections.jl. To make it more useful, we need to optimize methods related to the computation of the Fisher information matrix. I made a first attempt to improve performance by modifying the computation of kron(invη2, invη2). I believe further improvements are possible, but this serves as a starting point.

Any suggestions for additional optimizations to enhance the distribution's effectiveness are much welcome.

UPD: I added the piece of code that actually fixes the ReactiveBayes/ReactiveMP.jl#387

@albertpod albertpod requested a review from Nimrais August 21, 2024 10:14
@albertpod albertpod marked this pull request as draft August 21, 2024 10:15
@albertpod albertpod marked this pull request as ready for review August 21, 2024 11:45
@albertpod
Copy link
Member Author

ping @Nimrais

Copy link
Collaborator

@Nimrais Nimrais left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

First, great work on implementing the MvNormalMeanScalePrecision distribution and integrating it into the ExponentialFamily.jl.

return prod(BayesBase.default_prod_rule(wleft, wright), wleft, wright)
end

function BayesBase.rand(rng::AbstractRNG, dist::MvGaussianMeanScalePrecision{T}) where {T}
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

function BayesBase.rand(rng::AbstractRNG, dist::MvGaussianMeanScalePrecision{T}) where {T}
μ, γ = mean(dist), scale(dist)
return μ .+ (1 / γ) .* randn(rng, T, length(μ))
end

Avoid constructing the identity matrix I(length(μ)) and directly scale the random vector.
Use broadcasting with ., which is more efficient and avoids unnecessary allocations.


# FIXME: This is not the most efficient way to generate random samples within container
# it needs to work with scale method, not with std
function BayesBase.rand!(
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similiarly to rand

function BayesBase.rand!(rng::AbstractRNG, dist::MvGaussianMeanScalePrecision, container::AbstractArray{T}) where {T <: Real}
    μ, γ = mean(dist), scale(dist)
    randn!(rng, container)
    @. container = μ + (1 / γ) * container
    return container
end

Btw I think rand just need to re-use rand!

@Nimrais
Copy link
Collaborator

Nimrais commented Sep 25, 2024

@bvdmitri I think PR is ready for review, but I need some help with efficient implementation of the fisher. The only tests that are failing are once that checking that fisher in this parametrisation is really faster.

@Nimrais Nimrais self-requested a review September 25, 2024 12:54
@albertpod
Copy link
Member Author

Thanks for refactoring this, @Nimrais!

@bvdmitri
Copy link
Member

@Nimrais ping, what needs to be done to finish this PR?

@Nimrais
Copy link
Collaborator

Nimrais commented Oct 15, 2024

@bvdmitri The basic functionality is now implemented. The only remaining task is to implement the multiplication of the Fisher matrix with another matrix for the Arrowhead matrix structure in BayesBase. Currently, this multiplication is only done in our tests. There's already a generic method for all matrix structures using the getindex method, so implementing this method is the last thing we need to do.

@Nimrais
Copy link
Collaborator

Nimrais commented Oct 15, 2024

@bvdmitri basically we need this PR ReactiveBayes/BayesBase.jl#25 be merged to finish this one

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants