Bayesian Deep Learning with Multilevel Trace-class Neural Networks

TL;DR

This paper introduces a Bayesian inference framework for trace-class neural networks (TNNs), leveraging multilevel Monte Carlo methods to efficiently compute posterior expectations, with demonstrated applications in machine learning tasks.

Contribution

It develops an MLMC-based approach for Bayesian TNNs, providing a computationally efficient method with proven optimal complexity, and validates it through numerical experiments.

Findings

MLMC achieves optimal complexity for Bayesian TNNs

Numerical experiments confirm effectiveness in regression, classification, and reinforcement learning

TNN priors are identifiable and have desirable convergence properties

Abstract

In this article we consider Bayesian inference associated to deep neural networks (DNNs) and in particular, trace-class neural network (TNN) priors which can be preferable to traditional DNNs as (a) they are identifiable and (b) they 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 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 regression, classification, and reinforcement learning.