Exploring possibility of Maximum Entropy Constrained Priors

[lstagner]

I was playing around with RxInfer and I wanted to see if I can use it to create constrained priors. There is an example in Sivia's Data Analysis: A Bayesian Tutorial.

Suppose that a die, with the usual 6 faces, was rolled a very large number of times; if we were only told that the average result was 4.5, what probability should we assign for the various possible outcomes ${X_i}$ that the face on top had $i$ dots? The information $I$ provided can be written as a simple constraint:
$\sum_{i=1}^6 i, p(X_i|I) = 4.5$

The solution to this problem is to do a constrained optimization of the entropy. I figured that since RxInfer does free energy minimization which is related to entropy it should be able to figure it out. I couldn't find anything in the documentation that so I figured I would try to just run a model where the measurements are the mean and just repeat the "measurement" of 4.5.

using RxInfer

@model function constrained_dice(m)
    x ~ Dirichlet(6,1) # has the property that draws satisfies sum(x) = 1
    m ~ PointMass(sum(i*x[i] for i=1:6))
end

dataset = fill(4.5,1000);

result = infer(
    model = constrained_dice(),
    data  = (m = dataset, )
)

Unfortunately, this doesnt work for a number of reasons.

1. Dirichlet distribution doesn't support the two argument form. Changing that to Dirichlet(ones(6)) fixes it.
2. Doesn't recognize that x should be a vector. Just use .~ fixes that.
3. It doesn't like the sum function. It gets the error no method matching +(::GraphPPL.NodeLabel, ::GraphPPL.NodeLabel) OK just manually do the sum. x[1] + 2*x[2] + 3*x[3] + 4*x[4] + 5*x[5] + 6*x[6]
4. It doesn't like PointMass: MethodError: no method matching factornode(::ReactiveMP.UndefinedNodeFunctionalForm, ::Type{PointMass}, ::Vector{Tuple{Symbol, AbstractVariable}}, ::Tuple{Vector{Int64}, Vector{Int64}}) OK. just replace that with a normal with a small variance.

Here is the "corrected" model

using RxInfer
@model function constrained_dice(m)
    x .~ Dirichlet(ones(6))
    m ~ NormalMeanVariance(x[1] + 2*x[2] + 3*x[3] + 4*x[4] + 5*x[5] + 6*x[6], 1e-6)
end
dataset = fill(4.5,1000);

result = infer(
    model = constrained_dice(),
    data  = (m = dataset, )
)

and this throws out a rule error RuleMethodError: no method matching rule for the given arguments.

At this point I am stuck and curious if this is the right approach to create constrained priors.

[albertpod]

Hi @lstagner!

(UPD from @Nimrais )
Thanks for trying out RxInfer.jl.

There are multiple ways to handle this problem. I'll provide how I'd approach it without overthinking, though others might suggest more interesting solutions. To get things going, you could do the following. I'll be using the Projections method due to the non-conjugacy:

using RxInfer
using ExponentialFamilyProjection

@model function constrained_dice(m, dims)
    x ~ Dirichlet(ones(dims))
    c  = collect(1:dims)
    for i in eachindex(m)
        m[i] ~ NormalMeanVariance(dot(x, c), tiny)
    end
end
dataset = fill(4.5, 100);

@constraints function dice_constraints(dims)
    q(x) :: ProjectedTo(Dirichlet, dims...)
end

constraints = dice_constraints((6,))

result = infer(
    model = constrained_dice(dims=6),
    data  = (m = dataset, ),
    constraints = constraints,
    iterations = 5
)

@show result.posteriors[:x][end]

I haven't checked whether results makes sense, so I am leaving this to you.

Another (perhaps faster) alternative would be this tutorial

For undefined rule, I suggest this doc!

@Nimrais, for some reason I couldn't turn on free energy computations. Can you have a look?

[Nimrais]

