[avehtari]

Summary

For Gaussian submodels with the identity link, this PR computes the forward search and the evaluation in the LOO CV with validated search from cross-products of the model matrix. The cost per candidate submodel then does not depend on the number of observations. varsel() with forward search for a Gaussian model with 50 observations and 200 predictor terms gets about 6 times faster. cv_varsel(validate_search = TRUE) for a Gaussian model with 40 observations and 15 predictor terms gets about 11 times faster. The rankings are identical to master.

This PR was assisted by Claude.

Overhead

A forward search projects the reference model onto every candidate submodel in every step: about P × nterms_max projections, each with one fit per projected draw. cv_varsel(validate_search = TRUE) repeats the search for each left-out observation and then evaluates every submodel along that observation's path. For a Gaussian submodel with the identity link, each fit is a weighted ridge least-squares problem. The search needs only the solution and the part of the cross-entropy that differs between candidates; the evaluation also needs the projected dispersion and the predictions. All of these follow from cross-products of the model matrix, which projpred can compute once per search or evaluation.

Changes

The PR has five commits that change the code, numbered below. One more commit, between commits 3 and 4, only adds the PR number to NEWS.md.

1. Ranking the candidates from cross-products.

• search_forward() first calls the new search_forward_gauss(). For the columns J of a candidate and draw s, glm_ridge() solves (G[J, J] + Λ) β_s = C[J, s], with G = Xc' W Xc, C = Xc' W M and Λ the penalty on the original scale. search_forward_gauss() ranks the candidates by −Σ_s w_s β_s' (G[J, J] + 2Λ) β_s, which is the part of ce_gauss() that differs between candidates.
• In each step, it factors the columns of the chosen terms once and extends the factor for each candidate (block Cholesky).
• The candidates come from select_possible_terms_size(), so compound search terms such as "x1 + x2 + x1:x2" work as in search_forward().
• After the ranking, proj_to_submodl() computes the submodel fits along the ranking, which search_forward() stores as before.
• search_forward_gauss() returns NULL, and search_forward() continues as before, in these cases:
  • another family or link, or the latent or augmented-data projection;
  • a custom divergence minimizer or projpred.glm_fitter;
  • non-zero offsets;
  • a formal argument of glm_ridge() other than thresh in the divergence-minimizer arguments;
  • group-level or smooth terms, or a factor in an interaction term (its dummy coding depends on the other terms of a submodel);
  • missing values or a zero-variance column.
• The new global option projpred.search_gauss_crossprod = FALSE turns this off. I documented it with the other internal global options.

2. Evaluation in the LOO CV with validated search. The new file R/gauss_crossprod.R collects the condition check (gauss_crossprod_regul()), the cross-products (gauss_crossprod_design()) and the new perf_eval_gauss(). perf_eval() calls perf_eval_gauss() if refit_prj = TRUE and the projected draws are reweighted for one left-out observation. For each submodel size along the path, perf_eval_gauss() computes the coefficients, the projected dispersion (dis_gauss()), the cross-entropy (ce_gauss()) and the prediction for the left-out observation. The header of the file gives the identities. The submodels along a path are nested, so the Cholesky factor grows block by block. Other evaluations do not change. So varsel() and proj_linpred() keep using the same code for their test-data summaries, which a test compares at machine precision.

3. All one-column candidates at once. Most candidates add one model-matrix column. For these, the block Cholesky step reduces to scalars, and search_forward_gauss() computes the criterion for all of them with a few matrix products instead of one factorization per candidate.

4. No submodel fits along the ranking in the validated search when the evaluation refits. After the ranking, search_forward_gauss() computes the submodel fits along the ranking with proj_to_submodl(), because search_forward() stores them in search_path$outdmins. In cv_varsel(validate_search = TRUE) with refit_prj = TRUE, the evaluation of each fold projects again and does not use these fits. In a profile of this case, they took 23 % of the time.

• .select(), search_forward() and search_forward_gauss() get the new argument return_outdmins = TRUE. With FALSE, search_forward_gauss() does not compute the fits and returns outdmins = NULL. The draw-wise search in search_forward() gets the fits during the search, so it ignores the argument.
• loo_varsel() and kfold_varsel() set return_outdmins = !refit_prj for the search of each fold. All other calls keep the default, so varsel(), project() and refit_prj = FALSE work as before.

5. The latent projection. The latent projection projects onto Gaussian submodels with the identity link on the latent scale, with the same submodel fitter and the same cross-entropy as the traditional Gaussian projection. Commits 1 to 4 excluded it; commit 5 removes this exclusion. The evaluation passes the latent-scale predictions and dispersions to weighted_summary_means() with the latent family, as the draw-wise evaluation does, so the performance on the original response scale does not change. For a latent projection of a Bernoulli model with 100 observations and 60 predictor terms (varsel(), 20 search and 20 evaluation clusters, nterms_max = 25, one workstation), the time goes from 45.1 s to 3.9 s with an identical ranking.

