Variational Inference on the Final-Layer Output of Neural Networks
Traditional neural networks are simple to train but they typically produce
overconfident predictions. In contrast, Bayesian neural networks provide good
uncertainty quantification but optimizing them is time consuming due to the
large parameter space. This paper proposes to combine the advantages of both
approaches by performing Variational Inference in the Final layer Output space
(VIFO), because the output space is much smaller than the parameter space. We
use neural networks to learn the mean and the variance of the probabilistic
output. Like standard, non-Beyesian models, VIFO enjoys simple training and one
can use Rademacher complexity to provide risk bounds for the model. On the
other hand, using the Bayesian formulation we incorporate collapsed variational
inference with VIFO which significantly improves the performance in practice.
Experiments show that VIFO and ensembles of VIFO provide a good tradeoff in
terms of run time and uncertainty quantification, especially for out of
distribution data.
Yadi Wei, Roni Khardon
# Direct Uncertainty Quantification

###### Abstract

Traditional neural networks are simple to train but they produce overconfident predictions, while Bayesian neural networks provide good uncertainty quantification but optimizing them is time consuming. This paper introduces a new approach, direct uncertainty quantification (DirectUQ), that combines their advantages where the neural network directly models uncertainty in output space, and captures both aleatoric and epistemic uncertainty. DirectUQ can be derived as an alternative variational lower bound, and hence benefits from collapsed variational inference that provides improved regularizers. On the other hand, like non-probabilistic models, DirectUQ enjoys simple training and one can use Rademacher complexity to provide risk bounds for the model. Experiments show that DirectUQ and ensembles of DirectUQ provide a good tradeoff in terms of run time and uncertainty quantification, especially for out of distribution data.

Machine Learning, ICML

## 1 Introduction

With the development of training and representation methods for deep learning, models using neural networks provide excellent predictions. However, such models fall behind in terms of uncertainty quantification and their predictions are often overconfident (Guo et al., 2017). Bayesian methods provide a methodology for uncertainty quantification by placing a prior over parameters and computing a posterior given observed data, but the computation required for such methods is often infeasible. Variational inference (VI) is one of the most popular approaches for approximating the Bayesian outcome, e.g., (Blundell et al., 2015; Graves, 2011; Wu et al., 2019). By minimizing the KL divergence between the variational distribution and the true posterior and constructing an evidence lower bound (ELBO), one can find the best approximation to the intractable posterior. However, when applied to deep learning, VI requires sampling to compute the ELBO, and it suffers from both high computational cost and large variance in gradient estimation. Wu et al. (2019) have proposed a deterministic variational inference (DVI) approach to alleviate the latter problem. The idea of this approach relies on the central limit theorem, which implies that with sufficiently many hidden neurons, the distribution of each layer forms a multivariate Gaussian distribution. Thus we only need to compute the mean and covariance of each layer. However, DVI still suffers from high computational cost and complex optimization.

Inspired by DVI, we observe that the only aspect that affects the prediction is the distribution of the last layer in the neural network. We therefore propose to use a neural network to directly output the mean and the diagonal covariance of the last layer, which avoids calculating a distribution over network weights and saves intermediate computations. Like all Bayesian methods, our method quantifies both epistemic and aleatoric uncertainty of the output. However, it does so in an explicit manner, so we call it _Direct Uncertainty Quantification_ (DirectUQ). DirectUQ has a single set of parameters and thus enjoys simple optimization as in non-Bayesian methods yet it has the advantage of uncertainty quantification in predictions similar to Bayesian methods.

We can motivate DirectUQ from several theoretical perspectives. First, due to its simplicity, one can derive risk bounds for the model through Rademacher complexity. Second, DirectUQ can be derived as an alternative ELBO for neural networks which uses priors and posteriors on the output layer. We show that, for the linear case, with expressive priors DirectUQ can capture the same predictions as standard variational inference. With practical priors and deep networks the model can be seen as a simpler alternative. Third, through the interpretation as ELBO, we derive improved priors (or regularizers) for DirectUQ motivated by collapsed variational inference (Tomczak et al., 2021) and empirical bayes (Wu et al., 2019). The new regularizers greatly improve the performance of DirectUQ.

An experimental evaluation compares DirectUQ with VI and non-Bayesian neural networks optimized by stochastic gradient descent (SGD), as well as ensembles of these models. The results show that (1) DirectUQ is slightly slower than SGD but much faster than VI, and (2) Direc
tUQ achieves better uncertainty quantification on shifted and out-of-distribution data while preserving the quality of in-distribution predictions. Overall, ensembles of DirectUQ provide good tradeoff in terms of run time and uncertainty quantification especially for out-of-distribution data.

### Related Work

DirectUQ shares some similarities with Kendall and Gal (2017), where both use neural networks to output the mean and covariance of the last layer. However, unlike their method DirectUQ represents a variational lower bound, and there are differences in the loss function and approach for epistemic uncertainty which are elaborated further in Section 2. Another similar line of work (Sun et al., 2019; Tran et al., 2022) performs variational inference on the functional space. However, they focus on choosing a better prior in weight space which is induced from Gaussian Process priors on functional space, whereas DirectUQ directly puts a simple prior in function space.

DirectUQ differs from existing variational inference methods. Last-layer variational inference (Brosse et al., 2020) performs variational inference on the weights of the last layer, while we perform variational inference on the output of the last layer. The local reparametrization trick (Tomczak et al., 2020; Oleksiienko et al., 2022) performs two forward passes with the mean and variance to sample the output for each layer, while we only require one pass and sample the output of the last layer for prediction.

