Density Uncertainty Layers for Reliable Uncertainty Estimation

Yookoon Park David M. Blei

Department of Computer Science Columbia University New York, NY 10027, USA
<EMAIL_ADDRESS>Department of Computer Science, Statistics Columbia
University New York, NY 10027, USA<EMAIL_ADDRESS>
###### Abstract

Assessing the predictive uncertainty of deep neural networks is crucial for
safety-related applications of deep learning. Although Bayesian deep learning
offers a principled framework for estimating model uncertainty, the common
approaches that approximate the parameter posterior often fail to deliver
reliable estimates of predictive uncertainty. In this paper, we propose a
novel criterion for reliable predictive uncertainty: a model’s predictive
variance should be grounded in the empirical density of the input. That is,
the model should produce higher uncertainty for inputs that are improbable in
the training data and lower uncertainty for inputs that are more probable. To
operationalize this criterion, we develop the density uncertainty layer, a
stochastic neural network architecture that satisfies the density uncertain
criterion by design. We study density uncertainty layers on the UCI and
CIFAR-10/100 uncertainty benchmarks. Compared to existing approaches, density
uncertainty layers provide more reliable uncertainty estimates and robust out-
of-distribution detection performance.

##  1 Introduction

The success of deep learning models in a range of applications has spurred
significant interest in deploying the for real-world predictions. But in high-
stakes domains, such as healthcare, finance, and autonomous systems, both the
model’s prediction and its predictive uncertainty are crucial. A challenge is
that conventional neural networks lack a robust mechanism for quantifying
uncertainty, and they tend to produce overconfident predictions (Guo et al.,
2017; Ovadia et al., 2019). This paper is about how to produce reliable
estimates of predictive uncertainty with deep neural networks.

Why is this a problem? Bayesian deep learning offers a principled framework
for quantifying predictive uncertainty by incorporating uncertainty about the
model parameters (Graves, 2011; Welling and Teh, 2011; Neal, 2012; Blundell et
al., 2015). However, the common variational inference (VI) approaches that
approximate the parameter posterior in Bayesian deep learning (Graves, 2011;
Blundell et al., 2015; Kingma et al., 2015; Dusenberry et al., 2020) often
fall short of providing reliable predictive uncertainty (Foong et al., 2019b;
Ober and Rasmussen, 2019). For example, Figure 1a-c demonstrate that these
methods produce either collapsed or flat predictive uncertainty around the
origin, even though no training data were observed in that region.

In this work, we propose a novel criterion for reliable predictive uncertainty
and a new stochastic neural network architecture to satisfy the criterion. The
density uncertainty criterion posits that a model’s predictive uncertainty
should be grounded in the empirical density of the input. That is, we should
see higher uncertainty for improbable inputs and lower for more probable ones.
As motivation, we will show how Bayesian linear regression inherently adheres
to the criterion.

We then develop the density uncertainty layer, a stochastic neural network
architecture that is designed to satisfy the density uncertainty criterion.
The idea is to fit an energy-based model of input and then satisfy the density
uncertainty criterion via a constraint on the predictive variance produced by
the approximate posterior. Density uncertainty layers serve as a flexible
building block for density-aware deep neural networks. See Figure 1d for how
Density Uncertainty produces high predictive uncertainty in the low-density
region in the input space around the origin.

(a) MFVI

(b) Variational Dropout

(c) Rank 1 BNN

(d) Density Uncertainty

Figure 1: Predictive uncertainty of a two-layer MLP on a toy regression
problem. The red dots denote the training data and the blue shades mark the
95th and 99th percentiles of the predictive variance. All baselines other than
Density Uncertainty fail to account for the in-between uncertainty in the low-
density region around the origin

We study density uncertainty layers on the UCI and CIFAR-10/100 uncertainty
benchmarks. We compare the proposed method to popular uncertainty estimation
methods for neural networks (Blundell et al., 2015; Kingma et al., 2015; Gal
and Ghahramani, 2016; Dusenberry et al., 2020). We find that the proposed
density-aware neural networks provide more reliable predictive uncertainty
estimates and robust out-of-distribution (OOD) detection performance.

Contributions. In summary, this paper makes the following contributions.

  1. 1.

We propose the density uncertainty criterion for reliable predictive
uncertainty, asserting that a model’s predictive uncertainty should be
grounded in the training density of the input.

  2. 2.

We present the density uncertainty layer whose predictive uncertainty adheres
to the density uncertainty criterion and serves as a general building block
for uncertainty-aware neural networks.

  3. 3.

We study the proposed density uncertainty layers on the UCI and CIFAR-10/100
uncertainty estimation benchmarks. On both benchmarks, this method performs
better than existing approaches.

##  2 Density Uncertainty Layers

###  2.1 Motivation: Bayesian Linear Regression

We first illustrate how Bayesian uncertainty for linear regression is grounded
in the empirical density of the input. In the next section, this observation
will motivate a novel criterion for reliable predictive uncertainty.

Consider a Bayesian linear regression model with input
$\bm{\mathbf{X}}\in\mathbb{R}^{N\times D}$, target
$\bm{\mathbf{y}}\in\mathbb{R}^{N}$, and weight $w\in\mathbb{R}^{D}$:

$\displaystyle p(w)$ $\displaystyle=\mathcal{N}(w|0,\alpha^{-1}I),$ (1)
$\displaystyle p(\bm{\mathbf{y}}|\bm{\mathbf{X}},w)$
$\displaystyle=\prod_{i=1}^{N}\mathcal{N}(y_{i}|w^{T}x_{i},\beta^{-1}),$ (2)

The posterior distribution of the weight $w$ given the observations
$\mathcal{D}=\\{\bm{\mathbf{X}},\bm{\mathbf{y}}\\}$ is

$\displaystyle p(w|\mathcal{D})=\mathcal{N}(w|\mu,\Lambda^{-1}),$ (3)
$\displaystyle\text{ where
}\mu=\beta\Lambda^{-1}\bm{\mathbf{X}}^{T}\bm{\mathbf{y}}\text{ and
}\Lambda=\beta\bm{\mathbf{X}}^{T}\bm{\mathbf{X}}+\alpha I.$ (4)

The posterior predictive distribution for a test input $x_{*}$ is obtained by
marginalizing out the weight from the posterior joint

$\displaystyle p(y_{*}|x_{*},\mathcal{D})$ $\displaystyle=\int
p(y_{*},w|x_{*},\mathcal{D})dw$ (5)
$\displaystyle=\mathcal{N}(y_{*}|\mu^{T}x_{*},\beta^{-1}+x_{*}^{T}\Lambda^{-1}x_{*}).$
(6)

Equation 6 helps establish the connection between the Bayesian predictive
uncertainty for linear regression and the empirical input density. Rewriting
the predictive variance in Equation 6,

$\displaystyle\text{Var}[y_{*}|x_{*},\mathcal{D}]$
$\displaystyle=\frac{1}{\beta}+\frac{2}{\beta N}E(x_{*}),$ (7)
$\displaystyle\text{ where }E(x_{*})$
$\displaystyle=\frac{1}{2}x_{*}^{T}{\Sigma}^{-1}x_{*},$ (8)
$\displaystyle{\Sigma}$
$\displaystyle=\frac{1}{N}(\bm{\mathbf{X}}^{T}\bm{\mathbf{X}}+\frac{\alpha}{\beta}I).$
(9)

Here $E(x)$ is an energy function, an unnormalized negative log density of $x$
such that

$\displaystyle p_{\text{energy}}(x)\propto\exp(-E(x))$ (10)

Specifically, Equations 8 to 9 describe an instance of Gaussian energy model
$\mathcal{N}(0,{\Sigma})$, where ${\Sigma}$ is the MAP estimate of the input
covariance matrix given the observations $\bm{\mathbf{X}}$. This shows that
Bayesian linear model derives its predictive uncertainty from the Gaussian
energy model of the input that is fitted to the training data. Consequently,
predictive uncertainty will be high for test inputs that are improbable in the
training density and low for those that are more probable, providing intuitive
and reliable predictive uncertainty. In the following sections, we build upon
this principle.

###  2.2 Variational Inference for BNNs