Results

The rankings are identical to master in all checks:

• a Gaussian model with N = 50 and P = 200;
• a model with a factor, a numeric interaction and a poly() term, with forward and L1 search, 1, 5 or 20 clusters, cv_varsel(validate_search = FALSE) and project();
• cv_varsel(validate_search = TRUE).

Outside the validated LOO evaluation, all numbers are identical. In the validated LOO evaluation, the summaries differ from master by at most 7e-14 (relative), because the dispersion, the cross-entropy and the predictions come from other but equivalent computations. When two candidates have almost the same criterion, the two computations could select different candidates. This did not happen in the checks.

Gaussian, N = 50, P = 200	master	this PR
varsel(), forward search, 20 clusters, nterms_max = 65, 400 evaluation draws	809 s	128 s
forward search, 1 cluster	68 s	16 s
LOO evaluation, cv_varsel(validate_search = FALSE)	109 s	109 s
test-set evaluation, varsel(d_test = )	111 s	111 s

The two evaluations reuse the search, so they do not change. The benchmark (CSC Roihu, one node) also gives identical results (difference 0) for a model with a factor and an interaction and for a Bernoulli model.

Gaussian, cv_varsel(validate_search = TRUE), n = 40, P = 15 (with a factor and an interaction), nloo = 40, 5 clusters, 50 evaluation draws	time
master	30.7 s
this PR, commits 1 to 3	4.8 s
this PR, all commits	2.7 s

These times are medians of 3 runs on one Linux workstation.

Tests

The new test_search_gauss.R (about 8 seconds, no Stan fits) checks that:

• varsel() gives the same rankings, performances and stored submodel fits with the option on and off (1 and 5 clusters);
• cv_varsel(validate_search = TRUE) gives the same results with the option on and off (tolerance 1e-10);
• perf_eval_gauss() is used once per left-out observation;
• search_forward_gauss() applies to a model with compound search terms and returns NULL for a factor in an interaction term;
• search_forward_gauss(return_outdmins = FALSE) gives the same ranking and no submodel fits;
• in the LOO CV with validated search and nloo = 5, the search of each of the 5 left-out observations runs with return_outdmins = FALSE.

The tests that call search_forward_gauss() or perf_eval_gauss() fail on master, where these functions do not exist; the last two tests also fail without commit 4. The full suite passes with NOT_CRAN=true: 260 tests, 55 823 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 other PRs

• With Avoid repeated formula updates, term extraction and training-data model matrices #578, search_forward_gauss() can pass the precomputed term labels to select_possible_terms_size(). I will add this one-line change after Avoid repeated formula updates, term extraction and training-data model matrices #578 is merged.
• With Keep the data-dependent bases #573, the draw-wise path also handles poly() terms in the LOO CV with validated search. After Keep the data-dependent bases #573 is merged, I will add a poly() term to the validated-LOO test.

[avehtari]

More details with equations

The three subsections follow the three commits of the PR, so each one can be pasted as a PR comment.

1. Ranking the candidates from cross-products

For a Gaussian submodel with the identity link, glm_ridge() solves a weighted ridge least-squares problem. It centers and scales the columns of the model matrix with the observation weights and does not penalize the intercept. Let $X_c$ be the model matrix of all search terms, centered at the weighted column means, $W = \mathrm{diag}(w_1, \dots, w_n)$ the observation weights, $M$ the $n \times S$ matrix of fitted values of the projected reference draws, and $d_j$ the weighted standard deviation of column $j$. Scaling column $j$ by $1 / d_j$ turns the penalty $\lambda \lVert \theta \rVert^2$ on the standardized scale into $\beta^\top \Lambda \beta$ with $\Lambda = \lambda , \mathrm{diag}(d_1^2, \dots, d_p^2)$ on the original scale. For the columns $J$ of a candidate submodel and draw $s$, the coefficients solve

$$ (G_{JJ} + \Lambda_{JJ}) \beta_s = C_{Js} , \qquad G = X_c^\top W X_c , \qquad C = X_c^\top W M . $$

The fitted values are $\hat{y}s = \bar{y}s + X{c,J} \beta_s$, with $\bar{y}s$ the weighted mean of column $s$ of $M$. With the Cholesky factor $G{JJ} + \Lambda{JJ} = R^\top R$ and $u_s = R^{-\top} C_{Js}$, let

$$ q_s = \lVert u_s \rVert^2 + \beta_s^\top \Lambda_{JJ} \beta_s = \beta_s^\top (G_{JJ} + 2 \Lambda_{JJ}) \beta_s . $$

