Faster augmented-data projection for ordinal and categorical submodels

NEWS.md

@@ -7,6 +7,10 @@ If you read this from a place other than <https://mc-stan.org/projpred/news/inde
 
 ## Major changes
 
+* Faster projections onto ordinal submodels without group-level terms (`brms::cumulative()`, `rstanarm::stan_polr()`): one C++ call fits all projected draws, with results equal to those of `MASS::polr()` up to its convergence tolerance. For example, `cv_varsel()` with `validate_search = TRUE` is about 23 times faster for a model with 6 categories. (GitHub: #583)
+
+* Faster projections onto categorical submodels without group-level terms (`brms::categorical()`): one C++ call fits all projected draws, with results equal to those of `nnet::multinom()` up to its convergence tolerance. For example, `cv_varsel()` with `validate_search = TRUE` is about 8 times faster for a model with 4 categories. (GitHub: #583)
+
 ## Minor changes
 
 ## Bug fixes and maintenance

R/RcppExports.R

@@ -1,6 +1,14 @@
 # Generated by using Rcpp::compileAttributes() -> do not edit by hand
 # Generator token: 10BE3573-1514-4C36-9D1C-5A225CD40393
 
+categ_fit_multi_c <- function(x, w_all, ncat, icpt, thresh, maxit) {
+    .Call('_projpred_categ_fit_multi_c', PACKAGE = 'projpred', x, w_all, ncat, icpt, thresh, maxit)
+}
+
+cumul_fit_multi_c <- function(x, w_all, ncat, offset, link, thresh, maxit) {
+    .Call('_projpred_cumul_fit_multi_c', PACKAGE = 'projpred', x, w_all, ncat, offset, link, thresh, maxit)
+}
+
 glm_elnet_c <- function(x, pseudo_obs, lambda, alpha, intercept, penalty, thresh, qa_updates_max, pmax, pmax_strict, beta, beta0, w0, as_updates_max = 50L) {
     .Call('_projpred_glm_elnet_c', PACKAGE = 'projpred', x, pseudo_obs, lambda, alpha, intercept, penalty, thresh, qa_updates_max, pmax, pmax_strict, beta, beta0, w0, as_updates_max)
 }

R/divergence_minimizers.R

@@ -694,6 +694,30 @@ divmin_augdat <- function(
   projpred_ws_aug <- unclass(projpred_ws_aug)
   attr(projpred_ws_aug, "ndiscrete") <- NULL
 
+  # Cumulative-link submodels without group-level terms: all draws in one C++
+  # call (see cumul_native_multi()):
+  if (identical(sdivmin, fit_cumul)) {
+    outdmin <- cumul_native_multi(formula, data = data, family = family,
+                                  ws_aug = projpred_ws_aug,
+                                  dot_args = list(...))
+    if (!is.null(outdmin)) {
+      check_conv(outdmin, rep(0L, length(outdmin)), do_check = do_check_conv)
+      return(outdmin)
+    }
+  }
+
+  # Multinomial-logit submodels without group-level terms: all draws in one C++
+  # call (see categ_native_multi()):
+  if (identical(sdivmin, fit_categ)) {
+    outdmin <- categ_native_multi(formula, data = data, family = family,
+                                  ws_aug = projpred_ws_aug,
+                                  dot_args = list(...))
+    if (!is.null(outdmin)) {
+      check_conv(outdmin, rep(0L, length(outdmin)), do_check = do_check_conv)
+      return(outdmin)
+    }
+  }
+
   if (is.null(getOption("projpred.parallel_proj_trigger")) &&
       !is.null(getOption("projpred.prll_prj_trigger"))) {
     warning(
@@ -868,6 +892,98 @@ fit_cumul <- function(formula, data, family, weights, ...) {
   return(fitobj)
 }
 
+# Fits the cumulative-link submodels of all draws in one C++ call
+# (cumul_fit_multi_c()) with the objective of MASS::polr(), instead of one
+# fit_cumul() call per draw. The augmented data hold C copies of the n
+# observations (one per response category), so this builds the model matrix
+# once, and the augmented weights of each draw form an n x C matrix. Newton
+# iterations replace the BFGS iterations of polr(), so the estimates agree with
+# those of polr() up to its convergence tolerance. The returned `polr` objects
+# hold the elements that predict.polr() (with `newdata`), coef() and
+# check_conv_s() use, without `fitted.values` and `lp`. Returns `NULL` for other
+# cases; divmin_augdat() then calls fit_cumul() per draw. Global option
+# `projpred.augdat_native = FALSE` switches this off.
+cumul_native_multi <- function(formula, data, family, ws_aug, dot_args) {
+  if (!getOption("projpred.augdat_native", TRUE)) {
+    return(NULL)
+  }
+  # fit_cumul() passes these arguments on to MASS::polr(); they are not
+  # reproduced here:
+  if (length(intersect(names(dot_args),
+                       c(methods::formalArgs(MASS::polr),
+                         methods::formalArgs(stats::optim))))) {
+    return(NULL)
+  }
+  method <- family$link
+  if (method == "logit") {
+    method <- "logistic"
+  } else if (method == "probit_approx") {
+    method <- "probit"
+  }
+  link_code <- match(method, c("logistic", "probit", "loglog", "cloglog",
+                               "cauchit"))
+  if (is.na(link_code)) {
+    return(NULL)
+  }
+  # As in MASS::polr():
+  mfr <- model.frame(formula, data = data)
+  if (nrow(mfr) != nrow(data) || !is.null(model.offset(mfr))) {
+    return(NULL)
+  }
+  trms <- attr(mfr, "terms")
+  x <- model.matrix(trms, mfr)
+  cons <- attr(x, "contrasts")
+  xint <- match("(Intercept)", colnames(x), nomatch = 0L)
+  if (xint == 0L) {
+    return(NULL)
+  }
+  x <- x[, -xint, drop = FALSE]
+  y <- model.response(mfr)
+  if (!is.factor(y)) {
+    return(NULL)
+  }
+  lev <- levels(y)
+  ncat <- length(lev)
+  n <- nrow(x) %/% ncat
+  ws_aug <- as.matrix(ws_aug)
+  # The augmented data have to be C copies of the n observations, in the order
+  # of the response categories:
+  if (ncat < 3L || n * ncat != nrow(x) || nrow(ws_aug) != nrow(x) ||
+      !identical(as.integer(y), rep(seq_len(ncat), each = n)) ||
+      anyNA(ws_aug) || any(ws_aug < 0)) {
+    return(NULL)
+  }
+  x1 <- x[seq_len(n), , drop = FALSE]
+  if (ncol(x1) > 0 && !all(x == x1[rep(seq_len(n), ncat), , drop = FALSE])) {
+    return(NULL)
+  }
+  nprj <- ncol(ws_aug)
+  out <- cumul_fit_multi_c(
+    unname(x1), matrix(ws_aug, nrow = n), ncat, rep(0, n), link_code,
+    thresh = 1e-12, maxit = 100L
+  )
+  xlvls <- .getXlevels(trms, mfr)
+  zeta_nms <- paste(lev[-ncat], lev[-1L], sep = "|")
+  pc <- ncol(x1)
+  # As in polr(), also for a submodel without predictors:
+  beta_nms <- colnames(x1) %||% character()
+  return(lapply(seq_len(nprj), function(s) {
+    beta <- setNames(out$beta[, s], beta_nms)
+    zeta <- setNames(out$zeta[, s], zeta_nms)
+    wsum <- sum(ws_aug[, s])
+    structure(
+      list(coefficients = beta, zeta = zeta, deviance = 2 * out$loss[s],
+           lev = lev, terms = trms, df.residual = wsum - pc - (ncat - 1L),
+           edf = pc + ncat - 1L, n = wsum, nobs = wsum,
+           call = quote(projpred_cumul_fit()), method = method,
+           convergence = out$convergence[s],
+           niter = c(f.evals = out$niter[s], g.evals = out$niter[s]),
+           contrasts = cons, xlevels = xlvls),
+      class = "polr"
+    )
+  }))
+}
+
 # Use ordinal::clmm() to fit multilevel submodels for the brms::cumulative()
 # family:
 fit_cumul_mlvl <- function(formula, data, family, weights, ...) {
@@ -971,6 +1087,102 @@ fit_categ <- function(formula, data, family, weights, ...) {
   return(fitobj)
 }
 
+# Fits the multinomial-logit submodels of all draws in one C++ call
+# (categ_fit_multi_c()) with the objective of nnet::multinom(), instead of one
+# fit_categ() call per draw. Like cumul_native_multi(), it builds the model
+# matrix once and uses the augmented weights as an n x C matrix. Newton
+# iterations replace the BFGS iterations of nnet::multinom() (by default at
+# most 100), so the estimates agree with those of nnet::multinom() up to its
+# convergence tolerance. The returned objects are copies of one `multinom`
+# object (from nnet::multinom() with `maxit = 0`) with the weights, the
+# objective and the fitted values of each draw. Returns `NULL` for other cases;
+# divmin_augdat() then calls fit_categ() per draw. Global option
+# `projpred.augdat_native = FALSE` switches this off.
+categ_native_multi <- function(formula, data, family, ws_aug, dot_args) {
+  if (!getOption("projpred.augdat_native", TRUE) || family$link != "logit") {
+    return(NULL)
+  }
+  # fit_categ() passes these arguments on to nnet::multinom(); they are not
+  # reproduced here:
+  if (length(intersect(names(dot_args),
+                       c(methods::formalArgs(nnet::multinom),
+                         methods::formalArgs(nnet::nnet))))) {
+    return(NULL)
+  }
+  mfr <- model.frame(formula, data = data)
+  if (nrow(mfr) != nrow(data) || !is.null(model.offset(mfr))) {
+    return(NULL)
+  }
+  trms <- attr(mfr, "terms")
+  x <- model.matrix(trms, mfr)
+  y <- model.response(mfr)
+  if (!is.factor(y)) {
+    return(NULL)
+  }
+  lev <- levels(y)
+  ncat <- length(lev)
+  n <- nrow(x) %/% ncat
+  ws_aug <- as.matrix(ws_aug)
+  # The augmented data have to be C copies of the n observations, in the order
+  # of the response categories:
+  if (ncat < 3L || n * ncat != nrow(x) || nrow(ws_aug) != nrow(x) ||
+      !identical(as.integer(y), rep(seq_len(ncat), each = n)) ||
+      anyNA(ws_aug) || any(ws_aug < 0)) {
+    return(NULL)
+  }
+  x1 <- x[seq_len(n), , drop = FALSE]
+  if (!all(x == x1[rep(seq_len(n), ncat), , drop = FALSE])) {
+    return(NULL)
+  }
+  # The template object (no iterations, so its weights are the zero start
+  # values):
+  tmpl <- do.call(nnet::multinom, list(formula = formula, data = data,
+                                       weights = ws_aug[, 1], maxit = 0,
+                                       trace = FALSE))
+  if (!identical(tmpl$vcoefnames, colnames(x1)) || tmpl$n[3L] != ncat ||
+      length(tmpl$wts) != ncat * (ncol(x1) + 1L)) {
+    return(NULL)
+  }
+  icpt <- match("(Intercept)", colnames(x1), nomatch = 0L) - 1L
+  nprj <- ncol(ws_aug)
+  out <- categ_fit_multi_c(unname(x1), matrix(ws_aug, nrow = n), ncat, icpt,
+                           thresh = 1e-12, maxit = 100L)
+  y_ind <- diag(ncat)[as.integer(y), , drop = FALSE]
+  dimnames(y_ind) <- dimnames(tmpl$fitted.values)
+  idx_aug <- rep(seq_len(n), ncat)
+  idx_rows <- seq_len(n)
+  # A loop instead of lapply() with an anonymous function, which the byte
+  # compiler would compile again for each projection:
+  outdmin <- vector("list", nprj)
+  for (s in seq_len(nprj)) {
+    coef_s <- out$coef[, , s]
+    dim(coef_s) <- dim(out$coef)[1:2]
+    # Weights as nnet::multinom() stores them (see nnet:::coef.multinom()):
+    # one row per category, the first column and the first row (reference
+    # category) zero:
+    wts_mat <- matrix(0, ncat, ncol(x1) + 1L)
+    wts_mat[-1L, -1L] <- t(coef_s)
+    eta <- cbind(0, x1 %*% coef_s)
+    eta_max <- eta[cbind(idx_rows, max.col(eta, ties.method = "first"))]
+    prbs <- exp(eta - eta_max)
+    prbs <- prbs / rowSums(prbs)
+    fitted_s <- prbs[idx_aug, , drop = FALSE]
+    dimnames(fitted_s) <- dimnames(tmpl$fitted.values)
+    fit_s <- tmpl
+    fit_s$wts <- as.vector(t(wts_mat))
+    fit_s$value <- out$loss[s]
+    fit_s$deviance <- 2 * out$loss[s]
+    fit_s$AIC <- fit_s$deviance + 2 * fit_s$edf
+    fit_s$convergence <- out$convergence[s]
+    fit_s$weights <- ws_aug[, s]
+    fit_s$fitted.values <- fitted_s
+    fit_s$residuals <- y_ind - fitted_s
+    fit_s$call <- quote(projpred_categ_fit())
+    outdmin[[s]] <- fit_s
+  }
+  return(outdmin)
+}
+
 # Use mclogit::mblogit() to fit multilevel submodels for the brms::categorical()
 # family:
 fit_categ_mlvl <- function(formula, projpred_formula_no_random,
@@ -1214,6 +1426,12 @@ subprd <- function(fits, newdata) {
 }
 
 subprd_augdat <- function(fits, newdata) {
+  # `polr` and `multinom` fits of the same submodel: linear predictors from the
+  # coefficients, for all draws at once (see subprd_augdat_direct()):
+  prd_direct <- subprd_augdat_direct(fits, newdata)
+  if (!is.null(prd_direct)) {
+    return(prd_direct)
+  }
   if (inherits(fits[[1]], "clmm")) {
     y_nm <- extract_terms_response(formula(fits[[1]]))$response
   }
@@ -1269,6 +1487,89 @@ subprd_augdat <- function(fits, newdata) {
   return(do.call(abind::abind, list(prd_list, rev.along = 0)))
 }
 
+# Linear predictors (on the link scale used by subprd_augdat()) of `polr` or
+# `multinom` fits that share the terms, the factor levels and the contrasts (as
+# the draw-wise fits of one submodel do). The model matrix for `newdata` is
+# built once, as in predict.polr() and predict.multinom(). For `polr`, the
+# linear predictors of the cumulative probabilities are zeta_j - eta, which
+# subprd_augdat() computes from the predicted category probabilities; for
+# `multinom`, they are the log odds against the first category. This function
+# computes both directly, so they differ from those of subprd_augdat() only by
+# rounding. Returns `NULL` for other cases.
+subprd_augdat_direct <- function(fits, newdata) {
+  fit_1 <- fits[[1]]
+  if (inherits(fit_1, "polr")) {
+    is_polr <- TRUE
+  } else if (inherits(fit_1, "multinom") && is.matrix(coef(fit_1))) {
+    is_polr <- FALSE
+  } else {
+    return(NULL)
+  }
+  same_cls <- function(fit) {
+    if (is_polr) {
+      inherits(fit, "polr") && !inherits(fit, "clmm")
+    } else {
+      inherits(fit, "multinom") && is.matrix(coef(fit))
+    }
+  }
+  # The terms are compared without the environment of their formula:
+  terms_noenv <- function(trms) {
+    attr(trms, ".Environment") <- NULL
+    trms
+  }
+  trms_1 <- terms_noenv(fit_1$terms)
+  if (!all(vapply(fits, function(fit) {
+    same_cls(fit) && identical(terms_noenv(fit$terms), trms_1) &&
+      identical(fit$xlevels, fit_1$xlevels) &&
+      identical(fit$contrasts, fit_1$contrasts)
+  }, logical(1)))) {
+    return(NULL)
+  }
+  trms <- delete.response(fit_1$terms)
+  mfr <- model.frame(trms, newdata, na.action = function(x) x,
+                     xlev = fit_1$xlevels)
+  if (!is.null(model.offset(mfr))) {
+    return(NULL)
+  }
+  x <- model.matrix(trms, mfr, contrasts.arg = fit_1$contrasts)
+  n <- nrow(x)
+  nprj <- length(fits)
+  if (is_polr) {
+    xint <- match("(Intercept)", colnames(x), nomatch = 0L)
+    if (xint > 0L) {
+      x <- x[, -xint, drop = FALSE]
+    }
+    if (!identical(colnames(x) %||% character(),
+                   names(coef(fit_1)) %||% character())) {
+      return(NULL)
+    }
+    zetas <- do.call(cbind, lapply(fits, "[[", "zeta"))
+    if (ncol(x) > 0) {
+      etas <- x %*% do.call(cbind, lapply(fits, coef))
+    } else {
+      etas <- matrix(0, n, nprj)
+    }
+    ncat_m1 <- nrow(zetas)
+    out <- array(NA_real_, dim = c(n, ncat_m1, nprj))
+    for (s in seq_len(nprj)) {
+      out[, , s] <- outer(-etas[, s], zetas[, s], "+")
+    }
+    cat_nms <- utils::head(fit_1$lev, -1)
+  } else {
+    coef_1 <- coef(fit_1)
+    if (!identical(colnames(x), colnames(coef_1))) {
+      return(NULL)
+    }
+    out <- array(NA_real_, dim = c(n, nrow(coef_1), nprj))
+    for (s in seq_len(nprj)) {
+      out[, , s] <- x %*% t(coef(fits[[s]]))
+    }
+    cat_nms <- rownames(coef_1)
+  }
+  dimnames(out) <- list(rownames(newdata), cat_nms, NULL)
+  return(out)
+}
+
 subprd_augdat_binom <- function(fits, newdata) {
   augprd <- subprd(fits, newdata = newdata)
   return(array(augprd, dim = c(nrow(augprd), 1L, ncol(augprd))))