Improved uncertainty quantification for neural networks with Bayesian
  last layer
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.
Felix Fiedler, Sergio Lucia
# Improved uncertainty quantification for neural networks with Bayesian last layer

###### Abstract

Uncertainty quantification is an essential task in machine learning - a task in which neural networks (NNs) have traditionally not excelled. Bayesian neural networks (BNNs), in which parameters and predictions are probability distributions, can be a remedy for some applications, but often require expensive sampling for training and inference. NNs with Bayesian last layer (BLL) are simplified BNNs where only the weights in the last layer and the predictions follow a normal distribution. They are conceptually related to Bayesian linear regression (BLR) which has recently gained popularity in learning based-control under uncertainty. Both consider a non-linear feature space which is linearly mapped to the output, and hyperparameters, for example the noise variance, For NNs with BLL, these hyperparameters should include the deterministic weights of all other layers, as these impact the feature space and thus the predictive performance. Unfortunately, the marginal likelihood is expensive to evaluate in this setting and prohibits direct training through back-propagation.

In this work, we present a reformulation of the BLL log-marginal likelihood, which considers weights in previous layers as hyperparameters and allows for efficient training through back-propagation. Furthermore, we derive a simple method to improve the extrapolation uncertainty of NNs with BLL. In a multivariate toy example and in the case of a dynamic system identification task, we show that NNs with BLL, trained with our proposed algorithm, outperform standard BLR with NN features.

 Bayesian last layer, Bayesian neural network, uncertainty quantification, system identification

## I Introduction

Machine learning tries to capture real world patterns and trends through data. Both, the data and these underlying patterns are subject to uncertainty. For many applications, especially those where machine learning is applied to safety critical tasks, it is imperative to predict the pattern and its corresponding uncertainty. An important example is learning-based control under uncertainty, where probabilistic system models are identified from data and used for safe control decisions [1, 2, 3, 4, 5, 6].

Bayesian linear regression (BLR) [1, 2, 3] and Gaussian processes (GPs) [1, 4] are prominent methods for probabilistic system identification. Both assume that a nonlinear feature space can be mapped linearly to the outputs, with the main difference being that for BLR the nonlinear features are explicitly defined, while in GPs they are substituted through a kernel function [7]. This is a major advantage of GPs, as the most suitable features for BLR are often challenging to determine. On the other hand, GPs scale poorly with the number of data samples [8]. For big data problems they are typically approximated with sparse GPs [9], and can be further improved with deep kernel learning [10, 11].

Especially in recent years, neural networks (NNs) and deep learning have gained significant popularity for a vast variety of machine learning tasks [12]. NNs have also been successfully applied for control applications, often to infer the system model from data [13, 14, 15, 16]. Unfortunately, one of the main challenges with NNs is their tendency to overfit and their inability to express uncertainty [17]. Bayesian neural networks (BNNs), in which the weights and predictions are probability distributions, are a concept to tackle this shortcoming. Unfortunately, BNNs can be intractable to train and query and are often approximated in practice [17, 18]. Many of the approximate BNN techniques are sampling-based, e.g. Monte Carlo dropout [19], Markov chain Monte Carlo [20] and variational inference [21].

While sampling can approximate arbitrary predictive distributions, it can be a disadvantage in comparison to GPs and BLR which yield analytical results. A promising compromise between tractability and expressiveness are NNs with Bayesian

Fig. 1: Multivariate neural network with Bayesian last layer: Predicted mean and standard deviation for two outputs with different and unknown noise level. Training with the proposed Algorithm 1, which maximizes the log-marginal likelihood in (50), and yields the optimal parameter \(\alpha^{*}\). This value can then be adapted to improve the uncertainty quantification in the extrapolation regime by maximizing (42), yielding \(\alpha^{\max}\).
last layer (BLL) [22, 8, 23]. These networks can be seen as a simplified BNN, where only the weights of the output layer follow a Gaussian distribution, and the remainder of the layers contain deterministic weights. At the same time, they can be interpreted as a deep kernel learning approach with linear kernel. NNs with BLL are also stronlgy related to BLR in that they consider a nonlinear feature space which is mapped linearly onto the outputs. Similarly to GPs, BLR and other probabilistic models [7, 24], NNs with BLL are trained by maximizing the marginal likelihood with respect to the model hyperparameters. These hyperparameters include prior and noise variance and, importantly, the weights of the deterministic layers. While the log-marginal likelihood (LML) can be expressed analytically for NNs with BLL, it contains expressions such as the inverse of the precision matrix, making it unsuitable for direct gradient-based optimization. Previous works, especially for control applications, have therefore either assumed knowledge of all hyperparameters [2, 6, 16], including prior and noise variance, or have maximized the LML after training a NN with fixed features [5, 25]. In previous works that did include the deterministic weights as hyperparameters, maximizing the LML required sampling the surrogate posterior during training [23], or using an approximate precision matrix [10] to enable gradient-based optimization.

As a main contribution of this work, we suggest an approach to maximize the exact LML of a NN with BLL that does not require sampling and is suitable for gradient-based optimization. Most importantly, we avoid the matrix inverse in the LML by reintroducing the weights of the last layer, which were marginalized, as optimization variables. We show that our reformulation of the LML satisfies the conditions of optimality for the same solution as the original formulation. In this way, we provide a simpler training procedure, in comparison to training with variational inference, and can consistently outperform BLR with fixed features obtained from a NN. Our second main contribution is a simple algorithm to improve the BLL uncertainty quantification for extrapolation points. To this end, we relate the computation of the BLL covariance to an intuitive metric for extrapolation, which is inspired by the definition of extrapolation in [26]. Based on this relationship, the proposed algorithm adjusts a scalar parameter to improve the log-predictive density on additional validation data. Our proposed methods are derived for the multivariate-case, estimating individual noise-variances for each output with a first impression shown in Figure 1. Especially for the application of system identification, the multivariate case, using GPs and BLR, has previously been tackled by fitting an individual model to each output [1, 4, 5] or to assume i.i.d. noise for all outputs [2]. All of the presented algorithms and our results are available on online1.

