1 Introduction↩︎

Bayesian neural networks (BNNs) induce epistemic uncertainty over predictions by placing a distribution over the weights [1]–[4]. However, it is challenging to infer the weight posterior due to the high dimensionality and multi-modality of this distribution [5], [6]. Alternative BNN methods have been introduced to avoid the complexity of weight-space inference, which include combining multiple maximum-a-posteriori (MAP) solutions [7], performing inference in the function-space [8], or performing inference in a lower dimensional latent space [9]–[12].

A recent approach to simplify BNNs is node stochasticity, which assigns latent noise variables to hidden nodes of the network [9], [12]–[16]. By restricting inference to the node-based latent variables, node stochasticity greatly reduces the dimension of the posterior, as the number of nodes is orders of magnitude smaller than the number of weights in a neural network [12]. Within this framework, multiplying each hidden node with its own random variable has been shown to produce great predictive performance, while having dramatically smaller computational complexity compared to weight-space BNNs [12]–[14], [16].

In this paper, we focus on node-based BNNs, which represent epistemic uncertainty by inferring the posterior distribution of the multiplicative latent node variables while learning a point-estimate of the weight posterior [12], [17]. We show that node stochasticity simulates a set of implicit corruptions in the data space during training, and by learning in the presence of such corruptions, node-based BNNs achieve natural robustness against some real-world input corruptions. This is an important property because one of the key promises of BNNs is robustness under covariate shift [6], [18], defined as a change in the distribution of input features at test time with respect to that of the training data. Based on our findings, we derive an entropy regularization approach to improve out-of-distribution generalization for node-based BNNs.

In summary, our contributions are:

1. We demonstrate that node stochasticity simulates data-space corruptions during training. We show that the diversity of these corruptions corresponds to the entropy of the latent node variables, and training on more diverse generated corruptions produce node-based BNNs that are robust against a wider range of corruptions.

2. We derive an entropy-regularized variational inference formulation for node-based BNNs.

3. We demonstrate excellent empirical results in predictive uncertainty estimation under covariate shift due to corruptions compared to strong baselines on large-scale image classification tasks.

4. We show that, as a side effect, our approach provides robust learning in the presence of noisy training labels.

Our code is available at https://github.com/AaltoPML/node-BNN-covariate-shift.

2 Background↩︎

2.0.0.1 Neural networks.

We define a standard neural network \(\mathbf{f}(\mathbf{x})\) with \(L\) layers for an input \(\mathbf{x}\) as follows: \[\begin{align} \mathbf{f}^0(\mathbf{x}) &= \mathbf{x}\tag{1}\\ \mathbf{h}^\ell(\mathbf{x}) &= \mathbf{W}^\ell \mathbf{f}^{\ell-1}(\mathbf{x}) + \mathbf{b}^\ell\\ \mathbf{f}^\ell(\mathbf{x}) &= \sigma^\ell \big( \mathbf{h}^\ell(\mathbf{x}) \big), \qquad \forall \ell=1,\dots,L \tag{2}\\ \mathbf{f}(\mathbf{x}) &= \mathbf{f}^L(\mathbf{x}), \tag{3} \end{align}\] where the parameters \(\theta=\{\mathbf{W}^\ell, \mathbf{b}^\ell\}_{\ell=1}^L\) consist of the weights and biases, and the \(\{\sigma^\ell\}_{\ell=1}^L\) are the activation functions. For the \(\ell\)-th layer, \(\mathbf{h}^\ell\) and \(\mathbf{f}^\ell\) are the pre- and post-activations, respectively.

2.0.0.2 Node-based Bayesian neural networks.

Probabilistic neural networks constructed using node stochasticity have been studied by [9], [12]–[17], [19]. We focus on inducing node stochasticity by multiplying each hidden node with its own random latent variables, and follow the framework of [12] for optimization. A node-based BNN \(\mathbf{f}_\mathcal{Z}(\mathbf{x})\) is defined as: \[\begin{align} \mathbf{f}^0_\mathcal{Z}(\mathbf{x}) &= \mathbf{x}\tag{4} \\ \mathbf{h}^\ell_\mathcal{Z}(\mathbf{x}) &= (\mathbf{W}^\ell (\mathbf{f}^{\ell-1}_\mathcal{Z}(\mathbf{x}) \circ \mathbf{z}^\ell ) + \mathbf{b}^\ell) \circ \mathbf{s}^\ell \tag{5} \\ \mathbf{f}_\mathcal{Z}^\ell(\mathbf{x}) &= \sigma^\ell \big ( \mathbf{h}_\mathcal{Z}^\ell(\mathbf{x}) \big ), \qquad \forall \ell=1,\dots,L \tag{6} \\ \mathbf{f}_\mathcal{Z}(\mathbf{x}) &= \mathbf{f}^L_\mathcal{Z}(\mathbf{x}), \tag{7} \end{align}\] where \(\mathbf{z}^\ell\) and \(\mathbf{s}^\ell\) are the multiplicative latent random variables of the incoming and outgoing signal of the nodes of the \(\ell\)-th layer, and \(\circ\) denotes the Hadamard (element-wise) product. We collect all latent variables to \(\mathcal{Z}=\{\mathbf{z}^\ell, \mathbf{s}^\ell\}_{\ell=1}^L\).¹

