Improved uncertainty quantification for neural networks with Bayesian last layer

Felix Fiedler,Sergio Lucia
DOI: https://doi.org/10.1109/ACCESS.2023.3329685
2024-01-04
Abstract:Uncertainty quantification is an important task in machine learning - a task in which standardneural networks (NNs) have traditionally not excelled. This can be a limitation for safety-critical applications, where uncertainty-aware methods like Gaussian processes or Bayesian linear regression are often preferred. Bayesian neural networks are an approach to address this limitation. They assume probability distributions for all parameters and yield distributed predictions. However, training and inference are typically intractable and approximations must be employed. A promising approximation is NNs with Bayesian last layer (BLL). They assume distributed weights only in the linear output layer and yield a normally distributed prediction. To approximate the intractable Bayesian neural network, point estimates of the distributed weights in all but the last layer should be obtained by maximizing the marginal likelihood. This has previously been challenging, as the marginal likelihood is expensive to evaluate in this setting. We present a reformulation of the log-marginal likelihood of a NN with BLL which allows for efficient training using backpropagation. Furthermore, we address the challenge of uncertainty quantification for extrapolation points. We provide a metric to quantify the degree of extrapolation and derive a method to improve the uncertainty quantification for these points. Our methods are derived for the multivariate case and demonstrated in a simulation study. In comparison to Bayesian linear regression with fixed features, and a Bayesian neural network trained with variational inference, our proposed method achieves the highest log-predictive density on test data.
Machine Learning,Systems and Control
What problem does this paper attempt to address?
The paper attempts to address the problem of better uncertainty quantification in neural networks. Traditionally, standard neural networks perform poorly in uncertainty quantification, which is a significant limitation in applications requiring high reliability, such as autonomous driving and medical diagnosis. Bayesian Neural Networks (BNNs) are an improved approach, but their training and inference are often difficult to handle and require approximations. Specifically, the paper focuses on neural networks with a Bayesian Last Layer (BLL). This type of network assumes a probability distribution for the weights only in the output layer, resulting in normally distributed predictions. However, when training such networks, the computation of the marginal likelihood function is very complex, particularly requiring the inversion of the precision matrix, which is a challenge in practical applications. The main contributions of the paper include: 1. **Proposing a new method for maximizing the marginal likelihood function**: By reintroducing the marginalized last layer weights as optimization variables, it avoids the direct computation of the precision matrix inversion, making gradient-based optimization feasible. 2. **Improving uncertainty quantification for extrapolation points**: Proposing a metric based on affine cost to assess the degree of extrapolation points and providing an algorithm to adjust parameters to enhance the uncertainty quantification performance for these points. These methods not only simplify the training process but also outperform traditional Bayesian linear regression and variational inference-trained Bayesian neural networks in multivariate cases.