Footnote 1: [https://github.com/4flixt/2022_Paper_BLL_JML](https://github.com/4flixt/2022_Paper_BLL_JML).

This work is structured as follows. In Section II, we introduce NNs with BLL. The marginal likelihood, which is maximized during training, is discussed in Section III. In Section IV, we discuss interpolation and extrapolation for NNs with BLL and present an algorithm to improve the predictive distribution in the extrapolation regime. Finally, we present the application of BLL to a probabilistic system identification example in Section VI. The paper is concluded in Section VII.

## II Bayesian last layer

We investigate a dataset \(\mathcal{D}=(\mathbf{X},\mathbf{t})\) consisting of \(m\) data pairs of inputs \(\mathbf{x}\in\mathbb{R}^{n_{x}}\) and targets \(t\in\mathbb{R}\) from which the matrices \(\mathbf{X}=\left[\mathbf{x}_{1},\dots,\mathbf{x}_{m}\right]^{\top}\in\mathbb{R}^{m\times n _{x}}\) and \(\mathbf{t}=[t_{1},\dots,t_{m}]^{\top}\in\mathbb{R}^{m\times 1}\) are formed. We assume a scalar output for ease of notation and address the multivariate case in Section V. For the regression task, we introduce a feed-forward NN model with \(L\) hidden layers as

\[y=NN(\mathbf{x};\mathbb{W}_{L+1})=g_{L+1}\circ h_{L+1}\circ\dots\circ g_{1}\circ h _{1}(\mathbf{x}), \tag{1}\]

where \(\circ\) denotes function composition, \(g_{l}(\cdot),\forall l=1,\dots,L\) are activation functions and

\[\mathbf{h}_{l} =h_{l}(\mathbf{\phi}_{l-1})=\left[\mathbf{\phi}_{l-1}^{\top},\ 1\right]\mathbf{W}_{l} \tag{2}\] \[\mathbf{\phi}_{l} =g_{l}(\mathbf{h}_{l}), \tag{3}\]

are the operations with \(n_{\phi_{l}}\) neurons in layer \(l\). The set of weight matrices with cardinality \(L+1\) is denoted \(\mathbb{W}_{L+1}=\left\{\mathbf{W}_{1},\dots,\mathbf{W}_{L+1}\right\}\). As a requirement for BLL, we state the the following assumption.

**Assumption 1**.: _The NN (1) has a linear activation function in the output layer, i.e. \(g_{L+1}(h_{L+1})=h_{L+1}\)._

With the linear mapping, the output of the last internal layer will be of particular importance, such that we introduce the simplified notation:

\[\tilde{\mathbf{\phi}}=\mathbf{\phi}_{L},\ \text{and}\ \mathbf{\phi}=[\tilde{\mathbf{\phi}}^{ \top},\ 1]^{\top}, \tag{4}\]

where \(\tilde{\mathbf{\phi}}\in\mathbb{R}^{n_{\tilde{\phi}}}\) are referred to as _affine features_ and \(\mathbf{\phi}\in\mathbb{R}^{n_{\phi}}\) are the corresponding _linear features_ with \(n_{\phi}=n_{\tilde{\phi}}+1\). With slight abuse of notation, we use \(\tilde{\mathbf{\phi}}=\tilde{\mathbf{\phi}}(\mathbf{x};\mathbb{W}_{L})\) and \(\mathbf{\phi}=\mathbf{\phi}(\mathbf{x};\mathbb{W}_{L})\) as both, the values of the features, and the function, parameterized with \(\mathbb{W}_{L}\). We require two additional assumptions to state Lemma 1 which formally describes a NN with BLL.

**Assumption 2**.: _The NN (1) provides a feature space \(\mathbf{\phi}(\mathbf{x};\mathbb{W}_{L})\in\mathbb{R}^{n_{\phi}}\) from which the targets are obtained through a linear mapping, according to Assumption 1:_

\[\mathbf{t}=\mathbf{\phi}\left(\mathbf{X};\mathbb{W}_{L}\right)^{\top}\mathbf{w}+\epsilon=\mathbf{y }+\epsilon, \tag{5}\]

_where we introduce the set \(\mathbb{W}_{L}=\left\{\mathbf{W}_{1},\dots,\mathbf{W}_{L}\right\}\) for all weights until layer \(L\) and have \(\mathbf{w}=\mathbf{W}_{L+1}\) the weights of the output layer. The additive noise \(\mathbf{\epsilon}\in\mathbb{R}^{m}\) is zero-mean Gaussian normal distributed, that is, \(\epsilon\sim\mathcal{N}(0,\mathbf{\Sigma}_{E})\)._

**Assumption 3**.: _We have a prior belief for the weights of the output layer \(\mathbf{w}\sim\mathcal{N}(0,\mathbf{\Sigma}_{\mathbf{W}})\)._

This is exactly the setting of BLR [24], for which we introduce hyperparameters \(\mathbf{\Theta}=\{\mathbb{W}_{L},\mathbf{\Sigma}_{\mathbf{E}},\mathbf{\Sigma}_{\mathbf{W}}\}\) and state the posterior distribution of the parameters \(\mathbf{w}\) using Bayes' law:

\[p(\mathbf{w}|\mathcal{D},\mathbf{\Theta})=\frac{p(\mathcal{D}|\mathbf{w},\mathbf{ \Theta})p(\mathbf{w}|\mathbf{\Theta})}{p(\mathcal{D}|\mathbf{\Theta})} \tag{6}\]
This allows to formulate the Lemma of Bayesian last layer.

**Lemma 1** (Bayesian last layer [23]).: _If Assumptions 1-3 hold, we obtain a normal distribution for the predicted outputs:_

\[p(\mathbf{y}|\mathcal{D},\mathbf{\Theta},\mathbf{x}) =\mathcal{N}(\mathbf{\mu}_{y}^{\text{\tiny{N}}}(\mathbf{x}),\mathbf{\Sigma}_{y }^{\text{\tiny{N}}}(\mathbf{x})), \tag{7a}\] \[\mathbf{\mu}_{y}^{\text{\tiny{N}}}(\mathbf{x}) =\text{NN}(\mathbf{x};\mathbb{W}_{L+1}),\] (7b) \[\mathbf{\Sigma}_{y}^{\text{\tiny{N}}}(\mathbf{x}) =\mathbf{\phi}^{\top}\mathbf{\Lambda}_{\mathbf{\gamma}}^{-1}\mathbf{\phi},\] (7c) \[\mathbf{\Lambda}_{p}=\mathbf{\Sigma}_{p}^{-1} =\mathbf{\Phi}^{\top}\mathbf{\Sigma}_{E}^{-1}\mathbf{\Phi}+\mathbf{\Sigma}_{\mathbf{W }}^{-1}, \tag{7d}\]

_where \(\mathbf{\Phi}=\mathbf{\phi}(\mathbf{X};\mathbb{W}_{L})\) is the feature matrix for the training data, \(\mathbf{\phi}=\mathbf{\phi}(\mathbf{x};\mathbb{W}_{L})\) is the feature matrix test data, and \(\mathbf{\Lambda}_{p}\) is the precision matrix._

Proof.: The posterior \(p(\mathbf{w}|\mathcal{D},\mathbf{\Theta})\) in (6) yields a normal distribution in the weights:

\[\mathbf{w} \sim\mathcal{N}(\bar{\mathbf{w}},\mathbf{\Sigma}_{p}) \tag{8a}\] \[\text{with:}\quad\mathbf{\Sigma}_{p}^{-1}=\mathbf{\Lambda}_{p} =\mathbf{\Phi}^{\top}\mathbf{\Sigma}_{E}^{-1}\mathbf{\Phi}+\mathbf{\Sigma}_{W}^{-1},\] (8b) \[\bar{\mathbf{w}} =\mathbf{\Lambda}_{p}^{-1}\mathbf{\Phi}^{\top}\mathbf{\Sigma}_{E}^{-1}\mathbf{t}, \tag{8c}\]

as shown for the case of arbitrary features in [24]. Finally, the predicted output \(\mathbf{y}\) is a linear transformation of the random variable \(\mathbf{w}\), yielding the distribution in (7). 

We can also obtain a posterior distribution for the targets, for which we consider (5) and obtain

\[p(\mathbf{t}|\mathcal{D},\mathbf{\Theta},\mathbf{x}) =\mathcal{N}(\mathbf{\mu}_{y}^{\text{\tiny{N}}}(\mathbf{x}),\mathbf{\Sigma}_{t }^{\text{\tiny{N}}}(\mathbf{x})), \tag{9a}\] \[\mathbf{\Sigma}_{t}^{\text{\tiny{N}}}(\mathbf{x}) =\mathbf{\Sigma}_{y}^{\text{\tiny{N}}}(\mathbf{x})+\mathbf{\Sigma}_{\mathbf{E}}, \tag{9b}\]

where \(\mathbf{\mu}_{y}^{\text{\tiny{N}}}(\mathbf{x})\) and \(\mathbf{\Sigma}_{y}^{\text{\tiny{N}}}(\mathbf{x})\) stem from (7). We will revisit (9) again for the definition of performance metrics.

## III Marginal likelihood maximization

The hyperparameters \(\mathbf{\Theta}\) for NNs with BLL, BLR or GP models are typically trained by maximizing the LML in (6), which goes back to [27] and is coined _empirical Bayes_ or _type-2 maximum likelihood_. In BLR, that is, for a fixed feature space, the LML can be maximized as shown in [24]. Following this idea, the authors in [25] propose a straightforward BLL approach, where a fixed feature space is obtained through classical NN regression. However, a NN with BLL should approximate a full BNN where all weights are assumed to follow an unknown distribution function. It is therefore reasonable to include the set of weights \(\mathbb{W}_{L}\) until the last layer as hyperparameters and train them by maximizing the LML. Below we state the LML of a NN with BLL.

**Lemma 2** (Log-marginal likelihood).: _If Assumptions 1-3 hold, the negative LML of (6), denoted as \(J(\mathbf{\Theta};\mathcal{D})=-\log(p(\mathbf{y}=\mathbf{t}|\mathbf{X},\mathbf{ \Theta}))\), results in:_

\[J(\mathbf{\Theta};\mathcal{D})=\frac{m}{2}\log(2\pi)+\frac{1}{2} \log\det(\mathbf{\Sigma}_{\mathbf{W}})+\frac{1}{2}\log\det(\mathbf{\Sigma}_{E}) \tag{10}\] \[\qquad+\frac{1}{2}\log\det(\mathbf{\Lambda}_{p})+\frac{1}{2}\|\mathbf{t} -\mathbf{y}\|_{\mathbf{\Sigma}_{E}^{-1}}^{2}+\frac{1}{2}\|\bar{\mathbf{w}}\|_{\mathbf{\Sigma} _{\mathbf{W}}^{-1}}^{2}.\]

_In this formulation, we have \(\mathbf{\Lambda}_{p}\), \(\bar{\mathbf{w}}\) according to (8b), (8c), \(\mathbf{y}=\mathbf{\Phi}\bar{\mathbf{w}}\) and \(\mathbf{\Theta}=\{\mathbb{W}_{L},\mathbf{\Sigma}_{\mathbf{E}},\mathbf{\Sigma}_{\mathbf{W}}\}\)._

Proof.: The proof is analog to [24], where it is shown for fixed features. 

Unfortunately, it is typically intractable to determine the hyperparameters \(\mathbf{\Theta}\) by minimizing (10). In particular, we have that \(\mathbf{\Sigma}_{\mathbf{E}}\), which is a hyperparameter, and for which we need to compute the log-determinant, grows with the number of samples in the dataset. We therefore state a simplified formulation of (10) for which we require the following assumption.

**Assumption 4**.: _The additive noise introduced in Assumption 2 is i.i.d. for all samples \(m\), yielding \(\mathbf{\Sigma}_{E}=\sigma_{e}^{2}\mathbf{I}_{m}\). The weight prior introduced in Assumption 3 is i.i.d. but we assume a flat prior for the bias, which is incorporated by stating \(\mathbf{\Sigma}_{w}^{-1}=\sigma_{w}^{-2}\mathrm{diag}(\mathbf{I}_{n_{\hat{\sigma}}}, \ 0)\). We introduce \(\tilde{\mathbf{I}}_{n_{\phi}}=\mathrm{diag}(\mathbf{I}_{n_{\hat{\sigma}}^{\text{}}}, \ 0)\)._

The assumption of a flat prior for the bias term will have a negligible effect in practice but is important for the theoretical results presented in Section IV. For ease of notation, we will reuse \(J(\mathbf{\Theta};\mathcal{D})\) and \(\mathbf{\Theta}\) in following results.

**Result 1**.: _Let Assumptions 1-3 and Assumption 4 hold. The negative LML in (10) is:_

\[J(\mathbf{\Theta};\mathcal{D})=\frac{m}{2}\log(2\pi)+n_{\phi}\log (\sigma_{w})+m\log(\sigma_{e}) \tag{11}\] \[+\frac{1}{2}\log\det(\mathbf{\Lambda}_{p})+\frac{1}{2\sigma_{e}^{2}}\| \mathbf{t}-\mathbf{y}\|_{2}^{2}+\frac{1}{2\sigma_{w}^{2}}\|\bar{\mathbf{w}}\|_{2}^{2},\]

_where (8b) and (8c) are now:_

\[\mathbf{\Lambda}_{p} =\sigma_{e}^{-2}\mathbf{\Phi}^{\top}\mathbf{\Phi}+\frac{1}{\sigma_ {w}^{2}}\tilde{\mathbf{I}}_{n_{\phi}}, \tag{12a}\] \[\bar{\mathbf{w}} =\sigma_{e}^{-2}\mathbf{\Lambda}_{p}^{-1}\mathbf{\Phi}^{\top}\mathbf{t}. \tag{12b}\]

_In this setting, we denote \(\mathbf{\Theta}=\{\mathbb{W}_{L},\sigma_{e},\sigma_{w}\}\)._

### _Augmented log-marginal likelihood maximization_

To train a NN with BLL we seek to minimize the negative LML:

\[\min_{\mathbf{\Theta}}\quad J(\mathbf{\Theta};\mathcal{D}), \tag{13}\]

with \(J(\mathbf{\Theta};\mathcal{D})\) and \(\mathbf{\Theta}\) according to Result 1. This problem excludes the weights in the last layer of the NN as they have been marginalized. A result of this marginalization is the expression (12b) which computes these weights explicitly. Unfortunately, this creates a major challenge when iteratively solving problem (13) for the optimal values \(\mathbf{\Theta}^{*}\), as the computational graph contains unfavorable expressions such as the inverse of \(\mathbf{\Lambda}_{p}\), which are both numerically challenging and computationally expensive. Furthermore, (12b) requires leveraging the entire training dataset for the computation of the gradient \(\nabla_{\mathbf{\Theta}}J(\mathbf{\Theta};\mathcal{D})\).

The authors in [23] circumvent these issues by variational inference, replacing the LML objective by the evidence lower bound objective (ELBO). As the variational posterior they choose a Gaussian normal distribution, parameterized with mean and covariance. This is the obvious choice in the BLL setting where the true posterior is Gaussian as shown in (7) and, consequently, the ELBO objective is equivalent to the LML [23]. Variational inference comes at the cost, however, of introducing as optimization variables the parameters to describe the variational posterior. Furthermore, the variational inference training loop requires sampling this variational posterior, thus adding significant complexity to the algorithm.
In this work, we present an alternative approach, which simultaneously avoids parameterizing the variational posterior and yields a computational graph with lower complexity than (11) suitable for fast gradient-based optimization. To obtain this result we reformulate (13) as

\[\begin{split}\min_{\mathbf{\Theta},\bar{\mathbf{w}}}& J(\mathbf{\Theta},\bar{\mathbf{w}};\mathcal{D})\\ &\text{s.t.}\quad\bar{\mathbf{w}}=\frac{1}{\sigma_{e}^{2}}\mathbf{ \Lambda}_{p}^{-1}\mathbf{\Phi}^{\top}\mathbf{t},\end{split} \tag{14}\]

with \(\mathbf{\Theta}=\{\mathbb{W}_{L},\sigma_{e},\sigma_{w}\}\). Importantly, we introduce \(\bar{\mathbf{w}}\) as an optimization variable and add an equality constraint corresponding to (12b). The optimal solution \(\mathbf{\Theta}^{*}\) of (14) is thus identical to the optimal solution of (13). As a main contribution of this work, we state the following theorem.

**Theorem 1** (Augmented log-marginal likelihood maximization).: _The optimal solution \((\mathbf{\Theta}^{*},\bar{\mathbf{w}}^{*})\) of problem (14) is equivalent to the optimal solution obtained from the unconstrained problem:_

\[\min_{\mathbf{\Theta},\bar{\mathbf{w}}}\quad J(\mathbf{\Theta},\bar{\mathbf{w}};\mathcal{D}), \tag{15}\]

_where, as the only difference to (13), \(\bar{\mathbf{w}}\) is now an optimization variable and, in comparison to (14), the equality constraint has been dropped._

Proof.: The Lagrangian of problem (14) can be written as:

\[\begin{split}\mathcal{L}(\mathbf{\Theta},\bar{\mathbf{w}},\lambda; \mathcal{D})\\ =J(\mathbf{\Theta},\bar{\mathbf{w}};\mathcal{D})+\lambda^{\top}\left(\mathbf{ \Lambda}_{p}\bar{\mathbf{w}}-\frac{1}{\sigma_{e}^{2}}\mathbf{\Phi}^{\top}\mathbf{t}\right).\end{split} \tag{16}\]

We then state the first-order condition of optimality for the optimization variable \(\bar{\mathbf{w}}\):

\[\nabla_{\bar{\mathbf{w}}}\mathcal{L}(\mathbf{\Theta},\bar{\mathbf{w}},\lambda;\mathcal{D })=\nabla_{\bar{\mathbf{w}}}J(\mathbf{\Theta},\bar{\mathbf{w}};\mathcal{D})+\lambda^{\top }\mathbf{\Lambda}_{p}\overset{!}{=}0. \tag{17}\]

Considering (11) and (5), that is, \(\mathbf{y}=\mathbf{\Phi}\bar{\mathbf{w}}\), we obtain:

\[\nabla_{\bar{\mathbf{w}}} J(\mathbf{\Theta},\bar{\mathbf{w}};\mathcal{D}) \tag{18}\] \[=\nabla_{\bar{\mathbf{w}}}\left(\frac{1}{2\sigma_{e}^{2}}\|\mathbf{t}- \mathbf{\Phi}\bar{\mathbf{w}}\|_{2}^{2}+\frac{1}{2\sigma_{w}^{2}}\|\bar{\mathbf{w}}\|_{2}^ {2}\right)\] \[=\frac{2}{2\sigma_{e}^{2}}\left(\bar{\mathbf{w}}^{\top}\mathbf{\Phi}^{ \top}\mathbf{\Phi}-\mathbf{t}^{\top}\mathbf{\Phi}\right)+\frac{2}{2\sigma_{w}^{2}}\bar{\bm {w}}^{\top}\] (19) \[=\bar{\mathbf{w}}^{\top}\left(\frac{1}{\sigma_{e}^{2}}\mathbf{\Phi}^{\top }\mathbf{\Phi}+\frac{1}{\sigma_{w}^{2}}\tilde{\mathbf{I}}_{n_{\phi}}\right)-\frac{1}{ \sigma_{e}^{2}}\mathbf{t}^{\top}\mathbf{\Phi}\] (20) \[\overset{\eqref{eq:mpl}}{=}\bar{\mathbf{w}}^{\top}\mathbf{\Lambda}_{p}- \frac{1}{\sigma_{e}^{2}}\mathbf{t}^{\top}\mathbf{\Phi}. \tag{21}\]

We substitute (21) into (17) and obtain:

\[\lambda^{\top}\mathbf{\Lambda}_{p} =-\bar{\mathbf{w}}^{\top}\mathbf{\Lambda}_{p}+\frac{1}{\sigma_{e}^{2}} \mathbf{t}^{\top}\mathbf{\Phi}, \tag{22}\] \[\Leftrightarrow\lambda =-\mathbf{\Lambda}_{p}^{-1}\mathbf{\Lambda}_{p}\bar{\mathbf{w}}+\frac{1}{ \sigma_{e}^{2}}\mathbf{\Lambda}_{p}^{-1}\mathbf{\Phi}^{\top}\mathbf{t},\] (23) \[\Leftrightarrow\lambda =-\bar{\mathbf{w}}+\bar{\mathbf{w}}=0, \tag{24}\]

where in (23) we have substituted (12b). From (24), we obtain that \(\lambda=0\) and the Lagrangian of problem (14) thus simplifies to:

\[\mathcal{L}(\mathbf{\Theta},\bar{\mathbf{w}},\lambda;\mathcal{D})=J(\mathbf{\Theta},\bar{ \mathbf{w}};\mathcal{D}). \tag{25}\]

This is exactly the Lagrangian of problem (15). 

Theorem 1 therefore enables us to maximize the LML stated in (11) without having to explicitly compute \(\bar{\mathbf{w}}\) according to Equation (12b). This means, in particular, that we are not required to compute the inverse of \(\mathbf{\Lambda}_{p}\) to express the LML. Consequentially, we can use back-propagation and gradient-based optimization to maximize the LML, which significantly simplifies training NNs with BLL.

### _Change of variables and scaling_

With Theorem 1 we can state the LML as a function of \(\bar{\mathbf{w}}\) which is now included in the set of hyperparameters \(\mathbf{\Theta}=\{\mathbb{W}_{L},\bar{\mathbf{w}},a,b\}=\{\mathbb{W}_{L+1},a,b\}\). Furthermore, we propose a change of variables and scale the objective function to improve numerical stability. In particular, we introduce:

\[\alpha=\frac{\sigma_{w}^{2}}{\sigma_{e}^{2}}, \tag{26}\]

which can be interpreted as a signal-to-noise ratio and

\[a=\log(\alpha),\quad b=\log(\sigma_{e}). \tag{27}\]

to ensure that \(\sigma_{e}>0\) and \(\sigma_{w}>0\) without constraining the problem. These changes are formalized in the the following result.

**Result 2**.: _Considering the definition of \(\alpha\) in (26), we reformulate (12a):_

\[\begin{split}\mathbf{\Lambda}_{p}&=\sigma_{e}^{-2}\mathbf{ \Phi}^{\top}\mathbf{\Phi}+\sigma_{w}^{-2}\tilde{\mathbf{I}}_{n_{\phi}}\\ &=\sigma_{e}^{-2}\left(\mathbf{\Phi}^{\top}\mathbf{\Phi}+\alpha^{-1} \tilde{\mathbf{I}}_{n_{\phi}}\right)=\sigma_{e}^{-2}\tilde{\mathbf{\Lambda}}_{p},\end{split} \tag{28}\]

_where_

\[\tilde{\mathbf{\Lambda}}_{p}=\mathbf{\Phi}^{\top}\mathbf{\Phi}+\alpha^{-1}\tilde{\mathbf{I}}_{n _{\phi}}. \tag{29}\]

_The scaled negative LML from (11) can then be written as:_

\[\begin{split} J(\mathbf{\Theta};\mathcal{D})&=\frac{1}{ 2}\log(2\pi)+\frac{n_{\phi}}{2m}a+b\\ &+\frac{1}{2m}\log\det(\tilde{\mathbf{\Lambda}}_{p})+\frac{1}{2m}\exp( -2b)\|\mathbf{t}-\mathbf{y}\|_{2}^{2}\\ &+\frac{1}{2m}\exp\left(-a-2b\right)\|\bar{\mathbf{w}}\|_{2}^{2}\end{split} \tag{30}\]

_with \(\mathbf{\Theta}=\{\mathbb{W}_{L+1},a,b\}\)._

To train a NN with BLL and univariate output we consider in the following the LML and hyperparameters in the form of Result 2. Additionally, the newly introduced parameter \(\alpha\) in (26) will play an important role in the following discussion on interpolation and extrapolation.

## IV Interpolation and extrapolation with BLL

One of the main challenges of the BLL predictive distribution (7) is that for arbitrary extrapolation points the required assumptions for Lemma 1 will not hold. In this section, we discuss the behavior of the BLL predictive distribution in the interpolation and extrapolation regime and propose a method to improve the performance of BLL for extrapolation.

To formalize the notion of interpolation and extrapolation, we follow the authors [26] definition of interpolation (thus implicitly defining extrapolation) for which we also need to define the convex hull.
**Definition 1** (Convex hull).: _The convex hull of a set of samples \(\mathbf{X}\in\mathbb{R}^{m\times n_{x}}\) is defined as the set:_

\[\mathbb{C}_{\mathbf{X}}=\left\{\mathbf{X}^{\top}\mathbf{\nu}|\mathbf{\nu}\in\mathbb{R}^{m},\ \sum\mathbf{\nu}=1,\mathbf{\nu}\geq 0\right\}.\]

**Definition 2** (Interpolation).: _A sample \(\mathbf{x}\) is considered to be an interpolation point of a set of samples \(\mathbf{X}\), given a feature space \(\mathbf{\hat{\Phi}}(\mathbf{X};\mathbb{W}_{L})\) which satisfies Assumption 1 and 2, if:_

\[\mathbf{\hat{\phi}}(\mathbf{x})\in\mathbb{C}_{\mathbf{\hat{\Phi}}(\mathbf{X};\mathbb{W}_{L})}.\]

Interpolation is an attribute of the input space but its definition considers the learned feature space of the neural network. Considering the feature space in Definition 2 may seem counter-intuitive but applies well to reality, where nonlinear features for regression can show arbitrary behavior between data points, even for a univariate input. In this case, interpolation points in the input domain are rightly classified as extrapolation points by considering the feature domain.

Classifying a point as interpolation or extrapolation is a binary decision. In reality this is a shortcoming, as different degrees of extrapolation are possible. That is, a point "close" to the convex hull might still lead to a trustworthy prediction. The distance to a set, e.g. the convex hull, is defined below.

**Definition 3** (Distance).: _For a set \(\mathbb{X}\subset\mathbb{R}^{n_{x}}\) and a point \(\mathbf{x}\in\mathbb{R}^{n_{x}}\) we define the distance \(d(\mathbf{x},\mathbb{X})\) as:_

\[d(\mathbf{x},\mathbb{X})=\inf_{a\in\mathbb{X}}\|x-a\|_{2}^{2}. \tag{31}\]

### _Quantification of interpolation and extrapolation_

In the following, we seek to define a metric to quantify the degree of extrapolation for a nonlinear regression model with feature space. Intuitively, such a metric could be based on the distance, according to Definition 3, to the convex hull of the features. Unfortunately, the distance to the convex hull results in an optimization problem that scales with the number of samples and the feature dimension, which can be a limitation for practical applications. Instead, we propose an approximate metric for which we introduce the _affine cost_ in the following. But first, as related concepts of the affine cost, we define the well known _span_ and _affine hull_.

**Definition 4** (Span).: _The span of a set of samples \(\mathbf{X}\in\mathbb{R}^{m\times n_{x}}\) is defined as the set:_

\[\mathbb{S}_{\mathbf{X}}=\left\{\mathbf{X}^{\top}\mathbf{\nu}|\mathbf{\nu}\in\mathbb{R}^{m} \right\}.\]

**Definition 5** (Affine hull).: _The affine hull of a set of samples \(\mathbf{X}\in\mathbb{R}^{m\times n_{x}}\) is defined as the set:_

\[\mathbb{A}_{\mathbf{X}}=\left\{\mathbf{X}^{\top}\mathbf{\nu}|\mathbf{\nu}\in\mathbb{R}^{m},\ \sum\mathbf{\nu}=1\right\}.\]

**Definition 6** (Affine cost).: _The affine cost of a set of samples \(\mathbf{X}\in\mathbb{R}^{m\times n_{x}}\) and for a test point \(\mathbf{x}\in\mathbb{R}^{n_{x}}\) is defined as the optimal cost:_

\[c_{\mathbb{A}_{\mathbf{X}}}(\mathbf{x})=\operatorname*{minimize}_{\mathbf{ \nu},\mathbf{e}} \|\mathbf{\nu}\|_{2}^{2}+\gamma\|\mathbf{e}\|_{2}^{2} \tag{32a}\] \[\operatorname{subject\ to:} \mathbf{X}^{\top}\mathbf{\nu}+\mathbf{e}=\mathbf{x},\] (32b) \[\sum\mathbf{\nu}=1, \tag{32c}\]

_where spanning coefficient \(\mathbf{\nu}\in\mathbb{R}^{m}\) and residual variable \(\mathbf{e}\in\mathbb{R}^{n_{x}}\) are optimization variables and \(\gamma\in\mathbb{R}\) is a weighting factor for the residuals._

The affine cost naturally complements the definition of the affine hull by computing the norm of the respective spanning coefficients \(\mathbf{\nu}\) from Definition 5. Importantly, test values
\(\mathbb{A}_{\mathbf{X}}\) also have a value assigned, for which the \(\gamma\)-weighted norm of the residuals \(\mathbf{e}\) is considered.

The relationship of convex hull, span and affine hull as well as the affine cost can be inspected in Figure 2. In this figure, two exemplary sets of features \(\mathbf{\tilde{\Phi}}\in\mathbb{R}^{m\times n_{\hat{\phi}}}\), both with \(m=3\) and \(n_{\hat{\phi}}=2\) are presented and we compare the relationship of test to training samples.

By comparing, for example Figure 2 a) and d), we can see that the affine cost has a strong relationship with the convex hull on which we have based the notion of interpolation in Definition 2. In particular, we consider the level set \(\mathbb{L}_{\mathbb{A}_{\mathbf{\phi}}}=\{\mathbf{\tilde{\phi}}\in\mathbb{R}^{n_{\hat{ \phi}}}|c_{\mathbb{A}_{\hat{\mathbf{\phi}}}}(\mathbf{\tilde{\phi}})\leq l\}\) obtained with the affine cost and suitable level \(l\) as a soft approximation of the convex hull. Such level sets can be seen in Figure 2 d) and h) as the contour lines of \(c_{\mathbb{A}_{\hat{\mathbf{\phi}}}}\). It holds that for a test point \(\mathbf{\tilde{\phi}}\), the affine cost grows with the distance to the level set \(\mathbb{L}_{\mathbb{A}_{\mathbf{\phi}}}\). In this sense, we consider the distance \(d(\mathbf{\tilde{\phi}},\mathbb{L}_{\mathbb{A}_{\mathbf{\phi}}})\) and, more directly, the affine cost \(c_{\mathbb{A}_{\mathbf{\phi}}}(\mathbf{\tilde{\phi}})\) itself as the desired metric for the degree of extrapolation

