Skip to content

Faster projections onto multilevel and additive submodels - #584

Open
avehtari wants to merge 4 commits into
masterfrom
ml-add-submodels
Open

avehtari wants to merge 4 commits into
masterfrom
ml-add-submodels

Conversation

@avehtari

Copy link
Copy Markdown
Member

This PR was assisted by Claude.

Summary

For a multilevel submodel, divmin() fits each projected draw with lme4::lmer() (Gaussian) or lme4::glmer(), and init_submodl() then predicts each draw's fit for the training data with predict.merMod() and repair_re(). Each lmer() call builds the model frame, the fixed-effects model matrix and the random-effects structure (lme4::lFormula()), although only the response differs between the draws; and each prediction builds the fixed-effects matrix and the random-effects design again, twice. In a profile of cv_varsel(validate_search = TRUE) for a Gaussian multilevel model, the predictions took about 60 % of the time and lFormula() about 15 %.

This PR has three commits.

  1. Predictions (subprd_mer_direct(), called first in subprd()). If all fits are lmerMod or glmerMod fits of the same submodel and newdata has no new group levels (and the global option projpred.mlvl_pred_new is FALSE), repair_re() would add exactly zero and draw no random numbers. subprd_mer_direct() then skips it and computes the predictions $X \beta_s + Z b_s$ of all draws from one fixed-effects model matrix and one random-effects design (lme4::mkNewReTrms()), with the steps of predict.merMod(). It maps the conditional modes of each draw onto the rows of $Z$ by the names of the group levels, and checks the mapping against mkNewReTrms() for the first draw. In all other cases, the existing code runs, with the same random numbers. The predictions are identical to master (largest absolute difference 0 in the checks), including the dimension names.

  2. Shared lFormula() for lmer() submodels (lmer_shared_callback() in divmin(), Gaussian submodels with the identity link). lme4::lmer() is lFormula() followed by mkLmerDevfun(), optimizeLmer(), checkConv() and mkMerMod(). This commit calls lFormula() once per projection and repeats the other steps of lmer() for each draw. lmer_shared_fit() transcribes these steps from the body of lmer() in lme4 2.0-6 to 2.1-1 (lme4 is GPL-2 | GPL-3), with the response column replaced, fresh copies of the vectors that lme4's C++ objects modify in place (Lambdat@x, theta, Zt@x, X), and the same call and warning texts as lmer(). The fits are all.equal() (with tolerance = 0) to those of lmer(), with identical theta. (Two lmer() fits of the same data are not identical() either, because of lme4's reference-class objects.) Errors fall back to the draw-wise lmer() fit.

  3. Shared setup for mgcv::gam() submodels (gam_shared_callback() in divmin(), additive submodels without group-level terms). This commit builds the setup of gam() (model frame, smooth bases, penalties) once per projection with mgcv::gam(fit = FALSE) and fits each draw with mgcv::gam(G = ), the public mgcv interface for a prepared setup. It sets the elements that contain the response name and call as in a fit by fit_gam_callback(), so the fits are all.equal() (with tolerance = 0) to the draw-wise fits. subprd() predicts gam fits of one submodel from one linear-predictor matrix (predict.gam(type = "lpmatrix")) for all draws; the predictions are identical to the per-draw predict.gam(). The internal option projpred.gam_shared = FALSE switches the shared setup off.

In commits 2 and 3, shared_setup_callback() builds the shared setup from the data of the first draw-wise fitter call in divmin(). #580 removes the observations with weight 0 in divmin() before the draw-wise fitter calls. A setup built before this removal would include these observations. If the data of a later call differ from the data of the first call, divmin() fits that draw with the draw-wise fitter.

Commit 2 repeats steps of lmer(), so it has three guards:

  • At the first use in an R session, lmer_shared_check() (0.2 s, a fixed problem without random numbers: two grouping factors, a random slope, observation weights, 3 draws) compares lmer_shared_fit() with lmer(). If they differ, for example after a change in lme4, projpred uses one lmer() call per draw as before. (Without the fresh copies, the check returns FALSE.)
  • tests/testthat/test_lmer_shared.R runs on CRAN (0.75 s): the check itself and a comparison of divmin() with and without sharing. So a change in lmer() or in lme4's in-place behavior also makes projpred's checks fail, including the reverse-dependency checks of lme4.
  • The internal option projpred.lmer_shared = FALSE switches commit 2 off.