To learn the network parameters, we follow [12] and perform variational inference [20] over the weight parameters \(\theta\) and latent node variables \(\mathcal{Z}\). We begin by defining a prior \(p(\theta, \mathcal{Z}) = p(\theta)p(\mathcal{Z})\). We set a variational posterior approximation \(q_{\hat{\theta},\phi}(\theta,\mathcal{Z})=q_{\hat{\theta}}(\theta)q_\phi(\mathcal{Z})\), where \(q_{\hat{\theta}}(\theta)=\delta(\theta-\hat{\theta})\) is a Dirac delta distribution and \(q_\phi(\mathcal{Z})\) is a Gaussian or a mixture of Gaussians distribution. We infer the posterior by minimizing the Kullback-Leibler (KL) divergence between variational approximation \(q\) and true posterior \(p(\theta, \mathcal{Z}|\mathcal{D})\). This is equivalent to maximizing the evidence lower bound (ELBO): \[\begin{gather} \mathcal{L}(\hat{\theta}, \phi) = \mathbb{E}_{q_\phi(\mathcal{Z})} \Big[ \log p(\mathcal{D}| \hat{\theta},\mathcal{Z}) \Big] \\ - \mathrm{KL}\Big[q_\phi(\mathcal{Z}) \, || \, p(\mathcal{Z})\Big] + \log p(\hat{\theta}). \label{eq:elbo} \end{gather}\tag{8}\] In essence, we find a MAP solution for the more numerous weights \(\theta\), while inferring the posterior distribution of the latent variables \(\mathcal{Z}\). We refer the reader to 8 for detailed derivations.

2.0.0.3 Neural networks under covariate shift.

In this paper, we focus on covariate shift from input corruptions, following the setting of [21]. To simulate covariate shift, one can take an input \(\mathbf{x}\) assumed to come from the same distribution as the training samples and apply an input corruption \(\mathbf{g}^0\) to form a shifted version \(\mathbf{x}^c\) of \(\mathbf{x}\): \[\mathbf{x}^c = \mathbf{x}+ \mathbf{g}^0(\mathbf{x}).\] For instance, \(\mathbf{x}\) could be an image and \(\mathbf{g}^0\) can represent the shot noise corruption [21]. The input corruption \(\mathbf{g}^0(\mathbf{x})\) creates a shift in the output of each layer \(\mathbf{g}^\ell(\mathbf{x})\) (see 1). We can approximate these shifts by first-order Taylor expansion (see 10 for full derivation), \[\begin{align} \tag{9} \overbrace{\mathbf{g}^\ell(\mathbf{x})}^{\text{shift}} &= \overbrace{\mathbf{f}^\ell(\mathbf{x}^c)}^{\text{corrupted output}} - \overbrace{\mathbf{f}^\ell(\mathbf{x})}^{\text{clean output}} \\ &\approx \mathbf{J}_\sigma \big[ \mathbf{h}^\ell(\mathbf{x}) \big] \big(\mathbf{W}^{\ell}\mathbf{g}^{\ell-1}(\mathbf{x})\big), \tag{10} \end{align}\] where \(\mathbf{J}_{\sigma^\ell} = \partial \sigma^{\ell}/\partial \mathbf{h}^\ell\) denotes the (diagonal) Jacobian of the activation \(\sigma^\ell\) with respect to \(\mathbf{h}^\ell\). While \(\mathbf{g}^0\) causes activation shifts in every layer of the network, we focus on the shift in the final output layer \(\mathbf{g}^L\). The approximation in 10 shows that this shift depends on the input \(\mathbf{x}\), the network’s architecture (e.g., choice of activation functions) and parameters \(\theta\). We measure the robustness of a network with respect to a corruption \(\mathbf{g}^0(\cdot)\) on the dataset \(\mathcal{D}=\{\mathbf{x}_n, \mathbf{y}_n\}_{n=1}^N\) by the induced mean square shift, \[\begin{align} \mathrm{MSS}_g &= \frac{1}{N} \sum_{n=1}^N ||\mathbf{g}^L(\mathbf{x}_n)||^2_2, \label{eq:AS95g} \end{align}\tag{11}\] where \(\mathrm{MSS}_g\) is the average shift on the data. Ideally, we want \(\mathrm{MSS}_g\) to be small for the network to still provide nearly correct predictions given corrupted inputs. When the training data and the architecture are fixed, \(\mathrm{MSS}_g\) depends on the parameters \(\theta\). A direct approach to find \(\theta\) minimizing \(\mathrm{MSS}_g\) is to apply the input corruption \(\mathbf{g}^0\) to each input \(\mathbf{x}_n\) during training to teach the network to output the correct label \(\mathbf{y}_n\) given \(\mathbf{g}^0(\mathbf{x}_n)\). However, this approach is not domain-agnostic and requires defining a list of corruptions beforehand. In the next sections, we discuss the usage of multiplicative latent node variables as an implicit way to simulate covariate shifts during training.

3 Characterizing implicit corruptions↩︎

In this section, we demonstrate that multiplicative node variables correspond to implicit input corruptions. We show how to extract and visualize these new corruptions.

3.1 Relating input corruptions and multiplicative nodes↩︎