Various alternative Bayesian techniques have been proposed. One direction is to get samples from the true posterior, as in Markov chain Monte Carlo methods (Wenzel et al., 2020; Izmailov et al., 2021). Expectation propagation aims to minimize the reverse KL divergence to the true posterior (Teh et al., 2015; Li et al., 2015). These Bayesian methods, including variational inference, often suffer from high computational cost and therefore hybrid methods were proposed. Stochastic weight averaging Gaussian (Maddox et al., 2019) forms a Gaussian distribution over parameters from the stochastic gradient descent trajectory. Dropout (Gal and Ghahramani, 2016) randomly sets weights 0 to model the epistemic uncertainty. Deep ensembles (Lakshminarayanan et al., 2017) use ensembles of models learned with random initialization and shuffling of data points and then average the predictions. These methods implicitly perform approximate inference. In addition to these methods, there are also non-Bayesian methods to calibrate the overconfident predictions, for example, temperature scaling (Guo et al., 2017) introduces a temperature parameter to anneal the predictive distribution to avoid high confidence. DirectUQ strikes a balance between simplicity and modelling power to enable simple training and Bayesian uncertainty quantification.

## 2 Direct Uncertainty Quantification

In this section we describe our DirectUQ method in detail. Given a neural network parametrized by weights \(W\) and input \(x\), the output layer is \(z=f_{W}(x)\). In classification, \(z\) is a \(k\)-dimensional vector, where \(k\) is the number of classes. The probability of being class \(i\) is defined as

\[p(y=i|z)=\text{softmax}(z)_{i}=\frac{\exp z_{i}}{\sum_{j}\exp z_{j}}. \tag{1}\]

In regression, \(z=(m,l)\) is a 2-dimensional vector. We apply a function \(g\) on \(l\) that maps \(l\) to a positive real number. The probability of the output \(y\) is:

\[p(y|z)=\mathcal{N}(y|m,l)=\frac{1}{\sqrt{2\pi g(l)}}\exp\left(- \frac{(y-m)^{2}}{2g(l)}\right). \tag{2}\]

Traditional non-Bayesian deep learning learns weights by minimizing \(-\log p(y|z)\) using stochastic gradient descent. We refer to this as SGD.

By fixing the weights \(W\), SGD maps \(x\) to \(z\) deterministically, while Bayesian methods seek to map \(x\) to a distribution over \(z\). Variational inference puts a distribution over \(W\) and marginalizes out to get a distribution over \(z\). According to central limit theorem, with sufficiently wide neural network, the marginalization distribution of \(z\) is normal (Wu et al., 2019). DirectUQ pursues this in a direct manner. It has two sets of weights, \(W_{1}\) and \(W_{2}\) (with shared components), to model the mean and variance of \(z\). That is, \(\mu_{q}(x)=f_{W_{1}}(x)\), \(\sigma_{q}(x)=g(f_{W_{2}}(x))\), where \(g:\mathbb{R}\rightarrow\mathbb{R}^{+}\) maps the output to positive real numbers as the variance is positive. Thus, \(q(z|x)=\mathcal{N}(z|\mu_{q}(x),\text{diag}(\sigma_{q}^{2}(x)))\), where \(\mu_{q}(x),\sigma_{q}^{2}(x)\) are vectors of the corresponding dimension. We will call \(q(z|x)\) the _variational predictive distribution_.

Note the the prediction over \(y\) is given by \(\int q(z|x)p(y|z)dz\) which is different from the basic model. In regression, the basic model of Eq (2) is known as the mean-variance estimator (Kabir et al., 2018; Khosravi et al., 2011; Kendall and Gal, 2017). Applying DirectUQ on this model, we will have **four** outputs, two of which are the means of \(m\) and \(l\), and the other two are the variances of \(m\) and \(l\). The mean of \(l\) represents the heteroscedastic aleatoric uncertainty and the variances of \(m\) and \(l\) model the epistemic uncertainty as the variances come from marginalization of Bayesian posteriors. In classification, using Eq (1), the mean of \(z\) induces aleatoric uncertainty and the variance of \(z\) represents epistemic uncertainty.

The standard Bayesian approach puts a prior on the weights \(W\). Instead, since DirectUQ models the distribution over \(z\), we put a prior over \(z\). We consider two options, a conditional prior \(p(z|x)\) and a simpler prior \(p(z)\). Both of these choices
yield a valid ELBO using the same steps:

\[\log p(y|x) \geq\mathbb{E}_{q(z|x)}\left[\log\frac{p(y,z|x)}{q(z|x)}\right]\] \[=\mathbb{E}_{q(z|x)}[\log p(y|z)]-\text{KL}(q(z|x)||p(z|x)). \tag{3}\]

Eq (3) is defined for every \((x,y)\). For a dataset \(\mathcal{D}=\{(x,y)\}\), we optimize \(W_{1}\) and \(W_{2}\) such that:

\[\sum_{(x,y)\in\mathcal{D}}\left\{\mathbb{E}_{q(z|x)}[\log p(y|z)]-\text{KL}(q (z|x)||p(z|x))\right\}\]

