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

[spinkney]

Hi brms users, I’m hoping to get eyes on this pr Add fast sum-to-zero parameterizations for group effects by spinkney · Pull Request #1919 · paul-buerkner/brms · GitHub

Warnings, it is large and it is AI generated because I’m not familiar with the brms repo nor do I have the time to make this pr. Does that mean it sucks? I hope not. I gave it the math and my notes on the s2z parameterization for hierarchical models that is equivalent to the conventional Bayesian model for normal and student-t effects. This includes both independent and dependent cases of the varying/random effects. I also gave it my notes on centering/non-centering and data driven partial centering so users can specify.

What I need from folks to to test it out in two ways. The first is find the bugs or short-cuts chatgpt did. The second, and more interesting, is to run your hierarchical models and compare the fitting times and diagnostics to the conventional fit.

The syntax is

bf(y ~ x * z + (1 + x * z | gr(g, s2z = TRUE)))

# centered physical S2Z coordinates (default)
gr(g, s2z = TRUE)

# standardized/non-centered S2Z coordinates
gr(g, s2z = TRUE, center = FALSE)

# a fixed partial-centering fraction
gr(g, s2z = TRUE, center = 0.35)

# group- and coefficient-specific fractions derived from the fitted design
gr(g, s2z = TRUE, center = "auto")

If you wish to have varying standard deviations for the effects (possibly because of domain knowledge) this is supported through a scale = "varying" argument.

bf(y ~ x * z +
     (1 + x * z | gr(g, s2z = TRUE, scale = "varying")))

prior(normal(0, 0.20), class = "sdlog", group = "g", coef = "x")

[avehtari]

I tested with Park rule data (the case study shows the need for sum-to-zero, but had to be made with Stan code as this wasn’t yet in brms):

fit1 <- brm(y ~ male_name + white_name + n_skipped_full +
                (1 | item) +
                (1 | respondent),
             family = bernoulli(),
             data = data_park,
             init = 0.1,
             control = list(max_treedepth = 8))

fit1s2z <- brm(y ~ male_name + white_name + n_skipped_full +
                   (1 | gr(item, s2z = TRUE)) +
                   (1 | gr(respondent, s2z = TRUE)),
             family = bernoulli(),
             data = data_park,
             init = 0.1,
             control = list(max_treedepth = 8))

With sum-to-zero the sampling time went from 1488s to 511s and the smallest bulk-ESS went from 263 to 3924, that is, 43x improvement in ESS/s. This is awesome!

[avehtari]

Another big win with pupil data

data("df_pupil_complete", package = "bcogsci")

fit4 <- brm(bf(p_size ~ load + (load | subj),
               sigma ~ subj),
               data = df_pupil_complete)

fit4_s2z <- brm(bf(p_size ~ load + (load | gr(subj, s2z=TRUE)), 
                   sigma ~ subj),
                data = df_pupil_complete)

Sampling time goes from 301s to 60s, and the lowest bulk-ESS goes from 154 to 5266, that is, 171x improvement in ESS/s.

Overall I see the biggest wins with bigger data, when there are more than (.|.) term, or like in this pupil example the covariate dependend sigma also has posterior dependency with mean formula and adding sum-to-zero makes the posterior easier. So far I haven’t seen catastrophic slow-down with any example I’ve tried. center="auto" doesn’t seem to help with models/data I’ve tested (slightly worse than center=TRUE or no difference).

[avehtari]

Another one with pupil data, so that the s2z is used also for sigma

fit5 <- brm(bf(p_size ~ load + (load | subj),
               sigma ~ (1 | subj)),
            data = df_pupil_complete)

fit5s2z <- brm(bf(p_size ~ load + (load | gr(subj, s2z=TRUE)),
                  sigma ~ (1 | gr(subj, s2z=TRUE))),
               data = df_pupil_complete)

Sampling wall time goes from 608s to 105s and lowest bulk-ESS from 518 to 4228, that is, 47x improvement in ESS/s.

[paul.buerkner]

This is very impressive indeed! Excellent work @spinkney and thank you @avehtari for testing!

[wds15]

Under which circumstances is sum to zero actually worse? I wonder if s2z should become the default option whenever it applies. Why would I not want it? Sure, in brms a reason to not want it is when I want to intercept the brms code and I expect the brms code to define thinks in certain ways… but even then it should be possible to formulate the Stan model changes such that the “conventional” group specific means are available in the usual form, just derived from the more efficient way of coding it, no?

