Enforcing mean-field constraints in simple Gaussian linear 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?

[bvdmitri]

Hey! Thanks for the interest in RxInfer!
As I can see from the model specification, both z1 and z2 are conditionally independent in the model itself. In other words, they don't have a shared factor (e.g. something like x ~ SomeNode(z1, z2, ...)). IMO they are already factorized out and applying extra mean-field won't change anything.

Thus being said, you can double-check that as in your original code the specification inside the @constraints macro won't do anything. If you want to separate variables, you should write q(z1, z2) = q(z1)q(z2) instead. Alternatively, you can use constraints = MeanField() in the infer call to apply mean-field globally for the entire model.

[joacorapela]

Thanks for your quick reply @bvdmitri . I have unsuccesfuly tried different versions of the model to enforce the mean field approximation. Below is a version following your suggestions.

using RxInfer                                                                                                                                                                                                 
                                                                                                                                                                                                              
# Define the model generator                                                                                                                                                                                  
@model function lg(x, a)                                                                                                                                                                                      
    z ~ MvNormal(mean = [0.0, 0.0], cov = Matrix{Float64}(I, 2, 2))                                                                                                                                           
    x ~ MvNormal(mean = a * z, cov = Matrix{Float64}(I, 1, 1))                                                                                                                                                
end                                                                                                                                                                                                           
                                                                                                                                                                                                              
a = [2.0 3.0]                                                                                                                                                                                                 
x = [4.0]                                                                                                                                                                                                     
                                                                                                                                                                                                              
# Perform inference                                                                                                                                                                                           
bayesian_result = infer(                                                                                                                                                                                      
    model = lg(a=a),                                                                                                                                                                                          
    data = (x = x,),                                                                                                                                                                                          
)                                                                                                                                                                                                             
                                                                                                                                                                                                              
@constraints function mf_constraints()                                                                                                                                                                        
    q(z) = q(z[1])q(z[2])                                                                                                                                                                                     
end                                                                                                                                                                                                           
                                                                                                                                                                                                              
mf_result = infer(                                                                                                                                                                                            
    model = lg(a=a),                                                                                                                                                                                          
    data = (x = x,),                                                                                                                                                                                          
    constraints = mf_constraints(),                                                                                                                                                                           
)                                                                                                                                                                                                             
                                                                                                                                                                                                              
# Print results                                                                                                                                                                                               
                                                                                                                                                                                                              
println("*** Bayesian ***")                                                                                                                                                                                   
                                                                                                                                                                                                              
mu = mean(bayesian_result.posteriors[:z])                                                                                                                                                                     
covM = cov(bayesian_result.posteriors[:z])                                                                                                                                                                    
println("mu_1=$mu[1]")                                                                                                                                                                                        
println("mu_2=$mu[2]")                                                                                                                                                                                        
println("var_1=$(covM[1,1])")                                                                                                                                                                                 
println("var_2=$(covM[2,2])")                                                                                                                                                                                 
println("cov_12=$(covM[1,2])")                                                                                                                                                                                
                                                                                                                                                                                                              
println("*** Variational ***")                                                                                                                                                                                
                                                                                                                                                                                                              
mu = mean(mf_result.posteriors[:z])                                                                                                                                                                           
covM = cov(mf_result.posteriors[:z])                                                                                                                                                                          
println("mu_1=$mu[1]")                                                                                                                                                                                        
println("mu_2=$mu[2]")                                                                                                                                                                                        
println("var_1=$(covM[1,1])")                                                                                                                                                                                 
println("var_2=$(covM[2,2])")                                                                                                                                                                                 
println("cov_12=$(covM[1,2])") 

Output

*** Bayesian ***
mu_1=[0.5714285714285714, 0.8571428571428571][1]
mu_2=[0.5714285714285714, 0.8571428571428571][2]
var_1=0.7142857142857142
var_2=0.3571428571428571
cov_12=-0.4285714285714285
*** Variational ***
mu_1=[0.5714285714285714, 0.8571428571428571][1]
mu_2=[0.5714285714285714, 0.8571428571428571][2]
var_1=0.7142857142857142
var_2=0.3571428571428571
cov_12=-0.4285714285714285