is maximized. We regard the negation of the first term \(\mathbb{E}_{q(z|x)}[-\log p(y|z)]\) as the loss term and treat \(\text{KL}(q(z|x)||p(z))\) as a regularizer. In practice, we put a coefficient \(\eta\) on the regularizer, as is often done in variational approximations (Higgins et al., 2017). Then, from a regularized loss minimization view, our objective becomes:

\[\min_{W_{1},W_{2}}\sum_{(x,y)\in\mathcal{D}}\left\{\mathbb{E}_{q(z|x)}[-\log p (y|z)]+\eta\text{KL}(q(z|x)||p(z|x))\right\}. \tag{4}\]

In most cases, the loss term is intractable, so we use Monte Carlo samples to approximate it, where implementations can use reparametrization to reduce the variance in gradients:

\[\mathbb{E}_{q(z|x)}[-\log p(y|z)]\approx\frac{1}{M}\sum_{m=1}^{M}-\log p(y|z^{ (m)}), \tag{5}\]

where \(z^{(m)}\sim q(z|x)\).

We can now discuss the relation to Kendall & Gal (2017) in more details. This work also draws samples in the output space but replaces the loss term in Eq (5) with the cross entropy loss, \(-\log\frac{1}{M}\sum_{m}p(y|z^{(m)})\). Furthermore, they use dropout and their objective has no explicit regularization, whereas our formulation yields a valid ELBO.

## 3 Rademacher Complexity of DirectUQ

In this section we provide generalization bounds for DirectUQ through Rademacher Complexity. We need to make the following assumptions.

**Assumption 3.1**.: \(\log p(y|z)\) is \(L_{0}\)-Lipschitz in \(z\), i.e., \(\log p(y|z)-\log p(y|z^{\prime})\leq L_{0}\|z-z^{\prime}\|_{2}\).

**Assumption 3.2**.: The link function \(g\) is \(L_{1}\)-Lipschitz.

We show in the Appendix that these assumptions hold for classification and with a smoothed loss for regression. With the two assumptions, we derive the main result:

**Theorem 3.3**.: _Let \(\mathcal{H}\) be the set of functions that can be represented with neural networks with parameter space \(\mathcal{V}\), \(\mathcal{H}=\{f_{W}(\cdot)|W\in\mathcal{W}\}\). DirectUQ has two components, so the DirectUQ hypothesis class is \(\mathcal{H}\times\mathcal{H}\)\(=\{(f_{W_{1}}(\cdot),f_{W_{2}}(\cdot))\mid W=(W_{1},W_{2}),W_{1},W_{2} \in\mathcal{W}\}\). Let \(l\) be the loss function for DirectUQ, \(l(W,(x,y))=E_{q_{W}(z|x)}[-\log p(y|x)]\). Then the Rademacher complexity of DirectUQ is bounded as \(R(l\circ(\mathcal{H}\times\mathcal{H})\circ S)\leq(L_{0}(1+L_{1})K)\cdot R( \mathcal{H}\circ S)\), where \(K\) is the dimension of \(z\)._

Proof.: We show that the loss is Lipschitz in \(f_{W_{1}}(x)\) and \(f_{W_{2}}(x)\). Fix any \(x\), and \(W,W^{\prime}\). We denote the mean and standard deviation of \(q_{W}(z|x)\) by \(\mu\) and \(s\) and the same for \(q_{W^{\prime}}(z|x)\).

\[\mathbb{E}_{q_{W}(z|x)}[-\log p(y|z)]-\mathbb{E}_{q_{W^{\prime}} (z|x)}[-\log p(y|z)]\] \[= \mathbb{E}_{\epsilon\sim\mathcal{N}(0,I)}[\log p(y|\mu^{\prime}+ \epsilon\cdot s^{\prime})-\log p(y|\mu+\epsilon\cdot s)]\] \[\leq \mathbb{E}_{\epsilon\sim\mathcal{N}(0,I)}\left[L_{0}\|(\mu-\mu^{ \prime})+\epsilon\cdot(s-s^{\prime})\|_{2}\right]\] \[\leq L_{0}\|\mu-\mu^{\prime}\|_{2}+L_{0}\mathbb{E}_{\epsilon}[\sqrt{ \epsilon^{2}(s-s^{\prime})^{2}}]\] \[\leq L_{0}\|\mu-\mu^{\prime}\|_{2}+L_{0}\sqrt{\mathbb{E}_{\epsilon}[ \epsilon^{2}(s-s^{\prime})^{2}]}\] \[= L_{0}(\|\mu-\mu^{\prime}\|_{2}+\|s-s^{\prime}\|_{2})\] \[\leq L_{0}(\|\mu-\mu^{\prime}\|_{1}+\|s-s^{\prime}\|_{1})\]

where the third equation uses the Lipschitz assumption and the fifth uses Jensen's inequality. In the above equations the product and power are performed element-wise, and the derivation holds for multivariate \(z\). The loss function is Lipschitz in \(\mu\), which is exactly \(f_{W_{1}}(x)\). Further, \(s\) is \(L_{1}\)-Lipschitz in the logit \(f_{W_{2}}(x)\), thus, the loss function is \((L_{0}(1+L_{1}))\)-Lipschitz in \(f_{W_{1}}(x)\) and \(f_{W_{2}}(x)\). The theorem follows from Rademacher complexity bounds for bivariate functions (Lemma A.1) and for functions over vectors (Corollary A.2). 