In any case… this is super interesting.. I am building this trick into my own R packages right now.

[avehtari]

@wds15

Under which circumstances is sum to zero actually worse? I wonder if s2z should become the default option whenever it applies.

I’m copying my comment from the PR discussion

• s2z=TRUE helps most when the conditional posterior for the location is narrow, this can happen with a few groups with strong likelihood each or with many groups with weak likelihood each. Furthermore it helps more when there are many hierarchical components with each going from sum-to-non-zero to sum-to-zero
• center=TRUE helps most when each group has strong likelihood, and center=FALSE helps most when each group has weak likelihood.
• With strong group likelihoods, s2z=TRUE, center=TRUE makes sense, but I have already two examples were s2z=TRUE, center=FALSE has better performance, and at least one example where s2z=TRUE, center="auto" has the best performance, so I think it would be nice to have center exposed to users.
• The current implementation doesn’t allow s2z=FALSE, center=TRUE, so I didn’t test that.

At least in one example s2z has 3x sampling time without small drop in ESS. It seems s2z adds some overhead and can make the posterior more difficult.

EDIT: edited the last paragraph

[spinkney]

Thanks for testing @avehtari this is excellent and it’s the two things I think really benefit from the s2z. Both in terms of making larger size models fit quicker and making them fit better. I do think that users will want to reduce the warmup and sampling iterations because they should see large increases in ESS.

@wds15

Under which circumstances is sum to zero actually worse? I wonder if s2z should become the default option whenever it applies.

The main drawback I see is that it’s not universally equivalent to the conventional bayesian fit with priors which are not elliptically symmetric. I’ve limited the pr to just normal and student-t for now. However, most users are using the normal distribution to fit these models so it’s not a problem.

I do think the default should use s2z under normal distribution priors but I think it’s prudent to get this in as optional and have users test it to make it a mature enough feature for default usage. Just like Aki’s last example, I have also encountered models where s2z is not better from ESS and does increase runtime. This seems to happen in smaller data models and the good here is that you’re not getting worse estimates, it’s just slower.

[jroon]

Hi. So I’m trying it out on some data that I know is a bit difficult to fit (not entirely clear why). This is data from this study: https://www.tandfonline.com/doi/full/10.1080/21678421.2024.2359556#d1e675 To test the s2z I’ve thrown out the mutlivariate outcomes so I’m using only a single outcome. Here are my model specifications:

fit_model1 <- brm(
    bf(outfvc ~ 1 + days_from_baseline * sexsite + cohort + 
           (1 + days_from_baseline | gr(uin, s2z = FALSE) ) ),
    data = df2, chains=4, seed = 1234345,
    iter = 2000)
fit_model2 <- brm(
    bf(outfvc ~ 1 + days_from_baseline * sexsite + cohort + 
           (1 + days_from_baseline | gr(uin, s2z = TRUE) ) ),
    data = df2, chains=4, seed = 1234345,
    iter = 2000)
fit_model3 <- brm(
    bf(outfvc ~ 1 + days_from_baseline * sexsite + cohort + 
           (1 + days_from_baseline | gr(uin, s2z = TRUE, center = FALSE) ) ),
    data = df2, chains=4, seed = 1234345,
    iter = 2000)
fit_model4 <- brm(
    bf(outfvc ~ 1 + days_from_baseline * sexsite + cohort + 
           (1 + days_from_baseline | gr(uin, s2z = TRUE, center = "auto") ) ),
    data = df2, chains=4, seed = 1234345,
    iter = 2000)

The data consists of 836 observations from 269 individuals (uin is the subject id variable in the model). Here are the times for the model fit over 4 chains:

> rstan::get_elapsed_time(fit_model1$fit) 
        warmup sample
chain:1 24.064 11.000
chain:2 22.432 13.084
chain:3 22.054 11.368
chain:4 23.532 17.220
> rstan::get_elapsed_time(fit_model2$fit) 
        warmup sample
chain:1 33.012 16.576
chain:2 32.518 16.532
chain:3 32.536 16.508
chain:4 32.767 16.564
> rstan::get_elapsed_time(fit_model3$fit) 
        warmup sample
