Help testing brms pr for sum-to-zero and partial centering

Interfaces brms
21

Thank you @spinkney for responding to my request and for the update.

I have installed the branch, and it works fine. I’m now testing nonlinear models.

One issue I’ve run into is that I can’t use a prior that has been passed to the data block. For example, if I create:

stanvar(x = array(8.51), name = "prior_ult", block = "data")

and then try to use it in the prior via:

prior(normal(prior_ult, 0.25), nlpar = "ult")

it seems the prior must be a numeric value such as 8.51 and not a symbol like prior_ult. Could you advise on the intended way to pass such data-block priors into nonlinear parameters in this branch?

Here is the example code to reproduce the error:

library(brms)

form_loss <- bf(
  cum ~ exp(ult) *
    (1 - exp(-(dev / exp(theta))^exp(omega))),
  
  # ult, omega, and theta are now on log scales
  ult ~ 1 + (1 | gr(AY, s2z = TRUE, center = FALSE)),
  omega ~ 1 + (1 | gr(AY, s2z = TRUE, center = FALSE)),
  theta ~ 1 + (1 | gr(AY, s2z = TRUE, center = FALSE)),
  
  nl = TRUE
)

prior_loss <- c(
  prior(normal(prior_ult, 0.25), nlpar = "ult"),   # log(5000)
  # prior(normal(8.51719319141624, 0.25), nlpar = "ult"),   # log(5000)
  prior(normal(0, 0.5), nlpar = "omega"),                # log(1)
  prior(normal(3.80666248977032, 0.3), nlpar = "theta"), # log(45)
  
  prior(exponential(2), class = "sd", group = "AY",
        nlpar = "ult"),
  prior(exponential(2), class = "sd", group = "AY",
        nlpar = "omega"),
  prior(exponential(2), class = "sd", group = "AY",
        nlpar = "theta"),
  
  prior(normal(0, 500), class = "sigma")
)

fit_loss_s2z <- brm(
  form_loss,
  data = loss,
  family = gaussian(),
  stanvars = stanvar(x = array(8.51), name = "prior_ult", block = "data"),
  prior = prior_loss,
  backend = "cmdstanr",
  chains = 4,
  cores = 4,
  iter = 1000,
  control = list(
    adapt_delta = 0.8,
    max_treedepth = 10
  )
)

Error in `stop_re_s2z()`:
! All arguments of population-level priors used with the sum-to-zero parameterization must currently be numeric constants (coefficient 'Intercept').
S2Z capability 'active_prior_arguments' is unavailable for response 'cum', family 'gaussian', dpar 'mu', nlpar 'ult', group 'AY', ID '1', coefficient(s) 'Intercept', prior 'normal(prior_ult,0.25)'. Remedy: replace symbolic arguments with finite numeric constants.
Run `rlang::last_trace()` to see where the error occurred.
22

I encountered this as well but I thought it was a limitation with brms :), can you confirm that this capability exists for non-s2z?

23

Yes, for non-s2z it works fine. Please see the code below where I replaced s2z = TRUEwith s2z = FALSE

library(brms)

form_loss <- bf(
  cum ~ exp(ult) *
    (1 - exp(-(dev / exp(theta))^exp(omega))),
  
  # ult, omega, and theta are now on log scales
  ult ~ 1 + (1 | gr(AY, s2z = FALSE, center = FALSE)),
  omega ~ 1 + (1 | gr(AY, s2z = FALSE, center = FALSE)),
  theta ~ 1 + (1 | gr(AY, s2z = FALSE, center = FALSE)),
  
  nl = TRUE
)

prior_loss <- c(
  prior(normal(prior_ult, 0.25), nlpar = "ult"),   # log(5000)
  # prior(normal(8.51719319141624, 0.25), nlpar = "ult"),   # log(5000)
  prior(normal(0, 0.5), nlpar = "omega"),                # log(1)
  prior(normal(3.80666248977032, 0.3), nlpar = "theta"), # log(45)
  
  prior(exponential(2), class = "sd", group = "AY",
        nlpar = "ult"),
  prior(exponential(2), class = "sd", group = "AY",
        nlpar = "omega"),
  prior(exponential(2), class = "sd", group = "AY",
        nlpar = "theta"),
  
  prior(normal(0, 500), class = "sigma")
)

fit_loss_s2z <- brm(
  form_loss,
  data = loss,
  family = gaussian(),
  stanvars = stanvar(x = array(8.51), name = "prior_ult", block = "data"),
  prior = prior_loss,
  backend = "cmdstanr",
  chains = 4,
  cores = 4,
  iter = 1000,
  control = list(
    adapt_delta = 0.8,
    max_treedepth = 10
  )
)
24

Super cool. We should have StanCon more often—we get something great from @spinkney every year!

Could you share your notes? I think that’d make it easier to review and I’d also like to understand the specific way you’re partially non-centering and how you’re figuring out the weights.