Hence Rademacher complexity for DirectUQ is bounded through the Rademacher complexity of deterministic neural networks. This shows one advantage of DirectUQ which is more amenable to analysis than standard VI due to its simplicity. The Rademacher complexity for neural networks is \(O\left(\frac{B_{W}B_{x}}{\sqrt{N}}\right)\)(Golowich et al., 2018), where \(B_{W}\) bounds the norm of the weights and \(B_{x}\) bounds the input. The Rademacher complexity of DirectUQ is of the same order.

## 4 Comparison of DirectUQ and VI

DirectUQ is inspired by deterministic variational inference and it highly reduces the computational cost. We explore whether DirectUQ can produce exactly the same predictive distribution as VI. In this section, we show that this is the case for linear models but that for deep models DirectUQ is less powerful. We first introduce the setting of linear models. Let the parameter be \(\theta\), then the model is defined as:

\[y|x,\theta\sim p(y|\theta^{\top}x). \tag{6}\]

For example, \(p(y|\theta^{\top}x)=\mathcal{N}(y|\theta^{\top}x,\frac{1}{\beta})\) where \(\beta\) is a constant for Bayesian linear regression; and \(p(y=1|\theta^{\top}x)=\)
\(\frac{1}{1+\exp(-\theta^{\top}x)}\) for Bayesian binary classification.

The standard approach specifies the prior of \(\theta\) to be \(p(\theta)=\mathcal{N}(\theta|m_{0},S_{0})\), and uses \(q(\theta)=\mathcal{N}(\theta|m,S)\). Then the ELBO objective for linear models with a dataset \((X_{N},Y_{N})\) of size \(N\), where \(X_{N}=(x_{1},x_{2},\ldots,x_{N})\in\mathbb{R}^{d\times N}\) and \(Y_{N}=(y_{1},y_{2},\ldots,y_{N})\in\mathbb{R}^{N}\), is:

\[\sum_{i=1}^{N}\mathbb{E}_{q(\theta)}[\log p(y_{i}|\theta^{\top}x_ {i})]-\text{KL}(q(\theta)||p(\theta))\] \[= \sum_{i=1}^{N}\mathbb{E}_{q(\theta)}[\log p(y_{i}|\theta^{\top}x_ {i})]-\frac{1}{2}[\text{tr}(S_{0}^{-1}S)-\log|S_{0}^{-1}S|]\] \[-\frac{1}{2}(m-m_{0})^{\top}S_{0}^{-1}(m-m_{0})]. \tag{7}\]

As the following theorem shows, if we use a conditional correlated prior and a variational posterior that correlates data points, then in the linear case DirectUQ can recover the ELBO and the solution of VI.

**Theorem 4.1**.: _Let \(p(z|X_{N})=\mathcal{N}(z|m_{0}^{\top}X_{N},X_{N}^{\top}S_{0}X_{N})\) be a correlated and data-specific prior (which means that for different data \(x\), we have a different prior over \(z\)), and \(q(z|x)=\mathcal{N}(z|\theta_{1}^{\top}x,x^{\top}V_{2}x)\) be the variational predictive distribution of \(z\), where \(\theta_{1}\) and \(V_{2}\) are the parameters to be optimized. Then the DirectUQ objective is equivalent to the ELBO objective implying identical predictive distributions._

Proof.: The DirectUQ objective is:

\[\sum_{i=1}^{N}\left\{\mathbb{E}_{q(z|x_{i})}[\log p(y_{i}|z)] \right\}-\text{KL}(q(z|X_{n})||p(z|X_{N}))\] \[=\sum_{i=1}^{N}\left\{\mathbb{E}_{q(z|x_{i})}[\log p(y_{i}|z)]\right\}\] \[-\frac{1}{2}\text{tr}((X_{N}^{\top}S_{0}X_{N})^{-1}(X_{N}^{\top} V_{2}X_{N}))\] \[+ \frac{1}{2}\log|(X_{N}^{\top}S_{0}X_{N})^{-1}(X_{N}^{\top}V_{2}X_ {N})|\] \[- \frac{1}{2}(\theta_{1}^{\top}X_{N}-m_{0}^{\top}X_{N})^{\top}(X_{N }^{\top}S_{0}X_{N})^{-1}(\theta_{1}^{\top}X_{N}-m_{0}^{\top}X_{N}). \tag{8}\]

Assume \(N>d\). First consider the loss term. Let \(L\) be the Cholesky decomposition of \(V_{2}\), i.e. \(V_{2}=LL^{\top}\). By reparametrization, for \(\epsilon\sim\mathcal{N}(0,I_{d})\), \(\theta_{1}^{\top}x_{i}+x_{i}^{\top}L\epsilon\sim\mathcal{N}(\theta_{1}^{\top} x_{i},x_{i}^{\top}LL^{\top}x_{i})\) and thus

\[\mathbb{E}_{q(z|x_{i})}[\log p(y_{i}|z)] =\mathbb{E}_{\epsilon\sim\mathcal{N}(0,I_{d})}[\log p(y_{i}| \theta_{1}^{\top}x_{i}+x_{i}^{\top}L\epsilon)]\] \[=\mathbb{E}_{\epsilon\sim\mathcal{N}(0,I_{d})}[\log p(y_{i}|( \theta_{1}+L\epsilon)^{\top}x_{i})]\] \[=\mathbb{E}_{\theta\sim\mathcal{N}(\theta_{1},LL^{\top})}[\log p( y_{i}|\theta^{\top}x_{i})], \tag{9}\]