The node-based BNN of induces the predictive posterior \(p(\mathbf{f}^\ell_\mathcal{Z}(\mathbf{x}))\) over the \(\ell\)-th layer outputs by marginalizing over the variational latent parameter posterior \(q(\mathcal{Z}_{\leq \ell})\). Optimization of the variational objective in 8 enforces the model to achieve low loss on the training data despite each layer output being corrupted by noise from \(q(\mathcal{Z})\), represented by the expected log likelihood term of the ELBO. Let \(\hat{\mathbf{f}}^\ell(\mathbf{x})\) denote the mean predictive posterior, \[\begin{align} \hat{\mathbf{f}}^\ell(\mathbf{x}) &= \mathbb{E}_{q(\mathcal{Z})} \big[\mathbf{f}_\mathcal{Z}^\ell(\mathbf{x})\big], \quad \forall \ell=1,\ldots,L, \end{align}\] and where we denote the final output \(\hat{\mathbf{f}}(\mathbf{x}) = \hat{\mathbf{f}}^L(\mathbf{x})\). If the shifted output \(\hat{\mathbf{f}}^\ell(\mathbf{x}+\mathbf{g}^0(\mathbf{x})) = \hat{\mathbf{f}}^\ell(\mathbf{x}) + \mathbf{g}^\ell(\mathbf{x})\) caused by corrupting a training sample \(\mathbf{x}\) using \(\mathbf{g}^0\) lies within the predictive distribution of \(\mathbf{f}^\ell_\mathcal{Z}(\mathbf{x})\) (blue dot in 1), then the model can map this corrupted version of \(\mathbf{x}\) to its correct label. This implies robustness against the space of implicit corruptions generated by \(q(\mathcal{Z})\), which indirectly leads to robustness against real corruptions. However, standard variational inference will converge to a posterior whose entropy is calibrated for the variability in the training data, but does not necessarily account for corruptions caused by covariate shifts. Thus, the posterior might cover the corruption with low severity \(\mathbf{g}_0\) (blue dot in 1), but not the one with higher severity \(\mathbf{g}_1\) (green dot in 1). To promote predictive distributions that are more robust to perturbations, we propose to increase the entropy of \(p(\mathbf{f}_\mathcal{Z}^\ell(\mathbf{x}))\) by increasing the entropy of the variational posterior \(q(\mathcal{Z})\).

3.1.0.1 Empirical demonstration.

To illustrate our intuition, we present an example with two node-based BNNs, one with high entropy and one with lower entropy. We use the all-cnn-c architecture of [22] and cifar10 [23]. We initialize the standard deviations of \(q(\mathcal{Z})\) for the low-entropy model using the half-normal \(\mathcal{N}^+(0.16, 0.02)\), while we use \(\mathcal{N}^+(0.32, 0.02)\) for the high-entropy model. For brevity, we refer to the former model as \(\mathcal{M}_{16}\) and the latter model as \(\mathcal{M}_{32}\). In the left plot of 3, we show that, after training, \(\mathcal{M}_{32}\) retains higher variational posterior entropy than \(\mathcal{M}_{16}\) due to having higher initial standard deviations for \(q(\mathcal{Z})\).² We use principal component analysis (PCA) to visualize the samples from the output distribution \(p(\mathbf{f}^\ell_\mathcal{Z}(\mathbf{x}))\) of the \(\ell\)-th layer with respect to one input image \(\mathbf{x}\), as well as the output \(\{\hat{\mathbf{f}}^\ell(\mathbf{x}+\mathbf{g}_i(\mathbf{x}))\}_{i=1}^{95}\) under the real image corruptions \(\{\mathbf{g}_i\}_{i=1}^{95}\) from [21]. There are 19 corruption types with 5 levels of severity, totalling 95 corruption functions. 2 shows the activations of the last layer, projected into a two-dimensional subspace with PCA for visualization. From this figure, we can see that there is more overlap between samples from \(p(\mathbf{f}_\mathcal{Z}^\ell(\mathbf{x}))\) and the shifted outputs \(\{\hat{\mathbf{f}}^\ell(\mathbf{x}+\mathbf{g}_i(\mathbf{x}))\}_{i=1}^{95}\) for \(\mathcal{M}_{32}\) in 2 (d) than for \(\mathcal{M}_{16}\) in 2 (c). This indicates that during training the posterior of \(\mathcal{M}_{32}\) is able to simulate a larger number of implicit corruptions bearing resemblance to the real-world corruptions than the posterior of \(\mathcal{M}_{16}\), leading to better negative log-likelihood (NLL) accross all level of corruptions as well as on the clean test set in 3. This example supports our intuition that increasing the entropy of the latent variables \(\mathcal{Z}\) allows them to simulate more diverse implicit corruptions, thereby boosting the model’s robustness against a wider range of input corruptions.

3.1.0.2 Why latent variables at every layer?

In principle, we could have introduced latent variables only to the first layer of the network, as the shift simulated in the first layer will propagate to subsequent layers. However, modern NNs contain asymmetric activation functions such as ReLU or Softplus, which can attenuate the signal of the shift in the later layers. Thus, the latent variables in every layer (after the first one) maintain the strength of the shift throughout the network during the forward pass. Moreover, by using latent variables at every layer – as opposed to only the first layer – we can simulate a more diverse set of input corruptions, since we can map each sample \(\mathcal{Z}\) from \(q(\mathcal{Z})\) to an input corruption as shown in the following section.

3.2 Visualizing the implicit corruptions↩︎

Next, we show how to find the explicit image corruptions that correspond to the stochasticity of the predictive posterior. Let \(\mathcal{Z}\) be a sample drawn from \(q(\mathcal{Z})\). If we assume that \(\mathcal{Z}\) corresponds to an input corruption \(\mathbf{g}(\mathbf{x})\): \[\mathbf{f}_{\mathcal{Z}} (\mathbf{x}) = \hat{\mathbf{f}}(\mathbf{x}+ \mathbf{g}(\mathbf{x}))\] then we can approximately solve for \(\mathbf{g}(\mathbf{x}) = \mathbf{x}^c - \mathbf{x}\) by finding \(\mathbf{x}^c\) that minimizes \[\label{eq:xc95loss} \mathcal{L}(\mathbf{x}^c) = \frac{1}{2}\left|\left|\mathbf{f}_{\mathcal{Z}}(\mathbf{x})-\hat{\mathbf{f}}(\mathbf{x}^c)\right|\right|^2_2 + \frac{\lambda}{2}\big|\big| \mathbf{g}(\mathbf{x}) \big|\big|_2^2\tag{12}\] using gradient descent. The second term with a coefficent \(\lambda \geq 0\) regularizes the norm of \(\mathbf{g}(\mathbf{x})\). This approach is similar to the method of finding adversarial examples of [24]. 4 visualizes the corruptions generated by \(\mathcal{M}_{32}\) on a test image of cifar10 under different \(\lambda\). We can see that \(\lambda\) controls the severity of the corruptions, with smaller \(\lambda\) corresponding to higher severity.

