Linear Regression

Turing is powerful when applied to complex hierarchical models, but it can also be put to task at common statistical procedures, like linear regression. This tutorial covers how to implement a linear regression model in Turing.

Set Up

We begin by importing all the necessary libraries.

# Import Turing.
using Turing

# Package for loading the data set.
using RDatasets

# Package for visualization.
using StatsPlots

# Functionality for splitting the data.
using MLUtils: splitobs

# Functionality for constructing arrays with identical elements efficiently.
using FillArrays

# Functionality for normalizing the data and evaluating the model predictions.
using StatsBase

# Functionality for working with scaled identity matrices.
using LinearAlgebra

# Set a seed for reproducibility.
using Random
Random.seed!(0);

setprogress!(false)

We will use the mtcars dataset from the RDatasets package. mtcars contains a variety of statistics on different car models, including their miles per gallon, number of cylinders, and horsepower, among others.

We want to know if we can construct a Bayesian linear regression model to predict the miles per gallon of a car, given the other statistics it has. Let us take a look at the data we have.

# Load the dataset.
data = RDatasets.dataset("datasets", "mtcars")

# Show the first six rows of the dataset.
first(data, 6)

6×12 DataFrame

Row	Model	MPG	Cyl	Disp	HP	DRat	WT	QSec	VS	AM	Gear	Carb
	String31	Float64	Int64	Float64	Int64	Float64	Float64	Float64	Int64	Int64	Int64	Int64
1	Mazda RX4	21.0	6	160.0	110	3.9	2.62	16.46	0	1	4	4
2	Mazda RX4 Wag	21.0	6	160.0	110	3.9	2.875	17.02	0	1	4	4
3	Datsun 710	22.8	4	108.0	93	3.85	2.32	18.61	1	1	4	1
4	Hornet 4 Drive	21.4	6	258.0	110	3.08	3.215	19.44	1	0	3	1
5	Hornet Sportabout	18.7	8	360.0	175	3.15	3.44	17.02	0	0	3	2
6	Valiant	18.1	6	225.0	105	2.76	3.46	20.22	1	0	3	1

size(data)

(32, 12)

The next step is to get our data ready for testing. We’ll split the mtcars dataset into two subsets, one for training our model and one for evaluating our model. Then, we separate the targets we want to learn (MPG, in this case) and standardize the datasets by subtracting each column’s means and dividing by the standard deviation of that column. The resulting data is not very familiar looking, but this standardization process helps the sampler converge far easier.

# Remove the model column.
select!(data, Not(:Model))

# Split our dataset 70%/30% into training/test sets.
trainset, testset = map(DataFrame, splitobs(data; at=0.7, shuffle=true))

# Turing requires data in matrix form.
target = :MPG
train = Matrix(select(trainset, Not(target)))
test = Matrix(select(testset, Not(target)))
train_target = trainset[:, target]
test_target = testset[:, target]

# Standardize the features.
dt_features = fit(ZScoreTransform, train; dims=1)
StatsBase.transform!(dt_features, train)
StatsBase.transform!(dt_features, test)

# Standardize the targets.
dt_targets = fit(ZScoreTransform, train_target)
StatsBase.transform!(dt_targets, train_target)
StatsBase.transform!(dt_targets, test_target);

Model Specification

In a traditional frequentist model using OLS, our model might look like:

\[ \mathrm{MPG}_i = \alpha + \boldsymbol{\beta}^\mathsf{T}\boldsymbol{X_i} \]

where \(\boldsymbol{\beta}\) is a vector of coefficients and \(\boldsymbol{X}\) is a vector of inputs for observation \(i\). The Bayesian model we are more concerned with is the following:

\[ \mathrm{MPG}_i \sim \mathcal{N}(\alpha + \boldsymbol{\beta}^\mathsf{T}\boldsymbol{X_i}, \sigma^2) \]

where \(\alpha\) is an intercept term common to all observations, \(\boldsymbol{\beta}\) is a coefficient vector, \(\boldsymbol{X_i}\) is the observed data for car \(i\), and \(\sigma^2\) is a common variance term.

