Redundant computations with multilevel models and Slice sampler

Sorry it took a little bit, but here’s a notebook: colab notebook

The idea is just a simple recovery of 1-D gaussians; hopefully an easy task for the samplers. In summary: the notebook generates data from univariate Normal distributions, a different distribution for each of M groups. Then it tries to recover the parameters of those distributions using MCMC. It keeps track of a few statistics for each likelihood node: number of logp calls, number of unique distribution parameters seen, and number of repeat parameters seen. In order to do that, the Normal likelihood is a DensityDist wrapping an Aesara op for each node; and the op uses a cache to map parameter values to already-computed logps. It tries sampling with Slice and Metropolis on datasets generated from M=2^0,\dots,2^4 groups. I think it makes sense to keep the total number of observations fixed, so that’s what it does here.

And here’s a summary of those statistics for each run!

Interesting to note that the number of logp() calls to each likelihood node seems to grow linearly, as does the number of logp() calls with repeated parameter values. Also interesting to see that the number of unique parameter vals seen by each node doesn’t really change as M changes. But as the number of groups (and hence likelihood nodes) also grows linearly, computation time grows exponentially.