3.2.0.1 Is a model robust against its own corruptions?

We use both models \(\mathcal{M}_{16}\) and \(\mathcal{M}_{32}\) to generate corruptions on a subset of \(1024\) test images of cifar10. We generate \(8\) corruptions per test image. The left plot of 5 shows that each model is robust against its own implicit corruptions even when the corruption is severe, as evidenced by the small performance degradation under different \(\lambda\). By comparing the right plot to the left plot, we can see that each model is less resistant to the corruptions generated by the other model than its own corruptions. Crucially, however, the performance of \(\mathcal{M}_{32}\) under the corruptions generated by \(\mathcal{M}_{16}\) is better than the reverse. This example thus suggests that while each model is resistant to its own corruptions, the model with higher entropy shows better robustness against the corruptions created by the other model.³

4 Maximizing variational entropy↩︎

The previous sections motivated the usage of variational posteriors with high entropy from the perspective of simulating a diverse set of input corruptions. In this section, we discuss a simple method to increase the variational entropy.

4.1 The augmented ELBO↩︎

Our goal is to find posterior approximations that have high entropy. In the previous section, we considered a heuristic approach of initializing \(q(\mathcal{Z})\) with high entropy (3). However, if the initial entropy of \(q(\mathcal{Z})\) is too high, training will converge slowly due to high variance in the gradients.

Here we consider the approach of augmenting the original ELBO in 8 with an extra \(\gamma\)-weighted entropy term, adapting [25]. The augmented \(\gamma\)-ELBO is \[\begin{align} &\mathcal{L}_\gamma(\hat{\theta}, \phi) = \mathcal{L}(\hat{\theta}, \phi) + \gamma \mathbb{H}\big[q_\phi(\mathcal{Z})\big]\\ &= \underbrace{\mathbb{E}_{q_{\phi}(\mathcal{Z})}\Big[\log p(\mathcal{D}| \hat{\theta},\mathcal{Z})\Big]}_{\text{expected log-likelihood}} - \underbrace{\mathbb{H}\big[q_{\phi}(\mathcal{Z}), p(\mathcal{Z})\big]}_{\text{cross-entropy}} \\ & \qquad \qquad \quad+ \underbrace{(\gamma+1) \mathbb{H}\big[q_\phi(\mathcal{Z})\big]}_{\text{variational entropy}} + \underbrace{\log p(\hat{\theta})}_{\text{weight prior}}, \label{eq:entropy95elbo} \end{align}\tag{13}\] where we decompose the KL into its cross-entropy and entropy terms. \(\gamma \ge 0\) controls the amount of extra entropy, with \(\gamma=0\) reducing to the classic ELBO in 8 . We can interpret the terms in 13 as follows: the first term fits the variational parameters to the dataset; the second and fourth terms regularize \(\phi\) and \(\hat{\theta}\) respectively; the third term increases the entropy of the variational posterior.

4.2 Tempered posterior inference↩︎

One could also arrive at 13 by minimizing the KL divergence between the approximate posterior \(q_{\phi}(\hat{\theta},\mathcal{Z})\) and the tempered posterior \(p_\gamma(\theta, \mathcal{Z}|\mathcal{D})\) [25]: \[\begin{align} \label{eq:tempered95posterior} p_\gamma(\theta, \mathcal{Z}|\mathcal{D}) &= \frac{p(\mathcal{D}|\theta, \mathcal{Z})^{\tau}p(\mathcal{Z},\theta)^{\tau}}{p_\gamma(\mathcal{D})} \\ p_\gamma(\mathcal{D}) &= \int_{\theta}\int_{\mathcal{Z}} p(\mathcal{D}|\theta, \mathcal{Z})^{\tau}p(\mathcal{Z},\theta)^{\tau} d\mathcal{Z}d\theta, \end{align}\tag{14}\] where the temperature \(\tau = 1/(\gamma+1)\). The tempered posterior variational approximation \[\begin{align} \mathop{\mathrm{arg\,min}}_{\hat{\theta},\phi} \:\: \frac{1}{\tau} \, \mathrm{KL}\left[ q_{\phi}(\hat{\theta}, \mathcal{Z}) \, \left|\right| \, p_\gamma(\theta, \mathcal{Z}|\mathcal{D}) \right] \end{align}\] is equivalent to tempered ELBO maximization \[\begin{align} \mathop{\mathrm{arg\,max}}_{\hat{\theta}, \phi} \: \mathcal{L}_\gamma(\hat{\theta}, \phi) - \log p_\gamma(\mathcal{D})^{\frac{1}{\tau}}. \end{align}\] We refer the reader to 9 for detailed derivations. The entropy-regularized \(\gamma\)-ELBO thus corresponds to the family of tempered variational inference, and with positive \(\gamma > 0\), to ‘hot’ posteriors [5]. In the next section, we will demonstrate empirically the benefits of such hot posteriors in node-based BNNs.

5 Experiments↩︎

In this section, we present experimental results of node-based BNNs on image classification tasks. For the datasets, we use cifar [23] and tinyimagenet [26], which have corrupted versions of the test set provided by [21]. We use vgg16 [27], resnet18 [28] and preactresnet18 [29] for the architectures. We test three structures of latent variables: in, where we only use the input latent variables \(\{\mathbf{z}^\ell\}_{\ell=1}^L\); out, where we only use the output latent variables \(\{\mathbf{s}^\ell\}_{\ell=1}^L\); and both, where we use both \(\{\mathbf{z}^\ell\}_{\ell=1}^L\) and \(\{\mathbf{s}^\ell\}_{\ell=1}^L\). We use \(K \in \{1, 2, 4\}\) Gaussian component(s) in the variational posterior. For each result, we report the mean and standard deviation over multiple runs.