The fact that we observe a non-zero covariance between z_1 and z_2 suggests that RxInfer is not applying the constraint. What can I do to enforce it?

[bvdmitri]

Well, no, its not what I suggested. I suggested to either use

@constraints function mf_constraints()                                                                                                                                                                        
    q(z1, z2) = q(z1)q(z2)                                                                                                                                                                                           
end    

for your original model or simply constraints = MeanField() in the infer function.

[bvdmitri]

The second example of yours cannot apply a mean-field constraint within a single variable z. It is not supported to apply constraints on individual components of a single multivariate variable (an array of different variables is fine).

Thus being said, I will still reemphasize that in my opinion, in this model

@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   

result should not change with or without constraints = MeanField() because RxInfer uses so-called Bethe factorization assumption. The way I see it, the variables are already factorized out and there is no need to assign mean-field constraints in this particular model and even if you do, it should not change the result.

What are you trying to solve in particular?

[joacorapela]

In the simple linear Guassian model above, I am trying to get the variational mean-field approximation of p(z|x). With all the model specifications and inference methods that I have tried, I can only get the uncontrained posterior p(z|x). Can I get the variational mean-field approximation of p(z|x) in RxInfer?

Both the unconstrained posterior and the mean-field approximation can be solved analytically for this simple model (details here). With the parameter a and observation x in the above example, for the unconstrained posterior we should get

mean(z1) = 0.5714285714285714
mean(z2) = 0.8571428571428571
var(z1) = 0.7142857142857142
var(z2) = 0.3571428571428571
covar(z1, z2) = −0.42857142857142855

and for the mean-field approximation we should get

mean(z1) = 0.5714285714285714
mean(z2) = 0.8571428571428571
var(z1) = 0.2
var(z2) = 0.1

Above you said:

As I can see from the model specification, both z1 and z2 are conditionally independent in the model itself. In other words, they don't have a shared factor (e.g. something like x ~ SomeNode(z1, z2, ...)). IMO they are already factorized out and applying extra mean-field won't change anything.

In the model

@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

z1 and z2 are independent a priori, but observing x makes them dependent. Thus, applying the mean-field constraint should change the inferred values, as shown above.

How should I specify the model and the inference method in RxInfer to obtain the mean-field approximation?

mean(z1) = 0.5714285714285714
mean(z2) = 0.8571428571428571
var(z1) = 0.2
var(z2) = 0.1

[albertpod]

@joacorapela sorry this conversation stopped, didn't the the last message.

Look, if you really want to enforce mean-field, you'd have to change your model a bit:

@model function lg(x, a)                                                                                                                                                                                      
    z1 ~ Normal(mean=0.0, var=1.0)                                                                                                                                                                            
    z2 ~ Normal(mean=0.0, var=1.0)    
    tmp1 ~ a[1]*z1
    tmp2 ~ a[2]*z2                                                                                                                                                                   
    x ~ SoftAddition(tmp1, tmp2, 1.0)                                                                                                                                                               
end

You then would have to introduce a node and workout rules for variational message passing and belief propagation if you want to compare these thing, or you introduce only variational updates and compare the variances between this model and the previous model you had.

Additionally, I believe that something with projection methods might work here, but not sure yet.

[albertpod]

@joacorapela, I just realised that you can use softdot node here that supports full mean-field.

@model function lg(x, a)                                                                                                                                                                                      
    z1 ~ Normal(mean=0.0, var=1.0)                                                                                                                                                                            
    z2 ~ Normal(mean=0.0, var=1.0)
    c1 = [1.0, 0.0]
    c2 = [0.0, 1.0]
    z ~  z1*c1 + z2*c2                                                                                                                                            
    x ~ softdot(a, z, 1.0)                                                                                                                                                               
end