LOO with brms::add_criterion crashes the R session

Modeling
, ,
1

Hi folks,

I’m trying to compare 8 brms models (6 as of now, still waiting for two to finish) using loo with moment_match = TRUE but keep constantly running into crashes where the OS kills the R process. These models are:

M2 <- brms::brm(brms::brmsformula(mvbind(F00727, F00679, F00709) ~ state + (1|taxa)) + set_rescor(TRUE), data = df, chains = 8, cores = 8, threads = 4, iter = 5000, warmup = 2500, backend = "cmdstanr", save_pars = save_pars(all = TRUE))
M4 <- brms::brm(brms::brmsformula(mvbind(F00727, F00679, F00709) ~ state + (1|q|taxa)) + set_rescor(TRUE), data = df, chains = 8, cores = 8, threads = 4, iter = 5000, warmup = 2500, backend = "cmdstanr", save_pars = save_pars(all = TRUE))
M1 <- brms::brm(brms::brmsformula(mvbind(F00727, F00679, F00709) ~ state + (1|gr(binominal, cov = corrmat)) + (1|taxa)) + set_rescor(TRUE), data = df, data2 = list(corrmat = corrmat), chains = 8, cores = 8, threads = 4, iter = 5000, warmup = 2500, backend = "cmdstanr", save_pars = save_pars(all = TRUE))
M3 <- brms::brm(brms::brmsformula(mvbind(F00727, F00679, F00709) ~ state + (1|p|gr(binominal, cov = corrmat)) + (1|q|taxa)) + set_rescor(TRUE), data = df, data2 = list(corrmat = corrmat), chains = 8, cores = 8, threads = 4, iter = 5000, warmup = 2500, backend = "cmdstanr", save_pars = save_pars(all = TRUE))

N2 <- brms::brm(brms::brmsformula(mvbind(F00727, F00679, F00709) ~ 1 + (1|taxa)) + set_rescor(TRUE), data = fred4, chains = 8, cores = 8, threads = 4, iter = 5000, warmup = 2500, backend = "cmdstanr", save_pars = save_pars(all = TRUE))
N4 <- brms::brm(brms::brmsformula(mvbind(F00727, F00679, F00709) ~ 1 + (1|q|taxa)) + set_rescor(TRUE), data = fred4, chains = 8, cores = 8, threads = 4, iter = 5000, warmup = 2500, backend = "cmdstanr", save_pars = save_pars(all = TRUE))
N1 <- brms::brm(brms::brmsformula(mvbind(F00727, F00679, F00709) ~ 1 + (1|gr(binominal, cov = corrmat)) + (1|taxa)) + set_rescor(TRUE), data = fred4, data2 = list(corrmat = corrmat), chains = 8, cores = 8, threads = 4, iter = 5000, warmup = 2500, backend = "cmdstanr", save_pars = save_pars(all = TRUE))
N3 <- brms::brm(brms::brmsformula(mvbind(F00727, F00679, F00709) ~ 1 + (1|p|gr(binominal, cov = corrmat)) + (1|q|taxa)) + set_rescor(TRUE), data = fred4, data2 = list(corrmat = corrmat), chains = 8, cores = 8, threads = 4, iter = 5000, warmup = 2500, backend = "cmdstanr", save_pars = save_pars(all = TRUE))

The problem happens when adding the loo criterion to the fitted models (reloaded from serialized .Rds files).

M1 <- brms::add_criterion(readRDS("./ScratchData/brms_phylo.Rds"), criterion = "loo", moment_match = TRUE, save_psis = TRUE, cores = 8, pointwise = TRUE)

Every time, the R process gets killed by the OS because of it maxxes out the RAM (the server I’m using has 32 virtual cores and 126 GiB of RAM). This happens shortly after I see the output

Recompiling the model with 'rstan'
Recompilation done

Setting pointwise = TRUE didn’t help either, unfortunately!

Please see the screenshot for reference:

Screenshot (127)
Screenshot (127)1888×944 150 KB

The models I’m trying to add loo to are relatively big (1.2 GiB to 2.3 GiBs each on disk), so I partly understand the memory issues. However, brms::add_criterion works just fine when moment_match is set to FALSE. In fact I was able to add loo to all the six models and run loo::loo_compare in that case - but got the following warnings about pareto K estimates advising me to set moment_match to TRUE.

> M1 <- brms::add_criterion(M1, criterion = "loo", moment_match = FALSE, save_psis = TRUE, cores = 8, pointwise = FALSE)
Warning messages:
1: Ignoring relative efficiencies as some were NA. See argument 'r_eff' in ?loo::loo for more details. 
2: Found 267 observations with a pareto_k > 0.7 in model 'M1'. We recommend to set 'moment_match = TRUE' in order to perform moment matching for problematic observations.  

> M2 <- brms::add_criterion(M2, criterion = "loo", moment_match = FALSE, save_psis = TRUE, cores = 8, pointwise = FALSE)
Warning messages:
1: Ignoring relative efficiencies as some were NA. See argument 'r_eff' in ?loo::loo for more details. 
2: Found 696 observations with a pareto_k > 0.7 in model 'M2'. We recommend to set 'moment_match = TRUE' in order to perform moment matching for problematic observations.  