5.1 Effects of \(\gamma\) on covariate shift↩︎

In this section, we study the changes in performance of the model trained with the \(\gamma\)-ELBO objective as we increase \(\gamma\). We perform experiments with vgg16 on cifar100, and use the corrupted test set of cifar100 provided by [21]. In 6, we show the out model’s behaviour under a different number of Gaussian components \(K\). In 7, we show the results of a model with \(K=4\) components under the different latent variable structures in, out, and both.

These figures show that performance across different test sets improves as \(\gamma\) increases up until a threshold and then degrades afterward. The optimal \(\gamma\) for each set of test images correlates with the severity of the corruptions, where more severe corruptions can be handled by enforcing more diverse set of implicit corruptions during training. However, learning on a more diverse implicit corruptions requires higher capacity, and reduces the learning capacity needed to obtain good performance on the in-distribution (ID) data. The entropy coefficient \(\gamma\) thus controls the induced trade-off between ID performance and out-of-distribution (OOD) robustness.

6 shows that for ID data, the optimal performance of the model (at optimal \(\gamma\)) remains similar under different \(K\). On OOD data, however, higher \(K\) consistently produces better results as \(\gamma\) varies. The optimal \(\gamma\) is higher for variational distributions with more components. This finding is likely because with more mixture components, the variational posterior can approximate the true posterior more accurately, and thus it can better expand into the high-density region of the true posterior as its entropy increases.

7 shows the optimal performance on ID data is quite similar between different latent architectures. On OOD, the optimal performance of using both input and output latent variables is similar to using only output latent variables, while using only input latent variables produces slightly worse optimal performance. The optimal \(\gamma\) is lower when the model uses both types of latent variables (\(\mathbf{z},\mathbf{s}\)), because the entropy of the product of two latent variables increases rapidly as we increase the entropy of both latent variables.

We also observe these patterns in other architectures and datasets (see 16). In summary, from our experimental results we find that using only output latent variables with a sufficient number of components (e.g., \(K=4\)) achieves excellent results for node-based BNNs in our benchmark.

5.2 Effects of \(\gamma\) on robustness against noisy labels↩︎

Learning wrong labels amounts to memorizing random patterns, which requires more capacity from the model than learning generalizable patterns [30]. We hypothesize that if we corrupt wrongly labelled training samples with sufficiently diverse implicit corruptions, we overwhelm the neural network making it unable to memorize these spurious patterns during training. To test this intuition, we follow the experiment in [31], where we take a percentage of training samples in cifar10 and corrupt their labels. We thus split the training set into two parts: \(\mathcal{D}_1\) containing only samples with correct labels, and \(\mathcal{D}_2\) including those with wrong labels. We then track the final NLL of \(\mathcal{D}_1\) and \(\mathcal{D}_2\) under different \(\gamma\), and visualize the results in 8. This figure shows that as \(\gamma\) increases, the NLL of \(\mathcal{D}_2\) (noisy labels) increases much faster than that of \(\mathcal{D}_1\) (clean labels), indicating that the network fails to learn random patterns under simulated corruptions. As a consequence, the model generalizes better on the test set, as shown in 9.

5.3 Benchmark results↩︎

10 11 present the results of node-based BNNs and baselines on cifar10/cifar100 and tinyimagenet. We choose SWAG [32], cSG-HMC [33] and ASAM [34] as our baselines. These are strong baselines, as both SWAG and cSG-HMC have demonstrated state-of-the-art uncertainty estimation, while ASAM produce better MAP models than stochastic gradient descent by actively seeking wide loss valleys. We repeated each experiment 25 times with different random seeds. For each method, we also consider its ensemble version where we combine 5 models from different runs when making predictions. For the ensemble versions, each experiment is repeated 5 times. We use 30 Monte Carlo samples for node-based BNNs, SWAG, cSG-HMC and their ensemble versions to estimate the posterior predictive distribution. We use standard performance metrics of expected calibration error (ECE) [35], NLL and predictive error. We use resnet18 for cifar10/cifar100 and preactresnet18 for tinyimagenet. We also include the result of vgg16 on cifar10/cifar100 in 14. For evaluation, we use the corrupted test images provided by [21].

On cifar100, node-based BNNs outperform the baselines in NLL and error, however SWAG performs best on cifar10. Interestingly, in cifar100, node-based BNNs and their ensembles have worse ECE than the baselines on ID data, however as the test images become increasingly corrupted, the ECEs of the baselines degrade rapidly while the ECE of node-based BNNs remains below a threshold. Similar behaviors are observed on tinyimagenet, with the node-based BNNs produce the lowest NLL and error while not experiencing ECE degradation under corruptions.

6 Related works↩︎

6.0.0.1 Multiplicative latent node variables in BNNs.

There have been several earlier works that utilize multiplicative latent node variables, either as a primary source of predictive uncertainty such as MC-Dropout [14], Variational Dropout [13], Rank-1 BNNs [12] and Structured Dropout [16]; or to improve the flexibility of the mean-field Gaussian posterior in variational inference [19]. Here we study the contribution of these latent variables to robustness under covariate shift.

6.0.0.2 BNNs under covariate shift.

Previous works have evaluated the predictive uncertainty of BNNs under covariate shift [6], [18], with the recent work by [6] showing that standard BNNs with high-fidelity posteriors perform worse than MAP solutions under covariate shift. [36] attributed this phenomenon to the absence of posterior contraction on the null-space of the data manifold. This problem is avoided in node-based BNNs as they still maintain a point-estimate for the weights.

