Implementation of pseudo-marginal HMC

[jacob-hjortlund]

Hi all!

I'm currently working on a research project that involves a hierarchical model with (as far as I've concluded) complicated posterior density. See here for an outline of the model. Expansions of this model to higher dimensions excludes the use of emcee and means that we can't effectively marginalize over the gamma distributed latent parameters mentioned in Appendix A.

I've been working on sampling this model using HMC via numpyro and blackjax but, as I mentioned, I run into complications due to the geometry of the prior. I've been reading around for solutions and become pretty interested in pseudo-marginal HMC, as outlined in Alenlöv et al. 2016.

Have there been any attempts, or are there any plans, to implement such a sampler? If not, then I might be willing to give it a shot. I imagine it wouldn't be too hard to build off of the existing HMC kernel, assuming I don't implement the integrator mentioned in the article. Be great to hear your takes on it before I dig into it further.

[albcab]

Hi, thanks for your interest in blackjax!

As much as it'd be great to have pseudo-marginal HMC implemented on blackjax, I have to mention that usually you'll get a lot more out of marginalizing parameters, reparametrizing/refactoring, and modifying the hyperparameters of the priors (giving stronger prior information usually) when doing inference. An MCMC method will seldom solve modelling problems (but maybe it would work for you, never used pseudo-marginal MCMC/HMC). I've never worked with the types of models you are mentioning here, so let me know if I'm correct. When you say you cannot "marginalize over the gamma distributed latent parameters", you mean you are not solving the integrals in (A4) numerically and instead sampling the y_i? If so, that looks hard to do, specially because these y_i are probably (strongly) correlated to the rest of the parameters in the model. What about finding their MAP estimates, fixing them and then using the sampler the paper recommends?

That being said, there are no current plans of implementing pseudo-marginal HMC and it would be great if you would implement it. You can open an issue with your proposal for implementation, then after some discussion you can open a PR.