It’d also be great to get your notes on how to convert back to what you would have gotten without sum-to-zero. The outline for that was in the StanCon talk, but 20m wasn’t quite enough time for me to process the approach. Is the back-conversion in the PR?

I’m pretty much never laying down priors to such a degree of precision that the difference would matter.

25

I’m testing this on the air pollution dataset that I used for the BRMS tutorial at StanCon 2023 -

This is the dataset from the “Visualization in Bayesian Workflow”, Gabry et al 2019 paper.
There are 6003 observations. Following the paper, the StanCon 2023 notebook explores 3 groupings of the data, fits them in BRMS, and then runs loo to compare the fits.

The data is measurements of pm2.5 (particulate matter pollution at the 2.5 micron level)
The groups are very unbalanced.

> table(sites$cluster_log_region)  # Gabry does some clustering on log pm2.5
   1    2    3    4    5    6 
  42 1677 1555 1694  347  688 

> table(sites$cluster_region)  # Gabry does some clustering on pm2.5
   1    2    3    4    5    6 
  27 3534 1696  418  315   13 

> table(sites$super_region)  # straw man - WHO super-region names
   1    2    3    4    5    6    7 
3051  273  464  518  308 1312   77

If you fit all three models and run loo, clustering on log pm2.5 scores the best, so let’s start there:

> elapsed_time(fit_cluster_log_region)
  chain_id  warmup sampling   total
1        1 107.149  161.641 268.790
2        2 102.446  167.001 269.447
3        3 101.629  179.791 281.420
4        4 101.958  171.276 273.234
> elapsed_time(fit_cluster_log_s2z_centered)
  chain_id warmup sampling  total
1        1 35.706   38.051 73.757
2        2 35.006   39.566 74.572
3        3 33.735   40.846 74.581
4        4 34.637   33.568 68.205
> elapsed_time(fit_cluster_log_s2z_noncentered)
  chain_id warmup sampling  total
1        1 31.511   29.053 60.564
2        2 37.015   36.055 73.070
3        3 33.273   37.210 70.483
4        4 30.245   32.155 62.400

The fit_cluster_log_region run - using all BRMS defaults - resulted in poor adaption with stepsize = 0.006 (!) so 40% of the iterations hit max_treedepth.
The fit_cluster_log_region_s2z_noncentered run had stepsize=0.02.

s2z rules!! SP rules!!

I am continuing to play around with the most unbalanced grouping on cluster_region because I think that this is an interesting test case for the centered/auto/non-centered/fixed part of this PR.

Update: on the more unbalanced dataset the centered parameterization works better. go figure.

fit_cluster_region = brm(log_pm25 ~ log_sat + (log_sat | cluster_region), data=sites, cores=4)
fit_cluster_region_s2z_centered = brm(log_pm25 ~ log_sat + (log_sat | gr(cluster_region, s2z=TRUE, center=TRUE)), data=sites, cores=4)
fit_cluster_region_s2z_noncentered = brm(log_pm25 ~ log_sat + (log_sat | gr(cluster_region, s2z=TRUE, center=FALSE)), data=sites, cores=4)

> elapsed_time(fit_cluster_region_s2z_centered)
  chain_id warmup sampling   total
1        1 70.774   82.316 153.090
2        2 67.557   85.167 152.724
3        3 70.249   83.133 153.382
4        4 67.500   81.343 148.843
> elapsed_time(fit_cluster_region_s2z_noncentered)
  chain_id  warmup sampling   total
1        1 106.830  129.053 235.883
2        2 115.644  142.132 257.776
3        3 105.072  115.305 220.377
4        4 108.151  131.616 239.767
> elapsed_time(fit_cluster_region)
  chain_id  warmup sampling   total
1        1 112.310  167.090 279.400
2        2 126.595  177.202 303.797
3        3 123.077  169.912 292.989
4        4 119.293  158.070 277.363

Next up, the covid model with test sensitivity, specificity with a bunch of categorical predictors which requires has a custom likelihood in BRMS.

26

Thanks for including these types of models! I tested spinkney/brms@6ae3c9b.

I wanted to try a model with a bunch of group terms, like so:

mvbind(y1, y2) ~ 1 + 
  (1 | Acor | gr(group1, cov = A, s2z = ...)) +
  (1 | Bcor | gr(group2, cov = B, s2z = ...)) +
  (1 | Ecor | gr(group3, s2z = ...))

but this ran into:

A predictor participating in a cross-predictor sum-to-zero ID cannot yet contain another conventional sum-to-zero block.

So I fit a slightly simpler model, like so:

mvbind(y1, y2) ~ 1 +
  (1 | Acor | gr(group1, cov = A, s2z = ...)) +
  (1 | Ecor | group2)

So applying the sum-to-zero only on the group effect following a covariance matrix.

For context, this a quantitative genetics model, estimating the heritability and genetic correlations of two traits, measured in 14,100 observations of 3,229 fish. So the groups are additive-genetic (a dense pedigree relatedness matrix, cov = A) and permanent environment (repeated photos per fish, median 4).