For \(\sigma^2\), we assign a prior of truncated(Normal(0, 100); lower=0). This is consistent with Andrew Gelman’s recommendations on noninformative priors for variance. The intercept term (\(\alpha\)) is assumed to be normally distributed with a mean of zero and a variance of three. This represents our assumptions that miles per gallon can be explained mostly by our assorted variables, but a high variance term indicates our uncertainty about that. Each coefficient is assumed to be normally distributed with a mean of zero and a variance of 10. We do not know that our coefficients are different from zero, and we don’t know which ones are likely to be the most important, so the variance term is quite high. Lastly, each observation \(y_i\) is distributed according to the calculated mu term given by \(\alpha + \boldsymbol{\beta}^\mathsf{T}\boldsymbol{X_i}\).

# Bayesian linear regression.
@model function linear_regression(x, y)
    # Set variance prior.
    σ² ~ truncated(Normal(0, 100); lower=0)

    # Set intercept prior.
    intercept ~ Normal(0, sqrt(3))

    # Set the priors on our coefficients.
    nfeatures = size(x, 2)
    coefficients ~ MvNormal(Zeros(nfeatures), 10.0 * I)

    # Calculate all the mu terms.
    mu = intercept .+ x * coefficients
    return y ~ MvNormal(mu, σ² * I)
end

linear_regression (generic function with 2 methods)

With our model specified, we can call the sampler. We will use the No U-Turn Sampler (NUTS) here.

model = linear_regression(train, train_target)
chain = sample(model, NUTS(), 5_000)

┌ Info: Found initial step size
└   ϵ = 0.4

Chains MCMC chain (5000×24×1 Array{Float64, 3}):

Iterations        = 1001:1:6000
Number of chains  = 1
Samples per chain = 5000
Wall duration     = 11.66 seconds
Compute duration  = 11.66 seconds
parameters        = σ², intercept, coefficients[1], coefficients[2], coefficients[3], coefficients[4], coefficients[5], coefficients[6], coefficients[7], coefficients[8], coefficients[9], coefficients[10]
internals         = lp, n_steps, is_accept, acceptance_rate, log_density, hamiltonian_energy, hamiltonian_energy_error, max_hamiltonian_energy_error, tree_depth, numerical_error, step_size, nom_step_size

Summary Statistics
        parameters      mean       std      mcse    ess_bulk    ess_tail       ⋯
            Symbol   Float64   Float64   Float64     Float64     Float64   Flo ⋯

                σ²    0.2862    0.1649    0.0053    417.5864    168.9836    1. ⋯
         intercept    0.0001    0.1158    0.0019   3772.4461   2879.3778    1. ⋯
   coefficients[1]    0.2040    0.4602    0.0083   3050.5898   2824.1061    1. ⋯
   coefficients[2]    0.2158    0.4999    0.0104   2320.9800   2680.2657    0. ⋯
   coefficients[3]   -0.2398    0.3694    0.0071   2666.1863   2947.6896    1. ⋯
   coefficients[4]    0.1161    0.2628    0.0040   4243.9492   3435.8319    1. ⋯
   coefficients[5]   -0.8272    0.4415    0.0092   2304.6859   2330.2041    1. ⋯
   coefficients[6]    0.4129    0.3436    0.0067   2579.6572   2700.4003    1. ⋯
   coefficients[7]   -0.0325    0.2382    0.0042   3215.3818   2877.2899    1. ⋯
   coefficients[8]    0.2149    0.2766    0.0048   3318.0267   3003.2340    1. ⋯
   coefficients[9]    0.0014    0.3065    0.0052   3523.5073   2833.7321    0. ⋯
  coefficients[10]    0.0179    0.3225    0.0066   2384.0778   2961.1605    0. ⋯
                                                               2 columns omitted

Quantiles
        parameters      2.5%     25.0%     50.0%     75.0%     97.5%
            Symbol   Float64   Float64   Float64   Float64   Float64

                σ²    0.1030    0.1821    0.2456    0.3414    0.7082
         intercept   -0.2296   -0.0704    0.0015    0.0715    0.2268
   coefficients[1]   -0.7051   -0.0778    0.2034    0.4862    1.1212
   coefficients[2]   -0.7894   -0.0940    0.2079    0.5312    1.2000
   coefficients[3]   -0.9758   -0.4738   -0.2276   -0.0047    0.4934
   coefficients[4]   -0.4013   -0.0440    0.1165    0.2811    0.6299
   coefficients[5]   -1.7305   -1.1076   -0.8204   -0.5465    0.0660
   coefficients[6]   -0.2492    0.1895    0.4037    0.6328    1.1039
   coefficients[7]   -0.5110   -0.1851   -0.0314    0.1174    0.4404
   coefficients[8]   -0.3512    0.0425    0.2101    0.3863    0.7657
   coefficients[9]   -0.5992   -0.1941   -0.0047    0.1925    0.6234
  coefficients[10]   -0.6075   -0.1880    0.0157    0.2179    0.6688

