[avehtari]

Summary

This PR reduces overhead in projections for submodels without group-level or additive terms, and more for such submodels with the Gaussian family and the identity link. varsel() for a Gaussian model with 50 observations and 200 predictor terms gets about 7 times faster. In a Bernoulli example, varsel() is about 2 times faster. Memory use is also reduced.

This PR was assisted by Claude.

Overhead

1. For a submodel without group-level or additive terms, divmin() calls fit_glm_ridge_callback() once per projected draw (or cluster). The draw-wise formulas differ only in the response, but each call builds the model frame and the model matrix again, and subprd() builds the model matrix for new data once per draw.
2. For the Gaussian family with the identity link, each call then runs glm_ridge() (C++, with R callbacks for the pseudo-data), although the fit is one weighted ridge least-squares solve with the same left-hand side for all draws.

In a profile of varsel() for a Gaussian model with N = 50 and P = 200 (forward search, 20 clusters, nterms_max = 65, 400 draws for the evaluation), model.frame() took 43 % and model.matrix() took 40 % of the time.

Changes

The PR has two commits, which can be reviewed in order. Each commit changes some computations that master does S times, so that they are done only once.

1. Share one model matrix across the draw-wise fits.

• divmin() calls the new shared_model_matrix(), which builds the model matrix, the factor levels and the response names once. glm_ridge_callback_shared() gives them to fit_glm_ridge_callback() (new argument projpred_shared_x), which then takes only the response of the draw from data.
• subprd() calls the new subprd_subfit_shared() for several subfit objects with the same predictor terms. It builds one model matrix for the new data and computes the predictions of all draws with one matrix product.
• The model-matrix code of fit_glm_ridge_callback() and predict.subfit() moves without change into glm_model_matrix() and subfit_newdata_x().
• shared_model_matrix() returns NULL, and divmin() uses the draw-wise code as before, if the right-hand sides differ, if a response is not a plain column of data, or if there are missing values. A non-default projpred.glm_fitter or a custom div_minimizer also keeps the draw-wise code.

2. Solve the Gaussian identity-link submodels for all draws at once.

• The new glm_ridge_gauss_multi() standardizes the shared model matrix as glm_ridge() does, solves the normal equations with one right-hand side per draw, transforms the coefficients back to the original scale, and returns one subfit object per draw with the same elements as fit_glm_ridge_callback(). divmin() calls it before the draw-wise loop, also for a single draw.
• It returns NULL (draw-wise code) for other families and links, for a model matrix with a zero-variance column, for unusable observation weights, and if ... contains a formal argument of glm_ridge() other than family, weights or thresh (for example penalty or beta_init). It ignores the other arguments in ..., as fit_glm_ridge_callback() does. Thus varsel() and proj_linpred() use the same fit, and the test in test_varsel.R that compares them at .Machine$double.eps passes.

All draw-wise subfit objects refer to one model matrix in memory instead of S copies. saveRDS() and parallel workers still write one copy per fit.

Results

All rankings are identical to master. The largest absolute difference is 4.8e-10, in the test-set ELPD sums (about -1 800, so a relative difference of 3e-13). The ELPD summaries, the pointwise lppd and mu, the projected draws and proj_linpred() differ by at most 6e-14. The checks used a Gaussian model with N = 50 and P = 200, a Gaussian model with factor and interaction terms, and a Bernoulli model (commit 1 only, because commit 2 is Gaussian only).

Gaussian, N = 50, P = 200	master	commit 1	commit 2
varsel(), settings as above	818 s	268 s	109 s
LOO evaluation, cv_varsel(validate_search = FALSE)	112 s	26 s	5.0 s
test-set evaluation, varsel(d_test = )	113 s	26 s	6.7 s

A Bernoulli varsel() with N = 80 takes 1.3 s on master and 0.6 s with commit 1.

Peak memory (RSS) of one R process for a Gaussian reference model with P = 50 and 4 000 draws, commit 1:

case	master	commit 1
project(), n = 50 000, 4 000 projected draws	41.1 GB	19.2 GB
varsel(), n = 50 000, 4 000 projected draws	44.6 GB	34.8 GB
varsel(), n = 10 000, 4 000 projected draws	10.4 GB	6.5 GB

In 48 configurations (n from 100 to 50 000; 20 clusters, 400 or 4 000 projected draws; varsel() and project()), the peak memory and the run time were never higher than with master.

Tests

The new test_submodel_fit_shared.R checks that:

• glm_ridge_gauss_multi() gives the same subfit objects as the draw-wise fit_glm_ridge_callback() (coefficients to 1e-10; x, y, xlvls and the element names identical), for a model with a factor and an interaction and with observation weights;
• glm_ridge_gauss_multi() returns NULL for a non-identity link and for a non-Gaussian family;
• shared_model_matrix() returns NULL for different right-hand sides, for a response that is not a plain column, and for a missing predictor value;
• subprd_subfit_shared() and subprd() give the same predictions for new data as the draw-wise predict.subfit();
• project() with a Gaussian reference model uses glm_ridge_gauss_multi() (counted with trace()).

All five tests fail on master (functions not found). The full suite passes with NOT_CRAN=true: 258 tests, 55 857 expectations, 0 failures, the same 4 warnings as master. Tested with this change applied to master 73cf0ee (projpred 2.11.0.9000), R 4.6.1, brms 2.23.0, rstanarm 2.32.2, cmdstanr 0.9.0.9001, CmdStan 2.39.0, on CSC Roihu (Linux, AMD EPYC).

Interaction with #573

