Enforcing mean-field constraints in simple model?

[joacorapela]

I am new to RxInfer. Sorry if the question is naive.

It seems that RxInfer is not applying the mean field constraints, as the variance of the mean-field approximation should be smaler than that of the unconstrained posterior. Can I enforce the mean field constraints?

using LinearAlgebra                                                                                                                                                                                           
using RxInfer                                                                                                                                                                                                 
                                                                                                                                                                                                              
# Define the model generator                                                                                                                                                                                  
@model function lg(x, a)                                                                                                                                                                                      
    z1 ~ Normal(mean=0.0, var=1.0)                                                                                                                                                                            
    z2 ~ Normal(mean=0.0, var=1.0)                                                                                                                                                                            
    x ~ Normal(mean=a[1]*z1 + a[2]*z2, var=1.0)                                                                                                                                                               
end                                                                                                                                                                                                           
                                                                                                                                                                                                              
a = [2.0 3.0]                                                                                                                                                                                                 
x = 4.0                                                                                                                                                                                                       
                                                                                                                                                                                                              
# Perform inference                                                                                                                                                                                           
unconstrained_result = infer(                                                                                                                                                                                 
    model = lg(a=a),                                                                                                                                                                                          
    data = (x = x,),                                                                                                                                                                                          
)                                                                                                                                                                                                             
                                                                                                                                                                                                              
@constraints function mf_constraints()                                                                                                                                                                        
    q(z1) = q(z1)                                                                                                                                                                                             
    q(z2) = q(z2)                                                                                                                                                                                             
end                                                                                                                                                                                                           
                                                                                                                                                                                                              
mf_result = infer(                                                                                                                                                                                            
    model = lg(a=a),                                                                                                                                                                                          
    data = (x = x,),                                                                                                                                                                                          
    constraints = mf_constraints(),                                                                                                                                                                           
)                                                                                                                                                                                                             
                                                                                                                                                                                                              
# Print results                                                                                                                                                                                               
                                                                                                                                                                                                              
println("*** Unconstrained ***")                                                                                                                                                                              
                                                                                                                                                                                                              
println("mu_1=$(mean(unconstrained_result.posteriors[:z1]))")                                                                                                                                                 
println("mu_2=$(mean(unconstrained_result.posteriors[:z2]))")                                                                                                                                                 
println("var_1=$(var(unconstrained_result.posteriors[:z1]))")                                                                                                                                                 
println("var_2=$(var(unconstrained_result.posteriors[:z2]))")                                                                                                                                                 
                                                                                                                                                                                                              
println("*** Mean Field ***")                                                                                                                                                                                 
                                                                                                                                                                                                              
println("mu_1=$(mean(mf_result.posteriors[:z1]))")                                                                                                                                                            
println("mu_2=$(mean(mf_result.posteriors[:z2]))")                                                                                                                                                            
println("var_1=$(var(mf_result.posteriors[:z1]))")                                                                                                                                                            
println("var_2=$(var(mf_result.posteriors[:z2]))")                                                                                                                                                            

Output

*** Unconstrained ***
mu_1=0.5714285714285715
mu_2=0.8571428571428573
var_1=0.7142857142857143
var_2=0.35714285714285715
*** Mean Field ***
mu_1=0.5714285714285715
mu_2=0.8571428571428573
var_1=0.7142857142857143
var_2=0.35714285714285715

The mean-field variances should be smaller than the unconstrained ones, which suggests that RxInfer is not applying the mean-field constraints. Can I enforce them?

[joacorapela]

Address as issue #479