Frequentist Ridge Regression

One of the problems that arises from equation 4 is that we need to select a value for \(\lambda\), the parameter that controls how much we learn from our training data. But how should we select an optimal \(\lambda\)? Of course, we could just pick a value based on our intuition, but that would likely lead to non-optimal solutions in many cases. Therefore, frequentist ridge regression typically relies on cross-validation (CV) to select \(\lambda\) (see here for more details on hyper-parameter tuning/optimization). CV proceeds by:

1. generating a grid of values for our hyper-parameter (\(\lambda\))
2. splitting the training data into k equivalently sized training sets (k folds),
3. fitting a model to k-1 folds and testing prediction accuracy on the left-out fold,
4. repeat step 3 until each fold has been left out exactly once, and
5. iterating steps 2-4 for each hyper-parameter in the grid defined by step 1

When we are finished, we select the value of \(\lambda\) in the grid that minimizes the CV error (step 3). The R code below proceeds through the above steps, using the R glmnet package to fit the model:

# Penalty term values to test
lambdas <- 10^seq(3, -2, by = -.1)

# Fit and cross-validate
fit_ridge <- glmnet(x = train_scale[,-8], y = train_scale[,8], 
                    alpha = 0, lambda = lambdas, intercept = F)
cv_ridge <- cv.glmnet(x = train_scale[,-8], y = train_scale[,8], 
                      alpha = 0, lambda = lambdas, intercept = F, 
                      grouped = FALSE)

# Find optimal penalty term (lambda)
opt_lambda <- cv_ridge$lambda.min

# Generate predictions with opitmal lambda model from cross-validation
y_pred <- predict(fit_ridge, s = opt_lambda, newx = test_scale[,-8])

# Plot cor(predicted, actual)
qplot(x = y_pred[,1], y = test_scale[,8],
      main = paste0("Frequentist Ridge Regression:\nEstimating ", expression(lambda), 
                    " with CV\nr = ", round(cor(test_scale[,8], y_pred[,1]), 2))) +
  #geom_point(aes(x = , y = train_scale[,8]))
  xlab("Model Predicted Pr(Acceptance)") +
  ylab("Actual Pr(Acceptance)") +
  theme_minimal(base_size = 20)

Woah, much better than traditional linear regression! Before diving into these results, however, let’s first explore Bayesian ridge regression.

Bayesian Ridge Regression

Bayesian Ridge regression differs from the frequentist variant in only one way, and it is with how we think of the \(\lambda\) penalty term. In the frequentist perspective, we showed that \(\lambda\) effectively tells our model how much it is allowed to learn from the data. In the Bayesian world, we can capture such an effect in the form of a prior distribution over our \(\beta\) weights. To reveal the extraordinary power hiding behind this simple idea, let’s first discuss Bayesian linear regression.

Bayesian models view estimation as a problem of integrating prior information with information gained from data, which we formalize using probability distributions. This differs from the frequntist view, which treats regression as an opimization problem that results in a point estimate (e.g., minimizing squared error). A Bayesian regression model takes the form of:

\[y_{i} | \beta_{0},\boldsymbol{x_{i}},\boldsymbol{\beta},\sigma \sim \mathcal{N}(\beta_{0} + \sum_{j=1}^{p}x_{ij}\beta_{j}, \sigma) \tag{5}\] Importantly, Bayesian models require us to specify a prior distribution for each parameter we seek to estimate. Therefore, we need to specify a prior on the intercept (\(\beta_{0}\)), slopes (\(\boldsymbol{\beta}\)), and error variance (\(\sigma\)) in equation 5. Since we are standardizing all of our predictors and outcome variable(s), we will ignore the intercept term. Then, we are left with \(\boldsymbol{\beta}\) and \(\sigma\). Crucially, our choice of prior distribution on \(\boldsymbol{\beta}\) is what determines how much information we learn from the data, analagous to the penalty term \(\lambda\) used for frequentist regularization.

Like before, you can gain an intuition for this behavior by imagnining extreme cases. For example, suppose that we assumed that \(\boldsymbol{\beta} \sim \mathcal{U}(-\infty,+\infty)\). That is, we assume that \(\boldsymbol{\beta}\) can take on any real-valued number, and every value is equally likely. In this case, the mode of the posterior distribution on each \(\beta\) weight will be equivalent to the maximum likelihood estimate of the respective \(\beta\) weight, and by extention the OLS estimate under the normality assumption (see our prior post on parameter estimation for more details). If an unbounded uniform distribution on \(\boldsymbol{\beta}\) produces the same behavior as traditional linear regression, it follows from our discussions above that it also allows us to maximally learn from the data.