For the behavior of this metric we distinguish two important cases. In the first case, small values of the affine cost are achieved for test points that are within the affine hull, i.e. \(\mathbf{\tilde{\phi}}\in\mathbb{A}_{\mathbf{\tilde{\Phi}}}\), and for small distances to the convex hull. According to Definition 2, these include all interpolation points and what we consider mild extrapolation. Both test points in Figure 2 h) are examples for this case.

In the second case, a test point is not within the affine hull \(\mathbf{\tilde{\phi}}\notin\mathbb{A}_{\mathbf{\tilde{\Phi}}}\). The affine cost is now influenced primarily through the choice of the parameter \(\gamma\). This can be seen by inspection of (32), where the residuals \(\mathbf{e}\) must be used if \(\mathbf{\tilde{\phi}}\notin\mathbb{A}_{\mathbf{\tilde{\Phi}}}\). The cost may then be dominated by the term \(\|\mathbf{e}\|_{2}^{2}\) which is weighted with \(\gamma\). We reason that in the second case the test point can be considered an extrapolation point and \(\gamma\) can be interpreted as a penalty for extrapolation. This case be observed in Figure 2 d) for the test point with higher affine cost.

### _Relationship of affine cost and BLL_

As another main contribution of this work, we introduce Theorem 2 in this subsection to establish the relationship of the BLL covariance and the previously presented affine cost.

