JAGS vs Stan performance - how to compare 'workflow efficiency'?

Hi everyone,

I have a bunch of models in JAGS and have translated them to Stan. The Stan versions, however, are not identical to the JAGS versions (I had to introduce some smoothing for Stan to work). The posteriors of the shared parameters are comparable, as are the predictions made by the models, and also biological reasons for the smoothing introduced in the Stan models exist. I am happy to this point.

Now the reason to move to Stan was that JAGS was very slow for some more complex model structures. Stan is noticeably faster. I would like to somehow quantify this, but it seems tricky to find a good measure for this. I thought of things like minimum ESS of all parameters / time.

I also came across these older posts where similar issues are discussed:

In both cases, issues with the used methods are pointed out. Since these threads are already some years old, I wanted to ask if there is now a consensus on how to do this. Since my Stan models include smoothing that the JAGS models lack, I realize this isn’t a pure head-to-head algorithm benchmark on an identical target distribution, but rather a workflow efficiency comparison (i.e., time to reach valid, usable inference for the biological problem).

Thanks for any thoughts!

My thoughts would be to use a simulated data set with known values and see how well you recover values. Or, use a wide range of parameters and fit multiple parameters if you are interested in validating of the model rather than a single use application.

Hi Richard, thanks for the input! Parameter recovery with simulated data is definitely a great check for validating that the added smoothing doesn’t introduce bias. I did some checks in that direction, and I am happy with the models. However, my question is specifically about metrics for computational efficiency rather than model validation. I would like to know what is the best way to report / measure the speedup? Is minimum ESS/seconds a good approach? I know it is a bit apples to oranges because the models are not identical, but for this “real-world scenario” of adjusting a model slightly to be able to move to Stan, how could one quantify efficiency gains? Cheers, Friedrich

Have you seen this paper? It has some ideas

https://sites.stat.columbia.edu/gelman/research/published/nestedRhat.pdf

Thanks, it is new to me, I will read it.

You’re welcome.

Also I shared the wrong article. I am on a phone and away from my computer.

I think this is the one I wanted to share

https://sites.stat.columbia.edu/gelman/research/published/Vehtari_etal_2020_rhat_ess.pdf

Vehrari, German, and others have a new r hat method.

I’d also recommend Cole Monnahan et al.'s paper, " Faster estimation of Bayesian models in ecology using Hamiltonian Monte Carlo":

https://besjournals.onlinelibrary.wiley.com/doi/pdf/10.1111/2041-210x.12681

Could you say more about what this means? I’m particularly curious about JAGS models that can’t be translated to Stan.

This is standard, but noisy, because the ESS estimators are noisy. Plus, you can get ESS > 1 in cases where the ESS = 0 if you’re only subsampling a subspace of the posterior. So you need an independent measure from ESS to make sure you’re getting the correct answer.

If you want an example to verify this yourself, try to fit Neal’s funnel straight up in 100 dimensions in either JAGS or Stan. You will not get samples from the posterior—they’ll all be biased. But That and ESS are going to happily report positive values because it mixes well within the region that it does explore.

There are two ways to do this. One, you can measure coverage calibration, but that doesn’t really get at efficiency.

The other common measure is squared error against known expectation values. Even if you simulate data, if you do not have an analytic posterior, you will have to sample to get the values you need to recover. This is what posteriordb does, for example.

Definitely a better choice than nested R-hat (great paper with a lot of insight about R-hat, though, and totally worth reading, especially if you really want to understand the bias and variance decomposition of R-hat and the role of transient bias [before convergence]). Aki’s paper (cited above) is about how to measure R-hat and ESS more robustly in several cases where ordinary R-hat falls down (e.g., wide tailed distributions). So I would definitely recommend using this revised approach to R-hat estimation. It’s coded in posterior in R and arviz in Python. I don’t know how easy it would be to use posterior with JAGS if you’re going to do that.

Also, I’d recommend paying attention to everything @betanalpha had to say on those previous threads.

P.S. JAGS just released version 5.0.

Thanks both, and thanks Bob for the detailed answer—the Monnahan et al. paper is new to me, and I’ll work through it.

On the smoothing: The models are Bayesian wrappers around a deterministic crop-phenology simulator: a day-by-day development recursion driven by temperature and photoperiod, with observations recorded as development-stage scores. There are different versions of this model with different combinations of predictors. There is no analytic posterior. These models (and the JAGS versions) use multiple ifelse() statements for temperature boundaries and transitions between plant development stages. This is where I use sigmoids for smoothing. The steepnesses of the sigmoids are parameters too, so the Stan versions have more parameters to estimate than the JAGS versions. However, the steepness of the temperature gates is not identified; its posterior simply matches the prior, and predictions are invariant once it exceeds a low threshold, including in the hard-step limit. So, at the operating point, the sigmoid and step versions produce the same forward map, which is what I think makes the comparison meaningful rather than apples-to-oranges.

On the comparison itself: I should clarify that this is a “soft-quantitative workflow check”, something like “What does it cost me to get usable inference for this problem with each tool?”, rather than a strict sampler benchmark on a shared target. Same hardware, same chain/core layout, three seeds per model, and rank-normalized Rhat / ESS from the posterior package for both engines (feeding JAGS output through posterior works fine, for what it’s worth).

On your ESS caveat: I report ESS/time only for runs that pass the convergence criterion, and I use an independent check: regenerating out-of-sample predictions from every fit’s posterior and scoring them against held-out data. The more demanding models do not converge in JAGS within the time budget I set, so for the comparison with Stan I will just state that JAGS did not converge whereas Stan did converge. For the models that converge in both Stan and JAGS I report minimum ESS per time.

Headline numbers (with the above caveat attached): Minimum ESS/time ratios are roughly 20–50x in Stan’s favor (depending on predictors and model complexity), with JAGS receiving equal or greater wall-clock time than Stan in three of the five models. I will standardize the wall-clock time by running Stan for more iterations, determining the iteration reached at a fixed wall-clock threshold, and discarding samples beyond that point. For JAGS, I can set wall-clock time directly. I’m running additional simulations now to check the stability of these numbers. A few factors I haven’t checked yet include the impact of the warmup-to-iteration ratio in Stan and adaptation tuning in JAGS (beyond ensuring I don’t trigger insufficient adaptation warnings).