Vecchia Gaussian Process Ensembles on Internal Representations of Deep
  Neural Networks
For regression tasks, standard Gaussian processes (GPs) provide natural
uncertainty quantification, while deep neural networks (DNNs) excel at
representation learning. We propose to synergistically combine these two
approaches in a hybrid method consisting of an ensemble of GPs built on the
output of hidden layers of a DNN. GP scalability is achieved via Vecchia
approximations that exploit nearest-neighbor conditional independence. The
resulting deep Vecchia ensemble not only imbues the DNN with uncertainty
quantification but can also provide more accurate and robust predictions. We
demonstrate the utility of our model on several datasets and carry out
experiments to understand the inner workings of the proposed method.
Felix Jimenez, Matthias Katzfuss
# Vecchia Gaussian Process Ensembles on Internal Representations of Deep Neural Networks

###### Abstract

For regression tasks, standard Gaussian processes (GPs) provide natural uncertainty quantification, while deep neural networks (DNNs) excel at representation learning. We propose to synergistically combine these two approaches in a hybrid method consisting of an ensemble of GPs built on the output of hidden layers of a DNN. GP scalability is achieved via Vecchia approximations that exploit nearest-neighbor conditional independence. The resulting deep Vecchia ensemble not only imbues the DNN with uncertainty quantification but can also provide more accurate and robust predictions. We demonstrate the utility of our model on several datasets and carry out experiments to understand the inner workings of the proposed method.

nearest neighbors; Vecchia approximation; Gaussian process ensemble

## 1 Introduction

In recent years, deep neural networks (DNNs) have achieved remarkable success in various tasks such as image recognition, natural language processing, and speech recognition. However, despite their excellent performance, these models have certain limitations, such as their lack of uncertainty quantification (UQ). Much of UQ for DNNs is based on a Bayesian approach that models network weights as random variables [22] or has involved ensembles of networks [18].

Gaussian processes (GPs) provide natural UQ, but they lack the representation learning that makes DNNs successful. Standard GPs are known to scale poorly with large datasets, but GP approximations are plentiful and one such method is the Vecchia approximation [31, 16], which uses nearest-neighbor conditioning sets to exploit conditional independence among the data. In high dimensions, finding the appropriate conditioning sets is challenging, so a procedure that can discover the conditional independence can expand the applicability of the Vecchia approximation.

While the literature on Bayesian neural networks and deep ensembles is rich, less has been done on using the outputs of internal layers of DNNs to quantify uncertainty for the neural network. In this paper we propose the deep Vecchia ensemble (DVE), which uses the outputs of internal layers in conjunction with the Vecchia approximation to provide UQ for DNNs. By using the outputs of internal layers of the DNN, we inject the representation learning of DNNs into GPs. As a result, we are able to use the DNN to discover better conditioning sets for the Vecchia approximation. A high-level summary of this approach is given in Figure 1, where the outputs of internal layers are given in the bottom right three panels and the different conditioning sets extracted from the neural network are shown in the panel on the left.

We believe that the proposed methodology has the potential to improve the reliability and robustness of deep learning models in various applications while simultaneously improving Vecchia GPs. The remainder of the paper is organized as follows: We begin in Section 2 with a review of related concepts and define the intermediate representations upon which we later build. In Section 3, we contextualize our approach within existing work on uncertainty quantification
for neural networks and GP approximation. The proposed methodology is detailed in Section 4. We then apply our model to pretrained neural networks in Section 5. This is followed by an experiment in Section 6 that explores how our model works. In Section 7, we discuss potential application areas for our method, the limitations of our approach, and future research directions that could improve upon our work.

## 2 Preliminaries

In this section, we provide background information on GPs, including GP regression, as well as the Vecchia approximation to GPs. We also review GP ensembles, which combine multiple GPs to provide a distribution over the response variable. Finally, we define internal spaces as intermediate representations of the input data generated by the neural network.

### Modeling functions with Gaussian processes

Suppose we have a dataset \(\mathcal{D}=\{\mathbf{X},\mathbf{y}\}=\{\mathbf{x}_{i},y_{i}\}_{i=1}^{n}\), where \(\mathbf{x}_{i}\in\mathcal{X}\), \(y_{i}=f(\mathbf{x}_{i})+\epsilon_{i}\), \(\epsilon_{i}\stackrel{{ iid}}{{\sim}}\mathcal{N}(0,\sigma^{2})\), and \(f:\mathcal{X}\rightarrow\mathbb{R}\). Our goal is to make predictions at \(n_{p}\) test locations \(\mathbf{X}^{*}=\{\mathbf{x}_{j}^{*}\}_{j=1}^{n_{p}}\) in a probabilistic fashion. A common approach is to assign the latent function \(f\) a zero-mean Gaussian process prior with covariance function \(K_{\mathbf{\theta}}(\mathbf{x},\mathbf{x}^{\prime})\) parameterized by \(\mathbf{\theta}\). Letting \(\mathbf{y}^{*}=\{f(\mathbf{x}_{j}^{*})+\epsilon_{j}\}_{j=1}^{n_{p}},p(\mathbf{y}^{*},\mathbf{ y})\) will be jointly multivariate normal, which implies:

\[p(\mathbf{y}^{*}|\mathbf{y})=\mathcal{N}(\mathbf{\mu}^{*},\mathbf{K}_{\mathbf{\theta}^{*}}+\sigma^ {2}\mathbf{I}_{n_{p}}), \tag{1}\]

where \(\mathbf{\mu}^{*}\) and \(\mathbf{K}_{\mathbf{\theta}^{*}}\) depend on \(\mathcal{D}\) and \(\mathbf{\theta}\). To select \(\mathbf{\theta}^{*}\), we maximize the marginal log-likelihood:

\[p(\mathbf{y})=\int p(\mathbf{y}|\mathbf{f})p(\mathbf{f})d\mathbf{f}=\mathcal{N}(\mathbf{0},\mathbf{K}_{ \mathbf{\theta}}(\mathbf{X},\mathbf{X})+\sigma^{2}\mathbf{I}_{n}). \tag{2}\]

Both Equation (1) and (2) have time and space complexities of \(\mathcal{O}(n^{3})\) and \(\mathcal{O}(n^{2})\), respectively.

### Approximating Gaussian processes using Vecchia

The joint distribution of a collection of random variables, \(\mathbf{y}=\{y_{1},y_{2},...,y_{n}\}\), can be decomposed as \(p(\mathbf{y})=\prod_{i=1}^{n}p(y_{i}|\mathbf{y}_{h(i)})\), \(h(i)=\{1,...,i-1\}\). Vecchia's approximation [31] conditions the \(i^{th}\) observation on a set indexed by \(g(i)\), with \(g(i)\subset h(i)\). The approximate joint distribution is then,

\[p(\mathbf{y})\approx\prod_{i=1}^{n}p(y_{i}|\mathbf{y}_{g(i)}). \tag{3}\]

Figure 1: Nearest-neighbor conditioning sets for different input representations obtained via the layers of a DNN. Left: The red point has different nearest-neighbor conditioning sets based on the metric induced by the internal representations of the neural network. The brown, blue, and pink shaded areas denote the regions in input space that will be mapped to a hypersphere in the first, second, and third intermediate spaces, respectively. The conditioning sets derived from the different regions may overlap, as in the blue and brown region, or be disjoint from the others as in the pink region. Right: The data used to train the network is mapped through each layer and the output, or feature map, of each datum is stored. For a red test point, we asses uncertainty by asking: How many instances in the training data took similar paths through the network?
Letting \(\mathbf{\tilde{y}}=(\mathbf{y},\mathbf{y}^{*})\), where \(\mathbf{y}^{*}=\{y^{*}_{n+1},...,y^{*}_{n+n_{p}}\}\) are test points ordered after the observed data, the approximate conditional distribution for \(\mathbf{y}^{*}\) is:

\[p(\mathbf{y}^{*}|\mathbf{y})\approx\prod_{j=n+1}^{n+n_{p}}p(y^{*}_{j}|\mathbf{\tilde{y}}_{g (j)}). \tag{4}\]

For fixed \(g(i)\)'s, such that \(|g(i)|\leq m\) with \(m\ll n\), the approximation in (3) reduces the in-memory space complexity of Equation (2) to \(\mathcal{O}(n_{p}m)\). Using the approximation in Equation (4) reduces the time and space complexity of Equation (1) to \(\mathcal{O}(n_{p}m)\). An additional approximation utilizing mini-batches [14] of size \(n_{b}\) results in an in-memory space complexity of \(\mathcal{O}(n_{b}m^{2})\) for Equation (2) and a time complexity of \(\mathcal{O}(n_{b}m^{3})\).

For the \(i^{th}\) data point, the \(m\) ordered nearest neighbors \(g(i)\) are typically determined based on a distance metric, such as Euclidean distance, between the input variables \(\mathbf{x}\). To compute these ordered nearest neighbors, it is assumed that the data has been sorted according to some ordering procedure, and for each \(i\), the \(m\) nearest neighbors that precede \(i\) in the ordering are selected.

