Bayesian Deep Learning with Multilevel Trace-class Neural Networks

TL;DR

This paper develops a principled framework for Bayesian deep learning by coupling trace-class neural network priors with multilevel Monte Carlo, enabling Bayesian inference for deep nets with infinite width while achieving the canonical $O(\varepsilon^{-2})$ computational cost. The authors derive strong convergence rates for trace-class priors across multilevel widths, establish an MSE bound for the MLSMC estimator, and prove that the estimator converges to the true posterior under suitable conditions. They implement a multilevel sequential Monte Carlo sampler (MLSMC) using pCN-type mutations and level-wise width increments $n_l=2^l$, and validate the theory through extensive experiments in regression, classification (including MNIST and IMDb), and reinforcement learning, demonstrating substantial computational savings over single-level methods. The results highlight the practical viability of scalable Bayesian deep learning with uncertainty quantification, and point to future directions such as multi-index MLMC and integration with more advanced neural architectures.

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.

Figures (14)

• Figure 1: Increment 2nd moment vs. levels. The decay is $\mathcal{O}(2^{-3l})$. Therefore the variance decays faster than $1/$cost, which is the canonical regime. Left: activation function of $\textrm{ReLU}(z) =\max\{0,z\}$. Right: activation function of $\sigma(z) =\tanh(z)$.
• Figure 1: Regression problem: error vs cost plots for SMC and MLSMC using TNN priors. Top left: $\alpha=3$. Top right: $\alpha = 2.0$. Bottom left: $\alpha=1.9$. Bottom right: $\alpha=1.7$. Credible sets are provided in the thin blue and red curves.
• Figure 2: Increment 2nd moment vs. levels. The decay is $\mathcal{O}(2^{-l})$. Therefore the variance decays slower than $1/$cost, which is a sub-canonical regime. Left: activation function of $\textrm{ReLU}(z) =\max\{0,z\}$. Right: activation function of $\sigma(z) =\tanh(z)$.
• Figure 2: Regression problem: error vs cost plots for SMC and MLSMC using TNN priors. Left: $\alpha=1.4$. Right: $\alpha = 1.1$. Credible sets are provided by thin blue and red curves.
• Figure 3: Classification problem: Our data is generated as a 2D spiral with two classes, Class $k=1$ being in blue and Class $k=2$ in yellow.

Theorems & Definitions (15)

• Theorem 2.1: Giles MBG08
• Proposition 3.1
• Proposition 4.1
• Proof 1
• Remark 4.2: Strong Convergence
• Remark 4.3: Smoothness
• Corollary 4.4
• Remark 4.5
• Remark 4.6: Mutation kernel
• Remark 5.1: Parameter Tuning