A Bayesian Neural Network (BNN) $\Omega(x;\theta)$ treats its parameters
$\theta$ of neural networks as random variables from a prior $p(\theta)$.
Combined with a likelihood function $p(y|\Omega(x;\theta))$ and the
observations $\mathcal{D}=\\{(x_{1},y_{1}),...,(x_{N},y_{N})\\}$, the
posterior distribution of the parameters is defined by the Bayes’ rule:

$\displaystyle p(\theta|\mathcal{D})\propto
p(\theta)\prod_{i=1}^{N}p(y_{i}|\Omega(x_{i};\theta))$ (11)

However, the posterior is highly complex and intractable in BNNs. Variational
inference (VI) circumvents this problem by optimizing a tractable
approximation to the posterior. For example, the approximate distribution
$q(\theta)$ commonly takes the form of fully-factorized Gaussians whose
parameters are denoted by $\phi$. VI fits the approximate posterior by
maximizing the evidence lower-bound (ELBO) (Blundell et al., 2015):

$\displaystyle\mathcal{L}_{\phi}$
$\displaystyle=\mathbb{E}_{q(\theta)}\Big{[}\sum_{i=1}^{N}\log
p(y_{i}|\Omega(x_{i};\theta))\Big{]}-D(q(\theta)\|p(\theta)),$

where $D$ is the Kullback-Leibler divergence. At test time, the posterior
predictive distribution is approximated with $M$ Monte Carlo samples from
$q(\theta)$:

$\displaystyle
p(y_{*}|x_{*},\mathcal{D})\approx\frac{1}{M}\sum_{m=1}^{M}p(y_{*}|\Omega(x_{*};\theta_{m})),$
(12)

###  2.3 The Density Uncertainty Criterion

We have shown that Bayesian predictive uncertainty has desirable properties
for a linear model. However, VI-based BNNs often fails to provide reliable
uncertainty estimates (Foong et al., 2019a; Ober and Rasmussen, 2019) as
demonstrated in Figure 1 where the VI baselines fail to accurately capture the
density of the data in their predictive uncertainty.

To address this, we propose a novel uncertainty criterion that asserts model’s
predictive uncertainty should be grounded in the training density of the
input. To formalize the concept, we first introduce an auxiliary energy-based
generative model of the input:

$\displaystyle p_{\text{energy}}(x;\omega)\propto\exp(-E(x;\omega)),$ (13)

where the generative parameter $\omega$ is fitted to the training data
$x_{1},...,x_{N}$. We omit $\omega$ for simplicity hereafter.

We now establish the density uncertainty criterion.

######  Definition 1 (Density uncertainty criterion).

For a parameterized model $f(x;\theta)$ with a distribution on the parameters
$q(\theta)$ and an non-negative energy model $E(x)\geq 0$ fitted to the
training data, the model $f(x;\theta)$ adheres to the density uncertainty
criterion if

$\displaystyle\text{Var}_{q(\theta)}[f(x;\theta)]\propto E(x)\text{ for all
}x\in\mathcal{X}.$ (14)

The criterion posits that the predictive uncertainty should be proportional to
the energy of the input, producing high predictive uncertainty for inputs that
are improbable in the training distribution and low uncertainty for inputs
that are more probable. As illustrated in Section 2.1, Bayesian linear
regression inherently satisfies the density uncertainty criterion of Equation
14.

Incorporating the density uncertainty criterion into the ELBO objective yields

$\displaystyle\operatorname*{arg\,max}_{q\in\mathcal{Q}}\mathbb{E}_{q(\theta)}\Big{[}\sum_{i=1}^{N}\log
p(y_{i}|f(x_{i};\theta))\Big{]}-D(q(\theta)\|p(\theta))$ $\displaystyle\text{
s.t. }\text{Var}_{q(\theta)}[f(x;\theta)]\propto E(x)\text{ for all
}x\in\mathcal{X},$ (15)

where $D$ is the Kullback-Leibler divergence and $p$ is a prior. The
constraint in Equation 15 ensures the predictive uncertainty adheres to the
density uncertainty criterion, so that its predictive uncertainty is
consistently derived from the training density of the input, yielding reliable
uncertainty estimates.

Example: Bayesian linear regression Revisiting the Bayesian linear regression
example (Section 2.1), the family of parameter distribution $q(w)$ that
satisfies the density uncertainty criterion is

$\displaystyle q(w)$ $\displaystyle=\mathcal{N}(\mu,\gamma{\Sigma}^{-1}),$
(16) $\displaystyle\text{ where }{\Sigma}$
$\displaystyle=\frac{1}{N}(\bm{\mathbf{X}}^{T}\bm{\mathbf{X}}+\frac{\alpha}{\beta}I),$
(17)

where $\gamma$ is a scaling scalar and $\mu$ is a trainable parameter. The
posterior precision of the weight is now tied to the training covariance
estimate of the input. This results in the predictive variance of

$\displaystyle\text{Var}_{q(w)}[f(x_{*};w)]\propto
E(x_{*})=\frac{1}{2}x_{*}\widehat{\Sigma}^{-1}x_{*},$ (18)

satisfying the density uncertainty criterion for a Gaussian energy model $E$.
Notably, the constrained parameter posterior (Equation 16) recovers the true
posterior in Bayesian linear regression (Equation 3).

Reparametrized objective More generally, we introduce a reparametrized version
of the objective in Equation 15. Consider a stochastic function
$f(x,\epsilon;\phi)$ but now with a deterministic parameter $\phi$, exogenous
noise variable $\epsilon$, and a noise distribution $q(\epsilon)$:

$\displaystyle\operatorname*{arg\,max}_{\phi}\mathbb{E}_{q(\epsilon)}\Big{[}\sum_{i=1}^{N}\log
p(y_{i}|f(x_{i},\epsilon;\phi))\Big{]}-D(q(\epsilon)\|p(\epsilon))$
$\displaystyle\text{ s.t. }\text{Var}_{q(\epsilon)}[f(x,\epsilon;\phi)]\propto
E(x)\text{ for all }x\in\mathcal{X},$ (19)

A broad class of $q(\theta)$ such as normal distributions admits this kind of
reparametrization. The reparametrized objective offers more flexibility in how
we incorporate noise into the model while adhering to the density uncertainty
criterion. For example, we may choose to directly inject noise into the low-
dimensional function output instead of sampling the high-dimensional
parameters, thereby improving the computational efficiency and reducing the
gradient variance (e.g. Kingma et al. (2015)). Therefore, we default to this
form in the remainder of the paper.

###  2.4 The Density Uncertainty Layer

We now select an appropriate energy model $E(x)$ and the structure of a
stochastic function $\Omega(x,\epsilon;\phi)$ for deep neural networks. In
this work, we consider a popular form of residual network architecture

$\displaystyle a_{1}$ $\displaystyle=f_{1}(x,\epsilon_{1};\phi_{1}),$ (20)
$\displaystyle h_{\ell}$ $\displaystyle=\pi(a_{\ell}),$ (21) $\displaystyle
a_{\ell+1}$
$\displaystyle=a_{\ell}+f_{\ell}(h_{\ell},\epsilon_{\ell};\phi_{\ell}),\text{
for }l=1,...,L$ (22)

where $f_{l}$ is a stochastic linear layer with deterministic parameter
$\phi_{\ell}$ and exogenous noise variable $\epsilon_{\ell}$. $\pi$ is a
nonlinear activation function such as ReLU.

Section 2.1 shows that Bayesian linear regression derives its predictive
uncertainty from a Gaussian energy model. Motivated by this insight, we pair
each linear layer $f_{\ell}$ with a Gaussian energy model $E_{\ell}$ and make
the individual linear layers $f_{1},f_{2},...,f_{L}$ adhere to the density
uncertainty criterion layer-wise:

$\displaystyle\operatorname*{arg\,max}_{\phi}\mathbb{E}_{q}\Big{[}\sum_{i=1}^{N}\log
p(y_{i}|\Omega(x_{i},\bm{\mathbf{\epsilon}};\bm{\mathbf{\phi}}))\Big{]}-D(q(\bm{\mathbf{\epsilon}})\|p(\bm{\mathbf{\epsilon}}))$
$\displaystyle\text{ s.t.
}\text{Var}_{q(\epsilon_{\ell})}[f_{\ell}^{j}(h_{\ell},\epsilon_{\ell};\phi_{\ell})]\propto
E_{\ell}(h_{\ell})$ (23) $\displaystyle\text{ for all
}h_{\ell}\in\mathcal{H},\,j=1,...,D,\text{ and }\ell=1,...,L,$ (24)