> M4 <- brms::add_criterion(M4, criterion = "loo", moment_match = FALSE, save_psis = TRUE, cores = 8, pointwise = FALSE)
Warning messages:
1: Ignoring relative efficiencies as some were NA. See argument 'r_eff' in ?loo::loo for more details. 
2: Found 625 observations with a pareto_k > 0.7 in model 'M4'. We recommend to set 'moment_match = TRUE' in order to perform moment matching for problematic observations.  

> N1 <- brms::add_criterion(N1, criterion = "loo", moment_match = FALSE, save_psis = TRUE, cores = 8, pointwise = FALSE)
Warning messages:
1: Ignoring relative efficiencies as some were NA. See argument 'r_eff' in ?loo::loo for more details. 
2: Found 257 observations with a pareto_k > 0.7 in model 'N1'. We recommend to set 'moment_match = TRUE' in order to perform moment matching for problematic observations.  

> N2 <- brms::add_criterion(N2, criterion = "loo", moment_match = FALSE, save_psis = TRUE, cores = 8, pointwise = FALSE)
Warning messages:
1: Ignoring relative efficiencies as some were NA. See argument 'r_eff' in ?loo::loo for more details. 
2: Found 720 observations with a pareto_k > 0.7 in model 'N2'. We recommend to set 'moment_match = TRUE' in order to perform moment matching for problematic observations.  

> N4 <- brms::add_criterion(N4, criterion = "loo", moment_match = FALSE, save_psis = TRUE, cores = 8, pointwise = FALSE)
Warning messages:
1: Ignoring relative efficiencies as some were NA. See argument 'r_eff' in ?loo::loo for more details. 
2: Found 662 observations with a pareto_k > 0.7 in model 'N4'. We recommend to set 'moment_match = TRUE' in order to perform moment matching for problematic observations.  

> compres <- loo::loo_compare(M1, M2, M4, N1, N2, N4) 
Warning message:
Not all models have the same y variable. ('yhash' attributes do not match) 
> print(compres)
 model elpd_diff se_diff p_worse diag_diff        diag_elpd
    M4       0.0     0.0      NA           625 k_psis > 0.7
    N4     -21.6   431.0    0.52           662 k_psis > 0.7
    M1    -142.2    81.2    0.96           267 k_psis > 0.7
    N1    -144.5   416.2    0.64           257 k_psis > 0.7
    M2    -844.7    81.2    1.00           696 k_psis > 0.7
    N2    -865.7   409.7    0.98           720 k_psis > 0.7

This is why I decided to proceed with moment_match = TRUE but the process ends up being extremely demanding on the RAM. With options(mc.cores = 1) options(loo.cores = 1) I was able to do brms::add_criterion for one model but with more cores specified, the RAM use goes over the roof (each core ends up using about 22 GiBs of RAM) and the process gets killed by the OS. Doesn’t help having to load in these huge .Rds models into the process’s memory as well.

Unfortunately, I cannot afford to re-run the brms fits (I’ve already completed, serialized and backed up 6 fits and have two models at 20% needing 20 more days to complete). I’d appreciate if anyone can help me find a workaround to do a moment matched loo comparison? I don’t even mind the non moment matched loo if I ignoring the warnings is okay. Also don’t understand why it says Not all models have the same y variable. ('yhash' attributes do not match) even though all the Y variables are identical across all the brms models. Please feel free to ask for more info if needed. Thanks!

System specs:

Ubuntu 24.04.4 LTS (Noble Numbat)
R version - 4.6.1
brms version - 2.23.0
cmdstanr version - 0.9.0
CmdStan version - 2.39.0
rstan version - 2.32.7
Stan version - 2.32.2
StanHeaders version - 2.32.10
2

First, with that many high Pareto-k’s moment matching is not able to help. I’ve been testing limits of moment matching recently and will publish results in fall, but can already say that it is helpful only with a small number of high khats.

Second, if you still want to try, you could use cores = 1 to reduce the memory use.

Some of the models have data=df and some data=fred4. Are you sure Y is bit identical in these?

As all models have many high k_psis, there is additional uncertainty in the comparison. I would say that M2 and N2 are worse than M4 and N4. There is not much difference between M4 and N4. M1 and N1 are not much worse than M4 and N4. N-models have quite similar average performance as the corresponding M-models, but the very high se_diff indicates the individual predictions differ a lot. It is possible that with more accurate LOO-computation the differences would be more clear.

I have some ideas I could test. Do you have a smaller scale data you could share so I could do some experiments?

3

Hey, thank you so much for the reply. It’s unfortunate to hear that moment matching may not be of much use in this case. However, pp_check density plots showed good overlap between predictions and expected values (there were variations between models - some showed better overlap for certain traits but overall all models seemed decent). But I did notice the hit the maximum treedepth limit of 10.warnings in some models.