where the last equality uses reparametrization in a reverse order. By aligning \(\theta_{1}=m\) and \(V_{2}=LL^{\top}=S\), we recognize that Eq (9) is exactly the loss term in Eq (7). Thus the low-dimensional posterior on \(z\) yields the same loss term as the high-dimensional posterior over \(W\).

For the regularization, we use the pseudo inverse

\[(X_{N}^{\top}S_{0}X_{N})^{-1}=X_{N}^{\top}(X_{N}X_{N}^{\top})^{-1}S_{0}^{-1}(X_ {N}X_{N}^{\top})^{-1}X_{N}\]

and the same for \(V_{2}\). Then the regularization term in Eq (7) can be simplified to:

\[-\text{KL}(q(z|X_{N})||p(z|X_{N}))\] \[= -\frac{1}{2}\text{tr}(S_{0}^{-1}V_{2})+\frac{1}{2}\log|S_{0}^{-1} V_{2}|\] \[-\frac{1}{2}(\theta_{1}-m_{0})^{\top}S_{0}^{-1}(\theta_{1}-m_{0}). \tag{10}\]

By aligning \(\theta_{1}=m,V_{2}=S\), we can see that eq (10) is exactly the same as the regularizer in eq (7). 

However, as the next theorem shows, for the non-linear case even if we ignore the prior we may not be able to recover the loss term exactly.

**Theorem 4.2**.: _Given a neural network \(f_{W}\) parametrized by \(W\) and a mean-field Gaussian distribution \(q(W)\) over \(W\), there may not exist a set of parameters \(\widetilde{W}\) such that for all input \(x\) we have \(\mathbb{E}_{q(W)}[f_{W}(x)]=f_{\widetilde{W}}(x)\)._

Proof Sketch.: Consider the neural network with only one hidden unit and only one input. No matter what \(W\) we choose, there exists some region of \(x\) such that \(f_{\widetilde{W}}(x)=0\) due to the non-linear activation. However, \(\mathbb{E}_{q(W)}[f_{W}(x)]\) will include two segments of non-zero functions due to the properties of truncated normal distribution, thus, there is no \(\widetilde{W}\) that can perfectly recover \(\mathbb{E}_{q(W)}[f_{W}(x)]\). 

The results of this section point to the potential and limitations of DirectUQ in terms of expressiveness relative to standard VI. In practice, computing a correlated and data-specific prior \(p(z|x)\) is complex, and tuning its hyperparameters would be challenging. Hence, for a practical algorithm we propose to use a simple prior \(p(z)\) independent of \(x\). In addition, to reduce computational complexity, we do not learn a full covariance matrix and focus on the diagonal approximation. This limits expressive power but enables fast training of ensembles of DirectUQ.

## 5 Collapsed VI Applied to DUQ

Bayesian methods are often sensitive to the choice of prior parameters. To overcome this, Wu et al. (2019) used empirical Bayes (EB) to select the value of the prior parameters, and Tomczak et al. (2021) proposed collapsed variational inference, which defined a hierarchical model and performed inference on the prior parameters as well. Empirical Bayes
can be regarded as a special case of collapsed variational inference. In this section, we show how this idea is applicable in DirectUQ. In addition to \(z\), we model the prior mean \(\mu_{p}\) and variance \(\sigma_{p}^{2}\) as Bayesian parameters. Now the prior becomes \(p(z|\mu_{p},\sigma_{p}^{2})p(\mu_{p},\sigma_{p}^{2})\) and the variational distribution is \(q(z|x)q(\mu_{p},\sigma_{p}^{2})\). Then the objective becomes:

\[\log p(y|x)\geq\mathbb{E}_{q(z|x)q(\mu_{p},\sigma_{p}^{2})}\left[ \log\frac{p(y,z,\mu_{p},\sigma_{p}^{2}|x)}{q(z|x)q(\mu_{p},\sigma_{p}^{2})}\right]\] \[=\mathbb{E}_{q(z|x)}[\log p(y|z)]-\mathbb{E}_{q(\mu_{p},\sigma_{p }^{2})}[\text{KL}(q(z|x)||p(z|\mu_{p},\sigma_{p}^{2}))]\] \[-\text{KL}(q(\mu_{p},\sigma_{p}^{2})||p(\mu_{p},\sigma_{p}^{2})). \tag{11}\]

Similar to eq (4), we treat the first term as a loss term and the other two terms as a regularizer along with a coefficient \(\eta\) and aggregate over all data. Since the loss term does not contain \(\mu_{p}\) and \(\sigma_{p}^{2}\), we can get the optimal \(q^{*}(\mu_{p},\sigma_{p}^{2})\) by optimizing the regularizer and the choice of \(\eta\) will not affect \(q^{*}(\mu_{p},\sigma_{p}^{2})\). Then we can plug in the value of \(q^{*}\) into Eq (11). We next show how to compute \(q^{*}(\mu_{p},\sigma_{p}^{2})\) and the final collapsed variational inference objective. The derivations are similar to the ones in (Tomczak et al., 2021) but they are applied on \(z\) not on \(W\). In the following, we use \(K\) to denote the dimension of \(z\).

### Learn mean, fix variance

Let \(p(z|\mu_{p})=\mathcal{N}(z|\mu_{p},\gamma I)\), \(p(\mu_{p})=\mathcal{N}(\mu_{p}|0,\alpha I)\). Then \(q^{*}(\mu_{p}|x)\) is