where $j$ indexes the hidden units and $D$ is the width of the layers.
$\bm{\mathbf{\phi}}$ encompasses the deterministic weights of the neural
network and the posterior parameters of the noise distributions, and
$\bm{\mathbf{\epsilon}}$ involves all noise variables.

This design is based on several key observations: First, the Gaussian energy
facilitates efficient training and energy evaluation. Second, making the
individual layers stochastic fosters functional diversity and stochastic
regularization effect. Third, the complexity of the uncertainty landscape
grows naturally with the model complexity with size of in the neural network.
Lastly, a relaxed version of the density uncertainty criterion holds at the
network level which we describe next.

######  Proposition 1 (Density Uncertainty Criterion for Residual Networks).

Define the total energy
$E(x,h_{1},...,h_{\ell})=\sum_{\ell=1}^{L}E_{\ell}(h_{\ell})$. Assume the
followings:

  1. 1.

There exist $\alpha,\beta\in\mathbb{R}_{+}$ such that for all $j,\ell$

$\displaystyle\alpha
E_{\ell}(h_{\ell})\leq\text{Var}_{q(\epsilon_{\ell})}[f_{\ell}^{j}(h_{\ell},\epsilon_{\ell};\phi_{\ell})]\leq\beta
E_{\ell}(h_{\ell})$

  2. 2.

There exists $M\in\mathbb{R}_{+}$ such that $\|w_{\ell}^{j}\|_{2}^{2}\leq M$
for all $j,\ell$, where $w_{\ell}^{j}$ is the deterministic weight for $j$-th
hidden unit at layer $l$.

  3. 3.

$\phi$ is a 1-Lipschitz activation fucntion (e.g. ReLU).

Then the dimension-wise variance of the network output $a_{L+1}$ is bounded
from below by the expected total energy:

$\displaystyle\text{Var}_{q(\theta)}[a_{L+1}^{j}|x]\geq
C\,\mathbb{E}_{q(\theta)}[E(x,h_{1},...,h_{\ell})]$ (25)

for some constant $C$. Proof is in the appendix.

We now introduce the density uncertainty layer, a stochastic architecture that
by design satisfies the density uncertainty criterion layer-wise (Equation
23):

$\displaystyle
f_{\ell}^{j}(h_{\ell},\epsilon_{\ell};w_{\ell}^{j})=w_{\ell}^{j}\cdot
h_{\ell}+\epsilon_{\ell}^{j}\sqrt{E(h_{\ell})}+\eta_{\ell}^{j},$
$\displaystyle\text{ where
}E_{\ell}(h_{\ell})=\frac{1}{2}h_{\ell}^{T}\Sigma_{\ell}^{-1}h_{\ell},$ (26)
$\displaystyle
q(\epsilon_{\ell}^{j})=\mathcal{N}(0,\gamma_{\ell}^{j}),\,q(\eta_{\ell}^{j})=\mathcal{N}(0,\beta_{\ell}^{j}),$
(27)

where $\cdot$ denotes dot product, $w_{\ell}^{j}$ is the deterministic weight
for the $j$-th hidden unit at layer $\ell$, and $E_{\ell}(h_{\ell})$ is the
Gaussian energy model with the covariance parameter $\Sigma_{\ell}$. The layer
has two noise components $\epsilon_{\ell}^{j}$ and $\eta_{\ell}^{j}$, and the
predictive variance of

$\displaystyle\text{Var}_{q(\epsilon_{\ell})}[f_{\ell}^{j}(h_{\ell},\epsilon_{\ell};\phi_{\ell})]=\gamma_{\ell}^{j}E_{\ell}(h_{\ell})+\beta_{\ell}^{j},$
(28)

which satisfies the density uncertainty criterion as the bias term
$\beta_{\ell}^{j}$ can be always absorbed into the energy function. The
posterior noise variance parameters $\gamma_{\ell}^{j},\beta_{\ell}^{j}$ are
optimized using the ELBO, granting the model the flexibility to adjust the
noise scale for individual hidden units in the layer. While we assume Gaussian
noise for simplicity, we may potentially incorporate other noise distributions
such as heavy-tailed ones (Dusenberry et al., 2020) for enhanced robustness.

#### Gaussian energy

We adopt the LDL parametrization of the precision matrix
$\Sigma_{\ell}^{-1}=L_{\ell}D_{\ell}L_{\ell}^{T}$ for the layer-wise Gaussian
energy models, where $L_{\ell}$ is a lower unit triangular matrix, and
$D_{\ell}$ is a non-negative diagonal matrix. This admits efficient energy
evaluation

$\displaystyle
E_{\ell}(h_{\ell})=\frac{1}{2}h_{\ell}^{T}\Sigma_{\ell}^{-1}h_{\ell}=\frac{1}{2}\|D^{\frac{1}{2}}_{\ell}L_{\ell}^{T}h_{\ell}\|_{2}^{2},$

without computing the inverse of the covariance matrix and simplifies the log
determinant of the covariance matrix as $\log|\Sigma_{\ell}|=-\sum_{j}\log
D_{\ell}^{jj}$. For convolutional architectures, we replace the matrix-vector
product $L_{\ell}^{T}h_{\ell}$ with a masked convolution (Van den Oord et al.,
2016) to build a convolutional energy model. The main computational overhead
of density uncertainty layers comes from the matrix-vector product
$L_{\ell}^{T}h_{\ell}$ with $O(D^{2})$ complexity, same as regular linear
layers.

#### Mixture of rank-1 Gaussians

When the computational overhead is a major consideration, we propose to use a
mixture of rank-1 Gaussians as an alternative layer-wise energy model. More
specifically, the rank-1 Gaussian has the covariance matrix of the form:

$\displaystyle\Sigma=vv^{T}+D,$ (29)

where $v$ is the rank-1 factor and $D$ is a non-negative diagonal matrix.
Using the rank-1 construction, the inverse can be computed as

$\displaystyle\Sigma^{-1}=D^{-1}-\frac{D^{-1}vv^{T}D^{-1}}{1+v^{T}D^{-1}v},$
(30)

and the log determinant is also simplified as

$\displaystyle\log|\Sigma|=\log(1+v^{T}D^{-1}v)+\sum\log D_{jj},$ (31)

reducing the computational overhead to $O(D)$. However, as the rank-1
construction may be overly restrictive in practice, we choose a $K$ mixture of
rank-1 Gaussians as our layer-wise generative model, with the computational
complexity of $O(KD)$. We set $K$ as a fraction of the layer’s width and show
that it can be as small as $\approx 1\%$ with negligible performance drop.

Optimization The construction of the density uncertainty layer (Equation 27)
inherently satisfies the density uncertainty constraint, simplifying the
constrained optimization objective (Equation 23) to

$\displaystyle\mathcal{L}_{\bm{\mathbf{\phi}}}$
$\displaystyle=\mathbb{E}_{q(\bm{\mathbf{\epsilon}})}\Big{[}\sum_{i=1}^{N}\log
p(y_{i}|\Omega(x_{i},\bm{\mathbf{\epsilon}};\bm{\mathbf{\phi}}))\Big{]}+D(q(\bm{\mathbf{\epsilon}})\|p(\bm{\mathbf{\epsilon}}))$

We assume Gaussian priors with a shared variance for the noise variables.
Concurrently, the layer-wise energy models are optimized using the generative
objective:

$\displaystyle\mathcal{L}_{\bm{\mathbf{\omega}}}$
$\displaystyle=\sum_{i=1}^{N}\sum_{\ell=1}^{L}\mathbb{E}_{q(\bm{\mathbf{\epsilon}})}[\log
p_{\text{energy}}(h_{\ell};\omega_{\ell})]+\log p(\omega_{\ell}),$

where $p_{\text{energy}}(h_{\ell};\omega_{\ell})=\exp(-E(h_{\ell}))$ is the
Gaussian energy distribution for the input at layer $\ell$ and
$p(\omega_{\ell})$ is a prior on the generative parameter. The two objectives
are jointly optimized during training.

##  3 Related Work

This paper contributes to Bayesian deep learning and uncertainty estimation
for deep learning.

