Semi−separable Hamiltonian Monte Carlo for inference in Bayesian neural networks
Adam D Cobb‚ Atılım Güneş Baydin‚ Ivan Kiskin‚ Andrew Markham and Stephen Roberts
Abstract
We introduce a new method for performing inference in Bayesian neural networks (BNNs) using Hamiltonian Monte Carlo (HMC). We show how the previously introduced semi-separable HMC sampling scheme can be adapted to BNNs, which allows us to integrate over both the parameters and hyperparameters. We derive a suitable Riemannian metric for the BNN hyperparameters and show that it is positive definite. Our work is compared to both Monte Carlo dropout and a deterministic neural network, where our inference technique displays better calibrated uncertainties with comparable performance to current baselines. Our code is provided in a new open-source Python package, hamiltorch, which enables our method to scale to CNNs with over 400,000 parameters and take advantage of GPUs.