\[\operatorname*{arg\,min}_{q(\mu_{p})}\mathbb{E}_{q(\mu_{p})}[ \text{KL}(q(z|x)||p(z|\mu_{p}))]+\text{KL}(q(\mu_{p})||p(\mu_{p})),\]

and the optimal \(q^{*}(\mu_{p}|x)\) can be computed as:

\[\log q^{*}(\mu_{p}|x)\propto-\frac{(\mu_{q}(x)-\mu_{p})^{\top}( \mu_{q}(x)-\mu_{p})}{2\gamma}-\frac{\mu_{p}^{\top}\mu_{p}}{2\alpha},\]

and \(q^{*}(\mu_{p}|x)=\mathcal{N}(\mu_{p}|\frac{\alpha}{\alpha+\gamma}\mu_{q}(x), \frac{\alpha\gamma}{\alpha+\gamma})\). Notice that unlike the prior \(q^{*}(\mu_{p})\) depends on \(x\). If we put \(q^{*}\) back in the regularizer of eq (11), the regularizer becomes:

\[\frac{1}{2\gamma}\left[1^{\top}\sigma_{q}^{2}(x)+\frac{\gamma}{ \gamma+\alpha}\mu_{q}(x)^{\top}\mu_{q}(x)\right]-\frac{1}{2}1^{\top}\log \sigma_{q}^{2}(x)\] \[+\frac{K}{2}\log(\gamma+\alpha)-\frac{K}{2}. \tag{12}\]

As in (Tomczak et al., 2021), one thing to observe from eq (12) is that it puts a factor \(\frac{\gamma}{\gamma+\alpha}<1\) in front of \(\mu_{q}(x)^{\top}\mu_{q}(x)\), which weakens the regularization on \(\mu_{q}(x)\). We refer to this method as "duq-mean".

### Learn both mean and variance

Let \(p(z|\mu_{p},\sigma_{p}^{2})=\mathcal{N}(z|\mu_{p},\sigma_{p}^{2})\), \(p(\mu_{p})=\mathcal{N}(\mu_{p}|0,\frac{1}{t}\sigma_{p}^{2})\), \(p(\sigma_{p}^{2})=\mathcal{IG}(\sigma_{p}^{2}|\alpha,\beta)\), where \(\mathcal{IG}\) indicates the inverse Gamma distribution. The posterior of \(\mu_{p}\) is Gaussian and the posterior of \(\sigma_{p}^{2}\) is inverse Gamma. In this case we can show that \(q^{*}(\mu_{p}|x)=\mathcal{N}(\mu_{p}|\frac{1}{t+1}\mu_{q}(x),\frac{1}{(1+t)} \sigma_{p}^{2})\) and \(q^{*}(\sigma_{p}^{2}|x)=\mathcal{IG}(\sigma_{p}^{2}|(\alpha+\frac{1}{2})1,\beta +\frac{t}{2(t+1)}\mu_{q}(x)^{2}+\frac{1}{2}\sigma_{q}^{2}(x))\) and the regularizer becomes

\[(\alpha+\frac{1}{2})1^{\top}\log\left[\beta 1+\frac{t}{2(1+t)}\mu_{q}(x)^{2}+ \frac{1}{2}\sigma_{q}^{2}(x)\right] \tag{13}\] \[-\frac{1}{2}1^{\top}\log\sigma_{q}^{2}(x). \tag{14}\]

We refer to this method as "duq-mv".

### Empirical Bayes

Let \(p(\sigma_{p}^{2})=\mathcal{IG}(\sigma_{p}^{2}|\alpha,\beta)\), \(p(z|\sigma_{p}^{2})=\mathcal{N}(z|0,\sigma_{p}^{2})\) and let \(q(\sigma_{p}^{2})\) be a delta distribution \(\delta(s^{*}(x))\). Then the regularizer of eq (11) becomes:

\[\text{KL}(q(z|x)||p(z|\sigma_{p}^{2}))-\log p(\sigma_{p}^{2}). \tag{15}\]

By minimizing this objective we obtain the optimal value \(s^{*}(x)=\frac{\mu_{q}(x)^{\top}\mu_{q}(x)+1^{\top}\sigma_{q}^{2}(x)+2\beta}{K +2\alpha+2}\). Plugging this value into the KL term we obtain:

\[\frac{1}{2}\left[K\log\frac{\mu_{q}(x)^{\top}\mu_{q}(x)+1^{ \top}\sigma_{q}^{2}(x)+2\beta}{K+2\alpha+2}-1^{\top}\log|\sigma_{q}^{2}(x)| \right]\] \[-\frac{K}{2}+\frac{1}{2}\frac{(K+2\alpha+2)(\mu_{q}(x)^{\top}\mu_ {q}(x)+1^{\top}\sigma_{q}^{2}(x))}{\mu_{q}(x)^{\top}\mu_{q}(x)+1^{\top}\sigma_{q} ^{2}(x)+2\beta}. \tag{16}\]

Following (Wu et al., 2019), as shown in appendix B.3, including the prior term in the regularizer reduces its complexity and this harms performance in practice. Hence, we use Eq (16) as the regularizer. We call this method "duq-eb".

### L2 regularization