Results:

ordinary center = TRUE center = 0.5 center = FALSE center = “auto”
mean chain time (h) 26.6 28.0 28.1 28.0 17.3
slowest chain (h) 31.8 47.4 28.6 29.3 17.7
divergences 0 0 0 0 0
treedepth saturations 0 0 0 0 0
mean treedepth 8.25 8.50 8.00 8.00 8.00
max R-hat 1.078 1.081 1.049 1.038 1.063
min bulk ESS 59 56 117 78 94
min bulk ESS / h 1.86 1.19 4.10 2.67 5.32

image
image3000×840 106 KB

I found that center = "auto" works the best, and makes the model ~3x more efficient in terms of bulk ESS per hour.

Posteriors look pretty much identical.

Caveat

The R-hats are obv too high, and the ESS is very low, but this is a challenging and slow model to fit, so I only fit 4 chains, 500 warmup / 500 sampling.

27

Have you tried to use Pathinder to initialize MCMC sampling? It’s likely to improve warmup behavior.

Also if your data model is normal (family=gaussian()), then it would be worth testing to integrate out the latent values. I have examples and seen others reporting, too, up to 100x speed-ups.

28

Yes, I will. After Stancon I’ve had this pesky thing to do called my day job, it uses up a lot of time. After holiday weekend in the US I’ll see about putting these together. I’m thinking it makes more sense to just finish the paper and post it on arxiv.

Why everyone is getting the same posterior estimates is because the conversion is happening in the PR! The conversion does matter quite a bit. The differences, without knowing one could convert between them, are discussed at New Stan data type: zero_sum_vector.

29

❤

30

I’m glad it’s mostly working :)

I should’ve put the branch to test. That’s an older one.

install.packages("remotes")

remotes::install_github(
  "spinkney/brms",
  ref = "feature/s2z-group-effects",
  upgrade = "never"
)
31

The current version of the pr for auto centering is using pathfinder as an initial fit for the partial centering matrix. It’s probably similar to what Nicolas and your group was experimenting. We should probably discuss and see if we can make this better.

32

I added to RBesT sum to zero and it gave me about 5x more efficiency… adding now the auto centering thing (which I still need to digest, honestly) makes this into 10x more efficiency from my experiments. Is there something I can read on the auto centering trick?

Since I cannot use pathfinder in RBesT. as I want Stan 2.32 compatibility… Opus went ahead and crafted a version of pathfinder tailored to the varying intercept model… I still need to catch-up on that which is why the code is for now only local with me.

Anyway… saying it again - this is absolutely awesome. This model is used as the default to get informative meta-analytic-predictive priors enabling sample size reductions in clinical trials. It was already tuned a lot, but this makes the model far more stable with the priors we tend to use. So thanks a lot.

=====================================================
divergent-MAP stress test: legacy vs s2z (s2z, center, auto)
=====================================================
trials: 2000 | chains: 4 | J = 4 groups

Table: divergences, arms side by side

setting      trials   div% auto   div% center   div% legacy   div% s2z   div/fit auto   div/fit center   div/fit legacy   div/fit s2z
----------  -------  ----------  ------------  ------------  ---------  -------------  ---------------  ---------------  ------------
shipped        2000           0          99.6          19.3        0.5              0           100.25             0.39             0
matched99      2000           0          99.0          18.5        0.0              0            37.47             0.30             0
matched95      2000           0          99.4          75.6        0.4              0            95.92             3.64             0
tuned          2000           0          99.0          29.8        0.0              0            41.95             0.56             0
hard           2000           0          94.7           1.0        0.0              0            19.08             0.01             0

Table: cost, arms side by side

setting      trials   ESS/kgrad auto   ESS/kgrad center   ESS/kgrad legacy   ESS/kgrad s2z   sec auto   sec center   sec legacy   sec s2z
----------  -------  ---------------  -----------------  -----------------  --------------  ---------  -----------  -----------  --------
shipped        2000           106.04              42.41               9.56           42.83        0.8          1.4          1.3       1.0
matched99      2000            49.34              20.05               9.46           26.18        1.4          2.9          1.3       1.5
matched95      2000           105.22              42.53              14.51           42.88        0.8          1.4          0.9       1.0
tuned          2000            68.99              23.52              11.63           31.12        1.3          2.6          1.2       1.4
hard           2000            31.62               9.76               6.48           19.43        6.3          9.3          4.5       5.3
33

That is awesome! One of my goals, with the constraint parameterizations and this, is to make models with the current NUTS implementation as fast as possible. The pathfinder init is a bit of a cheat here but sampling is still occurring with NUTS.

34

I pushed a more stable version of the auto centering. It uses pathfinder and then a 2nd order taylor expansion around each s2z level log likelihood and the first and second order derivatives here. This approximates the posterior covariance under a normal assumption and then the s2z constraint is pushed through this. The remaining variance is the approximate weights that partially center the parameters.