Now, what can we do to restrict learning in a Bayesian model? Well, we can use a prior distribution that pulls the \(\beta\) weights toward 0 (unlike the unbounded uniform distribution), similar to the \(\lambda\) penalty parameter in frequentist regularization! Specifically, specifying the following prior distribution on \(\boldsymbol{\beta}\) is mathematically equivalent in expectation (see here) to using the ridge penalty in the frequentist model:

\[\boldsymbol{\beta} \sim \mathcal{N}(0, \sigma_{\beta})\tag{6}\]

There are two important points that the prior distribution in equation 6 conveys:

1. The normal distribution places very little prior probability on large-magnitude \(\beta\) weights (i.e. far from 0), while placing high prior probability on small-magnitude weights (i.e. near 0), and
2. \(\sigma_{\beta}\) controls how wide the normal distribution is, thus controlling the specific amount of prior probability placed on small- to large-magnitude \(\beta\) weights.

The Bayesian version differs most in that we jointly estimate the “penalization term” \(\sigma_{\beta}\) and \(\boldsymbol{\beta}\) in a single model, as opposed to using CV to select an optimal \(\lambda\) for frequentist regression. Because we are jointly estimating \(\sigma_{\beta}\) along with individual-level \(\beta\) weights, we can actually view Bayesian ridge regression as a simple hierarchical Bayesian model, where \(\sigma_{\beta}\) is interpreted as a group-level scaling parameter that is estimated from pooled information across individual \(\beta\) weights.

Below is the Stan code that specifies the Bayesian variant of ridge regression. Note that this code is based on Erp, Oberski, & Mulder (2019):

data{
    int N_train;             // "# training observations"
    int N_test;              // "# test observations"
    int N_pred;              // "# predictor variables"
    vector[N_train] y_train; // "training outcomes"
    matrix[N_train, N_pred] X_train; // "training data"
    matrix[N_test, N_pred] X_test;   // "testing data"
}
parameters{
    real<lower=0> sigma;   // "error SD"
    real<lower=0> sigma_B; // "hierarchical SD across betas"
    vector[N_pred] beta;   // "regression beta weights"
}
model{
  // "group-level (hierarchical) SD across betas"
  sigma_B ~ cauchy(0, 1);
  
  // "model error SD"
  sigma ~ normal(0, 1);
  
  // "beta prior (provides 'ridge' regularization)"
  beta ~ normal(0, sigma_B);
    
  // "model likelihood"
    y_train ~ normal(X_train*beta, sigma);
}
generated quantities{ 
    real y_test[N_test]; // "test data predictions"
    for(i in 1:N_test){
        y_test[i] = normal_rng(X_test[i,] * beta, sigma);
    }
}

The important parts of the above code are in the model{...} section. While the parameterization looks different from the frequentist model (equation 4), the behavior of the model is equivalent. In particular, we have a \(\sigma_{\beta}\) “penalty” term, can shrink the variation between \(\beta\) weights toward 0. Intuitively, as \(\sigma_{\beta} \rightarrow 0\), all (\(\beta\)) weights are shrunk toward 0 (no learning from data). Conversely, as \(\sigma_{\beta} \rightarrow \infty\), prior on \(\sigma_{\beta}\) becomes uniform. A uniform prior denotes no penalty at all, and we are left with traditional, non-regularized regression (albeit we get a joint posterior distribution as opposed to point estimates). It is worth noting that the mode of the resulting posterior distribution over \(\boldsymbol{\beta}\) will be equivalent to the maximum likelihood estimate when the prior on \(\boldsymbol{\beta}\) is uniform (for details, see MLE/MAP estimation described in a previous post).

Now, let’s try fitting the model!

# First, prepare data for Stan
stan_dat <- list(N_train = nrow(train_scale),
                 N_test  = nrow(test_scale),
                 N_pred  = ncol(train_scale)-1,
                 y_train = train_scale[,8],
                 X_train = train_scale[,-8],
                 X_test  = test_scale[,-8])

# Fit the model using Stan's NUTS HMC sampler
fit_bayes_ridge <- sampling(bayes_ridge, stan_dat, iter = 2000, 
                            warmup = 500, chains = 3, cores = 3)