The steps of lmer() are the same in lme4 2.0-6, 2.1-0 and 2.1-1.

Results

Predictions and fits are unchanged. In the bench below, the rankings, the suggested sizes and the CV ranking proportions are identical to master, and the largest relative difference in the ELPD estimates is 1.6e-16.

Bench: simulated data, n = 300, 6 predictors, grouping factors with 20 and 8 levels; rstanarm stan_lmer() (random intercepts and a random slope), stan_glmer() (Bernoulli) and stan_gamm4() (smooth term and random intercept) reference models; forward search, nterms_max = 6; cvt10 = cv_varsel(validate_search = TRUE, nloo = 10, nterms_max = 4, refit_prj = TRUE); 1 CPU per run, all versions in the same job on one node (CSC Roihu, AMD EPYC 9965), lme4 2.0-6. Times in seconds.

The speedup of a commit compares it with the column before it; the last column compares with master.

model call master commit 1 commits 1 + 2 speedup commit 1 speedup commit 2 speedup vs master
Gaussian, stan_lmer() vs 85.3 38.0 24.8 2.2 1.5 3.4
Gaussian, stan_lmer() cvt10 559.8 267.1 168.9 2.1 1.6 3.3
Gaussian, stan_gamm4() vs 105.1 83.1 74.1 1.3 1.1 1.4
Gaussian, stan_gamm4() cvt10 597.1 249.8 192.6 2.4 1.3 3.1
Bernoulli, stan_glmer() vs 393.8 331.7 (as commit 1) 1.2 - 1.2
Bernoulli, stan_glmer() cvt10 1852.2 1409.3 (as commit 1) 1.3 - 1.3

For the additive models (bench m1, same setup; commit 3 is measured against commits 1 + 2):

model call master commits 1 + 2 commits 1-3 speedup commit 3 speedup vs master
Gaussian, stan_gamm4() without group term vs 14.7 14.7 9.9 1.5 1.5
Gaussian, stan_gamm4() without group term cvt10 70.0 75.2 56.8 1.3 1.2
Gaussian, stan_gamm4() with group term cvt10 583.8 174.3 176.7 1.0 3.3

(The Bernoulli times are from a separate run with master and commit 1; commit 2 does not apply to glmer().) After this PR, lme4's optimizer and mkMerMod() take most of the time of the Gaussian runs.

Tests

The test suite passes with NOT_CRAN=true: 264 tests, 55 912 expectations, 0 failures, 0 errors, the same 4 warnings as master. Tested with this change applied to master 73cf0ee (projpred 2.11.0.9000), R 4.6.1, lme4 2.0-6, rstanarm 2.32.2, brms 2.23.0, on CSC Roihu (Linux, AMD EPYC).

New tests:

  • tests/testthat/test_mer_predict.R: subprd_mer_direct() and subprd() give the predictions of the draw-wise path, including the dimension names, for lmer fits (random slope and two grouping factors, random intercepts only, no fixed-effects predictors, a special formula with poly(), I() and a factor) and glmer fits, for the full data and a subset; the random number state is unchanged; NULL for a new group level, with projpred.mlvl_pred_new = TRUE, for other fits, without newdata and for fits of different submodels.

  • tests/testthat/test_lmer_shared.R (runs on CRAN): lmer_shared_check() is TRUE; divmin() gives fits that are all.equal(tolerance = 0) with identical theta and formulas with and without sharing (random slope, two grouping factors, poly(), a factor, weights), and also for data with one observation with weight 0; lmer_shared_callback() returns NULL with the option set to FALSE and for a missing response value.

  • tests/testthat/test_gam_shared.R: divmin() gives gam fits that are all.equal(tolerance = 0) with identical formulas with and without sharing (two smooth terms, a factor, weights), and also for data with one observation with weight 0; shared_setup_callback() builds the setup one time, from the data of its first call, and fits a call with different data with the draw-wise fitter; subprd() gives the predictions of predict.gam() for the full data and a subset, and subprd_gam_direct() returns NULL for fits of different submodels; gam_shared_callback() returns NULL with the option set to FALSE and for a missing response value.

The new tests fail on master (the functions do not exist there).

Overlap with other PRs

@avehtari avehtari added the performance Performance issues (i.e., high runtime or high memory usage). label Oct 26, 2026
@avehtari
avehtari requested a review from aloctavodia October 26, 2026 19:52

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

performance Performance issues (i.e., high runtime or high memory usage).

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant