Cross-Fit a Super Learner and Retain Fold-Specific Predictors

Cross-fitting proceeds by:

splitting the data into n_folds outer folds,
fitting a full super_learner() (with its own internal inner_n_folds-fold cross-validation to determine ensemble weights) on each outer training split,
predicting on the corresponding outer held-out fold,
retaining each fold-specific prediction function for later modified-data (e.g. intervention) prediction.

Usage

crossfit_super_learner(
  data,
  learners,
  formulas,
  y_variable = NULL,
  n_folds = 5,
  inner_n_folds = 5,
  determine_super_learner_weights = NULL,
  ensemble_or_discrete = c("ensemble", "discrete"),
  cv_schema = NULL,
  inner_cv_schema = NULL,
  outcome_type = c("continuous", "binary", "density", "multiclass"),
  extra_learner_args = NULL,
  cluster_ids = NULL,
  strata_ids = NULL,
  weights = NULL,
  rowids = NULL,
  loss_metric = NULL,
  use_complete_cases = FALSE
)

Arguments

data: Data to use in training a super_learner.
learners: A list of predictor/closure-returning-functions. See Details.
formulas: Either a single regression formula or a vector of regression formulas.
y_variable: Optional character name of the outcome variable; inferred from formulas when omitted.
n_folds: Number of outer cross-fitting folds.
inner_n_folds: Number of folds used by the inner super_learner() on each outer training split to estimate ensemble weights. Deliberately decoupled from n_folds.
determine_super_learner_weights: Function used by the inner super_learner() calls to determine ensemble weights. If NULL (default), inferred from outcome_type via default_determine_weights(), matching super_learner()'s own inference.
ensemble_or_discrete: Defaults to 'ensemble', but can be set to 'discrete'. Discrete super_learner() chooses only one of the candidate learners to have weight 1 in the resulting prediction algorithm, while ensemble super_learner() combines predictions from 1 or more candidate learners, with respective weights adding up to 1.
cv_schema: Function taking (data, n_folds) and returning a list with training_data and validation_data, each of length n_folds, used for the outer split. If omitted and cluster_ids or strata_ids are supplied, cv_origami_schema is used so the outer split respects them; otherwise cv_random_schema.
inner_cv_schema: Optional cv_schema for the inner super_learner() calls. If NULL (default), super_learner()'s own defaults apply (which respect cluster_ids/strata_ids, subset to each outer training fold).
outcome_type: One of 'continuous', 'binary', 'multiclass', or 'density'. outcome_type is used to infer the correct determine_super_learner_weights function if it is not explicitly passed.
extra_learner_args: A list of equal length to the learners with additional arguments to pass to each of the specified learners.
cluster_ids: (default: null) If specified, clusters will either be entirely assigned to training or validation (not both) in each cross-validation split.
strata_ids: (default: null) If specified, strata are balanced across training and validation splits so that strata appear in both the training and validation splits.
weights: If specified, (per observation) weights are used to indicate that risk minimization across models (i.e., the meta-learning step) should be targeted to higher weight observations.
rowids: (default: null) If specified, rowids are stored for use in out-of-fold prediction ($oof_predict()) and can also be used to validate/observe the behavior of the cross-validation folds construction.
loss_metric: Optional loss function used only to report cv_loss, the empirical loss of the cross-fitted out-of-fold predictions. If NULL (default), inferred from outcome_type via default_loss_metric().
use_complete_cases: (default: FALSE) If the data passed have any NA or NaN missing data, restrict the data to data[complete.cases(data),].

Value

An object of class "nadir_crossfit_sl", a list with:

$oof_predict(newdata, rowids = NULL): A prediction function that takes newdata with the same columns and rows as the training dataset and optionally rowids. If rowids are specified, the rows can be out of order compared to the original training data, but if rowids are not passed, then the ordering is assumed to be the same as the training data passed.
$oof_predictions: Numeric vector of out-of-fold predictions in the original row order of data (after any complete-case filtering; see $complete_rows). Rows never held out by cv_schema are NA.
$oof_predict_modified(modify): Full-length out-of-fold predictions after applying modify(newdata) to each held-out fold, e.g. modify = function(d) { d$treatment <- 1; d } to obtain cross-fitted $m_1$.
$oof_predict_fold(newdata_list = NULL, modify = NULL): List of per-fold prediction vectors; entry i is produced by the Super Learner trained without outer fold i.
$fold_assignments: Integer vector giving each row's outer validation fold (NA if never held out).
$fold_rows: List of row indices (into the analysis data) comprising each validation fold.
$sl_fits: List of the n_folds fitted super_learner() objects, e.g. for inspecting per-fold ensemble weights.
$cv_loss: Empirical loss of the cross-fitted predictions on held-out rows, or NA if it could not be computed.
$complete_rows: Indices into the originally supplied data retained after complete-case filtering (identity when no filtering occurred).
$warnings_from_fold_predictions: Present only if non-empty: warnings captured while predicting on the outer held-out folds.
$warnings_from_inner_super_learners: Present only if non-empty: a per-fold list of the warning fields captured by each inner super_learner() (see ?super_learner).
$warning_learners: Present only if non-empty: names of learners that signaled warnings in any fold.