#573 adds the element terms to subfit objects and uses it in predict.subfit(). This PR creates subfit objects in glm_ridge_gauss_multi() and moves the model-matrix code of predict.subfit() into subfit_newdata_x(). <After #573 is merged, I will merge master into this branch and add terms in the new code paths.>

[avehtari]

More details about the two changes

Notation (all sections)

symbol	meaning
$n$	number of observations
$p$	number of columns of the submodel model matrix, without intercept
$S$	number of projected draws or clusters
$X$	$n \times p$ submodel model matrix, without intercept column
$w_i$	observation weight of observation $i$
$y_{is}$	response of draw $s$: the reference-model fit $\mu_{is}$ for the Gaussian projection
$\lambda$	ridge penalty, regul (default $10^{-4}$)

Master

For a submodel without group-level or additive terms, divmin() calls fit_glm_ridge_callback() once per draw $s = 1, \dots, S$:

for s in 1..S:
  fr    <- model.frame(y_s ~ rhs, data)        # per draw
  X     <- model.matrix(rhs, fr)               # per draw, the same for all s
  fit_s <- glm_ridge(X, y_s, family, weights)  # C++ IRLS with R callbacks

validate_response_formula() splits the multi-response formula into one formula per draw, so the draw-wise formulas differ only in the response column. Thus $X$ is the same for all $s$, but master builds it $S$ times. For predictions for new data, subprd() also builds the new-data model matrix once per draw.

Commit 1: one model matrix

X <- model.matrix(rhs, model.frame(y_1 ~ rhs, data))   # once
for s in 1..S:
  fit_s <- glm_ridge(X, data[[y_s]], family, weights)

The shared matrix is used if all of these conditions hold:

• the right-hand sides of all draw-wise formulas are identical;
• each response is a plain column name of data;
• there are no missing values in the responses or in the predictors (the model frame has $n$ rows).

Otherwise the code is the same as on master. For predictions for new data, subprd_subfit_shared() builds one new-data matrix $\tilde{X}$ and computes all draws with one product:

$$ \hat{\eta} = [\mathbf{1}, \tilde{X}] \begin{bmatrix} \alpha_1 & \cdots & \alpha_S \ \beta_{1} & \cdots & \beta_{S} \end{bmatrix} , $$

where $\alpha_s$ is the intercept and $\beta_s$ the coefficient vector of draw $s$. All subfit objects refer to the same $X$ in memory. R copies an object only when it is changed, so $S$ references cost one copy.

Commit 2: one solve for all draws (Gaussian, identity link)

glm_ridge() standardizes the columns of $X$ with weighted means $m_j$ and weighted standard deviations $d_j$:

$$ m_j = \frac{\sum_i w_i x_{ij}}{\sum_i w_i} , \qquad z_{ij} = \frac{x_{ij} - m_j}{d_j} . $$

Let $Z = [\mathbf{1}, (z_{ij})]$ and $W = \mathrm{diag}(w_1, \dots, w_n)$. For draw $s$, the ridge problem on the standardized scale is

$$ \theta_s = \arg\min_{\theta} \sum_i w_i \left(y_{is} - z_i^\top \theta\right)^2 + \lambda \sum_{j \ge 1} \theta_j^2 . $$

For the Gaussian family with the identity link, the IRLS pseudo-data equal $y_{is}$ and the IRLS weights equal $w_i$. So IRLS stops after one exact step, and $\theta_s$ solves the normal equations

$$ \left(Z^\top W Z + R\right) \theta_s = Z^\top W y_s , \qquad R = \mathrm{diag}(0, \lambda, \dots, \lambda) . $$

The matrix on the left does not depend on $s$. Commit 2 stacks the responses into the $n \times S$ matrix $Y$ and solves once:

$$ \Theta = \left(Z^\top W Z + R\right)^{-1} Z^\top W Y . $$

Back-transform to the original scale:

$$ \beta_{js} = \theta_{js} / d_j , \qquad \alpha_s = \theta_{0s} - \sum_j m_j \beta_{js} . $$

if family is gaussian(identity) and X is shared:
  Zw    <- cbind(1, standardize(X, w)) * sqrt(w)
  A     <- crossprod(Zw) + R                          # (p+1) x (p+1), once
  Theta <- solve(A, crossprod(Zw, Y * sqrt(w)))       # S right-hand sides
  back-transform each column; one subfit object per draw
else:
  commit 1 path

Commit 2 uses the commit 1 path for other families or links, for a column with zero variance ($d_j = 0$), for missing or wrong-length weights, and if ... contains a glm_ridge() argument other than family, weights or thresh.

Cost

step	master	commit 1	commit 2
model frame, model matrix	$S$ R calls	1 R call	1 R call
$Z^\top W Z$	$S \cdot O(np^2)$	$S \cdot O(np^2)$	$O(np^2)$
factorization	$S \cdot O(p^3)$	$S \cdot O(p^3)$	$O(p^3)$
right-hand sides	$S \cdot O(np)$	$S \cdot O(np)$	$O(npS + p^2 S)$
IRLS loop, R callbacks	$S$ calls	$S$ calls	none

In the forward search with small $n$ (Gaussian, $n = 50$, 200 terms, nterms_max = 65, $S = 20$ clusters), the search projects $\sum_{k=0}^{64} (200 - k) = 10920$ candidate submodels, each with 20 draw-wise fits. The per-call R overhead of model.frame() and model.matrix() dominates there. Commit 1 removes it. Commit 2 removes the $S$ C++ calls and their R callbacks.