Bayesian Deep Learning with Multilevel Trace-Class Neural Networks

Source abstract

Abstract. In this article we consider Bayesian inference associated to deep neural networks (DNNs) and in particular, trace-class neural network (TNN) priors [T. Sell and S. S. Singh, Dimension-robust Function Space Mcmc with Neural Network Priors, 2020] which can be preferable to traditional DNNs because they (a) are identifiable and (b) possess desirable convergence properties. TNN priors are defined on functions with infinitely many hidden units, and have strongly convergent Karhunen–Loeve-type approximations with finitely many hidden units. A practical hurdle is that the Bayesian solution is computationally demanding, requiring simulation methods, so approaches to drive down the complexity are needed. In this paper, we leverage the strong convergence of TNN in order to apply multilevel Monte Carlo (MLMC) to these models. In particular, an MLMC method that was introduced by Beskos et al. [ SIAM/ASA J. Uncertain. Quantif., 6 (2018), pp. 762–786] is used to approximate posterior expectations of Bayesian TNN models with optimal computational complexity, and this is mathematically proved. The results are verified with several numerical experiments on model problems arising in machine learning, including some toy regression and classification models, MNIST image classification, and a challenging reinforcement learning problem. Furthermore, we illustrate the practical utility of the method on MNIST as well as IMDb sentiment classification.