# Extract posterior distribution (parameters and predictions)
post_ridge <- rstan::extract(fit_bayes_ridge)

# Compute mean of the posterior predictive distribution over test set predictors,
# which integrates out uncertainty in parameter estimates
y_pred_bayes <- apply(post_ridge$y_test, 2, mean)

# Plot correlation between posterior predicted mean and actual Pr(Acceptance)
qplot(x = y_pred_bayes, y = test_scale[,8],
      main = paste0("Bayesian Ridge Regression:\nEstimating ", expression(lambda), 
                    " Hierarchically\nr = ", round(cor(test_scale[,8], y_pred_bayes), 2))) +
    xlab("Model Predicted Pr(Acceptance)") +
    ylab("Actual Pr(Acceptance)") +
    theme_minimal(base_size = 20)

It is clear that the Bayesian ridge regression is giving results very similar (almost identical) to frequentist ridge regression as outlined above. Importantly, the Bayesian method also offers straightforward measures of uncertainty for both parameter estimates and predictions on the test data. The plot above integrates over the prediction uncertainty (i.e. integrating over the posterior predictive distribution). To really appreciate the prediction uncertainty, we compute the correlation between predicted and actual Pr(Acceptance) for each posterior sample, and then visualize the resulting distribution of correlations:

# Iterate through each posterior prediction and compute cor(predicted, actual),
# which preserves uncertainty in parameter estimates
y_pred_bayes2 <- apply(post_ridge$y_test, 1, function(x) cor(x, test_scale[,8]))

# Plot posterior predictive distribution of cor(predicted, actual)
qplot(x = y_pred_bayes2, geom = "histogram", bins = 50,
      xlab = "Correlation between\nPredicted and Actual Pr(Acceptance)") +
  theme_minimal(base_size = 20)

Clearly, there is good amount of uncertainty contained within the posterior predictive distribution, and the above plot shows it well. Importantly, this uncertainty reduce in proportion to how much data (i.e. observations) we have. Given that we only used 20 observations, these results are actually pretty good!

Another method is to visualize the scatterplot like above, but to include prediction intervals around the predicted mean estimates for each observation in the test set:

# `bayesplot` has many convenience functions for working with posteriors
color_scheme_set(scheme = "darkgray")
ppc_intervals(x = colMeans(post_ridge$y_test), y = test_scale[,8],
              yrep = post_ridge$y_test, prob = 0.95) +
  ggtitle("95% Posterior Prediction Intervals") +
  xlab("Model Predicted Pr(Acceptance)") +
  ylab("Actual Pr(Acceptance)") +
  theme_minimal(base_size = 20)

Looks great! We see that the 95% posterior prediction intervals contain the actual values in virtually all cases. Importantly, this visualization also makes it clear just how much uncertainty there is for each individual observation.

Next up, we will explore Traditional and Bayesian LASSO regressions, followed by a discussion of the benefits of Bayesian over frequentist regularization more broadly.

Frequentist LASSO Regression

The code below is very similar to what we had for traditional ridge regression, except that we will set \(\alpha = 1\) as opposed to \(\alpha = 0\):

# Scale training data (get rid of ID)
train_scale <- scale(train_dat[,2:9])

# Penalty term values to test
lambdas <- 10^seq(3, -2, by = -.1)

# Fit and cross-validate
fit_lasso <- glmnet(x = train_scale[,-8], y = train_scale[,8], 
                    alpha = 1, lambda = lambdas, intercept = F)
cv_lasso <- cv.glmnet(x = train_scale[,-8], y = train_scale[,8], 
                      alpha = 1, lambda = lambdas, intercept = F, 
                      grouped = FALSE)

# Find optimal penalty term (lambda)
opt_lambda_lasso <- cv_lasso$lambda.min

# Generate predictions on the test set
y_pred_lasso <- predict(fit_lasso, s = opt_lambda_lasso, newx = test_scale[,-8])

# Plot cor(predicted, actual)
qplot(x = y_pred_lasso[,1], y = test_scale[,8],
      main = paste0("r = ", round(cor(test_scale[,8], y_pred_lasso[,1]), 2))) +
  xlab("Model Predicted Pr(Acceptance)") +
  ylab("Actual Pr(Acceptance)") +
  theme_minimal(base_size = 20)