### Ensembling Gaussian processes

An ensemble of \(L\) GPs can be formed by partitioning the training data into disjoint subsets and fitting independent GPs to each subset of the training data. Assuming this independence structure, the hyperparameters, \(\mathbf{\theta}_{k}\), can be learned by maximizing the log-marginal likelihood, \(\sum_{k=1}^{L}\) log \(p_{k}(\mathbf{y}^{(k)}|\mathbf{\theta}_{k})\).

While in this work we allow each GP to have its own hyperparameters, product-of-expert GP ensembles [30; 3; 7; 26] typically share hyperparameters \(\mathbf{\theta}\) between the GPs. We are assuming \(\mathbf{\theta}_{k}\) contains all the hyperparameters in both the GP kernel and the likelihood.

The prediction of \(y^{*}\) at a location \(\mathbf{x}^{*}\) can be formed by using a generalized product of experts (gPoE) [3], an extension of product of experts (PoE) [13], which combines the predictions of \(L\) models by using a weighted product of the predictions. Letting \(p_{k}^{\beta_{k}(\mathbf{x}^{*})}(y^{*})\) denote the \(k^{th}\) model's predictive distribution for \(y^{*}\), the combined distribution for \(y^{*}\) is given by,

\[p(y^{*})\propto\prod_{k=1}^{L}p_{k}^{\beta_{k}(\mathbf{x}^{*})}(y^{*}). \tag{5}\]

The function \(\beta_{k}(\mathbf{x}^{*})\) is a weight for the \(k^{th}\) model's prediction based on the input \(\mathbf{x}^{*}\) such that \(\sum_{k=1}^{L}\beta_{k}(\mathbf{x}^{*})=1\). When each of the \(k\) predictions is a Gaussian distribution with mean \(\mu_{k}(\mathbf{x}^{*})\) and variance \(\mathbf{\sigma}_{k}(\mathbf{x}^{*})\), the combined distribution is also Gaussian with combined mean and variance given by,

\[\mu(\mathbf{x}^{*})=\sigma^{2}(\mathbf{x}^{*})\sum_{k=1}^{L}\beta_{k}(\mathbf{x}^{*}) \sigma_{k}^{-2}(\mathbf{x}^{*})\mu_{k}(\mathbf{x}^{*}),\qquad\quad\sigma^{2}(\mathbf{x}^{* })=(\sum_{k=1}^{L}\beta_{k}(\mathbf{x}^{*})\sigma_{k}^{-2}(\mathbf{x}^{*}))^{-1}. \tag{6}\]

For details on computing \(\beta_{k}\) see Appendix F.

Although we have focused on the combined predictive distribution for the observed \(y^{*}\), in the case of a Gaussian likelihood the same results hold for the combined distribution of the latent function \(f^{*}\) (with the variance terms changed to reflect we are working with the noiseless latent function). This implies we can choose to combine in either \(f\) or \(y\) space. For details on this see Appendix G.

### Extracting information from intermediate representations

Consider a finite sequence of functions \(f_{1},...,f_{L+1}\) such that \(f_{1}:\mathbb{R}^{d}\rightarrow\mathcal{A}_{1}\), \(f_{2}:\mathcal{A}_{1}\rightarrow\mathcal{A}_{2}\),...,\(f_{L}:\mathcal{A}_{L-1}\rightarrow\mathcal{A}_{L}\), \(f_{L+1}:\mathcal{A}_{L}\rightarrow\mathbb{R}\), where each \(\mathcal{A}_{i}\subset\mathbb{R}^{d_{i}}\).

The \(k^{th}\)**intermediate representation** of a point \(\mathbf{x}\in\mathbb{R}^{d}\) with respect to the sequence \(f_{1},...,f_{L+1}\) is defined to be \(\mathbf{e}_{k}(\mathbf{x}):=(f_{k}\circ f_{k-1}\circ...\circ f_{1})(\mathbf{x})\). The \(k^{th}\)**intermediate space**\(\mathcal{A}_{k}\) is the image of \((f_{k}\circ f_{k-1}\circ...\circ f_{1})(\cdot)\). A **composite model**\(\mathcal{M}\) is the composition of \(L\) such functions \(M(\mathbf{x}):=(f_{L+1}\circ f_{L-1}\circ...\circ f_{2}\circ f_{1})(\mathbf{x})\). A feed-forward neural network with \(L\) layers is an example of a composite model and the output of each layer in the network, excluding the final layer, results in an intermediate space.
Related work

We summarize work on improving scalability of GPs, quantifying uncertainty of neural networks, and using the outputs of internal layers of a neural network.