#### Bayesian Neural Networks and Uncertainty

Bayesian Neural Networks (BNNs) establish a principled framework for
estimating the uncertainty of neural networks by assuming their parameters are
latent variables that follow a prior distribution. Bayes’ rule, combined with
a likelihood function and observations, defines the posterior distribution of
the parameters. However, as exact posterior inference in BNNs is intractable,
the problem boils down to approximating the parameter posterior distribution.
For example, Markov Chain Monte Carlo (MCMC) (Welling and Teh, 2011; Chen et
al., 2014) simulates samples from the posterior distribution, using Langevin
(Welling and Teh, 2011) or Hamiltonian (Chen et al., 2014) dynamics. On the
other hand, the Laplace approximation (Ritter et al., 2018) applies a second-
order approximation at a mode of the posterior distribution.

Variational inference (VI) is a popular approach that reformulates inference
as an optimization problem. It seeks the best approximating distribution
within a distribution family that minimizes a discrepancy metric to the true
posterior, such as the Kullback-Leibler (KL) divergence. Graves (2011) adopt
fully-factorized Gaussian posteriors for the network’s parameters, and
Blundell et al. (2015) further incorporate the reparametrization trick (Kingma
and Welling, 2014; Rezende et al., 2014) to obtain unbiased, low-variance
gradient estimates with automatic differentiation. Louizos and Welling (2016)
enhance the expressiveness of the posteriors by utilizing matrix Gaussian
distributions for structured modeling of parameter correlations within each
layer. More recently, Ritter et al. (2021) propose sparse representations of
matrix Gaussian posteriors using inducing points (Snelson and Ghahramani,
2005; Titsias, 2009).

However, VI often fails to provide reliable uncertainty estimates for neural
networks in practice (Foong et al., 2019b; Ober and Rasmussen, 2019). This
failure might stem from the disconnect between Bayesian parameter uncertainty
and predictive uncertainty—while Bayesian methods focus on the posterior
parameter uncertainty, the practical interest often lies in estimating the
model’s predictive uncertainty. This gap may be further exacerbated by the the
restrictive independence assumptions (Trippe and Turner, 2017; Foong et al.,
2019a, 2020) and the mode-seeking behavior of the evidence lower-bound (ELBO)
(Bishop, 2006). Sun et al. (2019) tries to bridge the gap by performing
variational inference in the function space, albeit this requires additional
approximations to the intractable functional KL divergence. In contrast, the
density uncertainty criterion directly imposes a constraint on the model’s
predictive uncertainty, so that the predictive uncertainty is always grounded
in the training density of the input and provides reliable estimates of
predictive uncertainty.

#### Uncertainty Estimation for Deep Learning

This paper proposes a new methodology for estimating uncertainty in deep
learning models. Other alternatives for uncertainty estimation include Monte
Carlo dropout (Gal and Ghahramani, 2016), which interprets dropout
regularization as approximate Bayesian inference and estimates predictive
uncertainty by performing Monte Carlo integration using dropout at test time.
Variational Gaussian dropout Kingma et al. (2015), motivated by a continuous
approximation to dropout (Wang and Manning, 2013), applies multiplicative
noise to the preactivations. The authors show this corresponds to assuming a
degenerate parameter posterior distribution and propose a variational
inference method for adapting the dropout rates. Recently, Dusenberry et al.
(2020) propose to model only the uncertainty of rank-1 factors in the
network’s parameters by introducing dimension-wise multiplicative noise to
both the input and the output at each layer.

On the other hand, deep kernel learning methods (DKL) (Snoek et al., 2015;
Wilson et al., 2016; Liu et al., 2020; van Amersfoort et al., 2021) combine
the expressive power of neural networks with the uncertainty estimation
capability of Gaussian Processes (GPs) by using a deterministic neural network
as a feature extractor and applying the GP in the resulting feature space.
However, Ober et al. (2021) show that such construction is prone to
overfitting as it ignores any uncertainty associated with the neural network
feature extractor. To alleviate the problem, Liu et al. (2020); van Amersfoort
et al. (2021) limit the expressive power of the neural network feature
extractor using spectral normalization (Miyato et al., 2018) and apply
approximations to the GP posterior to circumvent the cubic complexity of GP
inference. While DKL methods provide interesting alternatives for uncertainty
estimation, in this work we primarily focus on the BNN and the related methods
(Blundell et al., 2015; Kingma et al., 2015; Gal and Ghahramani, 2016;
Dusenberry et al., 2020) that model the uncertainty of the neural network as a
whole.

##  4 Empirical Studies

Table 1: CIFAR-10 classification using the ResNet-14 architecture and 25 posterior samples. The average and the standard deviation over 3 random seeds are shown. Density Uncertainty significantly reduces the calibration error Method | Accuracy ($\uparrow$) | ECE ($\downarrow$) | NLL ($\downarrow$)
---|---|---|---
MFVI | 91.9 $\pm$ 0.1 | 0.080 $\pm$ 0.002 | 0.291 $\pm$ 0.004
MCDropout | 91.0 $\pm$ 0.1 | 0.016 $\pm$ 0.001 | 0.261 $\pm$ 0.001
VDropout | 92.2 $\pm$ 0.0 | 0.013 $\pm$ 0.001 | 0.233 $\pm$ 0.002
Rank-1 BNN | 92.2 $\pm$ 0.0 | 0.017 $\pm$ 0.002 | 0.231 $\pm$ 0.002
Density Uncertainty | 92.2 $\pm$ 0.3 | 0.004 $\pm$ 0.000 | 0.226 $\pm$ 0.003
Table 2: CIFAR-100 classification using the ResNet-14 architecture and 25 posterior samples. The average and the standard deviation over 3 random seeds are shown. Density Uncertainty improves all three metrics Method | Accuracy ($\uparrow$) | ECE ($\downarrow$) | NLL ($\downarrow$)
---|---|---|---
MFVI | 67.9 $\pm$ 0.4 | 0.132 $\pm$ 0.002 | 1.212 $\pm$ 0.017
MCDropout | 66.3 $\pm$ 0.3 | 0.066 $\pm$ 0.001 | 1.205 $\pm$ 0.005
VDropout | 67.8 $\pm$ 0.1 | 0.051 $\pm$ 0.002 | 1.167 $\pm$ 0.002
Rank-1 BNN | 68.0 $\pm$ 0.5 | 0.041 $\pm$ 0.003 | 1.132 $\pm$ 0.002
Density Uncertainty | 68.8 $\pm$ 0.2 | 0.011 $\pm$ 0.003 | 1.110 $\pm$ 0.006

In this section, we empirically demonstrate that the density uncertainty layer
delivers reliable predictive uncertainty estimates compared to the existing
approaches. The empirical studies are structured as follows:

  1. 1.

We visualize the predictive uncertainty landscape of different uncertainty
estimation methods and their failure modes, on a toy regression problem.

  2. 2.

We evaluate the uncertainty estimation performance on CIFAR-10/100
classification benchmarks (Krizhevsky and Hinton, 2009) using the ResNet-14
(He et al., 2016) and the Wide ResNet-28 (WRN-28) (Zagoruyko and Komodakis,
2016) architecture.

  3. 3.

We evaluate the out-of-distribution (OOD) detection performance on SVHN
(Netzer et al., 2011) using the models trained on CIFAR-10/100.

In the appendix, we include the results on the UCI regression benchmarks using
a MLP architecture.

We compare our method to the following popular uncertainty methods for deep
learning:

  1. 1.

Mean-field Variational Inference (MFVI) (Blundell et al., 2015) assumes fully
factorized normal posteriors on the neural network parameters.

  2. 2.

Monte Carlo Dropout (MCDropout) (Gal and Ghahramani, 2016) views dropout as
approximate Bayesian inference and applies dropout at test time in order to
estimate the predictive uncertainty.

  3. 3.

Variational Dropout (VDropout) (Kingma et al., 2015) applies Gaussian
multiplicative noise $\epsilon_{\ell}^{j}\sim\mathcal{N}(1,\alpha)$ to the
output of linear layers.

  4. 4.

Rank-1 BNN (Dusenberry et al., 2020) further extends Variational Dropout by
introducing multiplicative noise for both the layer’s input and output,
dimension-wise. We assume Gaussian noise.