Clearly, there is not much of a performance difference between frequentist ridge and LASSO regression in this dataset. However, it is worth noting that the LASSO is often more useful when we expect the solution to be sparse—that is, we expect that many of the predictors in our model are mostly noise, and are goal is to identify more robust effects. In such a case, the LASSO does a good job of setting \(\beta\) weights on noisy predictors to exactly 0.

Bayesian LASSO Regression

Recall that for Bayesian ridge regression, we only needed to specifiy a normal prior distribution to the \(\beta\) weights that we were aiming to regularize. For Bayesian LASSO regression, the only difference is in the form of the prior distribution. Specifically, setting a Laplace (i.e. double-exponential) prior on the \(\beta\) weights is mathematically equivalent in expectation to the frequentist LASSO penalty (see e.g., here):

\[\boldsymbol{\beta} \sim \text{double-exponential}(0, \tau_{\beta})\tag{8}\] Here, \(\tau_{\beta}~(0 < \tau_{\beta} < \infty)\) is a scale parameter that controls how peaked the prior distribution is around the center (0 in this case). As \(\tau_{\beta} \rightarrow 0\), then the model assigns infinite weight on 0-valued \(\beta\) weights (i.e. no learning from data). Conversely, as \(\tau_{\beta} \rightarrow \infty\), the prior reduces to a uniform prior, thus leading to no regularization at all (i.e traditional regression, albeit with posterior distributions rather than point estimates). It can be useful to visualize what this prior distribution looks like over a single \(\beta\) weight. For example, centered at 0 with \(\tau_{\beta} = 0.5\):

qplot(x = seq(-3, 3, .1), y = rmutil::dlaplace(seq(-3, 3, .1), 0, .5), geom = "line") +
  ggtitle("Laplace Prior with tau = 0.5") +
  xlab(expression(beta[1])) +
  ylab("Density") +
  theme_minimal(base_size = 20) 

## Registered S3 method overwritten by 'rmutil':
##   method         from
##   print.response httr

Compared to the ridge prior, which is a normal distribution, it is clear that the Laplace distribiution places much more probability mass directly on 0, which produces the variable selection effect specific to LASSO regression. Note also that such peakedness explains why there are sharp corners in the frequentist penalty function (see the LASSO contour plot above).

Below is the Stan code that specifies this Bayesian variant of LASSO regression. Like for the Bayesian ridge regression above, this code is loosely based on Erp, Oberski, & Mulder (2019):

data{
    int N_train;             // "# training observations"
    int N_test;              // "# test observations"
    int N_pred;              // "# predictor variables"
    vector[N_train] y_train; // "training outcomes"
    matrix[N_train, N_pred] X_train; // "training data"
    matrix[N_test, N_pred] X_test;   // "testing data"
}
parameters{
    real<lower=0> sigma;   // "error SD"
    real<lower=0> sigma_B; // "(hierarchical) SD across betas"
    vector[N_pred] beta;   // "regression beta weights"
}
model{
  // "group-level (hierarchical) SD across betas"
  sigma_B ~ cauchy(0, 1);
  
  // "Prior on SD"
  sigma ~ normal(0, 1);
  
  // "beta prior (Note this is the only change!)"
  beta ~ double_exponential(0, sigma_B); 
    
  // "model likelihood"
    y_train ~ normal(X_train*beta, sigma);
}
generated quantities{ 
    real y_test[N_test]; // "test data predictions"
    for(i in 1:N_test){
        y_test[i] = normal_rng(X_test[i,] * beta, sigma);
    }
}

Time to fit the model and visualize results:

# Fit the model using Stan's NUTS HMC sampler
fit_bayes_lasso <- sampling(bayes_lasso, stan_dat, iter = 2000, 
                            warmup = 500, chains = 3, cores = 3)

# Extract posterior distribution (parameters and predictions)
post_lasso <- rstan::extract(fit_bayes_lasso)

# Compute mean of the posterior predictive distribution over test set predictors,
# which integrates out uncertainty in parameter estimates
y_pred_bayes_lasso <- apply(post_lasso$y_test, 2, mean)

# Plot correlation between posterior predicted mean and actual Pr(Acceptance)
qplot(x = y_pred_bayes_lasso, y = test_scale[,8],
      main = paste0("r = ", round(cor(test_scale[,8], y_pred_bayes_lasso), 2))) +
    xlab("Model Predicted Pr(Acceptance)") +
    ylab("Actual Pr(Acceptance)") +
    theme_minimal(base_size = 20)

Like before, the Bayesian LASSO is producing very similar results compared to the frequentist version.