There are several common directions of research that have focused on scaling GPs. The most relevant for our work are the Vecchia approximation [31; 16], inducing-point GPs [12; 29], and ensembling methods [30; 7; 3]. The Vecchia approximation is based on an ordered conditional approximation of joint distributions, where conditioning sets are often computed using nearest neighbors, which is difficult in high input dimensions. Inducing point methods assume that the covariance matrix has a low-rank structure but they can result in overly smooth approximations. Finally, ensembling methods assign an independent GP to partitions of the data, and while this greatly improves scalability, much work has been done on improve and understanding the failure modes of these models [5; 26].

The work on uncertainty quantification (UQ) for deep learning most relevant to our paper are those involving ensembles of deep neural networks (DNNs) [18] and those using Bayesian neural networks (BNNs) [22]. Ensemble-based methods independently train DNNs on subsamples of the data. BNNs take the weights of the neural network to be random variables and use Bayes' rule to derive a posterior for the weights given the data. Each realization of the weights from the posterior implies a different network. By marginalizing over the posterior when forming predictions, the uncertainty in the weights is propagated to the uncertainty in the prediction. The posterior is rarely available in closed form and therefore we rely on approximate sampling or posterior approximation [22; 21; 4; 8]. Both BNNs and ensembles of DNNs consider multiple networks, but in this work we use a single network at a time. Therefore, our work can always be used within either framework and should not be taken as alternative to these methods, but simply as an extension.

In this work we heavily use the outputs of internal layers of DNNs, which has been done for explainable AI [2], deep kernel learning [32], deep Gaussian process [6] and within Resnet [11]. For explainable AI, the goal is to understand what features the network is learning to extract. Resnet's skip connections ensure, among other things, that adding new layers doesn't reduce the class of functions the network can model. However, neither of these two approaches use the outputs of the internal layers for UQ of DNNs or to improve GP approximations. While deep kernel learning is closest to our work, we use more than just the final layer which we believe is critical for uncertainty quantification and predictive performance (see the S-curve example in Section 6). Finally, unlike the deep Gaussian process, we are not feeding the output of one GP into the next, but instead feeding the output of network layers into Vecchia GPs.

## 4 Proposed methodology

This section outlines our proposed methodology, beginning with a concise model summary. We then describe our approach for extracting a sequence of datasets from a pretrained neural network, followed by a discussion of our conditioning set computation process. Next, we detail our method for building an ensemble using the collections of conditioning sets. Throughout, we explain the design choices that informed the development of our model.

### Model summary

We are provided with a dataset \(\mathcal{D}\) comprising \(n\) input-output pairs, \(\mathcal{D}=\{\mathbf{x}_{i},y_{i}\}_{i=1}^{n}\), and a composite model \(\mathcal{M}:\mathcal{X}\rightarrow\mathbb{R}\) that is trained to predict \(y_{i}\) given the corresponding \(\mathbf{x}_{i}\). Initially, we map the inputs \(\{\mathbf{x}_{i}\}_{i=1}^{n}\) to their respective intermediate representations, \(\{\mathbf{e}_{k}(\mathbf{x}_{i})\}_{i=1}^{n}\), where \(k=1,...,L\). We generate a collection of conditioning sets, \(\mathcal{G}_{k}\), for each set of intermediate representations, \(\{\mathbf{e}_{k}(\mathbf{x}_{i})\}_{i=1}^{n}\), by computing ordered nearest neighbors based on the Euclidean distance between the intermediate representations. The resulting tuple for each intermediate space \((\{\mathbf{e}_{k}(\mathbf{x}_{i}),\mathbf{y}_{i}\}_{i=1}^{n},\mathcal{G}_{k})\) is combined with an automatic relevance determination (ARD) kernel \(K_{\mathbf{\theta}_{k}}(\cdot,\cdot)\) of the appropriate dimension to form an ensemble of deep Vecchia Gaussian processes (GPs).

At test time, the test point will be mapped to its intermediate spaces and each Vecchia GP will provide a predictive mean and variance. The predictions are then combined via a product of experts to form a single mean and covariance estimate. The diagram in Figure 2 provides a high-level overview of the proposed model.

### Dataset extraction from a neural network

