HUIJZER.XYZ

Collinear Bayes

2021-11-17

In my post on Shapley values and multicollinearity, I looked into what happens when you fit a complex uninterpretable model on collinear or near-collinear data and try to figure out which features (variables) are important. The results were reasonable but not great. Luckily, there are still more things to try. Gelman et al. (2020) say that Bayesian models can do reasonably well on collinear data because they show high uncertainty in the estimated coefficients. Also, Bayesian models have a chance of fitting the data better as is beautifully shown in the Stan documentation. It can be quite tricky to implement though because a good parameterization is necessary (https://statmodeling.stat.columbia.edu/2019/07/07/collinearity-in-bayesian-models/).

Simulating data

Let's simulate some data with various columns are increasingly correlated with the outcome (and thus each other). Here, we assume that the data is centered around zero. This is easier for the Bayesian model to work with, but can often also make interpretation of the coefficients easier. There are various methods to rescale data, one is using MLDataUtils: rescale!. Note that rescale! bases the rescaling on the sample which is not recommended for small samples (Gelman, 2020). Instead, you can use knowledge that you have about the data such as the range of questionnaire scores or the weight of cars. Specifically, for example, it could be known for the data that the weight of a car is never below zero and unlikely to be above 3600 kg (8000 lbs); the weight of a Hummer H1.

begin
	using CairoMakie
	using CategoricalArrays: categorical
	using DataFrames: Not, DataFrame, select, stack, transform
	using GLM
	using Turing
	using Random: seed!
	using Statistics: rand, mean, cor
end
indexes = 1.0:150.0;
y_true(x) = x / last(indexes);
y_noise(x, corr_coefficient) = (corr_coefficient * y_true(x) - 0.5) + rand(Normal(0, 0.15));
df = let
    seed!(0)
    X = indexes
	A = y_noise.(indexes, 0)
    B = y_noise.(indexes, 0.05)
    C = y_noise.(indexes, 0.7)
    D = y_noise.(indexes, 1)
    E = y_noise.(indexes, 1)
    Y = y_noise.(indexes, 1)
    
    DataFrame(; X, A, B, C, D, E, Y)
end
X A B C D E Y
1.0 -0.398134 -0.646597 -0.287546 -0.619254 -0.638229 -0.39991
2.0 -0.375738 -0.613373 -0.559512 -0.557565 -0.485595 -0.263193
3.0 -0.552951 -0.345602 -0.584484 -0.633103 -0.457123 -0.687579
4.0 -0.520228 -0.623517 -0.557895 -0.365136 -0.430439 -0.610894
5.0 -0.412007 -0.714773 -0.549119 -0.703409 -0.41053 -0.543717
6.0 -0.4554 -0.695053 -0.467523 -0.206198 -0.493009 -0.34151
7.0 -0.490258 -0.475825 -0.538481 -0.638859 -0.180081 -0.411017
8.0 -0.516353 -0.300344 -0.282276 -0.53808 -0.370375 -0.419228
9.0 -0.577132 -0.712678 -0.550842 -0.379812 -0.406112 -0.434307
10.0 -0.26385 -0.599433 -0.557579 -0.333038 -0.4417 -0.427372
...
150.0 -0.554532 -0.428804 0.160114 0.616068 0.413887 0.688722

The data and the correlations look as follows:

Defining the model

This is a basic linear regression model similar to the one mentioned in the tutorials on https://turing.ml. The priors of this model are visualized below.

@model function linear_regression(X::Matrix, y)
	σ₂ ~ truncated(Normal(0, 100), 0, Inf)
	
	intercept ~ Normal(0, 0.4)

	n_features = size(X, 2)
    coef ~ MvNormal(n_features, 0.4)

	mu = intercept .+ X * coef
    y ~ MvNormal(mu, sqrt(σ₂))
end;
X = select(df, Not([:X, :Y]));	
model = let
	y = df.Y
	model = linear_regression(Matrix(X), y)
end;

Inspecting the prior

To verify that the priors are correctly set, we can use sample(model, Prior(), n_samples) from Turing.jl. This is shown below with the raw sample values on the left and the density plot for these values on the right.

n_samples = 1_000;

In this plot, everything looks good. On average, we expect our data to be zero (centered) and the variance looks reasonable. We expect the coefficients for the linear model to be between -0.5 and 0.5. Thanks to these priors, the sampler should have useful samples right from the start.

function mysample(model, sampler)
	n_chains = 3
	chns = sample(model, sampler, MCMCThreads(), n_samples, n_chains)
	return fix_names(chns)
end;

Estimating the parameters

When we fit the model, we have to decide on a sampler for this complex collinear case. NUTS is normally the best bet in Turing.jl, but let's first try HMC.

In the plots below, the different colors indicate different chains. All plots show good mixing and stationarity on the leftmost plots; the chains properly converged to the same outcome:

let
	chns = mysample(model, HMC(0.005, 10))
	plot_chain(chns)
end

Obtaining this outcome required setting the leapfrog size to a very low number. Normally, it is 0.05 or 0.1 which both did not work. What I mean by did not work is that the different chains did not converge, that is, gave different outcomes. Note that, thanks to the low leapfrog size, it took quite a few iterations for the chains to converge.

Let's try a more modern sampler, namely NUTS:

let
	chns = mysample(model, NUTS())
	plot_chain(chns)
end