Notice that all regularizations above are regularizing the variational predictive distribution \(q(z|x)\), not the parameters of the neural network. By analogy to MAP solutions of Bayesian neural networks, we can also regularize the weights directly by replacing the regularizers with \(\frac{\eta}{N}\|W\|_{2}\) in the objective, even though our posterior is on \(z\). We include this in our experiments and call it "duq-l2".

## 6 Experiments

In this section, we compare the empirical performance of DirectUQ with VI and hybrid methods related to stochastic gradient descent, as well as dropout (Gal and Ghahramani, 2016). VI candidates include the VI algorithm (Blundell et al., 2015) with fixed prior parameters (called vi-naive below), and other variations from collapsed variational inference (Tomczak et al., 2021) and empirical Bayes (Wu et al.,
2019). The hybrid methods include the non-Bayesian network optimized by stochastic gradient descent (marked as SGD below), as well as stochastic weight averaging (SWA, which uses the average of the SGD trajectory as the final weights) from Izmailov et al. (2018) and SWA-Gaussian (SWAG, which uses the SGD trajectory to form a Gaussian distribution over the neural network weight space) from Maddox et al. (2019). These methods are called hybrid methods because they do not have an explicit prior distribution over parameters but they can be interpreted as Bayesian methods. In addition to comparing the single-model performance of these methods, we compare their ensembles. The performance of each method is evaluated by the log loss on the test datasets and by two uncertainty quantification measures on shifted and out-of-distribution datasets.

### Experiment Details

We examine our methods on small datasets using a fully connected single layer neural network and on large image datasets using more complex neural network structures. DirectUQ is implemented such that \(W_{1}\) and \(W_{2}\) share all weights except the last layer. For the first set of experiments, we pick four regression datasets and four classification datasets from the UCI repository (Dua and Graff, 2017), and use a single-layer neural network with 50 hidden units. We split every dataset into 80%/10%/10% for training, validation and testing and perform standard normalization so that we do not need to tune hyperparameters for every dataset. For VI and DirectUQ, we select \(\eta\in\{0.1,0.3,0.5,0.7,0.9,1.0,3.0,5.0,7.5,10.0\}\) based on the validation log loss; for SGD, SWA and SWAG, we select weight decay in the same range as \(\eta\). Since the neural network is small, we do not include dropout in comparison. For the second set of experiments, we use four image datasets, CIFAR10, CIFAR100, SVHN1, STL10, together with two types of neural networks, AlexNet (Krizhevsky et al., 2012) and PreResNet20 (He et al., 2016), and fix \(\eta=0.1\) (except in "duq-l2" where we use \(\eta=5.0\)) and the dropout rate 0.1. Note that using \(\eta=1\) in VI yields poor predictions (see Appendix) and using \(\eta=0.1\) provides a stronger baseline. In all experiments, we use the Adam optimizer for VI and DirectUQ. For hybrid methods, we adapt the code from previous work (Izmailov et al., 2018; Maddox et al., 2019) and use the stochastic gradient descent optimizer for SGD, SWA and SWAG. Ensemble models use 5 runs with different random initializations and batch orders. Other details such as hyperparameters, learning rate and number of epochs are given in the Appendix.

Footnote 1: Dropout for SVHN on AlexNet fails and gives trivial uniform predictions, so we did not include it in the results.

### Run Time

Ignoring the data preprocessing time, we compare the run time of training 1 epoch of VI, DirectUQ and SGD. In Table

Figure 1: Uncertainty Quantification on AlexNet
1 we show the mean and standard deviation of 10 runs of these methods. Different regularizers have a small effect in terms of run time, so we only show that of vi-naive for VI and duq-mean for DirectUQ. As shown in Table 1, DirectUQ is much faster than VI and is slightly slower than SGD.

This is as expected because if we set the sample size to be \(M\), then SGD only needs 1 forward pass without sampling per batch, DirectUQ needs 1 forward pass and \(M\) samples of size \(K\) per batch, and VI needs \(M\) samples of the neural network size and \(M\) forward passes per batch. In addition, as the sample size \(M\) or the number of batches increases, the advantage of DirectUQ over VI expands as the increased forward passes will take more time, but the advantage of SGD over DirectUQ remains the same as sampling in a smaller space can be done in nearly constant time.

### Uncertainty Quantification for Shifted and Out-of-distribution Data

In this section we examine whether DirectUQ can capture the uncertainty in predictions for shifted and out-of-distribution (OOD) data. For uncertainty under data shift, STL10 and CIFAR10 can be treated as a shifted dataset for each other, as the figure size of STL10 is different from CIFAR10, and STL10 shares some classes with CIFAR10 so the labels are meaningful. Thus we can use expected calibration error (ECE, (Naeini et al., 2015)) to measure the uncertainty. That is, we separate data into bins of the same size according to the confidence level, calculate the difference between the accuracy and the averaged confidence in each bin and then average the differences among all bins, see Ovadia et al. (2019) for details. We select the number of bins to be 20. For uncertainty under OOD data, we choose the SVHN dataset as an OOD dataset for CIFAR10, as SVHN contains images of digits and the labels of SVHN are not meaningful in the context of CIFAR10. Thus, we cannot compute the accuracy. Instead, we use the entropy to measure the uncertainty. For the OOD data, we want our model to be as uncertain as possible and this implies high entropy in the predictive distribution. An alternative way to quantify the uncertainty is count-confidence plots, i.e., separate data into bins of the same confidence interval,