To begin we describe how we extract a sequence of datasets from a pretrained neural network and the data it was trained on. Suppose the first layer of our pretrained network is of the form \(f_{1}(\mathbf{x})=\sigma(\mathbf{x}^{T}\mathbf{W}^{(1)}+\mathbf{b}^{(1)})\), with, \(\mathbf{W}^{(1)}\in\mathbb{R}^{dxd_{1}}\), \(\mathbf{b}_{1}\in\mathbb{R}^{d_{1}}\) and \(\sigma_{1}:\mathbb{R}^{d_{1}}\rightarrow\mathbb{R}^{d_{1}}\). Then \(\mathbf{x}\)'s first intermediate representation, \(\mathbf{e}_{1}(\mathbf{x}):=f_{1}(\mathbf{x})\), will lie in a subset of \(\mathbb{R}^{d_{1}}\), that is \(\mathbf{e}_{1}(\mathbf{x})\in\mathcal{A}_{1}\subset\mathbb{R}^{d_{1}}\). Similarly, \(\mathbf{x}\)'s second intermediate representation is given by, \(\mathbf{e}_{2}(\mathbf{x}):=(f_{2}\circ f_{1})(\mathbf{x})\)
where \(f_{2}(\mathbf{e}_{2}(\mathbf{x}))=\sigma(\mathbf{e}_{2}(\mathbf{x})\mathbf{W}^{(2)}+\mathbf{b}^{(2)})\) with \(\mathbf{W}^{(2)}\in\mathbb{R}^{d_{1}xd_{2}}\), \(\mathbf{b}_{2}\in\mathbb{R}^{d_{2}}\) and \(\sigma_{2}:\mathbb{R}^{d_{2}}\rightarrow\mathbb{R}^{d_{2}}\). We can define the remaining \(L-2\) intermediate representations in a similar way and by combining the intermediate representations with the responses we can generate \(L\) datasets \(\{\mathcal{D}_{k}\}_{k=1}^{L}\) where \(\mathcal{D}_{k}=\{\mathbf{e}_{k}(\mathbf{x}_{i}),y_{i}\}_{i=1}^{n}\). We assume the same ordering of observations for all \(L\) datasets, so \(y_{i}\) will always refer to the same response for every \(\mathcal{D}_{k}\).

An example of generating \(L\) datasets is visualized in Figure 3 for the bike UCI dataset. Each of the panels contains one dataset to which we fit a Vecchia GP.

### Conditioning-set selection

Figure 3 suggests that the function \(\mathbf{e}_{k}(\mathbf{x})\mapsto y\) may be smooth and therefore we choose to estimate this function using a Gaussian process. For computational purposes, we approximate each of the GPs using Vecchia with nearest neighbors computed in the corresponding intermediate space. Doing so gives rise to several conditioning sets for each observation \(\mathbf{x}_{i}\).

For example, consider the \(i^{th}\) input \(\mathbf{x}_{i}\) and its previously ordered nearest neighbors. In Figure 3\(\mathbf{x}_{i}\) is given by the blue dot in the upper left most panel and the black crosses are \(\mathbf{x}_{i}\)'s previously ordered nearest neighbors. As \(\mathbf{x}_{i}\) and its nearest neighbors are mapped to their first intermediate representations, we see that \(\mathbf{x}_{i}\)'s new nearest neighbors (the magenta crosses) are different than the original conditioning set. The magenta crosses in each subsequent panel represent \(\mathbf{x}_{i}\)'s previously ordered nearest neighbors in the corresponding intermediate space. The magenta crosses changes from panel to panel, giving different conditioning sets as \(\mathbf{x}_{i}\) moves through the network.

Figure 2: The input data \(\mathbf{x}_{1},\mathbf{x}_{2},\mathbf{x}_{3},\ldots\) are fed into the network, and the output of each intermediate layer creates a new intermediate space (e.g., \(\mathcal{A}_{1}\)). Within each intermediate space, we can consider the nearest neighbors of any point in the space. For example, in the box labeled \(\mathcal{A}_{1}\), we see that the red test point \(x^{*}\) is closest to \(\mathbf{e}_{1},\mathbf{e}_{3},\mathbf{e}_{4}\) (in the first intermediate space), and in \(\mathcal{A}_{2}\), the closest points to \(x^{*}\) are \(\mathbf{e}_{2},\mathbf{e}_{3},\mathbf{e}_{4}\). We define \(L\) Vecchia GPs, denoted by \(\mathcal{V}_{1},\ldots,\mathcal{V}_{L}\), which estimate a distribution over the response using the nearest neighbors in their respective intermediate spaces. The resulting distributions are combined in a product-of-experts fashion to yield a single distribution parameterized by \(\hat{\mu}\) and \(\hat{\sigma}\). Although the neural network generates a point estimate \(\hat{y}^{*}\) for the value of \(y^{*}\), our method provides a distribution that better reflects the uncertainty in the prediction.
Letting \(g_{k}(i)\) denote the conditioning set for observation \(i\) from intermediate space \(k\), we can consider the **collection of conditioning sets**, which we denote by \(\mathcal{G}_{k}:=\{g_{k}(1),...,g_{k}(n)\}\). An \(L+1\) layer network will have \(L\) collections of conditioning sets, \(\mathcal{G}_{1},...,\mathcal{G}_{L}\). By using several collections of conditioning sets, we are utilizing the unique information present in the panels titled 'First Intermediate Space' through 'Fifth Intermediate Space'. This approach can lead to better predictions compared to using just one collection of conditioning sets.