chain:1 37.596 18.496
chain:2 36.973 17.200
chain:3 38.171 17.083
chain:4 36.774 17.052
> rstan::get_elapsed_time(fit_model4$fit) 
         warmup  sample
chain:1 188.411  87.301
chain:2 180.725  90.242
chain:3 194.140 106.509
chain:4 192.068 118.166

ESS:
Model1:

Model2:

Model3:

Model4:

So in this small data (relative to Aki’s datasets) turning on s2z slightly slowed down fit, but did improve ESS for the relevant variables - though had little effect on other variables. The non-centered version did worse, and the “auto” setting did really bad.
I cannot share any data but if you want me to try other stuff let me know please.

Edit: note removing the poorly sampled variables - i.e. * sexsite + cohort actually slows down the fit for all models and the overall pattern does not change!

[satpal]

Thanks @spinkney for introducing this great feature.

I wonder whether it also works with correlated random effects in a nonlinear model.

The following model fails:

fit_loss <- brm(
  bf(cum ~ ult * (1 - exp(-(dev/theta)^omega)),
     ult   ~ 1 + (1 | z | gr(AY, s2z=TRUE)), 
     omega ~ 1 + (1 | z | gr(AY, s2z=TRUE)), 
     theta ~ 1 + (1 | z | gr(AY, s2z=TRUE)), 
     nl = TRUE),
  data = loss, family = gaussian(),
  prior = c(
    prior(normal(5000, 1000), nlpar = "ult"),
    prior(normal(1, 2), nlpar = "omega"),
    prior(normal(45, 10), nlpar = "theta")
  ),
  control = list(adapt_delta = 0.9),
  iter = 1000
)


Error in `validate_re_s2z()`:! A sum-to-zero group-level ID cannot span multiple linear predictors. Run `rlang::last_trace()` to see where the error occurred.

whereas the model with uncorrelated random effects works.

fit_loss <- brm(
  bf(cum ~ ult * (1 - exp(-(dev/theta)^omega)),
     ult   ~ 1 + (1 || gr(AY, s2z=TRUE)), 
     omega ~ 1 + (1 || gr(AY, s2z=TRUE)), 
     theta ~ 1 + (1 || gr(AY, s2z=TRUE)), 
     nl = TRUE),
  data = loss, family = gaussian(),
  prior = c(
    prior(normal(5000, 1000), nlpar = "ult"),
    prior(normal(1, 2), nlpar = "omega"),
    prior(normal(45, 10), nlpar = "theta")
  ),
  control = list(adapt_delta = 0.9),
  iter = 1000
)

[spinkney]

Wow, that ESS on model 2 for s2z parameters! See how close to 0 those estimates are, this would be difficult to sample otherwise. I wonder if letting the other parameters vary across cohort would be better. Is there a reason not to group by cohort and have it vary by sexsite like in the below?

fit_model2 <- brm(
    bf(outfvc ~ 1 + days_from_baseline * sexsite + cohort + 
           (1 + sexsite | gr(cohort, s2z = TRUE) +
           (1 + days_from_baseline | gr(uin, s2z = TRUE) ) ),
    data = df2, chains=4, seed = 1234345,
    iter = 2000)

[spinkney]

That looks like a insurance model :). In your model the s2z can work but I wanted to get the linear model working first.

[jroon]

Hmm no I don’t see a reason not to group that way. So I just tried…interesting… now I’m getting divergences and tree dept warnings:

Model1: 224 divs, 2849 tree depth warnings
Model2: 16 divs, 3253 tree depth warnings
Model3: 6 divs, 3994 tree depth warnings
Model4: 1016 divs, 2984 tree depth warnings

I’ll take a look at some pairs plots later see what I can see

[wds15]

I went ahead and implemented (ok… Opus did with me complaining…) the sum to zero into the RBesTpackage I maintain on CRAN. This is a super simple model - varying intercept… that’s it. The crux is that priors are not the strongest on the heterogeneity parameter \\tau which makes this model difficult to sample in some circumstances. I have written a while ago a fake data example benchmark emulating a case I worked on… and back then I had to make the sampler do excessive warmup to get stable estimates without divergences (the model had to run on many data sets). The new s2z parametrisation appears to solve the matter elegantly. Here is a brief comparison of legacy vs new s2z model run:

=====================================================
divergent-MAP stress test: legacy vs s2z
=====================================================
trials: 1000 | chains: 4 | J = 3 groups