In addition, we include the last-layer GP methods of SNGP (Liu et al., 2020)
and DUE (van Amersfoort et al., 2021) for the WRN-28 experiments. We use the
DUE authors’ implementation for SNGP and DUE.

Experimental details for CIFAR-10/100 For the CIFAR-10/100 experiments, we use
the standard convolutional ResNet-14 architecture (He et al., 2016) and also
the significantly larger WRN-28 architecture (Zagoruyko and Komodakis, 2016)
with $2\times$ depth and $10\times$ width. We default to the full-rank
Gaussian energy models using LDL parametrization but also include the
efficient rank-1 mixture Gaussians for the larger WRN-28 experiments. We use
the ADAM optimizer with learning rate $0.1$ with batch size of 128 except for
MFVI where the learning rate is reduced to $0.01$ as higher learning rate led
to divergence. We train the models for 200 epochs with cosine learning rate
schedule without restarts (Loshchilov and Hutter, 2017). During training, we
apply random cropping and padding, and horizontal flipping data augmentations.
The input pixel values are normalized using the training pixel means and
standard deviations, channel-wise. We do not employ KL annealing or posterior
tempering but initialize the posterior standard deviation to a sufficiently
small value (e.g. $10^{-3}$) to stabilize training (Dusenberry et al., 2020).
The weight decay is set to $0.0001$ for CIFAR-10 and $0.0002$ for CIFAR-100
experiments. For WRN-28, we use weight decay of $0.0005$, learning rate of
$0.05$, and batch size of $64$.

Model-specific hyperparameters are searched on a grid on a randomly sampled
held-out set of CIFAR-10. The dropout rate for Monte Carlo Dropout is searched
over $\\{0.1,0.2,0.3\\}$ and set to $0.1$. For Variational Dropout, we find
that adapting the noise variance leads to under-fitting and thus fix the noise
variance. The multiplicative noise variance is searched over
$\\{0.1,0.25,0.5\\}$ and set to $0.1$. For Rank-1 BNN, the prior standard
deviation for the multiplicative noise distributions is searched over
$\\{0.01,0.1,1\\}$ and set to $0.1$. For Density Uncertainty, the prior noise
standard deviation is searched over $\\{0.1,1,10\\}$ and set to $1$.

All experiments are implemented in PyTorch (Paszke et al., 2019) and executed
on a single Titan X GPU. The code is available at
https://github.com/yookoon/density_uncertainty_layers

###  4.1 Toy Regression

We generate a toy regression problem in 1-D to illustrate the predictive
uncertainty landscapes of different uncertainty estimation methods. We
uniformly sample $x_{i}$ from $[-4,-2]\cup[2,4]$ and generate the target as
$y_{i}=x_{i}^{3}+\epsilon$. This leaves a in-between low density region in
$[-2,2]$. We normalize the input and the target to have zero-mean and unit
variance. We use a MLP with one hidden layer of width 50 as the base
architecture.

The predictive uncertainty landscapes on a toy regression problem are
visualized in Figure 1. The baselines fail to produce reliable in-between
uncertainty in the low-density region around the origin. For example, MFVI
gives collapsed in-between uncertainty while Variational Dropout and Rank 1
BNNs produce flat uncertainty in the low-density region. Theses two methods
both apply input-independent noise to the layers and lack a robust mechanism
for adjusting their uncertainty depending on the empirical density of the
input. In contrast, Density Uncertainty captures the density of the training
data and produces higher predictive uncertainty in the low-density regions.
This is because it derives its uncertainty from the energy model of the input,
adhering to the density uncertainty criterion.

Table 3: CIFAR-10 classification using the WRN28 architecture and 25 posterior samples. The average and the standard deviation over 3 random seeds are shown. Density Uncertainty improves all three metrics Method | Accuracy ($\uparrow$) | ECE ($\downarrow$) | NLL ($\downarrow$)
---|---|---|---
MFVI | 89.3 $\pm$ 1.4 | 0.310 $\pm$ 0.011 | 0.668 $\pm$ 0.024
MCDropout | 96.0 $\pm$ 0.2 | 0.013 $\pm$ 0.001 | 0.135 $\pm$ 0.002
VDropout | 96.2 $\pm$ 0.2 | 0.012 $\pm$ 0.001 | 0.128 $\pm$ 0.002
Rank-1 BNN | 95.8 $\pm$ 0.1 | 0.011 $\pm$ 0.002 | 0.143 $\pm$ 0.004
SNGP | 95.9 $\pm$ 0.1 | 0.015 $\pm$ 0.001 | 0.143 $\pm$ 0.002
DUE | 95.6 $\pm$ 0.1 | 0.014 $\pm$ 0.006 | 0.170 $\pm$ 0.005
Density Uncertainty |  |  |
Full-rank LDL | 96.4 $\pm$ 0.1 | 0.010 $\pm$ 0.001 | 0.119 $\pm$ 0.003
Rank-1 Mixture ($1.25\%$) | 96.5 $\pm$ 0.0 | 0.011 $\pm$ 0.001 | 0.118 $\pm$ 0.001
Table 4: CIFAR-100 classification using the WRN28 architecture and 25 posterior samples. The average and the standard deviation over 3 random seeds are shown. Density Uncertainty improves all three metrics Method | Accuracy ($\uparrow$) | ECE ($\downarrow$) | NLL ($\downarrow$)
---|---|---|---
MFVI | 74.2 $\pm$ 0.8 | 0.260 $\pm$ 0.006 | 1.197 $\pm$ 0.023
MCDropout | 80.2 $\pm$ 0.3 | 0.033 $\pm$ 0.003 | 0.766 $\pm$ 0.005
VDropout | 81.1 $\pm$ 0.3 | 0.036 $\pm$ 0.002 | 0.756 $\pm$ 0.011
Rank-1 BNN | 79.7 $\pm$ 0.2 | 0.032 $\pm$ 0.001 | 0.815 $\pm$ 0.008
SNGP | 80.5 $\pm$ 0.3 | 0.035 $\pm$ 0.006 | 0.782 $\pm$ 0.017
DUE | Out of memory error
Density Uncertainty |  |  |
Full-rank LDL | 82.3 $\pm$ 0.2 | 0.029 $\pm$ 0.003 | 0.684 $\pm$ 0.005
Rank-1 Mixture ($1.25\%$) | 82.3 $\pm$ 0.0 | 0.032 $\pm$ 0.003 | 0.692 $\pm$ 0.001

###  4.2 CIFAR-10 and CIFAR-100

For the CIFAR-10/100 experiments, we report accuracy, Expected Calibration
Error (ECE) (Naeini et al., 2015) and Negative Log-Likelihood (NLL). ECE
assesses the quality of the model’s predictive uncertainty estimates by
measuring how well-calibrated the model’s predictions are. NLL is a proper
scoring rule (Gneiting and Raftery, 2007) which favors predictions that are
both accurate and well-calibrated.

The results using ResNet-14 are summarized in Table 1 and Table 2. On
CIFAR-10, Density Uncertainty significantly reduces ECE compared to the
baselines while maintaining comparable accuracy and slightly improving NLL.
This indicates that Density Uncertainty provides more precise uncertainty
estimates with lower calibration error. On CIFAR-100, Density Uncertainty
improves the performance on metrics. Notably, Density Uncertainty again
reduces ECE significantly. These results show that Density Uncertainty
delivers the most reliable uncertainty estimates among the baselines.

Table 3 and Table 4 present the results using the larger WRN-28 architecture,
including the additional last layer GP baselines of SNGP (Liu et al., 2020)
and DUE (van Amersfoort et al., 2021). For Density Uncertainty, we also
experiment with the efficient rank-1 mixture Gaussians energy model to test
the scalability of the method in addition to the default LDL parametrization.
We set the number of mixture components to only $1.25\%$ of the layers’ width,
adding only negligible amount of parameter and computational overhead to the
neural network. Density Uncertainty brings the best performance in all three
metrics. Surprisingly, the rank-1 mixture with only $1.25\%$ number of
components achieves performance comparable with the full-rank energy model,
demonstrating the effectiveness of low-rank approximations in neural networks
(Maddox et al., 2019; Dusenberry et al., 2020) and the scalability of the
method. Table 5 shows the results on CIFAR-10 classification using the WRN28
architecture with varying number of rank-1 mixture components. We find that
the performance is robust for a wide range of values, ranging from $1.25\%$ to
$20\%$, and is comparable to the full-rank LDL model.

