Generalized Variational Inference in Function Spaces: Gaussian Measures meet Bayesian Deep Learning

We develop a framework for generalized variational inference in infinitedimensional function spaces and use it to construct a method termed Gaussian Wasserstein inference (GWI). GWI leverages the Wasserstein distance between Gaussian measures on the Hilbert space of square-integrable functions in order to determine a variational posterior using a tractable optimization criterion. It avoids pathologies arising in standard variational function space inference. An exciting application of GWI is the ability to use deep neural networks in the variational parametrization of GWI, combining their superior predictive performance with the principled uncertainty quantification analogous to that of Gaussian processes. The proposed method obtains state-of-the-art performance on several benchmark datasets. * equal contribution, order decided by coinflip † Work primarily done at the University of Oxford and finished at Amazon. 36th Conference on Neural Information Processing Systems (NeurIPS 2022).