The part of ce_gauss() that differs between candidates is $-\sum_s w_s q_s / \sum_i w_i$, with the weights $w_s$ of the projected draws. The search picks the candidate with the smallest value, as master does with the cross-entropy. $G$ and $C$ are computed once per search, so the cost per candidate depends on the number of columns, not on $n$.

In one step, all candidates share the columns $c$ of the chosen terms. The factor $R_c$ and $u_c = R_c^{-\top} C_c$ are computed once per step. A candidate with new columns $k$ extends them as a block Cholesky factor:

R_ck <- solve(t(R_c), A[c, k])              # A = G + Lambda
R_k  <- chol(A[k, k] - t(R_ck) %*% R_ck)
u_k  <- solve(t(R_k), C[k, ] - t(R_ck) %*% u_c)
beta_k <- solve(R_k, u_k)
beta_c <- solve(R_c, u_c - R_ck %*% beta_k)
q <- colSums(u_c^2) + colSums(u_k^2) + colSums(lambda_c * beta_c^2) + colSums(lambda_k * beta_k^2)

After the search, proj_to_submodl() fits the submodels along the ranking, as master stores them. The search returns NULL, and master's search runs, for other families or links, the latent or augmented-data projection, a custom divergence minimizer or submodel fitter, non-zero offsets, glm_ridge() arguments other than thresh, group-level or smooth terms, a factor inside an interaction term, missing values, or a zero-variance column.

2. Evaluation in the LOO CV with validated search

cv_varsel(validate_search = TRUE) evaluates every submodel along the path of each left-out observation. For Gaussian submodels, the projected dispersion and the cross-entropy follow from the same quantities, without the $n$ fitted values:

$$ \sum_i w_i (M_{is} - \hat{y}_{is})^2 = \sum_i w_i (M_{is} - \bar{y}_s)^2 - q_s , \qquad \sum_i w_i \left( -2 M_{is} \hat{y}_{is} + \hat{y}_{is}^2 \right) = -\bar{y}_s^2 \sum_i w_i - q_s . $$

The first identity follows from $\sum_i w_i (M_{is} - \bar{y}s)(\hat{y}{is} - \bar{y}s) = \beta_s^\top C{Js}$ and $\beta_s^\top C_{Js} = \beta_s^\top (G_{JJ} + \Lambda_{JJ}) \beta_s$. The prediction for the left-out observation is $\bar{y}s + (x{\mathrm{test}} - m)_J^\top \beta_s$, with $m$ the weighted column means of the training data. The submodels along a path are nested, so the factor $R$ grows by one block per size, as in the search.

This applies only in the LOO CV with validated search. varsel() and proj_linpred() keep their code, because a test compares their results at machine precision.

3. All one-column candidates at once

Most candidates add one column $j$. Then the block Cholesky step reduces to scalars, and all such candidates of a step can be computed with matrix operations. With $r_j = R_c^{-\top} A_{cj}$, $a_j = A_{jj} - \lVert r_j \rVert^2$, $b_0 = R_c^{-1} u_c$ and $v_j = R_c^{-1} r_j$:

$$ u_{js} = \frac{C_{js} - r_j^\top u_{cs}}{\sqrt{a_j}} , \qquad \beta_{js} = \frac{u_{js}}{\sqrt{a_j}} , \qquad \beta_{cs} = b_{0s} - v_j \beta_{js} . $$

Expanding $\beta_{cs}^\top \Lambda_{cc} \beta_{cs}$ gives

$$ q_{js} = u_{js}^2 + \lambda_j \beta_{js}^2 + \lVert u_{cs} \rVert^2 + b_{0s}^\top \Lambda_{cc} b_{0s} - 2 \beta_{js} , v_j^\top \Lambda_{cc} b_{0s} + \beta_{js}^2 , v_j^\top \Lambda_{cc} v_j , $$

where all terms for all candidates $j$ and draws $s$ come from a few matrix products. Candidates with several new columns (compound search terms, factors, poly() terms) use the block step of subsection 1

4. No submodel fits along the ranking in the validated search when the evaluation refits

After the ranking, search_forward_gauss() computes the submodel fits along the ranking with proj_to_submodl(): $K + 1$ projections for $K$ ranked terms, each with one fit per projected draw. In the draw-wise search of master, these fits are a by-product of the search. In cv_varsel(validate_search = TRUE) with refit_prj = TRUE, the evaluation of each fold projects again with its own draws, so the fits of the fold searches are not used:

for each fold:
  ranking <- search(training data of the fold)     # with return_outdmins = !refit_prj
  for each size along ranking:
    if refit_prj: fit <- proj_to_submodl(...)      # new fit; search fits not used
    else:         fit <- search_path$outdmins[...]  # search fits used

So loo_varsel() and kfold_varsel() pass return_outdmins = !refit_prj to the search, and search_forward_gauss() skips the fits when it is FALSE. In the validated LOO example of the PR, this takes the time from 4.8 s to 2.7 s.