Skip to content

Implement split-Rhat diagnostic to detect non-stationary MCMC chains

Vicentini Filippo requested to merge split-Rhat into master

Created by: femtobit

This PR introduces a flag to compute the split-R_hat diagnostic instead of the plain R_hat currently used. Split-R_hat is able to additionally detect non-stationarity within single chains.

See, e.g., Vehtari et al. (arXiv:1903.08008) for a modern discussion of this diagnostic (that paper points out some failure types that split-Rhat does not cover either and proposes an improved version, but let's do one step at a time here).

Here is a very simple example featuring two (identical) linearly increasing and thus non-stationary chains:

# stats_example.py
import netket as nk
import numpy as np

r = np.arange(100, dtype=float)
x = np.array([r, r])
print(nk.stats.statistics(x))

which gives these results (split-Rhat is currently enabled by the NETKET_USE_SPLIT_RHAT flag in this PR):

$ python3 stats_example.py
49.5 ± 3.5 [σ²=833.2, R̂=0.9950]
$ NETKET_USE_SPLIT_RHAT=1 python3 stats_example.py
49.5 ± 3.5 [σ²=833.2, R̂=1.3191]

The non-split Rhat is happy (with R^ ≈ 1.0) because the chains do have identical mean and variance, whereas the split-chain version correctly identifies the failed MC convergence (R^ >> 1.01).

There are some open questions regarding this PR:

  • Should split-Rhat be enabled by default? It makes Rhat strictly more useful, I would say, but this change will alter the vaules of Rhat for basically any VMC run then.
  • If we decide it should be enabled by default, do we need the flag to optionally restore the old behavior?
  • Do we care about the overhead of one more MPI call? (Its probably not a performance issue, but I haven't done any benchmarking yet.)
  • In cases where there is an odd number of samples per chain, I discard one sample per chain to get an even split (which makes the code much simpler). I think this is fine, unless in pathological cases where there are very few samples per chain (in which case the stationarity of the chains isn't a sensible property anyways).

Merge request reports