6.0.0.3 Dropout as data augmentation.

Similar to our study, a previous work by [37] studied Dropout from the data augmentation perspective. Here we study latent variables with more flexible posterior (mixture of Gaussians) and focus on simulating input corruptions for OOD robustness.

6.0.0.4 Adversarial robustness via feature perturbations.

Data-space perturbations have been investigated as a means to defend neural networks against adversarial attacks [38]–[40].

6.0.0.5 Tempered posteriors.

Tempered posteriors have been used in variational inference to obtain better variational posterior approximations [25]. A recent study put the focus on the cold posterior effect of weight-based BNNs [5]. We have shown that our approach of regularizing the variational entropy is equivalent to performing variational inference with a hot posterior as the target distribution. Tempered posteriors have also been studied in Bayesian statistics as a means to defend against model misspecification [41]–[44]. Covariate shift is a form of model misspecification, as model mismatch arises from using a model trained under different assumptions about the statistics of the data.

7 Conclusion↩︎

We analyzed node-based BNNs from the perspective of using latent node variables for simulating input corruptions. We showed that by regularizing the entropy of the latent variables, we increase the diversity of the implicit corruptions, and thus improve performance of node-based BNNs under covariate shift. Across cifar10, cifar100 and tinyimagenet, entropy regularized node-based BNNs produce excellent results in uncertainty metrics on OOD data.

In this study, we focused on variational inference, leaving the study of implicit corruptions under other approximate inference methods as future work. Furthermore, our work shows the benefits of hot posteriors and argues for an inherent trade-off between ID and OOD performance in node-based BNNs. It is an interesting future direction to study these questions in weight-based BNNs. Finally, our work presented entropy as a surprisingly useful summary statistic that can partially explain the complex connection between the variational posterior and corruption robustness. One important research direction is to develop more informative statistics that can better encapsulate this connection.

Acknowledgement↩︎

This work was supported by the Academy of Finland (Flagship programme: Finnish Center for Artificial Intelligence FCAI and grants no. 292334, 294238, 319264, 328400) and UKRI Turing AI World-Leading Researcher Fellowship, EP/W002973/1. We acknowledge the computational resources provided by Aalto Science-IT project and CSC–IT Center for Science, Finland.

8 Original ELBO derivation↩︎

Here we provide a detail derivation of the ELBO in 8 . We assume a prior \(p(\theta, \mathcal{Z})=p(\theta)p(\mathcal{Z})\) for the parameters \(\theta\) and latent variables \(\mathcal{Z}\), and we assume a variational posterior \(q_{\phi, \hat{\theta}}(\theta, \mathcal{Z})=\delta(\theta-\hat{\theta})q_\phi(\mathcal{Z})\) where \(\delta(.)\) is a Dirac delta distribution. We arrive at the ELBO in 8 by minimizing the KL divergence between the variational approximation and the true posterior with respect to the variational parameters \((\hat{\theta}, \phi)\): \[\begin{align} &\mathop{\mathrm{arg\,min}}_{\phi, \hat{\theta}} \mathrm{KL}\left[q_{\phi, \hat{\theta}}(\theta, \mathcal{Z})||p(\theta, \mathcal{Z}|\mathcal{D})\right] \\ &= \mathop{\mathrm{arg\,min}}_{\phi, \hat{\theta}} \mathbb{E}_{q_{\phi, \hat{\theta}}(\theta, \mathcal{Z})}\left[\log q_{\phi, \hat{\theta}}(\theta, \mathcal{Z}) - \log p(\mathcal{D}|\theta,\mathcal{Z}) - \log p(\theta, \mathcal{Z}) + \log p(\mathcal{D})\right] \\ &= \mathop{\mathrm{arg\,min}}_{\phi, \hat{\theta}} \mathbb{E}_{q_{\phi}(\mathcal{Z})}\left[-\log p(\mathcal{D}|\hat{\theta},\mathcal{Z})\right] + \mathrm{KL}\left[q_{\phi}(\mathcal{Z})||p(\mathcal{Z})\right] - \log p(\hat{\theta}) + \log p(\mathcal{D}) \\ &= \mathop{\mathrm{arg\,min}}_{\phi, \hat{\theta}} -\mathcal{L}(\hat{\theta}, \phi) \end{align}\]

9 Tempered ELBO derivation↩︎

Here we show a connection between the tempered posterior with temperature \(\tau=1/(\gamma+1)\) in 14 and the augmented ELBO in Section 4.1: \[\begin{align} &\mathop{\mathrm{arg\,min}}_{\phi, \hat{\theta}} \frac{1}{\tau} \mathrm{KL}\left[q_{\phi, \hat{\theta}}(\theta, \mathcal{Z}) \, || \, p_\gamma(\theta, \mathcal{Z}|\mathcal{D})\right] \\ &= \mathop{\mathrm{arg\,min}}_{\phi, \hat{\theta}} \frac{1}{\tau} \mathbb{E}_{q_{\phi, \hat{\theta}}(\theta, \mathcal{Z})}\left[\log q_{\phi, \hat{\theta}}(\theta, \mathcal{Z}) - \tau\log p(\mathcal{D}|\theta,\mathcal{Z}) - \tau\log p(\theta, \mathcal{Z}) + \log p_\gamma(\mathcal{D})\right] \\ &= \mathop{\mathrm{arg\,min}}_{\phi, \hat{\theta}} \mathbb{E}_{q_{\phi, \hat{\theta}}(z, \theta)}\left[\frac{1}{\tau}\log q_{\phi, \hat{\theta}}(\theta, \mathcal{Z}) - \log p(\mathcal{D}|\theta,\mathcal{Z}) - \log p(\theta) - \log p(\mathcal{Z}) \right] + \frac{1}{\tau}\log p_\gamma(\mathcal{D}) \\ &= \mathop{\mathrm{arg\,min}}_{\phi, \hat{\theta}} -\mathbb{E}_{q_{\phi}(\mathcal{Z})}\left[\log p(\mathcal{D}|\hat{\theta},\mathcal{Z})\right] + \mathrm{KL}\left[q_{\phi}(\mathcal{Z})||p(\mathcal{Z})\right] - \gamma \mathbb{H}\left[q_\phi(\mathcal{Z})\right] - \log p(\hat{\theta}) + \frac{1}{\tau}\log p_\gamma(\mathcal{D}) \\ &= \mathop{\mathrm{arg\,min}}_{\phi, \hat{\theta}} -\mathcal{L}_\gamma(\hat{\theta}, \phi) + \log p_\gamma(\mathcal{D})^{\frac{1}{\tau}} \end{align}\]