### Ensemble building from conditioning sets

Each \(\mathcal{G}_{k}\) represents a hypothesis about the conditional independence structure of the data. For instance, with time series data produced by an AR(2) process, the \(i^{th}\) observation should depend on the two preceding observations in time, while for an AR(1) process, the conditioning set is limited to the immediately preceding data point. However, in general, it is unclear which \(\mathcal{G}_{k}\) is most appropriate.

Since we do not know which \(\mathcal{G}_{k}\) is most suitable, we propose an ensemble of Vecchia GPs, each of which is based on one of the collections of conditioning sets. With each Vecchia GP independently fit to its corresponding dataset (e.g. \(\mathcal{D}_{k}\)), we combine the predictions from each member of the ensemble in a product-of-experts fashion. Details on training and prediction are provided in Appendix A.

## 5 Application to pretrained models

We report results using networks that were pretrained on a collection of UCI regression tasks and a network pretrained for a face age estimation task. In each case, we compare our results to those obtained with the original network, and, where appropriate, we also compare against other GP-based methods, such as SVI-GP. All GPs use an ARD Matern 3/2 kernel. We assess performance using root mean square error (RMSE) and negative log-likelihood (NLL) on the standardized test data for the UCI regression tasks. We use the mean absolute error (MAE) on the original scale for face age estimation. We combine our predictions in \(y\)-space and details on this choice are given in Appendix G.

### UCI benchmarks

We trained a feed-forward neural network on three train-validation-test splits of UCI regression tasks. After training, we constructed the deep Vecchia ensemble as detailed in Section 4 and made predictions for the test set using the

Figure 3: The 2D TSNE projection of input points in the original space of the bike UCI dataset (top left panel) and intermediate spaces (remaining five panels) with the color denoting the response value corresponding to each point. The blue dot is a fixed test point that has been propagated through the network. The black crosses denote the ordered nearest neighbors of the blue dot in the original input space, and their corresponding intermediate representations. The magenta crosses denote the ordered points nearest to the blue point within each intermediate space.
method described in Appendix A. Additionally, we trained the stochastic variational inference GP (SVI-GP) of Hensman, Fusi, and Lawrence [12] and the scaled Vecchia GP of Katzfuss, Guinness, and Lawrence [17] separately on the training data for prediction on the test set.

Table 1 compares NLL (left column under each dataset) and RMSE (right column under each dataset) for different methods on UCI datasets. The bold value in each column is the best performing (lower is better). The value in parentheses is the sample standard deviation of the mean of the metric over the three different data splits. The last line of the table specifies the sample size (\(n\)) and dimension (\(d\)) for each dataset. Our method is given by "DVE", "Neural Net" represents the original neural network, "Vecchia" stands for the scaled Vecchia model of Katzfuss, Guinness, and Lawrence [17], and SVI-GP is the model from Hensman, Fusi, and Lawrence [12].

We can see that DVE performs the best among all given models in terms of NLL and is second only to the neural network in RMSE on two datasets. Of course the DVE estimates are built upon the network, but it is remarkable how the DVE is able to provide uncertainty estimates to the neural network's predictions, providing much better performance than existing GP approximations.

### Face age estimation

We now use our deep Vecchia ensemble with a pretrained network for face age estimation. Face age estimation involves predicting the age of a subject in a photograph. The moving window regression (MWR) of Shin, Lee, and Kim [27] is a CNN-based approach that uses ordinal regression to solve this problem.

The MWR procedure uses a CNN to extract features of a test image as well as two references images. These features are concatenated together and fed into a feedforward neural network, called a \(\rho\)-regressor, to predict the age of the test image relative to the reference images. We use the pretrained global \(\rho\)-regressor for the UTKFace dataset from Shin, Lee, and Kim [27] to form the intermediate representations of a deep Vecchia ensemble as described in Section 4.1. For more details on the MWR procedure, the UTKFace dataset, our processing and comparisons, see Appendix D.