My dataset is also pretty big (5219 rows with repeated species measurements), would you say that’s still a big % of high Pareto-K s?

I copy pasted these two sets of model formulae from two different R scripts, hence the different names, sorry. But can confirm that the data was identical.

Here’s the link to the dataset and the phylogeny. The columns binominal and taxa are identical. If a smaller dataset is preferred, you can average the three columns (F00679, F00727 & F00709) by the column binominal or use a subset of the dataset - in that case I suppose the phylogeny will need to be trimmed before being passed to ape::vcv.phylo - happy to do that myself and provide a link if needed. The links given are for the full dataset and the complete 1301 species phylogeny.

Also including links to the .Rout files from the batch runs of M and N models - if those could be of any help.

Thank you!

4

732 of 1301 species have only one observation which makes this impossible for PSIS. You can find similar example in a case study Leave-one-out cross validation model checking and comparison: Roaches – Bayesian Workflow book. In this case, we should use integrated PSIS so that we integrate out the group (sepcies) specific parameters as shown in that case study.

In addition your models with phylo component had serious MCMC convergence issues and the results should not be trusted at all.

As priors and data model are normal, the integration can be made analytically and it is quite fast. By integrating out both species specific latent values and phylo effects, we get much better sampling performance both in wall time (depending on the model from 1.5min to 20min) and ESS (thousands), and only a few khat>0.7. I included also log-normal models as the target was clearly skewed and here is much more reliable comparison table

model elpd_diff se_diff k>0.7 max k
M3_log 0.0 0.0 0 0.68
M4_log −480.8 41.7 1 0.72
M1_log −558.6 59.7 0 0.55
M2_log −1371.3 83.7 1 0.72
M3_raw −8390.5 467.0 4 2.18
M1_raw −8508.9 467.1 3 2.41
M4_raw −8795.0 470.1 3 2.21
M2_raw −9011.5 469.7 4 2.00

M3 with log-normal is the clear winner. However, that one is slightly misspecified, too, as shown by LOO-PIT-ECDF-difference plots (see the case study and other case studies for other examples)

image
image937×653 40.4 KB

Both models have in general over-dispersed predictive distribution (S-shape with excess of LOO-PIT values near 0.5) but have also some outliers (which cause the overdispersion, red dots at the edges). The red color indicates where the observed ECDF differs too much from the expected ECDF for a well calibrated model. More thick tailed data model would be appropriate (but that would make the integration much more challenging).

pp_check is not good with overparameterized models, as illustrated in Figure 14 in the case study I mentioned above, and it misses the misspecification completely.

As some groups (species) have more than one observation, integrated LOO is using a trick I haven’t seen used before (I’ll publish a paper with that trick in fall). In addition for integrating out the group (species) specific and phylo specific variables, I did get help from Claude, which was able to recognize two important shortcuts: 1) for species the slowest part of the computation is enough to be made for 38 distinct numbers of observations per group, and 2) for phylo using Felsenstein pruning. It is likely that Claude performed well here as I did have a lot of related code and material for integrated LOO and LOGO. Although this was quite fast for me with help from Claude, unfortunately in general it’s not trivial to change brms to directly produce integrated version, but maybe some special cases could be included.

5

Thank you for taking the time to look into this and the really helpful comments. It’s appreciated. It took me a while to understand the details and had to look up a lot of stuff (I have a biology background - sorry). You were right about the species distribution being skewed, and I understand how that can become a problem. I read through chapter 24 that you linked and it was very informative.

If I understand correctly, the integrated PSIS can be visualized using pp_check(model, type = "loo_pit_ecdf", moment_match = TRUE, method = "correlated"), but in order to actually perform an integrated loo before comparing the models will require me to use Stan directly? That will also require refitting all the models again? I believe I can get the Stan code for the brms formulae using this function. If I did that for all the 8 models - I’ll have to look into generated quantities section of the model to do this, correct? If I get a hang of how to implement that section, I’ll finally be able to do the intergated loo, is that so? But I’ll have to refit all the models again - that’s inevitable? If that’s the way to proceed, I’ll sacrifice some time and do that anyways. Will you be ablt to show me the Stan code you used in your reply - It’d be really helpful to use as a starting point?

I do understand your concerns about these models being overparametrized - but unfortunately these 8 formulae are essential to test the hypothesis I have - I intend to make the conclusion based on which model turns out to the the best fit. The only thing I care about is picking the best model supported by the data and then interpreting its biological significance. I don’t plan to use the model’s coefficients to make conclusions - if that’s what you meant when you said “In addition your models with phylo component had serious MCMC convergence issues and the results should not be trusted at all”

Once again, thank you for your time, It’s appreciated.

6

I’ll share the Stan and R code for these 4 models. I’m just cleaning the code and verifying the cleaned code runs. Note that the code is specialized for the data stucture you shared and if you extend the models the intergation part may need to be changed, too. Rerunning MCMC is required, but much faster with the integrated code.

The overparameterization is problem only for PSIS-LOO and in phylo case also for the convergence diagnostics (possible multimodality), but analytic integration solves all these problem.