0.632+ Bootstrap Prediction Intervals in R from a caret-trained model

Question

I'm trying to write a function in R that calculates a central prediction and upper and lower prediction intervals from a trained caret model (i.e., a "train" object) using the 0.632+ Bootstrap approach.

In this effort, I'm attempting to follow a Python example (https://www.saattrupdan.com/posts/2020-03-01-bootstrap-prediction) as a guide. However, I'm having trouble replicating it in R. Any guidance would be appreciated.

My function is supposed to take a trained caret model, the training data, and new data as input and return prediction intervals. However, at present, my prediction interval values are not correct.

As highlighted in a comment by Mark Rieke, one issue is that the entire 0.632+ procedure needs to be done for every bootstrap split, but my current code fails to do this.

Here's my current code:

library(caret)

# Set the random seed for reproducibility
set.seed(123)

# Generate data
n <- 100
explainer <- runif(n)
y <- 1 + 0.2 * explainer + rnorm(n)
data <- data.frame(explainer, y)

# Fit linear regression models
fit_simple <- lm(y ~ explainer) # A plain old linear model
fit_caret <- train(
  y = y,
  x = data.frame(explainer),
  method = "lm"
) # An identical model, but fit using caret

new_data <- data.frame(explainer = runif(15, min = -10, max = 10))

# Function to calculate prediction intervals using 0.632+ Bootstrap
calculate_prediction_intervals <- function(model, new_data, alpha = 0.05) {
  # Extract training data and outcomes from the model
  X_train <- base::subset(model$trainingData, select = -c(.outcome))
  y_train <- as.numeric(model$trainingData$.outcome)
  n <- nrow(X_train)
  nbootstraps <- as.integer(sqrt(n))
  
  # Initialize matrices to store bootstrap predictions and validation residuals
  bootstrap_preds <- matrix(0, nrow(new_data), nbootstraps)
  val_residuals <- matrix(0, n, nbootstraps)
  
  for (b in 1:nbootstraps) {
    train_idxs <- sample(1:n, n, replace = TRUE)
    val_idxs <- setdiff(1:n, train_idxs)
    
    # Fit a bootstrap sample of the model
    fit_b <- train(
      y = y_train[train_idxs],
      x = X_train[train_idxs, , drop = FALSE],
      method = model$method,
      tuneGrid = model$bestTune,
      trControl = trainControl(method = "none", savePredictions = FALSE)
    )
    
    # Compute validation set predictions and residuals
    preds_val <- predict(fit_b, newdata = X_train[val_idxs, , drop = FALSE])
    val_residuals[val_idxs, b] <- y_train[val_idxs] - preds_val
    
    # Compute bootstrap predictions on new data
    preds_new <- predict(fit_b, newdata = new_data)
    bootstrap_preds[, b] <- preds_new
  }
  
  # Center the bootstrap predictions and residuals
  bootstrap_preds <- bootstrap_preds - colMeans(bootstrap_preds)
  val_residuals <- val_residuals - colMeans(val_residuals)
  
  # Fit the original model to the full training data
  fit <- train(
    y = y_train,
    x = X_train,
    method = model$method,
    tuneGrid = model$bestTune,
    trControl = trainControl(method = "none", savePredictions = FALSE)
  )
  
  preds <- predict(fit, newdata = X_train)
  train_residuals <- y_train - preds
  
  # Calculate various values needed for 0.632+ Bootstrap
  no_information_error <- mean(abs(sample(y_train) - sample(preds)))
  generalization <- abs(colMeans(val_residuals) - mean(train_residuals))
  no_information_val <- abs(no_information_error - train_residuals)
  relative_overfitting_rate <- mean(generalization / no_information_val)
  weight <- 0.632 / (1 - 0.368 * relative_overfitting_rate)
  
  # Calculate prediction residuals
  residuals <- (1 - weight) * train_residuals + weight * colMeans(val_residuals)
  
  # Calculate prediction percentiles
  percentiles <- apply(bootstrap_preds, 1, function(x) {
    quantile(x + residuals, probs = c(alpha / 2, 1 - alpha / 2))
  })
  
  # Create a data frame with predictions, lower, and upper limits
  result <- data.frame(
    fit = predict(fit, newdata = new_data),
    lwr = percentiles[1, ],
    upr = percentiles[2, ]
  )
  
  return(result)
}

My code fails to ~reproduce the expected prediction intervals for a linear model. Increasing the number of bootstrap resamples doesn't help this. Can you help me find where I went wrong?

> calculate_prediction_intervals(fit_caret, new_data)
           fit        lwr       upr
1   1.18302967 -0.2597420 1.1699486
2   2.07894173 -1.4669930 7.0949444
3   0.71611677 -2.1804343 0.4431974
4   1.37767478 -0.6438284 2.5235400
5   1.68312227 -0.9393278 4.4294951
6   1.71845385 -1.0413210 4.8058089
7   0.06639059 -6.7192473 1.1929259
8   0.58836348 -3.2036975 0.7598031
9   1.55414870 -0.7131324 3.5583779
10  0.04536204 -6.8536552 1.2401264
11  1.76387322 -1.0177667 5.0307556
12 -0.01836307 -7.4146538 1.4246235
13  1.29583653 -0.4646119 2.0345750
14  0.18768121 -5.8312821 1.0571434
15  1.33552830 -0.4831878 2.0921489
> predict(fit_simple, newdata =  new_data, interval= "prediction")
           fit        lwr      upr