In Table 2, we compare the results of the original global model (Global \(\rho\)-regressor), the deep Vecchia ensemble (DVE), a scaled Vecchia GP [17] fit to the output of the second to last layer (Vecchia) and the SVI-GP of Hensman, Fusi, and Lawrence [12] (SVI-GP) also fit to the outputs of the second-to-last layer. The latter two approaches can be viewed as variants of deep-kernel GPs.

By using the deep Vecchia ensemble, we were able to decrease the MAE for the global model and additionally provide an estimate of the model's uncertainty. The NLL for the deep ensemble also outperformed the models that provided estimates of uncertainty. The difference in NLL and MAE between the Vecchia GP on the second-to-last layer and the deep ensemble Vecchia supports our claim that there is value in using the information in the intermediate layers for both predictions and uncertainty quantification.

## 6 Exploring intermediate spaces

In this section we investigate the intermediate spaces of a small network to understand how the deep Vecchia ensemble is using different layers. We fit a feed forward network to samples from the S-curve dataset [28]. In the bottom row of Figure 4, we visualize the internal representations of the network. In the top row of Figure 4, we visualize the original

\begin{table}
\begin{tabular}{l c c c c c c c c c c c c c c c c c c c c c c c c c c} \hline \hline Method & \multicolumn{3}{c}{Kish-i0K} & \multicolumn{3}{c}{Elevators} & \multicolumn{3}{c}{KEGG} & \multicolumn{3}{c}{bike} & \multicolumn{3}{c}{Protein} \\ \hline DVE & **-1.445 (0.011)** & **0.057 (0.001)** & **0.424 (0.010)** & 0.373 (0.007) & **-1.021 (0.030)** & 0.087 (0.002) & **-3.154 (0.109)** & **0.016 (0.002)** & **0.646 (0.029)** & **0.494 (0.014)** \\ Neural Net & N/A & 0.069 (0.000) & N/A & **0.343 (0.006)** & N/A & **0.085 (0.002)** & N/A & 0.018 (0.001) & N/A & 0.014 (0.002) \\ Vecchia & -0.028 (0.015) & 0.209 (0.015) & 0.525 (0.015) & 0.414 (0.006) & -0.497 (0.026) & 0.098 (0.002) & 0.229 (0.010) & 0.183 (0.001) & 0.581 (0.003) & 0.571 (0.002) \\ SVM-GP & -0.093 (0.008) & 0.068 (0.005) & 0.455 (0.017) & 0.381 (0.007) & -0.912 (0.033) & 0.098 (0.003) & -1.400 (0.047) & 0.073 (0.004) & 1.085 (0.001) & 0.713 (0.001) \\ \hline \((n\_d)\) & \multicolumn{3}{c}{(400K,8)} & \multicolumn{3}{c}{(16.6K,18)} & \multicolumn{3}{c}{(48.8K,20)} & \multicolumn{3}{c}{(17.4K, 17)} & \multicolumn{3}{c}{(45.7K,89)} \\ \hline \hline \end{tabular}
\end{table}
Table 1: NLL (left columns) and RMSE (right columns) on UCI benchmarks

\begin{table}
\begin{tabular}{l c c} \hline \hline  & MAE & NLL \\ \cline{2-4} Global \(\rho\)-Regressor & 4.650 & N/A \\ DVE & 4.636 & 1.085 \\ SVI-GP & 4.661 & 2.450 \\ Vecchia & 4.693 & 1.789 \\ \hline \hline \end{tabular}
\end{table}
Table 2: MAE and NLL for UTKFace dataset
data and the predictions formed by the deep Vecchia ensemble and the neural network. For more details, see Appendix E.

The first internal representation had more separation between yellow and green points compared to the final representation, indicating the earlier layer held information distilled away by the final layer. Deep Vecchia ensemble preserved this information, while the neural network discarded it. The deep Vecchia prediction closely resembled the original dataset, supported by lower RMSE values compared to the neural network's prediction. We believe a similar phenomenon is occurring with the models used in Section 5.

To check whether the observed behavior held more generally we repeated the experiment using different networks and for every network we checked the results was the same. For example, with four hidden layers with (10, 5, 2, 2) units, the RMSE dropped from 0.18 to 4.2e-5. Changing the number of units per layer to (16,8,4,4) showed the RMSE drop from 0.007 to 3.9e-5.

### Interpreting internal representations

Our model can be seen from two perspectives: an ensembling of Vecchia GPs, and a method for quantifying epistemic uncertainty in neural networks. Below we relate both view points to the results in this section.