Table 5: Varying the number of components for rank-1 Gaussian mixture energy on CIFAR-10 classification using the WRN28 architecture. The average and the standard deviation over 3 random seeds are shown. The performance of rank-1 Gaussian mixure is robust for a wide range of values Method | Accuracy ($\uparrow$) | ECE ($\downarrow$) | NLL ($\downarrow$)
---|---|---|---
Rank-1 Mixture ($1.25\%$) | 96.5 $\pm$ 0.0 | 0.011 $\pm$ 0.001 | 0.118 $\pm$ 0.001
Rank-1 Mixture ($2.5\%$) | 96.4 $\pm$ 0.1 | 0.012 $\pm$ 0.001 | 0.124 $\pm$ 0.002
Rank-1 Mixture ($5\%$) | 96.4 $\pm$ 0.1 | 0.012 $\pm$ 0.001 | 0.122 $\pm$ 0.001
Rank-1 Mixture ($10\%$) | 96.5$\pm$ 0.0 | 0.012 $\pm$ 0.001 | 0.122 $\pm$ 0.001
Rank-1 Mixture ($20\%$) | 96.5$\pm$ 0.0 | 0.013 $\pm$ 0.000 | 0.121 $\pm$ 0.003
Full-rank LDL | 96.4 $\pm$ 0.1 | 0.010 $\pm$ 0.001 | 0.119 $\pm$ 0.003
Table 6: Out-of-distribution detection performance on SVHN using ResNet-14 models trained on CIFAR-10/100. The average and the standard deviation over 3 random seeds are shown. Using the energy statistics, Density Uncertainty gives the most robust out-of-distribution detection performance | CIFAR-10 $\rightarrow$ SVHN | CIFAR-100 $\rightarrow$ SVHN
---|---|---
Method | AUPRC ($\uparrow$) | AUROC ($\uparrow$) | AUPRC ($\uparrow$) | AUROC ($\uparrow$)
MFVI | 0.903 $\pm$ 0.009 | 0.830 $\pm$ 0.014 | 0.803 $\pm$ 0.008 | 0.640 $\pm$ 0.017
MCDropout | 0.892 $\pm$ 0.006 | 0.832 $\pm$ 0.010 | 0.817 $\pm$ 0.005 | 0.666 $\pm$ 0.009
VDropout | 0.917 $\pm$ 0.011 | 0.866 $\pm$ 0.018 | 0.822 $\pm$ 0.028 | 0.677 $\pm$ 0.052
Rank-1 BNN | 0.925 $\pm$ 0.003 | 0.880 $\pm$ 0.001 | 0.822 $\pm$ 0.027 | 0.681 $\pm$ 0.047
Density Uncertainty | 0.952 $\pm$ 0.026 | 0.893 $\pm$ 0.056 | 0.908 $\pm$ 0.015 | 0.800 $\pm$ 0.024

###  4.3 Out-of-Distribution Detection on SVHN

Can we deploy the energy models of the density uncertainty layers for
detecting out-of-distribution (OOD) inputs? Table 6 summarizes the the OOD
detection performance on SVHN, using the models trained on CIFAR-10/100. We
report the area under the precision-recall curve (AUPRC) and the receiver
operator characteristic (AUROC). Following the previous work (Ritter et al.,
2021), the baselines use the maximum predicted probability heuristic for OOD
detection. On the other hand, Density Uncertainty, equipped with layer-wise
generative energy models, can inherently perform OOD detection using the
energy statistics. Specifically, the energy at layer $\ell$ can be represented
as the squared sum of random variables:

$\displaystyle E_{\ell}(h_{\ell})=\sum_{j=1}^{D}(z_{\ell}^{j})^{2}\text{ where
}z_{\ell}=\Sigma_{\ell}^{-\frac{1}{2}}h_{\ell},$ (32)

where $z_{\ell}$ can be thought of as whitened input with
$\Sigma_{\ell}^{-1/2}$ decorrelating the layer’s input $h_{\ell}$. Based on
the NN-GP equivalence in the infinite width limit, we expect the energy to be
approximately normally-distributed. Therefore, we use the squared deviation of
energy from the in-distribution average

$\displaystyle|E_{\ell}(h_{\ell})-\mu_{\ell}|^{2}$ (33)

as a test statistic for OOD detection where the in-distribution average energy
$\mu_{\ell}$ is estimated using a in-distribution held-out set. We use the
energy of the last convolutional layer as it reflects the most high-level,
semantic information of the input. Table 6 demonstrates that Density
Uncertainty can most robustly detect OOD input among the baselines.

##  5 Conclusion

We proposed a novel density criterion for reliable uncertainty estimation,
asserting that the predictive uncertainty of a model should be grounded in the
empirical density of the input. A model that adheres to the criterion will
produce higher uncertainty for inputs that are improbable in the training
data, and lower uncertainty for those inputs that are more probable. We
formalized the concept as a constraint on the predictive variance of a
stochastic function and developed the density uncertainty layer as a flexible
building block for uncertainty-aware deep learning. Through the empirical
studies, we demonstrated that the proposed method provides the most reliable
uncertainty estimates and robust out-of-distribution detection performance
among the baselines. This could have practical applications in various fields
where robust uncertainty estimation is crucial, such as medical diagnosis,
autonomous driving, and financial forecasting.