Figure 3: Test log loss of four image datasets on AlexNet. Each dot with error bar represents the mean and standard deviation of 5 independent runs and the triangle represents the ensemble of each method that aggregate these 5 runs for prediction (same for Figure 8). The standard deviations for some VI and DUQ methods are very small.

Figure 2: Test log loss on UCI regression datasets. Each dot with the error bar shows the mean and standard deviation of 5 independent runs. The standard deviations for some methods are very small.
and then plot the number of data points of each bin as a function of the averaged confidence score. We include the count-confidence plots in the Appendix, as it is not easy to compare multiple methods visually in this manner. Instead, we summarize each count-confidence plot as a single number using the averaged entropy over the entire dataset.

Figure 1(a) and 1(d) show the ECE of each method under data shift. In some cases single-model DirectUQ performs less well than VI. But ensembles of DirectUQ are better than VI, ensembles of SGD and Dropout, and they are even competitive with ensembles of VI. The remaining plots in Figure 1 show the entropy of predictive distribution of OOD data for all methods. Here single-model DirectUQ outperforms SGD and Dropout and ensembles of DirectUQ outperform VI and are even competitive with ensembles of VI. Results for PreResNet20 are similar and they are included in the Appendix (Figure 4).

### In-distribution Log Loss

Figure 2 and Figure 7 (in Appendix) compare VI, DirectUQ and hybrid methods on small datasets. We omit the naive version of DirectUQ as it produces poor results. For regression, duq-mean and duq-mv both perform well and they are better than VI and the hybrid methods in three of four datasets. For classification, duq-l2, i.e., DirectUQ with L2 regularization performs the best among all DirectUQ methods and it is comparable to all VI methods and is better than hybrid methods in three out of four datasets.

Figure 3 and Figure 8 (in Appendix) compare all methods on large datasets. First we observe that using collapsed variational inference in VI does not significantly improve the performance as shown in Tomczak et al. (2021). This is because we use \(\eta=0.1\) that yields better performance while Tomczak et al. (2021) use \(\eta=1\). Second, we observe that duq-mean is the best among all DirectUQ methods and it is worse than VI but better than SGD. SWAG is slightly better in 2 cases but it is significantly worse in other cases. Dropout is not stable across different datasets. Third, we observe that for all methods, ensembles perform better than single models. The ensemble of duq-mean is competitive with VI. Thus, ensembles of DirectUQ can help alleviate the limitation of the expressiveness of DirectUQ. Ensembles of dropout also perform well. Last but not least, we observe that the ensemble of VI performs the best although optimizing each model in the ensemble is slow. Ensembles of VI have not been well studied in the literature and our experiments highlight their potential strong performance.

Overall, ensembles of DirectUQ are slightly worse but competitive with single model VI and ensembles of Dropout in terms of in-distribution log loss. On the other hand they provide better uncertainty quantification for out of distribution data, which is competitive with ensembles of VI. Among variants, DirectUQ with learn-mean regularizer (duq-mean and ens-duq-mean) provide the best overall performance.

## 7 Conclusion

In Bayesian neural networks, the distribution of the last layer directly affects the predictive distribution. Motivated by this fact, we proposed direct uncertainty quantification, DirectUQ, that uses a neural network to directly learn the mean and variance of the last layer. We showed that DirectUQ can match the expressive power of VI in linear cases with a strong prior but that in general it provides a less expressive model. On the other hand the simplicity of the model enables fast training and facilitates convergence analysis through Rademacher bounds. In addition, DirectUQ can be derived as a non-standard variational lower bound, which provides an approximation for the last layer. This connection allowed us to derive better regularizations for DirectUQ by using collapsed variational inference over a hierarchical prior. Empirical evaluation highlighted the strong performance of ensembles of VI (when using weak regularization \(\eta=0.1\) instead of 1) albeit at the cost of high run time. Ensembles of DirectUQ are competitive with VI and other methods in terms of in-distribution loss, they outperform VI for out of distribution data, and are even competitive with ensembles of VI. Hence DirectUQ gives a new attractive approach for approximate inference in Bayesian models. The efficiency of DirectUQ also means faster test time predictions which can be important when deploying Bayesian models for real-time applications. In future work it would be interesting to explore the complexity performance tradeoff provided by DirectUQ, and the connections to variational inference in functional space that induces complex priors.

\begin{table}
\begin{tabular}{|c|c c c c|} \hline dataset & CIFAR10 & CIFAR100 & SVHN & STL10 \\ \hline size & 50000 & 50000 & 73257 & 500 \\ \hline VI & \(8.51\pm 0.41\) & \(8.27\pm 0.40\) & \(11.56\pm 0.39\) & \(1.75\pm 0.41\) \\ DirectUQ & \(2.18\pm 0.39\) & \(2.17\pm 0.43\) & \(2.72\pm 0.38\) & \(1.16\pm 0.40\) \\ SGD & \(1.97\pm 0.41\) & \(1.99\pm 0.43\) & \(2.46\pm 0.40\) & \(1.12\pm 0.38\) \\ \hline \end{tabular}
\end{table}
Table 1: Running time for training 1 epoch with batch size 512, AlexNet
## Acknowledgements

This work was partly supported by NSF under grant IIS-1906694. Some of the experiments in this paper were run on the Big Red computing system at Indiana University, supported in part by Lilly Endowment, Inc., through its support for the Indiana University Pervasive Technology Institute.
