Free Energy for model selection?

[langestefan]

I was wondering if the free energy returned by rxinfer can be used as a model selection criterion.

Variational free energy provides an approximation to Bayesian model evidence so I would assume so.

Do you have any input or thoughts on this? Thanks! 🙏

[albertpod]

Hi @langestefan. It depends which type of models you compare. Generally speaking, you can't do that, see section 2.2 from Blei https://arxiv.org/pdf/1601.00670.

[wmkouw]

Hi @langestefan, I don't think the picture and your results match: you're doing linear regression and your picture shows curve fits. How did you generate this?

[wmkouw]

Ah! You're doing polynomial regression, i.e., performing a basis expansion on X and then fitting coefficients. Is that correct?

I suspect that your FE results are counter-intuitive because of your choice of prior parameters. You may be constraining your model too much. Try increasing the scale of the covariance matrix.

Also, it's not clear that the model is over-fitting. Over-fitting is of course the difference between training and test error. But remember that we're Bayesians and that we do not measure error in terms of distance, but in terms of probability of the observation under the posterior predictive distribution. In short, you have to take the uncertainty over your predictions into account. Try plotting the posterior uncertainty for your curve fits. We suspect that the uncertainty over M=6 grows much more rapidly than for the smaller model sizes.

[langestefan]

Hi @wmkouw, yes that is correct! Apologies for the confusion.

I think you are right about the uncertainty:

For the ribbon (+- 1 std) I have used here (Bishop 3.59):

$$\begin{align} p(\eta_{\mathrm{COP}}\,|\, \Delta \mathrm{T}, \mathcal{D}) &= \mathcal{N}\left(\eta_\mathrm{COP}\,|\, \boldsymbol{\mu}_{\boldsymbol{\theta}}^T \phi(\Delta \mathrm{T}), \sigma^2(\Delta \mathrm{T}) \right)\\\ \sigma^2(\Delta \mathrm{T}) &= \tau + \phi(\Delta \mathrm{T})^T \Sigma_{\boldsymbol{\theta}} \phi(\Delta \mathrm{T}) \end{align}$$

With:

$$\phi(x) = [1, x, x^2, \dots, x^M]$$

I have also increased the scale of the covariance matrix, and these are now the results:

model_order	FE	MAE_train	MAE_test
1	42650.2	0.433105	0.440047
2	41483.8	0.422798	0.430248
3	41480.3	0.424406	0.431945
4	41496.2	0.424657	0.432177
5	41341.5	0.418893	0.426726
6	41342.2	0.417627	0.425447

So using FE for selection seems inconclusive in this case?

Here is the full code for reference:

Model code

function polynomial(x::Vector{Float64}; order::Int64=1)
    N = length(x)
    return [1; [x[i]^j for i in 1:N, j in 1:order][:]]
end


@model function linear_regression(params, X, y)

    # define priors
    θ ~ MvNormal(params[:μ_θ], params[:Σ_θ])
    τ ~ Gamma(1.0, 1.0)

    # define likelihood
    y .~ Normal(mean=dot.(X, θ), precision=τ)
end


function linear_regression_simple(df::DataFrame,
    x_cols::Array{Symbol},
    y_col::Symbol;
    g::Function,
    n_iter::Int64=10,
    showprogress::Bool=false
)

    # build design matrix
    X = mapslices(x -> g(x), Matrix(df[!, x_cols]); dims=2)
    n_cols = size(X, 2)

    # dependent variable
    y = df[!, y_col]

    # create dataset
    X_data = [collect(i) for i in eachrow(X)]
    y_data = collect(y)

    prior_params = Dict(
        :μ_θ => zeros(n_cols),
        :Σ_θ => 100.0 * diageye(n_cols),
    )

    # define model initialization
    init = @initialization begin
        μ(θ) = MvNormal(zeros(n_cols), 100.0 * diageye(n_cols))
        q(τ) = vague(InverseGamma)
    end

    constraints = @constraints begin
        q(θ, τ) = q(θ)q(τ)
    end

    # infer model parameters
    result = infer(
        model=linear_regression(params=prior_params),
        data=(X=X_data, y=y_data,),
        initialization=init,
        returnvars=(θ=KeepLast(), τ=KeepLast()),
        iterations=n_iter,
        constraints=constraints,
        free_energy=true,
        options=(limit_stack_depth=500,),
        showprogress=showprogress
    )

    # posterior mean
    θ = mode(result.posteriors[:θ])
    f(x) = dot(θ, g(x))

    return (f=f, θ_MAP=θ, result=result)
end

[wmkouw]

Your uncertainty calculation for the posterior predictive is correct. Are you using \alpha / \beta for \tau?

Nice to see it visualized like that.

So using FE for selection seems inconclusive in this case?

Well, it is conclusive, namely M=5 is the best. It nearly matches the minimizer for MAE test, but I would still argue that MAE is not an appropriate metric because it does not take uncertainty into account.

You might be worried about the fact that model selection depends on your choice of prior parameter. But that is always going to be the case in Bayesian inference. Your options are then:

1. Pick a specific prior parameter based on reasoning. For example, you could pick \Sigma_\theta = 1e4 I because you think that the coefficients you need are within the range of [-3e2 3e2] with probability 99.7% (68 - 95 - 99.7 rule for sd).
2. Increase your the scale for your prior covariance matrix to a really large value, such as 1e15 I. With such a level of uncertainty for your prior, its effect on the posterior will be negligible.

Does this help you? If so, can I close the discussion?

[langestefan]

Yes! Thanks a lot for explaining.