## References

  * Bishop (2006) Christopher M. Bishop.  _Pattern recognition and machine learning_.  Springer-Verlag New York, Inc., Secaucus, NJ, 2006.
  * Blundell et al. (2015) Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra.  Weight uncertainty in neural network.  In _ICML_ , 2015.
  * Chen et al. (2014) Tianqi Chen, Emily Fox, and Carlos Guestrin.  Stochastic gradient hamiltonian monte carlo.  In _ICML_ , 2014.
  * Dua and Graff (2017) Dheeru Dua and Casey Graff.  UCI machine learning repository, 2017.  URL http://archive.ics.uci.edu/ml.
  * Dusenberry et al. (2020) Michael Dusenberry, Ghassen Jerfel, Yeming Wen, Yian Ma, Jasper Snoek, Katherine Heller, Balaji Lakshminarayanan, and Dustin Tran.  Efficient and scalable bayesian neural nets with rank-1 factors.  In _ICML_ , 2020.
  * Foong et al. (2019a) Andrew YK Foong, David R Burt, Yingzhen Li, and Richard E Turner.  Pathologies of factorised gaussian and mc dropout posteriors in bayesian neural networks.  In _NeurIPS 4th Workshop on Bayesian Deep Learning_ , 2019a.
  * Foong et al. (2019b) Andrew YK Foong, Yingzhen Li, José Miguel Hernández-Lobato, and Richard E Turner.  ’in-between’uncertainty in bayesian neural networks.  In _ICML Uncertainty and Robustness in Deep Learning Workshop_ , 2019b.
  * Foong et al. (2020) Andrew YK Foong, David R Burt, Yingzhen Li, and Richard E Turner.  On the expressiveness of approximate inference in bayesian neural networks.  In _NeurIPS_ , 2020.
  * Gal and Ghahramani (2016) Yarin Gal and Zoubin Ghahramani.  Dropout as a bayesian approximation: Representing model uncertainty in deep learning.  In _ICML_ , 2016.
  * Gneiting and Raftery (2007) Tilmann Gneiting and Adrian E Raftery.  Strictly proper scoring rules, prediction, and estimation.  _Journal of the American statistical Association_ , 102(477):359–378, 2007.
  * Graves (2011) Alex Graves.  Practical variational inference for neural networks.  In _NeurIPS_ , 2011.
  * Guo et al. (2017) Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger.  On calibration of modern neural networks.  In _ICML_ , 2017.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _CVPR_ , 2016.
  * Kingma and Welling (2014) Diederik P. Kingma and Max Welling.  Auto-encoding variational bayes.  In _ICLR_ , 2014.
  * Kingma et al. (2015) Diederik P Kingma, Tim Salimans, and Max Welling.  Variational dropout and the local reparameterization trick.  In _NeurIPS_ , 2015.
  * Krizhevsky and Hinton (2009) Alex Krizhevsky and Geoffrey Hinton.  Learning multiple layers of features from tiny images.  Technical report, Computer Science Department, University of Toronto, 2009\.
  * Liu et al. (2020) Jeremiah Liu, Zi Lin, Shreyas Padhy, Dustin Tran, Tania Bedrax Weiss, and Balaji Lakshminarayanan.  Simple and principled uncertainty estimation with deterministic deep learning via distance awareness.  In _NeurIPS_ , 2020.
  * Loshchilov and Hutter (2017) Ilya Loshchilov and Frank Hutter.  Sgdr: Stochastic gradient descent with warm restarts.  In _ICLR_ , 2017.
  * Louizos and Welling (2016) Christos Louizos and Max Welling.  Structured and efficient variational deep learning with matrix gaussian posteriors.  In _ICML_ , 2016.
  * Maddox et al. (2019) Wesley J Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson.  A simple baseline for bayesian uncertainty in deep learning.  2019\.
  * Miyato et al. (2018) Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida.  Spectral normalization for generative adversarial networks.  2018\.
  * Naeini et al. (2015) Mahdi Pakdaman Naeini, Gregory Cooper, and Milos Hauskrecht.  Obtaining well calibrated probabilities using bayesian binning.  In _AAAI_ , 2015.
  * Neal (2012) Radford M Neal.  _Bayesian learning for neural networks_ , volume 118.  Springer Science & Business Media, 2012.
  * Netzer et al. (2011) Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng.  Reading digits in natural images with unsupervised feature learning.  In _NeurIPS Workshop on Deep Learning and Unsupervised Feature Learning_ , 2011.
  * Ober and Rasmussen (2019) Sebastian W Ober and Carl Edward Rasmussen.  Benchmarking the neural linear model for regression.  In _AABI_ , 2019.
  * Ober et al. (2021) Sebastian W Ober, Carl E Rasmussen, and Mark van der Wilk.  The promises and pitfalls of deep kernel learning.  In _UAI_ , 2021.
  * Ovadia et al. (2019) Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, David Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanan, and Jasper Snoek.  Can you trust your model’s uncertainty? evaluating predictive uncertainty under dataset shift.  In _NeurIPS_ , 2019.
  * Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.  Pytorch: An imperative style, high-performance deep learning library.  In _NeurIPS_ , 2019.
  * Rezende et al. (2014) Danilo J. Rezende, Shakir Mohamed, and Daan Wierstra.  Stochastic backpropagation and approximate inference in deep generative models.  In _ICML_ , 2014.
  * Ritter et al. (2018) Hippolyt Ritter, Aleksandar Botev, and David Barber.  A scalable laplace approximation for neural networks.  In _ICLR_ , 2018.
  * Ritter et al. (2021) Hippolyt Ritter, Martin Kukla, Cheng Zhang, and Yingzhen Li.  Sparse uncertainty representation in deep learning with inducing weights.  In _NeurIPS_ , 2021.
  * Snelson and Ghahramani (2005) Edward Snelson and Zoubin Ghahramani.  Sparse gaussian processes using pseudo-inputs.  In _NeurIPS_ , 2005.
  * Snoek et al. (2015) Jasper Snoek, Oren Rippel, Kevin Swersky, Ryan Kiros, Nadathur Satish, Narayanan Sundaram, Mostofa Patwary, Mr Prabhat, and Ryan Adams.  Scalable bayesian optimization using deep neural networks.  In _ICML_ , 2015.
  * Sun et al. (2019) Shengyang Sun, Guodong Zhang, Jiaxin Shi, and Roger Grosse.  Functional variational bayesian neural networks.  In _ICLR_ , 2019.
  * Titsias (2009) Michalis Titsias.  Variational learning of inducing variables in sparse gaussian processes.  In _AISTATS_ , 2009.
  * Trippe and Turner (2017) Brian Trippe and Richard Turner.  Overpruning in variational bayesian neural networks.  In _NeurIPS Approximate Bayesian Inference Workshop_ , 2017.
  * van Amersfoort et al. (2021) Joost van Amersfoort, Lewis Smith, Andrew Jesson, Oscar Key, and Yarin Gal.  On feature collapse and deep kernel learning for single forward pass uncertainty.  _arXiv preprint arXiv:2102.11409_ , 2021.
  * Van den Oord et al. (2016) Aaron Van den Oord, Nal Kalchbrenner, Lasse Espeholt, Oriol Vinyals, Alex Graves, et al.  Conditional image generation with pixelcnn decoders.  In _NeurIPS_ , 2016.
  * Wang and Manning (2013) Sida Wang and Christopher Manning.  Fast dropout training.  In _ICML_ , 2013.
  * Welling and Teh (2011) Max Welling and Yee W Teh.  Bayesian learning via stochastic gradient langevin dynamics.  In _ICML_ , 2011.
  * Wilson et al. (2016) Andrew Gordon Wilson, Zhiting Hu, Ruslan Salakhutdinov, and Eric P Xing.  Deep kernel learning.  In _AISTATS_ , 2016.
  * Zagoruyko and Komodakis (2016) Sergey Zagoruyko and Nikos Komodakis.  Wide residual networks.  _arXiv preprint arXiv:1605.07146_ , 2016.

Table 7: Test NLL on UCI Regression. Lower is better. The average and the standard deviation over the 20 random train-test splits are shown Dataset | MCDropout | VDropout | Rank1-BNN | Density Uncertainty
---|---|---|---|---
Boston Housing | 2.794 $\pm$ 0.141 | 2.815 $\pm$ 0.092 | 2.588 $\pm$ 0.236 |  2.523 $\pm$ 0.205
Concrete Strength | 3.533 $\pm$ 0.037 | 3.497 $\pm$ 0.050 | 3.112 $\pm$ 0.081 |  3.093 $\pm$ 0.116
Energy Efficiency | 2.604 $\pm$ 0.065 | 2.525 $\pm$ 0.057 | 2.044 $\pm$ 0.099 |  2.034 $\pm$ 0.087
Kin8nm | -0.590 $\pm$ 0.012 | -0.792 $\pm$ 0.015 | -1.194 $\pm$ 0.028 |  -1.234 $\pm$ 0.036
Naval Propulsion | -3.437 $\pm$ 0.048 | -4.019 $\pm$ 0.008 | -4.956 $\pm$ 0.033 |  -5.274 $\pm$ 0.036
Protein Structure | 2.935 $\pm$ 0.006 | 2.852 $\pm$ 0.005 |  2.771 $\pm$ 0.010 | 2.821 $\pm$ 0.001
Wine Quality Red | 0.952 $\pm$ 0.062 |  0.950 $\pm$ 0.067 | 0.954 $\pm$ 0.108 | 0.981 $\pm$ 0.109
Yacht Hydrodynamics | 3.205 $\pm$ 0.062 | 3.206 $\pm$ 0.071 | 2.712 $\pm$ 0.084 |  2.593 $\pm$ 0.067
Year Prediction MSD | 3.568 $\pm$ NA |  3.558 $\pm$ NA | 3.566 $\pm$ NA | 3.570 $\pm$ NA
Table 8: Test RMSE on UCI Regression. Lower is better. The average and the standard deviation over the 20 random train-test splits are shown Dataset | MCDropout | VDropout | Rank1-BNN | Density Uncertainty
---|---|---|---|---
Boston Housing | 3.715 $\pm$ 0.952 | 3.710 $\pm$ 0.729 | 3.212 $\pm$ 0.758 |  2.957 $\pm$ 0.606
Concrete Strength | 7.589 $\pm$ 0.630 | 7.154 $\pm$ 0.672 | 5.348 $\pm$ 0.552 |  5.290 $\pm$ 0.639
Energy Efficiency | 3.073 $\pm$ 0.360 | 3.209 $\pm$ 0.429 | 1.718 $\pm$ 0.246 |  1.690 $\pm$ 0.268
Kin8nm | 0.121 $\pm$ 0.003 | 0.094 $\pm$ 0.003 | 0.073 $\pm$ 0.003 |  0.070 $\pm$ 0.002
Naval Propulsion | 0.007 $\pm$ 0.000 | 0.003 $\pm$ 0.000 | 0.010 $\pm$ 0.000 |  0.001 $\pm$ 0.000
Protein Structure | 4.549 $\pm$ 0.033 | 4.123 $\pm$ 0.028 |  3.894 $\pm$ 0.046 | 4.077 $\pm$ 0.043
Wine Quality Red | 0.626 $\pm$ 0.046 | 0.626 $\pm$ 0.049 |  0.623 $\pm$ 0.051 | 0.630 $\pm$ 0.050
Yacht Hydrodynamics | 5.060 $\pm$ 1.328 | 4.820 $\pm$ 1.175 | 3.070 $\pm$ 0.082 |  2.505 $\pm$ 0.060
Year Prediction MSD | 8.726 $\pm$ NA |  8.700 $\pm$ NA | 8.712 $\pm$ NA | 8.710 $\pm$ NA
Table 9: The size and the dimensionality of the UCI regression datasets Dataset | Size | Dimension
---|---|---
Boston Housing | 506 | 13
Concrete Strength | 1,030 | 8
Energy Efficiency | 768 | 8
Kin8nm | 8,192 | 8
Naval Propulsion | 11,934 | 16
Protein Structure | 45,730 | 9
Wine Quality Red | 1,599 | 11
Yacht Hydrodynamics | 308 | 6
Year Prediction MSD | 515,345 | 90