Table: divergences and cost, arms side by side

setting      trials   delta legacy   delta s2z   div% legacy   div% s2z   div/fit legacy   div/fit s2z   ESS/kgrad legacy   ESS/kgrad s2z   sec legacy   sec s2z
----------  -------  -------------  ----------  ------------  ---------  ---------------  ------------  -----------------  --------------  -----------  --------
shipped        1000          0.990       0.950          39.0        0.3             1.01             0               7.00           37.19          1.3       0.9
matched99      1000          0.990       0.990          37.4        0.0             0.85             0               7.04           23.43          1.3       1.3
matched95      1000          0.950       0.950          92.7        0.5             8.45             0              10.27           37.50          0.8       0.9
tuned          1000          0.990       0.990          50.1        0.0             1.30             0               8.46           27.71          1.2       1.2
hard           1000          0.999       0.999           1.4        0.0             0.05             0               4.93           17.03          4.0       4.4

yuck… that is not rendered nicely.. the headline is that s2z solves the matter (surprise).

If you like to run the thing yourself, there is a vibe coded benchmark for this on the PR which I am working on for this:

next release (#64)

I will probably continue to work on it a bit more… but so far I like what I see in terms of improvements.

Anything I can cite when I include this now?

… the idea is so simple that it feels odd to have overlooked it all the time… and these are the best innovations…

[spinkney]

I guess cite the Stancon26 talk. @NikVetr and I will have a paper on arxiv soon to cite. It wasn’t so obvious, we were messing with a bunch of ways to get it to work. In a simple model it seemed like the exact uncertainty of the s2z we needed to add was a single prior draw to add that 1-extra dimension back. Once we saw that all the math falls out.

[wds15]

Oh, I did spend some time on it to get my head around it… don’t get me wrong .. I do not say it is trivial here.. no. Simple after having understood it.. very different.

[Ax3man]

I tried running some models, but it’s not very clear to me what is or isn’t supported. E.g.:

phylo <- ape::read.nexus('https://paul-buerkner.github.io/data/phylo.nex')
A <- ape::vcv.phylo(phylo)
d <- read.table('https://paul-buerkner.github.io/data/data_simple.txt', header = TRUE)

# "Arguments 'by', 'cov', and 'pw' are not yet supported together with gr(..., s2z = TRUE)."
stancode(phen ~ cofactor + (1 | gr(phylo, cov = A, s2z = TRUE)), d, data2 = list(A = A))

# "A sum-to-zero group-level ID cannot span multiple linear predictors."
stancode(bf(mvbind(phen, cofactor) ~ 1 + (1 | q | gr(phylo, s2z = TRUE))), d)

I now see that the cov restriction is documented in ?gr, although the text for s2z there is pretty impenetrable to me. I don’t see the second restriction in the docs.

[satpal]

Dear @spinkney
If possible, could you please implement support for modelling correlated random effects, as illustrated in the example model? This example is an insurance model from the brms manual, which I included only for reference.

My actual use case involves a nonlinear mixed-effects growth model for analysing longitudinal height data, which has a similar group-level structure. Given that nonlinear mixed-effects models are typically slow to fit and often yield low effective sample sizes (ESS), I believe that the s2z transformation would be particularly beneficial in such cases.

Thanks

[spinkney]

Test out the recent push, you’ll have to reinstall that branch of brms. It allows nonlinear models.

However, I noticed that the sampler can easily get into difficult regions with s2z if the parameters are sufficiently far from 0. For example, I found that the in the loss data case. Trying to fit the model as you wrote doesn’t work well because a s2z parameterization on the group effects for ult, omega, and theta get pushed in weird negative areas when the left hand side variables should not really ever be 0. brms puts a wide prior on ult but the s2z constraint on the group means that some of the effects must be negative to account for the constraint and we get into low density modes.

In regular brms you might just add a lower bound on those but we can’t do that with s2z so the way to do it is to reparameterize

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(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(),
  prior = prior_loss,
  backend = "cmdstanr",
  chains = 4,
  cores = 4,
  iter = 1000,
  control = list(
    adapt_delta = 0.8,
    max_treedepth = 10
  )
)

[spinkney]

Both of these examples now work. This is a work-in-progress pr so don’t expect the documentation to help much yet. I appreciate the feedback that this will be necessary to help users.