Abstract:We present new Bayesian Last Layer models in the setting of multivariate regression under heteroscedastic noise, and propose an optimization algorithm for parameter learning. Bayesian Last Layer combines Bayesian modelling of the predictive distribution with neural networks for parameterization of the prior, and has the attractive property of uncertainty quantification with a single forward pass. The proposed framework is capable of disentangling the aleatoric and epistemic uncertainty, and can be used to transfer a canonically trained deep neural network to new data domains with uncertainty-aware capability.

What problem does this paper attempt to address?

This paper proposes a new Multivariate Bayesian Last Layer (BLL) model for addressing regression problems under heteroscedastic noise. In deep neural networks (DNNs), uncertainty quantification (UQ) of predictions is becoming increasingly important as it involves two types of uncertainties: aleatoric uncertainty and epistemic uncertainty. The focus of the paper is to develop a framework that can differentiate between these two uncertainties and enable the trained DNN to have uncertainty-aware capability in new data domains. The paper first reviews traditional linear regression models and then extends them to the case of multivariate heteroscedasticity. It introduces a new BLL model where the neural network is used to parameterize the prior distribution, allowing for one-shot forward propagation to quantify uncertainty. Through theoretical analysis of matrix variable distributions, the paper demonstrates how to decouple aleatoric and epistemic uncertainties in high dimensions, similar to the Deep Evidential Regression (DER) approach. The paper also discusses the issue of parameter estimation, particularly in the optimization of evidence (or marginal likelihood), where gradient estimates may be inaccurate due to the non-independence of batch data. To address this problem, they propose an algorithm called ELBO-EM (Evidence Lower Bound Expectation Maximization), which combines variational inference and expectation maximization (EM) to efficiently handle batch training and transfer learning. The experimental section showcases the applications of BLL in toy models, transfer learning, and multivariate time series prediction, demonstrating the effectiveness of the approach, especially in uncertainty quantification and model generalization capability. In summary, the main contribution of this paper is the development of a new BLL framework for handling uncertainty in multivariate regression, with particular emphasis on heteroscedastic noise, and providing an efficient optimization algorithm for batch training.

Multivariate Bayesian Last Layer for Regression: Uncertainty Quantification and Disentanglement

Improved uncertainty quantification for neural networks with Bayesian last layer

Variational Bayesian Last Layers

On Last-Layer Algorithms for Classification: Decoupling Representation from Uncertainty Estimation

Bayesian Layers: A Module for Neural Network Uncertainty

Flexible Bayesian Last Layer Models Using Implicit Priors and Diffusion Posterior Sampling

Bayesian Neural Networks Uncertainty Quantification with Cubature Rules

Density regression and uncertainty quantification with Bayesian deep noise neural networks

Bayesian Deep Convolutional Encoder-Decoder Networks for Surrogate Modeling and Uncertainty Quantification

Combining Model and Parameter Uncertainty in Bayesian Neural Networks

A Kernel Framework to Quantify a Model's Local Predictive Uncertainty under Data Distributional Shifts

Understanding Uncertainty in Bayesian Deep Learning

Uncertainty Quantification for nonparametric regression using Empirical Bayesian neural networks

Sparsifying Bayesian neural networks with latent binary variables and normalizing flows

High Accuracy Uncertainty-Aware Interatomic Force Modeling with Equivariant Bayesian Neural Networks

Variational Inference on the Final-Layer Output of Neural Networks

Bayesian meta learning for trustworthy uncertainty quantification

Functional Bayesian Neural Networks for Model Uncertainty Quantification

Bayesian Neural Networks: Essentials

Scalable Bayesian Inference in the Era of Deep Learning: From Gaussian Processes to Deep Neural Networks

Density Uncertainty Layers for Reliable Uncertainty Estimation