1   1.18302967 -0.9262779 3.292337
2   2.07894173 -4.5686088 8.726492
3   0.71611677 -2.0877607 3.519994
4   1.37767478 -1.4345098 4.189859
5   1.68312227 -2.6904110 6.056656
6   1.71845385 -2.8512314 6.288139
7   0.06639059 -6.2672902 6.400071
8   0.58836348 -2.8285939 4.005321
9   1.55414870 -2.1238365 5.232134
10  0.04536204 -6.4117391 6.502463
11  1.76387322 -3.0606644 6.588411
12 -0.01836307 -6.8508475 6.814121
13  1.29583653 -1.1747848 3.766458
14  0.18768121 -5.4394392 5.814802
15  1.33552830 -1.2942424 3.965299 

I am aware that alternatives to the method I am trying to replicate exist, e.g., conformal inference or even simply adding the raw residuals to predictions, but I'm hoping for a specific application here. The approach I am after should generally replicate the methods of https://arxiv.org/abs/2201.11676, similar to other approaches that have used tidymodels, e.g., https://www.bryanshalloway.com/2021/04/05/simulating-prediction-intervals/ and the workboots package (https://markjrieke.github.io/workboots/).

I plan to use this function on more complicated models (i.e., many predictors, not just linear models) from caret trained with the x and y data specified. I'm not using the formula method in caret. Due to this complexity, approaches that only work for linear models won't do the trick, either.

Answer 1

Following the approach from the Workboots package, with only a few adjustments to work with the caret objects, we can get all the bootstrapped predictions (with the corrected residuals added), the quantiles of predictions for a given alpha, and the fit on new data using the following code.

Note: This is slightly different from the original Python effort in formulation, though it's the same in effect.

# Function to generate prediction intervals for a caret model using bootstrapping
predict_caret_boots <-
  function(model,
           n = 2000,
           alpha = 0.05,
           new_data) {
    # Extract training data and outcomes from the model
    X_train <- base::subset(model$trainingData, select = -c(.outcome))
    y_train <- as.numeric(model$trainingData$.outcome)
    
    # Initialize a list to store predictions
    preds_list <- list()
    
    # Loop through n bootstrap resamples
    for (i in 1:n) {
      # Create a bootstrap sample
      train_idxs <- sample(length(y_train), replace = TRUE)
      boot_X_train <- X_train[train_idxs, , drop = FALSE]
      boot_y_train <- y_train[train_idxs]
      boot_X_oob <- X_train[-train_idxs, , drop = FALSE]
      boot_y_oob <- y_train[-train_idxs]
      
      # Fit a model on the bootstrap sample
      fit_b <- train(
        y = boot_y_train,
        x = boot_X_train,
        method = model$method,
        tuneGrid = model$bestTune,
        trControl = trainControl(method = "none", savePredictions = FALSE)
      )
      
      # Make predictions on the new data
      preds <- predict(fit_b, newdata = new_data)
      
      # Make predictions on training data
      preds_train <- predict(fit_b, newdata = boot_X_train)
      
      # Make predictions on OOB data
      preds_oob <- predict(fit_b, newdata = boot_X_oob)
      
      # Calculate training residuals
      resids_train <- boot_y_train - preds_train
      resids_train <- resids_train - mean(resids_train)
      
      # Calculate OOB residuals
      resids_oob <- boot_y_oob - preds_oob
      resids_oob <- resids_oob - mean(resids_oob)
      
      # Calculate no-information error rate (rmse_ni) with RMSE as the loss function
      combos <- tidyr::crossing(boot_y_train, preds_train)
      rmse_ni <- caret::RMSE(combos$preds_train, combos$boot_y_train)
      
      # Calculate overfit rate
      rmse_oob <- caret::RMSE(boot_y_oob, preds_oob)
      rmse_train <- caret::RMSE(boot_y_train, preds_train)
      overfit <- (rmse_oob - rmse_train) / (rmse_ni - rmse_train)
      
      # Calculate weight (if overfit = 0, weight = .632 & residual used will just be .632)
      # Use the actual proportion of distinct training/OOB samples, rather than the average of 0.632/0.368
      prop_368 <- length(boot_y_oob) / length(boot_y_train)
      prop_632 <- 1 - prop_368
      weight <- prop_632 / (1 - (prop_368 * overfit))
      
      # Determine residual std.dev based on weight
      sd_oob <- stats::sd(resids_oob)
      sd_train <- stats::sd(resids_train)
      sd_resid <- weight * sd_oob + (1 - weight) * sd_train
      
      # Add residuals to predictions
      preds <- preds + stats::rnorm(length(preds), 0, sd_resid)
      
      # Create a data frame with predictions and add it to the list
      preds_df <- data.frame(fit = preds)
      preds_list[[i]] <- preds_df
    }
    
    # Calculate quantiles for each row of preds_list
    
    preds_list <- data.frame(preds_list)
    
    quantiles <-
      apply(preds_list, 1, function(row)
        quantile(row, probs = c(alpha / 2, 1 - alpha / 2)))
    
    # Get the central fit, too
    fit_new <- predict(model, new_data)
    
    
    result <- list(
      preds = data.frame(preds_list),
      quantiles = t(data.frame(quantiles)),
      fit = data.frame(fit_new)
    )
    
    return(result)
  }

A little adjustment to this function could help it explicitly handle preprocessing options from caret, etc. But for now, this appears to do the trick beautifully!