From the Vecchia ensemble stance, we can see the intermediate representations in Figure 4 as providing three different distance metrics that can be used to compute possibly unique conditioning sets. The requirements to make this connection are formalized in Proposition 1. The proof of which can be found in Appendix B:

**Proposition 1**.: _Consider a sequence of injective functions \(f_{1},..,f_{L},f_{L+1}\) and a sequence of metric spaces \(\{(M_{i},d_{i})\}_{i=1}^{L+1}\) such that \(f_{i}:\mathcal{M}_{i}\rightarrow\mathcal{M}_{i+1}\) for \(i=1,...,L\). Then \(\{(M_{i},\tilde{d}_{i})\}_{i=2}^{L+1}\) defines a sequence of metric spaces, where \(\tilde{d}_{i}(x,y):=d_{i}(\mathbf{e}_{i}(\mathbf{x}),\mathbf{e}_{i}(\mathbf{x}^{\prime}))\) for any \(\mathbf{x},\mathbf{x}^{\prime}\in M_{1}\), with \(\mathbf{e}_{i}(\mathbf{x}):=f_{i}\circ f_{i-1}\circ...\circ f_{1}(\mathbf{x})\)._

Proposition 1's requirement that the function be injective ensures unique internal representations for unique input points, but continuity is not required. Discontinuous layers are useful when modeling functions with high Lipschitz constants or discontinuities. In such cases, predicting values using Euclidean nearest neighbors in \(\mathbb{R}^{d}\) can be inaccurate, but the function may be continuous with respect to a different distance metric on \(\mathbb{R}^{d}\). The difference in continuity

Figure 4: This figure shows the original data from the S-curve dataset in the top left panel. The top middle panel depicts the predictions generated by the deep Vecchia Ensemble, while the top right panel shows the predictions from the neural network. The three panels in the bottom row show the internal representations of the data from the three hidden layers of the network. The color of the points in each of these panels corresponds to the response value. The arrows pointing to the deep Vecchia prediction emphasize that the deep Vecchia ensemble builds predictions from all three layers. In contrast, the final network prediction relies on just the final layer, as indicated by the single arrow pointing to it.
is reminiscent of what we observed in Figure 3, where the function became progressively smoother for later intermediate spaces. While a similar goal is accomplished using deep kernels, by using multiple intermediate spaces we allow for competing hypothesis about what distance metric is best to be used.

From the perspective of uncertainty in neural networks, we can see the sequence of distance metrics as measuring how far a test point is away from instances in the training data. Of course, if a test point is close to existing training data with respect to the entire sequence of distance metrics, then the ensemble of Vecchia GPs will provide an estimate with small uncertainty estimates. If the test point is out of distribution, then the internal representations will not match existing data and we will return something close to the GP prior.

## 7 Conclusion

We presented an approach for UQ of pretrained neural networks. Our model has the added benefit of improving the Vecchia approximation and is the first approach to ensemble over the conditioning sets of the Vecchia approximation. The effectiveness of the model was shown on benchmark regression tasks and face age estimation. The study on internal representations provided insight into the inner workings of the approach.

We believe that our approach can find success in latent space Bayesian optimization [9] procedures that use a neural network to map latent representations to the response of interest, especially if the map between the latent space and the objective is not smooth. Our approach may reduce the need to provide additional structure to the latent space to allow for GP regression such as in Grosnit et al. [10]. We also believe that our method would work well in any application where nearest neighbors cannot be computed using Euclidean distance between inputs alone, such as high-dimensional regression.

Our proposed model has four key limitations. First, we need the data used to train the network which may not always be available. One possible solution is to simultaneously fit the network and an ensemble inducing-point GPs or to use a smaller validation set to build the deep Vecchia ensemble. The second limitation is that we combined predictions in \(y\)-space, so we are limited to using a Gaussian likelihood. This can be remedied by combining in \(f\)-space and defining a procedure to combine the noise terms in each of the Vecchia GP's likelihoods. The third limitation is that we use a DNN with fixed weights but many BNN methods assume the weights are random variables. To incorporate the DVE into BNNs, each intermediate representation can be expressed by a distribution rather than a point estimate. Then an error-in-variables Vecchia GP can be made by combining the kernel from Bachoc et al. [1] with the correlation Vecchia strategy of Kang and Katzfuss [15]. The final limitation is that our exploration was limited to feed-forward layers. To extend our model to more general architectures, a distance metric between layer outputs can be defined to make our approach applicable.

Given our promising results and a clear path towards solving our current limitations, we believe the deep Vecchia ensemble should be adopted more broadly, and we would be excited to see others improve upon our work.
