Hamiltonian Markov chain Monte Carlo for partitioned sample spaces with application to Bayesian deep neural nets

Allocating computation over multiple chains to reduce sampling time in MCMC is crucial in making MCMC more applicable in the state of the art models such as deep neural networks. One of the parallelization schemes for MCMC is partitioning the sample space to run different MCMC chains in each component of the partition (VanDerwerken and Schmidler in Parallel Markov chain Monte Carlo. arXiv:1312.7479, 2013; Basse et al. in Artificial intelligence and statistics, pp 1318-1327, 2016). In this work, we take Basse et al. (2016)'s bridge sampling approach and apply constrained Hamiltonian Monte Carlo on partitioned sample spaces. We propose a random dimension partition scheme that combines well with the constrained HMC. We empirically show that this approach can expedite MCMC sampling for any unnormalized target distribution such as Bayesian neural network in a high dimensional setting. Furthermore, in the presence of multi-modality, this algorithm is expected to be more efficient in mixing MCMC chains when proper partition elements are chosen.