10 Derivation of layer-wise activation shifts due to input corruptions↩︎

Here we explain in detail the approximation of layer-wise activation shifts in 10 . To simulate covariate shift, one can take an input \(\mathbf{x}\) assumed to come from the same distribution as the training samples and apply a corruption \(\mathbf{g}^0\) to form a shifted version \(\mathbf{x}^c\) of \(\mathbf{x}\): \[\mathbf{x}^c \triangleq \mathbf{x}+ \mathbf{g}^0(\mathbf{x})\] For instance, \(\mathbf{x}\) could be an image and \(\mathbf{g}^0\) can represent the shot noise corruption as seen in [21]. The corruption \(\mathbf{g}^0(\mathbf{x})\) creates a shift in the activation of the first layer \(\mathbf{f}^1\) which can be approximated using the first-order Taylor expansion: \[\begin{align} \mathbf{g}^1(\mathbf{x}) &= \mathbf{f}^1(\mathbf{x}^c) - \mathbf{f}^1(\mathbf{x}) \\ &= \sigma\left(\mathbf{W}^1(\mathbf{x}+\mathbf{g}^0(\mathbf{x})) + \mathbf{b}^1\right) - \sigma\left(\mathbf{W}^1\mathbf{x}+ \mathbf{b}^1\right) \\ &\approx \mathbf{J}_\sigma\left[ \mathbf{h}^1(\mathbf{x}) \right]\left(\mathbf{W}^1\mathbf{g}^0(\mathbf{x})\right) \end{align}\] where \(\mathbf{J}_\sigma=\partial \sigma/\partial \mathbf{h}\) denotes the Jacobian of the activation \(\sigma\) with respect to pre-activation outputs \(\mathbf{h}\). Similarly, the approximation of the activation shift in the second layer is: \[\begin{align} \mathbf{g}^2(\mathbf{x}) &= \mathbf{f}^2(\mathbf{x}^c) - \mathbf{f}^2(\mathbf{x}) \\ &= \sigma\left(\mathbf{W}^2\mathbf{f}^1(\mathbf{x}^c) + \mathbf{b}^2\right) - \sigma\left(\mathbf{W}^2\mathbf{f}^1(\mathbf{x})+ \mathbf{b}^2\right) \\ &= \sigma\left(\mathbf{W}^2(\mathbf{f}^1(\mathbf{x})+\mathbf{g}^1(\mathbf{x})) + \mathbf{b}^2\right) - \sigma\left(\mathbf{W}^2\mathbf{f}^1(\mathbf{x})+ \mathbf{b}^2\right) \\ &\approx \mathbf{J}_\sigma\left[ \mathbf{h}^2(\mathbf{x}) \right ]\left(\mathbf{W}^2\mathbf{g}^1(\mathbf{x})\right) \end{align}\] Generally, one can approximate the shift in the output of the \(\ell\)-th layer caused by \(\mathbf{g}(\mathbf{x})\) as: \[\mathbf{g}^\ell(\mathbf{x}) = \mathbf{f}^\ell(\mathbf{x}^c) - \mathbf{f}^\ell(\mathbf{x}) \approx \mathbf{J}_\sigma\left[ \mathbf{h}^\ell (\mathbf{x}) \right] \left(\mathbf{W}^{\ell}\mathbf{g}^{\ell-1}(\mathbf{x})\right)\]

11 Details on small-scale experiments↩︎

For the small-scale experiments in Section 3, we use the all-cnn-c architecture from [22]. We describe this architecture in Table 1. We train the model for \(90\) epochs, and only use the output latent variables and a posterior with \(1\) Gaussian component for this experiment

12 Additional visualization of outputs at each layer↩︎

In Section 3, we provide a PCA visualization of the outputs from the last layer of a node-based all-cnn-c BNN on one sample of cifar10. Here we also provide the same visualizations for the first two and the last two layers of the network. We use the same input image as 2.

13 Additional details on the experiments and hyperparameters↩︎

13.1 Approximation for the KL divergence with mixture variational posterior↩︎

We use a mixture of Gaussians (MoG) distribution with \(K\) equally-weighted components to provide a flexible approximation of the true posterior in the latent space: \[\begin{align} q(\mathcal{Z}) = \frac{1}{K} \sum_{k=1}^K q_k(\mathcal{Z}) \\ q_k(\mathcal{Z}) = \prod_{\ell=1}^L q_{k,\ell}(\mathcal{Z}^\ell) \\ q_{k,\ell}(\mathcal{Z}^\ell) = \mathcal{N}(\boldsymbol{\mu}_{k,\ell}, \mathrm{diag}\, \boldsymbol{\sigma}_{k,\ell}^2). \end{align}\] where \(L\) is the number of layers. We use a Gaussian prior with global scalar variance \(s^2\) for the latent prior, \[p(\mathcal{Z}) = \mathcal{N}(\mathbf{1}, s^2 I).\]

