The high cost of split R-hat

This post is by Bob.

I’ve been thinking a lot lately about R-hat given that I’m using it for online converging monitoring in our new Walnuts implementation. In that setting, where I use Welford accumulators to update R-hat estimates every iteration, I can’t use split R-hat without way too much buffering. So I’ve been thinking about the effect of splitting, too, and whether we need it. I asked Andrew and he said Kenny Shirley once produced an example where split R-hat diagnosed non-convergence that regular R-hat didn’t, but that example is lost to time and we’ve never seen this kind of behavior with NUTS as far as I know (please give us an example in the comments or via email to Andrew if you have).

Relating R-hat and ESS

My intuition was that we could set a low enough R-hat threshold that it would ensure a high enough effective sample size (ESS) when we crossed it. The relation’s a little tighter than I thought, with

    Rhat^2 ≈ 1 + M / ESS,

where M is the number of chains and ESS is the effective sample size of all chains combined. There’s a multivariate proof in Vats and Knudson, 2021, Revisitng the Gelman-Rubin diagnostic, Statistical Science, page 2 and section 5 for details, but it’s pretty straightforward to get the intuition when you reduce R-hat^2 to (N-1)/N + var(chain-means) / man(chain-variances) as Charles Margossian did in his nested R-hat paper. Vats and Knudson disapprove of Andrew and Aki’s suggested threshold of 1.1 from BDA3, because it is satisfied with a combined ESS of 20 across Andrew’s default 4 chains.

Being me, I tried to validate my intuition with simulations rather than linear algebra. Also, I like to see that things work in practice that theory entails to make sure I’ve understood all the assumptions baked into the theory (one can’t prove anything without assumptions!). When asked to code a simulation using ArviZ, Claude inserted a (2 * M) in the numerator in place of the M. Where did that come from, I asked? It told me it needed the factor of 2 because ArviZ uses split Rhat. D’oh! Of course it does, because we’ve doubled M without increasing ESS.

A worked example

Suppose we have 4 chains with a combined ESS of 400. Then sqrt(1 + 4/400) ≈ 1.005 and sqrt(1 + (2 * 4) / 400) ≈ 1.01. We’ve effectively doubled the number after the 1 by splitting. Unlike Vats and Knudson, I usually don’t need an ESS >> 100, so the 400 required for split R-hat < 1.01 is perhaps a bit too conservative for my tastes. On the other hand, we face a practical problem estimating ESS reliably with fewer than 50 or so ESS per chain. Estimation is challenging because it relies on autocorrelation estimates from the chains themselves, which become much noisier when based on shorter chains. (Side question: Do we not combine autocorrelation estimates across chains to reduce standard error because some chains might not be mixing?) Also, we know this algebra wasn't a coincidence of 4 chains and 400 draws. The Taylor expansion of sqrt(1 + x) is the convergent sequence

    sqrt(1 + x) = 1 + x/2 - x^2 / 8 + x^3 / 16 + ...

When x < 0.1, the first-order approximation, sqrt(1 + x) = 1 + x / 2, is good.

The bottom line for practitioners

We need around twice as many draws to get below a fixed threshold with split R-hat than with the original R-hat.