@albertpod There is a problem with how ReactiveMP runs free energy computation around the dot node. It needs to compute joint marginal for the x and c variables for some reason - I think @bvdmitri knows this part better.

I re-write your model a bit using the softdot node, a note for you @lstagner that softdot is exactly N(dot(θ, x), γ^(-1)).

I chose the precision 1 for the softdot - I tried with bigger values for the precision but the inference starts to diverge.

using RxInfer
using ExponentialFamilyProjection
using Plots

@model function constrained_dice(m, dims)
    x ~ Dirichlet(ones(dims))
    c  = collect(1:dims)
    for i in eachindex(m)
        m[i] ~ softdot(x, c, 1)
    end
end
dataset = fill(4.5, 100);

@constraints function dice_constraints(dims)
    q(x) :: ProjectedTo(Dirichlet, dims...)
end

constraints = dice_constraints((6,))

result = infer(
    model = constrained_dice(dims=6),
    data  = (m = dataset, ),
    constraints = constraints,
    iterations = 10,
    free_energy=true
)

@show result.posteriors[:x][end]

@show mean(result.posteriors[:x][end])

plot(result.free_energy)

[lstagner]

Thanks, for the help to get the code running unfortunetly it gives the wrong answer. Its not even close.

This problem can be solved by using Lagrange Multipliers.

The entropy has the form

$S = - \sum_{i=1}^6 x_i log(x_i)$

We have two constraints. The sum of the probabilities must sum to one and the second is that the mean is 4.5 (or whatever mean we want).

$c_1 = \sum_{i=1}^6 x_i = 1$
$c_2 = \sum_{i=1}^6 i x_i = 4.5$

As mentioned this is a lagrange multiplier problem. So we differentiate the entropy and constraints with respect to $x_i$ and multiply the constraints by their lagrange multiplier.

$\frac{dS}{dx_i} = \alpha \frac{dc_1}{dx_i} + \beta\frac{dc_2}{dx_i} \rightarrow -log(x_i) - 1 = \alpha + \beta i$

solving for $x_i$

$x_i = \exp(-(\alpha + 1 + \beta i))$

By enforcing the first constraint we can eliminate $\alpha$

$x_i = \frac{\exp(-\beta i)}{\sum_{i=1}^6 \exp(-\beta i)}$

Satisfying the mean constraint yields

$\frac{\sum_{i=1}^6 i \exp(-\beta i)}{\sum_{i=1}^6 \exp(-\beta i)} = 4.5$

You can find the value of $\beta$ numerically

using Roots
function maxent_constrained_dice(m,dims)
    f = beta -> sum(i*exp(-beta*i) for i=1:dims)/sum(exp(-beta*i) for i=1:dims) - m
    beta = fzero(f, 0.0)
    N = sum(exp(-beta*i) for i=1:dims)
    x = [exp(-beta*i)/N for i=1:dims]
    return beta, x
end

beta, x = maxent_constrained_dice(4.5,6)
(-0.3710489380810333, [0.05435316782649152, 0.07877154563305352, 0.11415997722944056, 
                       0.16544680311005336, 0.2397744404269, 0.347494065774061])

With that we can now compare the results generated by RxInfer and the true solution.

As you can see the true solution is nowhere near the mode of the posterior. I expect that the way I formulated the model is flawed in some way that is not immediately appararent or this is a limitation of the current version RxInfer.

My motivation for doing this is because I had an idea of creating priors for solutions of PDE's. Specifically the Grad-Shafranov Equation which is used my my field. I would represent the function as a Gaussian Process and then maximize the entropy with respect to the boundary conditions and the PDE itself. I expect that since the GP is a part of the exponential family it would be easier for RxInfer to find the MaxEnt prior.

[lstagner]

Interestingly enough. I get a much better answer if I only use one data point and set the precision to 10. If I increase the precision or the number of data points the results get worse.

EDIT:
For reference here is the result from Turing using only one data point. Like RxInfer Turing chains don't behave well when you try to have more than a few measurements.