Variable Depth Bayesian Neural Networks Using Reversible Jumps

Bayesian approaches to Neural Architecture Search (NAS) offer several benefits when compared to typical grid-search methods. Marginalization of a joint posterior distribution of network parameters as well as network architecture specification allows for a principled analysis of predictive and model uncertainty. In this paper, we propose a method of applying Reversible Jump Markov Chain Monte Carlo (RJMCMC) as inference over a distribution of neural network depths. We refer to these models as Reversible Jump Bayesian Neural Networks (RJBNN), and show that a tenable acceptance probability for across-dimension proposals can be achieved through the use of a composite sampler featuring the Hamiltonian Monte Carlo (HMC) based No U-Turn Sampler (NUTS) algorithm. We compare the results of an RJBNN model with standard Bayesian Neural Networks (BNN) of comparable architecture sizes on a regression task for the Boston Housing dataset.