We can also check the densities and traces of the parameters visually using the plot functionality.

plot(chain)

It looks like all parameters have converged.

AssertionError: AssertionError("Minimum ESS: 417.5864082260438 - not > 700")
AssertionError: Minimum ESS: 417.5864082260438 - not > 700
Stacktrace:
 [1] top-level scope
   @ ~/work/docs/docs/tutorials/05-linear-regression/index.qmd:201

Comparing to OLS

A satisfactory test of our model is to evaluate how well it predicts. Importantly, we want to compare our model to existing tools like OLS. The code below uses the GLM.jl package to generate a traditional OLS multiple regression model on the same data as our probabilistic model.

# Import the GLM package.
using GLM

# Perform multiple regression OLS.
train_with_intercept = hcat(ones(size(train, 1)), train)
ols = lm(train_with_intercept, train_target)

# Compute predictions on the training data set and unstandardize them.
train_prediction_ols = GLM.predict(ols)
StatsBase.reconstruct!(dt_targets, train_prediction_ols)

# Compute predictions on the test data set and unstandardize them.
test_with_intercept = hcat(ones(size(test, 1)), test)
test_prediction_ols = GLM.predict(ols, test_with_intercept)
StatsBase.reconstruct!(dt_targets, test_prediction_ols);

The function below accepts a chain and an input matrix and calculates predictions. We use the samples of the model parameters in the chain starting with sample 200.

# Make a prediction given an input vector.
function prediction(chain, x)
    p = get_params(chain[200:end, :, :])
    targets = p.intercept' .+ x * reduce(hcat, p.coefficients)'
    return vec(mean(targets; dims=2))
end

prediction (generic function with 1 method)

When we make predictions, we unstandardize them so they are more understandable.

# Calculate the predictions for the training and testing sets and unstandardize them.
train_prediction_bayes = prediction(chain, train)
StatsBase.reconstruct!(dt_targets, train_prediction_bayes)
test_prediction_bayes = prediction(chain, test)
StatsBase.reconstruct!(dt_targets, test_prediction_bayes)

# Show the predictions on the test data set.
DataFrame(; MPG=testset[!, target], Bayes=test_prediction_bayes, OLS=test_prediction_ols)

10×3 DataFrame

Row	MPG	Bayes	OLS
	Float64	Float64	Float64
1	30.4	26.1903	26.2335
2	17.8	18.8883	18.9211
3	32.4	27.0053	27.039
4	15.2	16.7462	16.6786
5	14.3	14.8959	14.8345
6	15.5	17.2442	17.235
7	24.4	20.7575	20.7586
8	22.8	24.859	24.8722
9	13.3	13.5221	13.4289
10	14.7	9.2232	9.17225

Now let’s evaluate the loss for each method, and each prediction set. We will use the mean squared error to evaluate loss, given by \[ \mathrm{MSE} = \frac{1}{n} \sum_{i=1}^n {(y_i - \hat{y_i})^2} \] where \(y_i\) is the actual value (true MPG) and \(\hat{y_i}\) is the predicted value using either OLS or Bayesian linear regression. A lower SSE indicates a closer fit to the data.

println(
    "Training set:",
    "\n\tBayes loss: ",
    msd(train_prediction_bayes, trainset[!, target]),
    "\n\tOLS loss: ",
    msd(train_prediction_ols, trainset[!, target]),
)

println(
    "Test set:",
    "\n\tBayes loss: ",
    msd(test_prediction_bayes, testset[!, target]),
    "\n\tOLS loss: ",
    msd(test_prediction_ols, testset[!, target]),
)

Training set:
    Bayes loss: 3.1350481645200405
    OLS loss: 3.1329887954334197
Test set:
    Bayes loss: 10.134886475620734
    OLS loss: 10.096483928907265

As we can see above, OLS and our Bayesian model fit our training and test data set about the same.