plus y_variable, outcome_type, n_folds, inner_n_folds, training_data, and validation_data.

Details

Each observation's out-of-fold prediction therefore comes from an ensemble whose base learners and ensemble weights were estimated entirely without that observation, which is the property required of cross-fitted nuisance estimators in semiparametric workflows (AIPW, TMLE, DML). Typical use is to obtain out-of-fold fits of an outcome model $m$, its interventional counterparts $m_1, m_0$ (via $predict_modified()), and a propensity model $g$, each from one crossfit_super_learner() call per nuisance.

Parallelism

Make sure to load the future package first using library(future).

Fitting is parallelized over outer folds through future.apply::future_lapply(). With plan(multisession) only the outer folds run in parallel and each inner super_learner() runs sequentially inside its worker. To also parallelize within each outer fold, declare a nested topology, e.g. plan(list(tweak(multisession, workers = 2), tweak(multisession, workers = I(4)))); the product across levels should not exceed future::availableCores().

Examples

# \donttest{
cf <- crossfit_super_learner(
  data = mtcars,
  formulas = mpg ~ disp + hp + am,
  learners = list(mean = lnr_mean, lm = lnr_lm, rf = lnr_rf),
  n_folds = 5,
  rowids = seq_len(nrow(mtcars))
)

# basic usage
coef(cf)
#>        mean        lm        rf
#> fold_1    0 0.6813384 0.3186616
#> fold_2    0 0.4418869 0.5581131
#> fold_3    0 0.1793374 0.8206626
#> fold_4    0 0.2721632 0.7278368
#> fold_5    0 0.3290426 0.6709574
head(fitted(cf))
#> [1] 21.98114 24.06128 26.18548 19.42724 15.65984 19.93665
head(residuals(cf))
#> [1] -0.9811405 -3.0612757 -3.3854840  1.9727577  3.0401641 -1.8366513
formula(cf)
#> mpg ~ disp + hp + am
#> <environment: 0x32dc9beb0>
nobs(cf)
#> [1] 32
summary(cf)
#> Summary of Cross-fitted Super Learner
#>   outcome: mpg (continuous);  n = 32
#>   outer folds: 5;  inner CV folds: 5
#> 
#> Held-out loss by outer fold:
#>  fold n_validation   loss
#>     1            6  7.402
#>     2            7  5.448
#>     3            6  4.279
#>     4            7  9.886
#>     5            6 18.670
#> Overall cross-fitted loss: 9.046
#> 
#> Ensemble weight stability across outer folds:
#>  learner mean_weight sd_weight min_weight max_weight n_folds_present
#>       rf      0.6192    0.1931     0.3187     0.8207               5
#>       lm      0.3808    0.1931     0.1793     0.6813               5
#>     mean      0.0000    0.0000     0.0000     0.0000               5

# but what cf is really for is out of fold predictions:
cf$oof_predict(mtcars, rowids = seq_len(nrow(mtcars)))
#>  [1] 21.98114 24.06128 26.18548 19.42724 15.65984 19.93665 13.56866 22.78791
#>  [9] 21.19578 20.09891 19.81215 17.34134 16.16197 17.34134 15.82786 15.15485
#> [17] 11.98064 28.30233 27.17797 27.05730 21.87528 17.05990 16.84988 13.63369
#> [25] 15.55110 28.29871 23.61663 23.32871 17.92060 22.38309 17.52743 24.02416

cf$oof_predictions # cross-fitted \hat m(X_i) prediction models
#>  [1] 21.98114 24.06128 26.18548 19.42724 15.65984 19.93665 13.56866 22.78791
#>  [9] 21.19578 20.09891 19.81215 17.34134 16.16197 17.34134 15.82786 15.15485
#> [17] 11.98064 28.30233 27.17797 27.05730 21.87528 17.05990 16.84988 13.63369
#> [25] 15.55110 28.29871 23.61663 23.32871 17.92060 22.38309 17.52743 24.02416

m1 <- cf$oof_predict_modified(function(d) {
  d$am <- 1
  d
}) # \hat m(1, X_i)
m0 <- cf$oof_predict_modified(function(d) {
  d$am <- 0
  d
}) # \hat m(0, X_i)

cf$cv_loss # honest empirical loss
#> [1] 9.045843
lapply(cf$sl_fits, `[[`, "learner_weights") # per-fold ensemble weights
#> [[1]]
#>      mean        lm        rf 
#> 0.0000000 0.6813384 0.3186616 
#> 
#> [[2]]
#>      mean        lm        rf 
#> 0.0000000 0.4418869 0.5581131 
#> 
#> [[3]]
#>      mean        lm        rf 
#> 0.0000000 0.1793374 0.8206626 
#> 
#> [[4]]
#>      mean        lm        rf 
#> 0.0000000 0.2721632 0.7278368 
#> 
#> [[5]]
#>      mean        lm        rf 
#> 0.0000000 0.3290426 0.6709574 
#> 
# }