The KL divergence decomposes into cross-entropy and entropy terms, \[\begin{align} \mathrm{KL}[q(\mathcal{Z}) \, || \, p(\mathcal{Z})] &= \mathbb{H}[q,p] - \mathbb{H}[q] = \frac{1}{K} \sum_{k=1}^K \mathbb{H}[q_k,p] - \mathbb{H}[q], \end{align}\] where the cross-entropy reduces into tractable terms \(\mathbb{H}[q_k,p]\) for Gaussians. The mixture entropy \(\mathbb{H}[q]\) remains intractable, but admits a lower bound [45], \[\begin{align} \label{eq:lb95entropy} \mathbb{H}[q] &\ge \frac{1}{K} \sum_{k=1}^K \mathbb{H}[q_k] - \frac{1}{K} \sum_{k=1}^K \log \left( \frac{1}{K} \sum_{r=1}^K \mathrm{BC}(q_k, q_r)\right) \triangleq \widehat{\mathbb{H}}[q] \end{align}\tag{15}\] where \[\mathrm{BC}(q,q') = \int \sqrt{q(\mathbf{z})} \sqrt{q'(\mathbf{z})} d\mathbf{z}\qquad \le 1\] is the Bhattacharyya kernel of overlap between two distributions [46], [47], and has a closed form solution for a pair of Gaussians \(q,q'\). The Bhattacharyya kernel has the convenient normalization property \(\mathrm{BC}(q,q) = 1\). The lower bound considers unary and pairwise component entropies.

13.2 Experimental details and hyperparameters↩︎

We actually maximizes the following objective to train the node-based BNNs on large-scale experiments: \[\mathcal{L}_{\gamma, \beta}(\hat{\theta}, \phi)=\mathbb{E}_{q_{\phi}(\mathcal{Z})}\Big[\log p(\mathcal{D}| \hat{\theta},\mathcal{Z})\Big] + \log p(\hat{\theta}) + \beta\left(-\mathbb{H}\big[q_{\phi}(\mathcal{Z}), p(\mathcal{Z})\big] + (\gamma+1) \widehat{\mathbb{H}}\big[q_\phi(\mathcal{Z})\big] \right)\] which is the augmented ELBO in 13 with additional coefficient \(\beta\) for the cross-entropy and variational entropy term. We also replace the intractable mixture entropy \(\mathbb{H}[q]\) with its tractable lower bound \(\widehat{\mathbb{H}}[q]\) presented in 15 . During training, we will anneal \(\beta\) from \(0\) to \(1\). We found this to have ease optimization and produce better final results. For all experiments, we estimate the expected log-likelihood in the loss function using 4 samples.

For all the experiments on cifar10/cifar100, we run each experiment for 300 epochs, where we increase \(\beta\) from \(0\) to \(1\) for the first 200 epochs. We use SGD as our optimizer, and we use a weight decay of \(0.0005\) for the parameters \(\theta\). We use a batch size of 128. For all the experiments on tinyimagenet, we run each experiment for 150 epochs, where we increase \(\beta\) from \(0\) to \(1\) for the first 100 epochs. We use a batch size of 256. Bellow, we use \(\lambda_1\) and \(\lambda_2\) to denote the learning rate of the parameters \(\theta\) and \(\phi\) respectively.

For vgg16, we set the initial learning rate \(\lambda_1=\lambda_2=0.05\), and we decrease \(\lambda_1\) linearly from \(0.05\) to \(0.0005\) from epoch 150 to epoch 270, while keeping \(\lambda_2\) fixed throughout training. We initialize the standard deviations with \(\mathcal{N}^+(0.30, 0.02)\) and set the standard deviation of the prior to \(0.30\).

For resnet18, we set the initial learning rate \(\lambda_1=\lambda_2=0.10\), and we decrease \(\lambda_1\) linearly from \(0.10\) to \(0.001\) from epoch 150 to epoch 270, while keeping \(\lambda_2\) fixed throughout training. We initialize the standard deviations with \(\mathcal{N}^+(0.40, 0.02)\) and set the standard deviation of the prior to \(0.40\).

For preactresnet18, we set the initial learning rate \(\lambda_1=\lambda_2=0.10\), and we decrease \(\lambda_1\) linearly from \(0.10\) to \(0.001\) from epoch 75 to epoch 135, while keeping \(\lambda_2\) fixed throughout training. We initialize the standard deviations with \(\mathcal{N}^+(0.30, 0.02)\) and set the standard deviation of the prior to \(0.30\).

13.3 Runtime↩︎

We report the average running times of different methods in Table 2. We used similar number of epochs for all methods in each experiment. All experiment were performed on one Tesla V100 GPU. Overall, node BNNs took 4 times longer to train than SWAG since we use 4 Monte Carlo samples per training sample to estimate the expected log-likelihood in the \(\gamma\)-ELBO. ASAM took 2 times longer to train than SWAG since they require two forward-backward passes per minibatch.

14 Additional benchmark results↩︎

Here we include the benchmark results of vgg16 on cifar10 and cifar100 in 13. We also include 14 and 15 as larger versions of 10 and 11.

15 The evolution of variational entropy during training↩︎

We visualize the progression of the variational entropy when trained using the original ELBO (without the \(\gamma\)-entropy term) under different settings in 16 17 18 19. We can observe the typical behaviour of variational inference that it tends to reduce the entropy of the variational posterior over time.

16 Additional results on the effect of \(\gamma\) on performance of node-based BNNs↩︎

Here we include 20 21 22 23 24 to show the effect of \(\gamma\) on performance of node-based BNNs under different architectures and datasets.

References↩︎