**Theorem 2** (BLL affine cost).: _If Assumption 4 holds, and with \(\gamma=\alpha\), the affine cost \(c_{\mathbb{A}_{\mathbf{\tilde{\Phi}}}}(\mathbf{\tilde{\phi}}(\mathbf{x}))\), according to Definition 6, is equivalent to the scaled BLL covariance (7c):_

\[c_{\mathbb{A}_{\mathbf{\tilde{\Phi}}}}(\mathbf{\tilde{\phi}}(\mathbf{x}))=\sigma_{e}^{-2} \mathbf{\Sigma}_{\mathbf{\tilde{\phi}}}^{\text{\tiny{\rm{\rm{\rm{\rm{\tiny{\tiny{ \tiny{\tiny{\tiny{\tiny{\tiny{\tiny{\tiny{\tiny{\tiny{\tiny{\tiny{\tiny \leftleftleftleftleftleftleft({ \left({ \left(   \left( \bm
test points with \(\tilde{\mathbf{\phi}}(\mathbf{x})\notin\mathbb{A}_{\tilde{\mathbf{\Phi}}}\), clearly indicating extrapolation through high values of the affine cost. The effect can be seen in Figure 2 by comparing subplot d) and h). The application of \(\log\det\)-regularization to obtain matrices with low rank is well known [28] and also applied in other fields such as compressed sensing [29].

The second important implication of Theorem 2 is the interpretation of the parameter \(\alpha\) (or \(\gamma\) respectively) in the context of the affine cost from Definition 6. The parameter directly controls the affine cost and thus the variance for test points \(\tilde{\mathbf{\phi}}(\mathbf{x})\notin\mathbb{A}_{\tilde{\mathbf{\Phi}}}\). This causes a dilemma: Naturally, we have the situation that all training samples \(\tilde{\mathbf{\Phi}}\) are within the affine hull of themselves. Therefore, extrapolation in the sense of \(\tilde{\mathbf{\phi}}(\mathbf{x})\notin\mathbb{A}_{\tilde{\mathbf{\Phi}}}\) does not occur during training and \(\alpha^{*}\), which maximizes the LML, might not yield desirable results.

To illustrate the issue we present a simple regression problem with \(n_{x}=n_{y}=1\) and \(m=3\) samples. We investigate a NN with \(n_{\hat{\phi}}=2\) which allows for a graphical representation of the feature space. The NN is trained by maximizing the LML (30), yielding the optimal parameters \(\mathbf{\Theta}^{*}\), with \(a^{*}=\log\alpha^{*}\), as introduced in (27). The predicted mean and standard deviation for the trained model using \(\alpha^{*}\) can be seen in Figure 3 a). The predicted mean of the NN, in light of the sparse training data, is suitable. However, the covariance shows that the prediction is overconfident. We show in Figure 3 b) and c) that this overconfidence can be tackled simply by increasing \(\alpha\) relative to the optimal value \(\alpha^{*}\).

The reason for this effect of \(\alpha\) on the extrapolation uncertainty can be seen by inspecting Figure 3 d)-f), where the features (recall \(n_{\hat{\phi}}=2\)) for the training (\(\tilde{\mathbf{\phi}}\)) and the test points (\(\tilde{\mathbf{\Phi}}\)) are displayed.

As desired, we have that \(\mathbb{A}_{\tilde{\mathbf{\Phi}}}\subset\mathbb{R}^{n_{\hat{\phi}}}\), that is, the training features lie in a subspace of \(\mathbb{R}^{2}\), which can be seen in Figure 3 where the training samples in the feature space could be connected by a straight line. This effect can be attributed to the \(\log\det\)-regularization in the LML, Extrapolation thus occurs for \(\tilde{\mathbf{\phi}}\notin\mathbb{A}_{\tilde{\mathbf{\Phi}}}\) and for these points the extrapolative uncertainty grows with increasing \(\alpha\). Importantly, by considering Definition 2, we also have extrapolation for test points that are within the convex hull of the input space, i.e. \(\mathbf{x}\in\mathbb{C}_{\mathbf{X}}\). In this example, the only true interpolation points are the training samples for which increasing \(\alpha\) has no significant effect on the predicted covariance.

In Figure 4, we further investigate the effect of increasing \(\alpha\) for the same regression problem and NN architecture as displayed in Figure 3. To quantify the quality of the predictive distribution, we use the mean log-predictive density (LPD) [30]:

\[\log\bar{p}(\mathbf{t}^{\text{test}})=\frac{1}{m_{\text{test}}}\sum_{i=1}^{m_{ \text{test}}}\log p(\mathfrak{t}=t_{i}^{\text{test}}|\mathcal{D},\mathbf{\Theta}, \mathbf{x}^{\text{test}}) \tag{42}\]

which evaluates the logarithm of the posterior distribution of the targets (9) for all test values and computes the average thereof. Importantly, we assume that test data is affected by the same noise as the training data and therefore compute the distribution of the targets. While the output distribution \(\mathbf{y}\) is only affected by the posterior uncertainty of the weights, the target distribution \(\mathbf{t}\) in (5) also consider the additive noise.

In Figure 4, we display log-predictive density (42) and the negative LML (30) for the optimal parameters \(\mathbf{\Theta}^{*}\) and as a function of \(\alpha\). As expected, \(\alpha^{*}\) minimizes the negative LML. We see, however, that the optimal value is almost at a plateau and that increasing \(\alpha\) has only a minor effect on the LML. We also see in the top diagram of Figure 4 that the log-predictive

Fig. 3: NN with BLL: Predicted mean and standard deviation (7) and feature space with \(n_{\hat{\phi}}=2\) for \(m=3\) training samples. The effect of parameter \(\alpha\) on the extrapolation uncertainty is shown by comparing the optimal \(\alpha^{*}\) (maximization of LML (30)) with suggested improved \(\alpha^{\text{max}}\).
density for the training points remains independent of \(\alpha\). On the other hand, we see that the log-predictive density for test points can be increased significantly by setting \(\alpha=\alpha^{\text{max}}\). Further increasing \(\alpha\) has a negative effect as the predicted posterior distribution becomes increasingly flat.

As a consequence of this investigation we propose Algorithm 1 to obtain a NN with BLL, trained with LML maximization, and enhanced extrapolation uncertainty.

```
0:\(\mathcal{D}^{\text{train}}\), \(\mathcal{D}^{\text{val}}\)
0: NN structure (\(L\), \(n_{\phi t},g_{l}(\cdot)\)\(\forall=1\ldots,L\)) \(\boldsymbol{\Theta}^{*}\leftarrow\arg\min_{\boldsymbol{\Theta}}J( \boldsymbol{\Theta};\mathcal{D}^{\text{train}})\)\(\triangleright\) (30) \(\alpha^{\text{max}}\leftarrow\arg\max_{\alpha}\log\tilde{p}(\boldsymbol{t}^{ \text{val}})\)\(\triangleright\) (42)
```

**Algorithm 1** LML optimization with adapted extrapolation penalty \(\alpha\) for enhanced BLL.

Optimizing for the single parameter \(\alpha\) in Algorithm 1 is often trivial and we obtain inexact solutions with a grid search.

## V The multivariate case

In general, we seek to investigate multivariate problems where, in contrast to Section II, we now have that \(\boldsymbol{t}\in\mathbb{R}^{n_{y}}\). As before, we consider a dataset \(\mathcal{D}=\left(\boldsymbol{X}_{\downarrow}\boldsymbol{T}\right)\) consisting of \(m\) samples and introduce \(\boldsymbol{T}=\left[\boldsymbol{t}_{1},\ldots,\boldsymbol{t}_{m}\right]^{+} \in\mathbb{R}^{m\times n_{y}}\).

To obtain the setting in which Lemma 1 holds, we augment Equation (5) for the multivariate case:

\[\boldsymbol{T}=\boldsymbol{\Phi}\boldsymbol{W}+\boldsymbol{E}, \tag{43}\]

where \(\boldsymbol{E}\in\mathbb{R}^{m\times n_{y}}\) is the matrix of residuals of the regression model. We use the \(\operatorname{vec}(\cdot)\) operation as defined in [31, Defn. 11.5] and vectorize (43):

\[\operatorname{vec}(\boldsymbol{T})=\operatorname{vec}(\boldsymbol{\Phi} \boldsymbol{W})+\operatorname{vec}(\boldsymbol{E}). \tag{44}\]

Introducing \(\otimes\) as the Kronecker product, this equation can be reformulated as [31, Prop. 11.16(b)]:

\[\operatorname{vec}(\boldsymbol{T})=(\boldsymbol{I}_{n_{y}}\otimes \boldsymbol{\Phi})\operatorname{vec}(\boldsymbol{W})+\operatorname{vec}( \boldsymbol{E}), \tag{45}\]

and expressed as:

\[\hat{\boldsymbol{t}}=\hat{\boldsymbol{\Phi}}\hat{\boldsymbol{w}}+\hat{ \boldsymbol{\epsilon}}. \tag{46}\]

In the form of (46), Lemma 1 directly applies to the multivariate case where the prior and noise covariances from Assumption 2 and 3 now refer to:

\[\hat{\boldsymbol{\epsilon}} \sim\mathcal{N}(0,\hat{\boldsymbol{\Sigma}}_{E}),\] \[\hat{\boldsymbol{w}} \sim\mathcal{N}(0,\hat{\boldsymbol{\Sigma}}_{\boldsymbol{W}}),\]

and for which we need to consider the multivariate feature matrix \(\hat{\boldsymbol{\Phi}}\).

### _Simplified training_

The multivariate settings adds significant complexity to the marginal likelihood maximization introduced in Section III. To reduce the computational complexity, we present Result 3 and two required assumptions in the following.

**Assumption 5**.: _Noise and prior, as introduced in Assumption 2 and 3, are uncorrelated for all \(n_{y}\) outputs and identical for all \(m\) samples and all \(n_{\phi}\) features, respectively. We denote \(\boldsymbol{\sigma}_{e}=\left[\sigma_{e,1},\ \ldots,\ \sigma_{e,n_{y}}\right]\) and \(\boldsymbol{\sigma}_{w}=\left[\sigma_{w,1},\ \ldots,\ \sigma_{w,n_{y}}\right]\) and obtain:_

\[\hat{\boldsymbol{\Sigma}}_{E} =\operatorname{diag}(\boldsymbol{\sigma}_{e})\otimes\boldsymbol{ I}_{m}, \tag{47}\] \[\hat{\boldsymbol{\Sigma}}_{W} =\operatorname{diag}(\boldsymbol{\sigma}_{w})\otimes\boldsymbol {I}_{n_{\phi}}. \tag{48}\]

We introduce \(\boldsymbol{\alpha}=\frac{\boldsymbol{\sigma}_{w}^{2}}{\boldsymbol{\sigma}_{ e}^{2}}\), similarly to (26), and assume the following.

**Assumption 6**.: _For all predicted outputs \(y_{i}\), we have the same parameter \(\alpha\), i.e.:_

\[\alpha_{1}=\cdots=\alpha_{n_{y}}\stackrel{{!}}{{=}}\alpha. \tag{49}\]

**Result 3**.: _Let Assumption 5 and 6 hold. The scaled negative LML in (10) for the vectorized multivariate case (46) is:_

\[J(\boldsymbol{\Theta};\mathcal{D}) =\frac{n_{y}}{2m}\left(m\log(2\pi)+n_{\phi}a+\log\det(\bar{ \boldsymbol{\Lambda}}_{p})\right) \tag{50}\] \[+\sum_{i=1}^{n_{y}}\left(b_{i}+\frac{1}{2m}\exp(-2b_{i})\| \boldsymbol{t}_{i}-\boldsymbol{y}_{i}\|_{2}^{2}\right)\] \[+\sum_{i=1}^{n_{y}}\left(\frac{1}{2m}\exp\left(-a-2b_{i}\right)\| \bar{\boldsymbol{w}}_{i}\|_{2}^{2}\right)\]

_with \(\boldsymbol{\Theta}=\left\{\mathbb{W}_{L+1},a,b_{1},\ldots,b_{n_{y}}\right\}\) and \(\bar{\boldsymbol{\Lambda}}_{p}\) according to (29). The full precision matrix \(\hat{\boldsymbol{\Lambda}}_{p}\) for the multivariate case can be obtained as:_

\[\hat{\boldsymbol{\Lambda}}_{p}=\operatorname{diag}(\boldsymbol{\sigma}_{e}^{-2 })\otimes\bar{\boldsymbol{\Lambda}}_{p}. \tag{51}\]

Proof.: The result follows directly from the properties of the Kronecker product [31] applied to the log-marginal likelihood in (10) with (46), as well as Assumption 5 and 6. 

Result 3 shows that the LML can be easily expressed for the multivariate case, allowing for fast and efficient NN training. This is largely due to Assumption 5 and 6. While Assumption 5 is a natural extension of Assumption 4 for the multivariate case, Assumption 6 might be questioned. We argue again with our interpretation of \(\alpha\) as an extrapolation

Fig. 4: Effect of \(\log(\alpha)\) on the LML (train points) for a trained NN and the mean log-predictive density (42) (test points). The same regression problem and NN as in Figure 3 is considered.
penalty weight, as discussed in Section IV. Importantly, extrapolation, as defined in Definition 2, is a property of the feature space and occurs regardless of the number of outputs.

Apart from simplifying the LML in (50), Assumption 5 and 6 also yield a simplified computation of the predictive distribution. The outputs are uncorrelated due to Assumption 5 and we can obtain independent covariance matrices:

\[\mathbf{\Sigma}_{\hat{y}_{i}}^{\text{\tiny{max}}}(\mathbf{x})=\mathbf{\phi}^{\top}\mathbf{ \Lambda}_{p,i}^{-1}\mathbf{\phi},\quad\text{with:}\ \mathbf{\Lambda}_{p,i}=\sigma_{e,i}^{-2}\bar{\mathbf{\Lambda}}_{p}, \tag{52}\]

where we have the same \(\bar{\mathbf{\Lambda}}_{p}\) for all outputs due to Assumption 6. Therefore, the complexity of evaluating the predictive distribution scales negligibly with number of outputs. This is a major advantage of NNs with BLL, in comparison to GPs, where it is common to fit an independent model for each output,

### _Toy example_

To showcase the multivariate NN with BLL we present a toy example in Figure 1. In the example, data is sampled from a function with \(n_{x}=1\) input and \(n_{y}=2\) outputs where both outputs are jointly learned by a single NN. Both outputs exhibit strong non-linear behavior with different noise-levels, in particular \(\sigma_{1}=0.05\) and \(\sigma_{2}=0.2\). We want to emphasize that these variances are assumed to be unknown.

The network is trained following Algorithm 1, using the multivariate LML from Result 3. For the toy example, a suitable structure with \(L=2\), \(n_{\hat{\phi}_{1}}=n_{\hat{\phi}_{2}}=n_{\hat{\phi}_{3}}=20\) and \(g_{1}(\cdot)=g_{2}(\cdot)=g_{3}(\cdot)=\tanh(\cdot)\) is determined using trial and error. To maximize the LML, the Adam [32] optimizer is used. To avoid overfitting, we use early-stopping for which we monitor the validation loss (the LML of the validation data). Finally, we compute the standard deviation of the predicted target distribution according to (9). This distribution is visually easier to interpret, as it should contain the noise disturbed train and test samples with a high probability.

In Figure 1 we see that the predictive distribution suitably describes both outputs. In particular, the predicted mean and variance is accurate in the interpolation regime where the true noise variance is \(\sigma_{1}=0.05\) and \(\sigma_{2}=0.2\) and the estimated variance \(\sigma_{1}=0.051\) and \(\sigma_{2}=0.17\). Furthermore, the extrapolation regime results in a variance that indicates extrapolation and approximates the true error.

In Figure 1 we also compare the extrapolation variance obtained with the optimal value \(\alpha^{*}\) and \(\alpha^{\text{max}}\). Only the latter is suitable to describe the extrapolation regime. To quantify the improvement, we compute the (negative) LML, the log-predictive density and the mean-squared error for \(\alpha^{*}\) and \(\alpha^{\text{max}}\). These performance indicators are presented in Table I. For comparison, we train the identical NN structure by minimizing the mean-squared-error and then use the learned features for BLR as in previous works [5, 25]. The second step of Algorithm 1, that is updating \(\alpha\) subsequently with validation data, can also be applied in this setting and the resulting performance metrics are also shown in Table I.

We have two important conclusions and a comment regarding the results in Table I. First, we see that updating \(\alpha\), as proposed in Algorithm 1, significantly improves the log-predictive density (42), both for BLL and BLR. Second, we see that the NN with BLL, which is trained according to Algorithm 1, outperforms BLR with respect to the log-predictive density. Finally, we want to comment on the NLML for the test data in Table I. This metric is displayed for the sake of completeness but has limited expressiveness when judging the models performance on the test data. The test data is chosen deliberately to contain extrapolation points for which the squared error between targets and predictions in (50) can be arbitrarily high. In contrast to the LPD in (42), where this error is weighted with the predicted variance of the targets, the NLML uses the estimated variance of the additive noise as weighting factors. This can lead to very poor results for the NLML if the error exceeds the magnitude of the noise variance, which is to be expected for extrapolation points. Consequentially, we propose to consider the LPD to evaluate the quality of uncertainty quantification.

## VI Bayesian last layer for system identification

An important and promising application of NNs with BLL is the probabilistic identification of non-linear dynamic systems. In discrete-time formulation, a dynamic system can be described as:

\[\mathbf{s}(k+1) =f_{\text{sys}}(\mathbf{s}(k),\mathbf{a}(k))+\mathbf{\epsilon}_{s}(k), \tag{53a}\] \[\mathbf{r}(k) =h_{\text{sys}}(\mathbf{s}(k),\mathbf{a}(k))+\mathbf{\epsilon}_{r}(k), \tag{53b}\]

with states \(\mathbf{s}\in\mathbb{R}^{n_{x}}\), actions \(\mathbf{a}\in\mathbb{R}^{n_{a}}\) and measurements (readings) \(\mathbf{r}\in\mathbb{R}^{n_{r}}\). We denote with \(k\) the discrete time-step \(t_{k}\) and assume unknown additive process noise \(\mathbf{\epsilon}_{s}\in\mathbb{R}^{n_{x}}\) and measurement noise \(\mathbf{\epsilon}_{r}\in\mathbb{R}^{n_{r}}\). For the system identification task, we will generally only have access to the actions and measurements of the system.
Assuming that system (53) is observable, it can be expressed as an nonlinear autoregressive model with exogenous inputs (NARX) [33]:

\[\Delta\mathbf{r}(k+1)= \tag{54}\] \[f_{\textsc{NARX}}(\mathbf{r}(k),\ldots,\mathbf{r}(k-l),\mathbf{a}(k),\ldots, \mathbf{a}(k-l))+\mathbf{\epsilon},\]

with parameter \(l\) and additive noise \(\mathbf{\epsilon}\). In (54) we predict \(\Delta\mathbf{r}(k+1)=\mathbf{r}(k+1)-\mathbf{r}(k)\), which is then used to obtain \(\mathbf{r}(k+1)\). In theory, this is equivalent to predicting \(\mathbf{r}(k+1)\) directly, but often shows better predictive performance in practice [34]. We also introduce

\[\hat{\mathbf{s}}=\left[\mathbf{r}(k),\ldots,\mathbf{r}(k-l),\mathbf{a}(k-1),\ldots,\mathbf{a}(k-l )\right], \tag{55}\]

as the NARX state, which is used to represent the system analogous to (53) as

\[\hat{\mathbf{s}}(k+1)=\hat{f}_{\textsc{NARX}}(\hat{\mathbf{s}}(k),\mathbf{a}(k)). \tag{56}\]

It is common to use NNs [35, 36] to approximate the NARX model \(f_{\textsc{NARX}}(\cdot)\). Learning a probabilistic NARX model with NN and BLL also estimates the covariance of the additive noise \(\mathbf{\epsilon}\) in (54) to approximate the effect of process and measurement noise in (53).

### _Linear system example_

We investigate the triple-mass-spring system described in [37] with the provided resources from Github2. The triple-mass-spring system is linear system with \(n_{s}=8\) states, \(n_{a}=2\) input actions and \(n_{r}=3\) measurements (the disc angles) with unknown and different additive noise. The motivation to investigate a linear system stems from the conviction that non-linear methods must first proof themselves in this simple setting.

Footnote 2: [https://github.com/4flix/DecPC_Perspective](https://github.com/4flix/DecPC_Perspective)

Data from a single simulation using a sequence of persistently exciting inputs is used to construct a dataset \(\mathcal{D}\) of inputs and targets for the NARX model (54) with parameter \(l=3\). This results in \(m=797\) samples used for training with \(n_{x}=15\) input features. We investigate a NN with \(L=2\) layers, \(n_{\widehat{\varphi}_{1}}=40\), \(n_{\widehat{\varphi}_{2}}=40\) and \(g_{1}(\cdot)=g_{2}(\cdot)=\tanh(\cdot)\) activation function. We again optimize with Adam and use early stopping based on the validation error to avoid overfitting.

In Figure 5 we showcase the results of the NN with BLL (trained with Algorithm 1) by displaying parity plots for the predicted vs. the true output including the standard deviation of the predictions. The predictive distribution is again computed using (9), to consider the additive noise of the test data. We compare the results for validation data (same distribution as training data) and test data which is created using a greater maximum amplitude of actions. The results in Figure 5 show that for both, validation and test data, suitable predictions are obtained for all outputs. As expected, the test data often yields larger deviations from the true values but, importantly, this error is suitably described by the increased variance.

We again compare the NN with BLL with BLR based on features from a trained (mean-squared-error) NN with the same architecture. These results are shown in Table II. We can see that the NN with BLL again outperforms the BLR

Fig. 5: Parity plots for predicted \(\Delta\mathbf{r}(k+1)\) of the triple-mass-spring system. Prediction and standard deviation (shown as \(\pm 3\sigma\)-error bars) from NN with BLL. Comparison of validation data vs. test data (obtained with double the input amplitude to enforce extrapolation.)
approach in terms of the LPD (42) with and without adapted \(\alpha\). Interestingly, we see that the LPD obtained with BLR benefits insignificantly from adapting \(\alpha\). We find that this is a result of an unsuitable feature space obtained from the trained (mean-squared-error) NN. In particular, this feature space has full rank, in which case the effect of \(\alpha\) on the predictive uncertainty is negligible, as discussed in Section IV.

As a final result of this section, we perform an open-loop simulation of the identified NARX models (NN with BLL and BLR with NN features) and compare it to the true system response. All trajectories are created with the same initial state and input trajectories. For the the open-loop simulation, the probabilistic NARX model is evaluated recursively yielding a normal distribution for the predicted next outputs at each step. We propagate uncertainty to successive outputs, similarly to an extended Kalman filter, by linearizing the NARX system as the current state:

\[\mathbf{A}(k)=\left.\frac{\partial\hat{f}_{\text{sys}}(\hat{\mathbf{s}},\mathbf{a})}{ \partial\hat{\mathbf{s}}}\right|_{\hat{\mathbf{s}}(k),\mathbf{a}(k)}. \tag{57}\]

The results are presented in Figure 6. We see that the estimated system response and uncertainty of the NN with BLL is an excellent match for the true system response. The NARX model with BLR based on NN features performs worse in terms of the mean trajectories and unfortunately leads to an unsuitable uncertainty prediction. This is mostly due to the linearization of the system model which yields an unstable system matrix in this case. This can also be seen in the Figure, where we display the spectral norm of \(\mathbf{A}\).

## VII Conclusions

Neural networks with Bayesian last layer are an attractive compromise between tractability and expressiveness in the field of Bayesian neural networks. They are strongly related to Bayesian linear regression in that they consider a learned nonlinear feature space and yield a Gaussian normal distribution as their prediction. As the main difference to Bayesian linear regression, a neural network with Bayesian last layer should be trained by maximizing the marginal likelihood which, importantly, considers all deterministic weights of the neural network as hyperparameters. This training has previously been difficult, as it requires sampling-based variational inference. To this end, our main contribution is a reformulation of the Bayesian last layer marginal likelihood to enable straightforward and fast gradient-based optimization. In comparison to Bayesian linear regression with features from a previously trained neural network, we find that training on the log-marginal likelihood also shows significant advantages in practice. Our second contribution, a simple algorithm to tune the extrapolative uncertainty also shows excellent results in our two presented simulation studies. Our presented methods are formulated for the multivariate case and investigated especially for the application of probabilistic system identification, where the predicted uncertainties suitably describe the approximation error. In future work, we seek to incorporate the probabilistic system model in a stochastic model predictive control framework.