##  Appendix A UCI Regression Results

We perform additional experiments on the UCI regression benchmarks [Dua and
Graff, 2017], one of the standard uncertainty estimation benchmarks for
Bayesian neural networks. We follow the experiment protocol of Gal and
Ghahramani [2016]. The benchmark includes regression datasets of varying size
ranging from 300 to 515K, and input dimensions ranging from 4 to 90, as
summarized in Table 9.

Experiment details We follow the experiment protocol of Gal and Ghahramani
[2016]. We report the negative log-likelihood (NLL) and the root mean squared
error (RMSE) on the test split using 10 posterior samples. For each dataset,
the results are averaged over 20 random train-test splits of the data (except
for Protein which uses 5 splits and Year which uses a single split). We
normalize the input using the mean and the standard deviation of the training
split. We use MLPs with two hidden layers of width 50. We increase the width
to 100 for the larger datasets of Protein and Year. All models are trained for
100 epochs with learning rate of 0.01 using momentum optimizer. We use batch
size of 128 and weight decay of 0.0001. We treat the output variance as a
trainable parameter. In order to minimize hyperparameter tuning, we use the
same hyperparameters used in the CIFAR-10/100 experiments except that we
initialize the posterior standard deviations to a higher value of 0.1 as lower
values led to overfitting.

Results Table 7 and Table 8 present the test NLL and the test RMSE on the
datasets. Overall, Density Uncertainty gives the best results, outperforming
the baselines on 6 of 9 datasets in both NLL and RMSE. These results
demonstrate that density uncertainty layers can also bring benefits for
regression problems potentially in low-data regimes besides the natural image
classification problems studied in the paper. However, a caveat is that the
variance of results are high on some datasets due to their small sizes and the
performance on these datasets can be sensitive to the hyperparameters.

##  Appendix B Proof of Proposition 1

For brevity, we omit the noise variable $\epsilon_{\ell}$ and the parameter
$\phi_{\ell}$ from the layer $f_{\ell}(h_{\ell};\epsilon_{\ell};\phi_{\ell})$.
First decompose the network output variance as

$\displaystyle\text{Var}[a_{L+1}^{j}|x]$ (34)
$\displaystyle=\text{Var}[a_{L}^{j}+f_{L}^{j}(h_{L})|x]$ (35)
$\displaystyle=\text{Var}[a_{L}^{j}|x]+\text{Var}[f_{L}^{j}(h_{L})|x]+2\text{Cov}[a_{L}^{j},f_{L}^{j}(h_{L})|x]$

The second term can be decomposed as

$\displaystyle\text{Var}[f_{L}^{j}(h_{L})|x]$ (36)
$\displaystyle=\mathbb{E}[\text{Var}[f_{L}^{j}(h_{L})|h_{L}]]+\text{Var}[\mathbb{E}[f_{L}^{j}(h_{L})|h_{L}]].$
(37)

Each of these components are bounded as

$\displaystyle\alpha\mathbb{E}[E_{L}(h_{L})]\leq\mathbb{E}[\text{Var}[f_{L}^{j}(h_{L})|h_{L}]]\leq\beta\mathbb{E}[E_{L}(h_{L})],$

by assumption, and

$\displaystyle 0\leq\text{Var}[\mathbb{E}[f_{L}^{j}(h_{L})|h_{L}]]\leq
MD\max_{k}\text{Var}[a_{L}^{k}|x]$ (38)

On the other hand, the covariance term is bounded as

$\displaystyle\Big{|}\text{Cov}[a_{L}^{j},f_{L}^{j}(h_{L})|x]\Big{|}\leq\sqrt{M}D\max_{k}\text{Var}[a_{L}^{k}|x]$
(39)

If there exist $C_{L},C_{L}^{\prime}$ such that for all $j$,

$\displaystyle
C_{L}\mathbb{E}[E(x_{1},...,h_{L-1})]\leq\text{Var}[a_{L}^{j}|x]$ (40)

and

$\displaystyle\text{Var}[a_{L}^{j}|x]\leq
C_{L}^{\prime}\,\mathbb{E}[E(x_{1},...,h_{L-1})].$ (41)

Combining the above, we have

$\displaystyle\text{Var}[a_{L+1}^{j}|x]\geq
C\,\mathbb{E}[E(x_{1},...,h_{L})],$ (42)

for some constant $C$ by induction.

##  Appendix C Bayesian Uncertainty in Linear Classification

In the paper, we demonstrated that the Bayesian uncertainty for linear
regression is grounded in the density estimate of input. But how about in
classification? We show that the Bayesian uncertainty in classification is
also based on a density estimate but with a weighted generative objective.

Consider a logistic classification problem, with the input
$\bm{\mathbf{X}}\in\mathbb{R}^{N\times D}$, the target
$\bm{\mathbf{y}}\in\mathbb{R}^{N}$, and the weight $w\in\mathbb{R}^{D}$:

$\displaystyle p(w)$ $\displaystyle=\mathcal{N}(w|0,\alpha^{-1}I),$ (43)
$\displaystyle p(\bm{\mathbf{y}}|\bm{\mathbf{X}},w)$
$\displaystyle=\prod_{i=1}^{N}y_{i}^{\sigma(w^{T}x_{i})}(1-y_{i})^{1-\sigma(w^{T}x_{i})},$
(44)

where $\sigma$ is the logistic sigmoid function: $\sigma(x)=1/(1+e^{-x})$.
Although the exact posterior distribution is intractable in this case, we can
obtain a Gaussian approximation using the Laplace’s approximation [Bishop,
2006]:

$\displaystyle q(w)=\mathcal{N}(w|\mu_{\text{MAP}},\Lambda^{-1}),$ (45)
$\displaystyle\Lambda=\sum_{i=1}^{N}\sigma(w^{T}x_{i})(1-\sigma(w^{T}x_{i}))x_{i}x_{i}^{T}+\alpha
I,$ (46)

and $\mu_{\text{MAP}}$ is the MAP estimate of the weight. Comparing the
posterior precision (Equation 46) to that of the regression in the paper
$\Lambda=\beta\sum_{i}x_{i}x_{i}^{T}+\alpha I$, we find that the precision
matrix in classification is a weighted estimate of the input covariance.
Noting that the weight $\sigma(w^{T}x_{i})(1-\sigma(w^{T}x_{i}))$ is higher
for inputs with more uncertain predictions (i.e., $\sigma(w^{T}x_{i})$ is
closer to $0.5$), Bayesian classification prioritizes inputs that give higher
prediction uncertainty, in contrast to the regression case where all inputs
were weighted uniformly. Recall that the posterior precision essentially
serves as a Gaussian density estimate of the input. Equation 46 shows that
Bayesian logistic classification performs weighted generative modeling of the
input, and by prioritizing inputs that are more informative, it can
potentially better utilize the capacity of the Gaussian energy model.

Despite this finding, in this work we choose to use the unweighted generative
objective for density uncertainty layers because (1) the weighted objective
can lead to an unfaithful estimate of the input density and (2) the capacity
of the generative model is not a major concern for density uncertainty layers,
as the complexity of the uncertainty landscape naturally grows with the number
of layers in the network.