An Overview of Uncertainty Quantification Methods
for Infinite Neural Networks

One fundamental tool used in machine learning to quantify uncertainty is Gaussian processes (GPs). They are primarily encountered in regression problems, but they can be extended to classification or clustering tasks [Kapoor et al., 2010, Kim and Lee, 2007].

Simply put, a GP works as follows: (i) for any input $\bm{x}\in\mathbb{R}^{n}$, it yields a probability distribution for any possible outcome $y\in\mathbb{R}$, and (ii) any finite collection of input points follows a multivariate Gaussian distribution $\mathcal{N}(\underline{\mu},\underline{\Sigma})$ with mean vector $\underline{\mu}\in\mathbb{R}^{n}$ and covariance matrix $\underline{\Sigma}\in\mathbb{R}^{n\times n}$ where one usually assumes the data to be centered ($\underline{\mu}=0$). The matrix $\underline{\Sigma}$ is determined by a kernel function $\Sigma$ that allows to incorporate prior knowledge about the shape and characteristics of the data distribution. We will denote a GP with kernel function $\Sigma$ applied on centered data as $\mathcal{GP}(0,\Sigma)$. The most commonly used kernel function is the squared exponential or radial basis function (RBF):

where the variance $\sigma$ and length $l$ serve as hyperparameters to encode domain knowledge. For a new test point $\bm{x}$, the prior distribution over $y$ is given by:

Then, the dataset $\mathcal{D}$ can be used as part of the Bayesian framework to update prior beliefs and obtain a posterior distribution $P$ by conditioning on the training points, as follows:

The most common tools used to quantify uncertainty specifically in neural networks come from the distinction between two fundamental approaches. On the one hand, we can use the variance of predictions of multiple neural networks that follow from different random initializations. We refer to this frequentist type of ensemble method as deep ensemble (DE). On the other hand, a Bayesian neural network (BNN) is the natural result of applying Bayesian inference to neural networks. In a BNN, we place priors – typically independent Gaussians – on the network weights and biases. As a result, the posterior predictive distribution gives a prediction interval instead of a single-point estimate.

BNNs used in practice have posterior distributions that are highly complex and multi-modal, challenging even the most sophisticated Markov chain Monte Carlo (MCMC) sampling methods. For that reason, practitioners oftentimes use mean-field variational inference (MFVI) to approximate the posterior with a factorized Gaussian, i.e. assume the posterior distributions over the network’s parameters to be independent.

An alternative to MFVI that scales better to large datasets than BNNs and GPs are neural linear models (NLMs).

First introduced by Lázaro-Gredilla and Figueiras-Vidal [2010] as Marginalized Neural Networks and later re-introduced by Snoek et al. [2015], NLMs consist of a Bayesian linear regression performed on the last layer of a deep neural network. They resemble BNNs in structure, except only the final linear layer has priors on its coefficients. An example is shown in Figure 1.

The first layers whose weights $\bm{\theta}$ are simple scalars can be viewed as a feature map $\Phi_{\bm{\theta}}$ that transforms input $\bm{x}$ into features $\Phi_{\bm{\theta}}(\bm{x})$, which are then fed into the Bayesian linear regression. Because of this particular structure, training and inference with NLMs differs from usual procedures and usually consist of two steps: training the deterministic layers first using a point estimate of the output layer weights, and computing the posterior of that Bayesian layer second. In [Thakur et al., 2021], the authors propose an instance of this procedure that has the benefit of not underestimating uncertainty in data-scarce regions.

When all priors on the output weights $\mathbf{w}$ are independent and identical Gaussian distributions, the posterior is a multivariate Gaussian whose mean $\underline{\mu}_{NLM}$ and covariance matrix $\underline{\Sigma}_{NLM}$ can be expressed in closed-form:

where $I$ is the identity matrix, and $\alpha=\frac{1}{{\sigma_{\epsilon}}^{2}}$ and $\beta=\frac{1}{{\sigma_{\mathbf{w}}}^{2}}$ are prior hyperparameters equal to the inverse of the variance of the regression noise and weight variables, respectively. These expressions reveal that the main computational effort of inference with NLMs lies in the inversion of the matrix $\alpha{\Phi_{\bm{\theta}}(\mathcal{X})}^{T}\Phi_{\bm{\theta}}(\mathcal{X})+\beta I$, whose dimension is the number of neurons in the last hidden layer. Compared with GPs who scale with the size of the training data set, as explained later in subsubsection 3.2.1, NLMs scale with the width of the neural network’s last layer, which is often of lower dimensionality. This is where NLMs gain their advantage in tractability and computational efficiency, which is why NLMs have gained popularity.

The last piece of background to understand is the theory describing the training process of a neural network with gradient descent. A helpful tool for that is the neural tangent kernel (NTK) introduced by Jacot et al. [2018]:

In general, this kernel function depends on the random initialization $\bm{\theta_{0}}$ and changes over time. However, Jacot et al. [2018] show that in the infinite-width limit, the NTK at initialization $\hat{\Theta}_{0}$ becomes deterministic and only depends on the activation function and network architecture. In this case, the NTK also stays constant during training for small enough learning rates: $\hat{\Theta}_{t}=\hat{\Theta}_{0}$. We call this deterministic and constant kernel $\Theta$. Yang [2020] extends these convergence results from multilayer perceptrons (MLPs) to other architectures including convolutional and batch-normalization layers.

Now, we investigate how the NTK describes change in the output function of a NN during training. Jacot et al. [2018] and Lee et al. [2019] show that for small enough learning rates, the NN prediction $f_{t}(\bm{x}):=f(\bm{x},\bm{\theta}_{t})$ at training time $t$ follows an ordinary differential equation (ODE) that is characterized by the NTK:

Here, $\nu$ denotes the learning rate and $\nabla_{f_{t}(\mathcal{X})}\mathcal{L}$ describes the gradient of the loss with respect to model’s output on the training set. Using an MSE loss for $\mathcal{L}$ and the fact that the NTK stays constant over time, Equation 8 simplifies to a linear ordinary differential equation with known exponentially decaying solution. So, the output for the training dataset at time $t$ is

This shows that the NN achieves zero training loss because the NTK is a positive definite matrix, hence the exponential term converges to zero.

Equivalently, based on the results of [Lee et al., 2019], He et al. [2020] show that $f_{t}$ can be expressed based on the random initialization $f_{0}$. For $t\rightarrow\infty$ this gives:

Under the condition that the network parameters follow independent and identical (iid) distributions with zero mean and finite variance and the activation function of the last layer is bounded, the prior predictive distribution converges to a GP with zero mean and finite variance [Neal, 1996]. The intuition is as follows: for a fixed input $\bm{x}$, every hidden layer output $\phi_{i}$ is a random variable with finite mean and variance as the activation function is bounded. When multiplied by the final layer weights, $\phi_{i}w_{i}$ are iid random variables with zero mean and finite variance. According to the central limit theorem, their sum converges to a Gaussian distribution. If we rescale the variances according to the network width, the network output becomes a standard normal distribution $f(\bm{x})\sim\mathcal{N}(0,1)$. As this holds true for every $\bm{x}$, the neural network is equivalent to a GP called the neural network Gaussian process (NNGP). A sampled function from this NNGP is equivalent to the output function of a randomly initialized NN and to a prior predictive sample of a BNN.

While Neal [1996] gives a proof of existence, i.e. that a NN converges to a GP in the infinite-width limit, Neal does not provide any analytic form of the covariance function $\mathcal{K}$. Williams [1997] builds on that work and provides closed-form kernel functions for different sets of activation functions and prior distributions on the network parameters.

So far we have not considered any data but made predictions based on randomly drawn initial parameters. In the following sections, we discuss three approaches to incorporate the training data. First, we update the GP prior kernel by conditioning on the training data to obtain a GP posterior kernel. Second, we use gradient-based methods to train the NN. Third, we make the neural network Bayesian and approximate the BNN posterior.