ReactiveMP.jl icon indicating copy to clipboard operation
ReactiveMP.jl copied to clipboard

Feature request: Broadcast distributions / mean-field MvNormal?

Open svilupp opened this issue 3 years ago • 0 comments

At the moment it's challenging to define a mean-field Multivariate Normal distribution (ie, with a diagonal covariance matrix).

Eg, assuming I have a regression y=beta * X + eps, where beta is MvNormal of dimension dim_x At the moment I would can easily use MvNormal and fit the full covariance matrix

beta  ~ MvNormalMeanPrecision(zeros(dim_x), diageye(dim_x))

for i in 1:n
    y[i] ~ NormalMeanPrecision(dot(X[i,:],beta), 1.)
end

If I wanted a mean-field MvNormal instead (eg, to have faster fitting), I would have to jump through a lot of hoops (eg, unroll dot product etc), which would ultimately give me no benefits

# create a RV with dimension `dim_x`
beta  ~ randomvar(dim_x)

# since I cannot use dot() with randomvar(), I would need to unroll it
# I would need some temporary variable 
# because I cannot accumulate (+=), I would need to save intermediate results into a helper dimension
temp ~ randomvar(n,dim_x)

for i in 1:n
    # unsafe! just to show the needed steps
    # assumes dim_x>1
    temp[i,1]=X[i,1] * beta[1]
    for j in 2:dim_x
        temp[i,j] ~ temp[i,j-1] + X[i,j] * beta[j]
    end
    
    # and the use it for the normal likelihood
    y[i] ~ NormalMeanPrecision(loc[i,dim_x], 1.)
end

As a user, I'd like to define something like

beta ~ randomvar(Normal(0,1),n) 

# or ala Turing when all dists are the same
beta ~ filldist(Normal(0,1),n)
# or if they are different
beta ~ arraydist([Normal(mu,1) for mu in mu_array])

Ideally, this object would then support all the operators like MvNormal (eg, dot product)

It would be helpful when dim_x gets large (eg, 1000s)

svilupp avatar Jun 17 '22 20:06 svilupp