Metropolis-CVAE: Bootstrapping Labels for Bayesian Inference via Semi-Supervised Conditional Variational Autoencoders

Jonathan Doucette, Christian Kames, Alexander Rauscher

Keywords: Metropolis-Hastings, CVAE, Bayesian Inference, MCMC, Semi-Supervised

Abstract: In Bayesian parameter estimation, models make simplifying assumptions to make parameter inference feasible. If learned inference methods are trained using data simulated by models, however, distributional differences between simulated and observed data may lead to biased inference results on the observed data. In this work, we introduce a semi-supervised learned Bayesian inference method which makes use of both simulated data -- for which the underlying parameters are known by construction -- and unlabeled data, which may depend on nuissance parameters not captured by the simulation procedure. A conditional variational autoencoder (CVAE) is trained to perform approximate inference simultaneously on the sets of labeled simulated data and unlabeled data, where the unlabeled data is initialized with arbitrary pseudo labels. At each training iteration, new candidate pseudo labels are drawn from the CVAE posterior and the pseudo labels are updated using the Metropolis-Hastings algorithm. This process results in a Markov chain of bootstrapped pseudo labels for each unlabeled datum, effectively performing online Markov chain Monte Carlo (MCMC) inference wherein the proposal distribution is a CVAE informed by labeled simulated data, producing proposals which are increasingly likely to be accepted as training proceeds. The resulting CVAE is able to efficiently produce samples from the posterior distributions of both the simulated and unlabeled data, implicitly marginalizing over nuissance parameters in the unlabeled data. We demonstrate the effectiveness of this method in magnetic resonance imaging (MRI) where MCMC is computationally impractical to due the (3+1)D nature of the images, showing improvement against traditional MCMC inference in both speed and posterior quality.

Code Of Conduct: I certify that all co-authors of this work have read and commit to adhering to the NeurIPS Statement on Ethics, Fairness, Inclusivity, and Code of Conduct.

TL;DR: Metropolis-CVAEs perform approximate Bayesian inference on unlabeled measured data by utilizing simulated data - which may be out-of-distribution from the unlabeled data - to bootstrap pseudo labels onto the unlabeled data during training.