Incorporating Unlabelled Data into Bayesian Neural Networks
Conventional Bayesian Neural Networks (BNNs) are unable to leverage
unlabelled data to improve their predictions. To overcome this limitation, we
introduce Self-Supervised Bayesian Neural Networks, which use unlabelled data
to learn models with suitable prior predictive distributions. This is achieved
by leveraging contrastive pretraining techniques and optimising a variational
lower bound. We then show that the prior predictive distributions of
self-supervised BNNs capture problem semantics better than conventional BNN
priors. In turn, our approach offers improved predictive performance over
conventional BNNs, especially in low-budget regimes.
Mrinank Sharma, Tom Rainforth, Yee Whye Teh, Vincent Fortuin
# Incorporating Unlabelled Data into

###### Abstract

Conventional Bayesian Neural Networks (BNNs) cannot leverage unlabelled data to improve their predictions. To overcome this limitation, we introduce _Self-Supervised Bayesian Neural Networks_, which use unlabelled data to learn improved prior predictive distributions by maximising an evidence lower bound during an unsupervised pre-training step. With a novel methodology developed to better understand prior predictive distributions, we then show that self-supervised prior predictives capture image semantics better than conventional BNN priors. In our empirical evaluations, we see that self-supervised BNNs offer the label efficiency of self-supervised methods and the uncertainty estimates of Bayesian methods, particularly outperforming conventional BNNs in low-to-medium data regimes.

## 1 Introduction

Bayesian Neural Networks (BNNs) are powerful probabilistic models that combine the flexibility of deep neural networks with the theoretical underpinning of Bayesian methods (Mackay, 1992; Neal, 1995). Indeed, as they place priors over their parameters and perform posterior inference, BNN advocates consider them to be a principled approach for uncertainty estimation (Wilson and Izmailov, 2020; Abdar et al., 2021), which can be helpful for label-efficient learning (Gal et al., 2017).

However, despite the prevalence of unlabelled data and the rise of unsupervised learning, conventional BNNs cannot harness unlabelled data for improved uncertainty estimates and label efficiency. Rather, practitioners have focused on improving label efficiency and predictive performance by imparting information with priors over network parameters or predictive functions (e.g., Louizos et al., 2017; Tran et al., 2020; Matsubara et al., 2021; Fortuin et al., 2021). But it stands to reason that the vast store of information contained in unlabelled data should be incorporated into BNNs, and that the potential benefit of doing so likely exceeds the benefit of designing better, but ultimately human-specified, priors over parameters or functions. Unfortunately, as standard BNNs are explicitly only models for supervised prediction, they cannot leverage such unlabelled data by conditioning on it.

To overcome this limitation, **we introduce _Self-Supervised Bayesian Neural Networks_** (SS3), which condition on pseudo-labelled tasks generated using unlabelled data. We show this can be seen as **incorporating unlabelled data into the prior predictive distribution**. That is, our method improves the prior predictive by using unlabelled label to _learn_ it, not by using a different but ultimately _human-specified_ prior over parameters or functions. Self-supervised BNNs first train a deterministic encoder to **maximise a lower bound of a log-marginal likelihood** derived from unlabelled data in the unsupervised pre-training step. Although the encoder is deterministic, it induces a prior predictive distribution over the downstream labels such that augmented image pairs likely have the same label, while distinct images likely have different labels (Fig. 0(a)). Following the pre-training step, we then condition a subset of the network parameters on the downstream labelled data to make predictions.
As a further contribution, **we develop a methodology to better understand prior predictive distributions** of functions with high-dimensional inputs (SS4). Since it is often easy to reason about the semantic relationship between points, our approach investigates the _joint_ prior predictive for input pairs with known semantic relation. **We then demonstrate that self-supervised BNN prior predictives reflect input pair semantic similarity better than normal BNN priors**. In particular, while normal BNN priors struggle to distinguish same-class input pairs and different-class input pairs (Fig. 0(c)), self-supervised BNN priors do so much better (Fig. 0(b)).

Finally, **we demonstrate that the improved prior predictives of self-supervised BNNs are helpful in practice** (SS5). In our experiments, we find that self-supervised BNNs combine the label efficiency of self-supervised approaches with the uncertainty estimates of Bayesian methods. In particular, self-supervised BNNs outperform conventional BNNs at low-to-medium data regimes in semi-supervised and active learning settings, whilst also offering better-calibrated predictions than SimCLR (Chen et al., 2020), which is a particularly popular self-supervised learning algorithm.

## 2 Background: Bayesian Neural Networks

Let \(f_{\theta}(x)\) be a neural network with parameters \(\theta\) and \(\mathcal{D}=\{(x_{i},y_{i})\}_{i=1}^{N}\) be an observed dataset where we want to predict \(y\) from \(x\). A Bayesian Neural Network (BNN) specifies a prior over parameters, \(p(\theta)\), and a likelihood, \(p(y|f_{\theta}(x))\), which in turn define the posterior \(p(\theta|\mathcal{D})\propto p(\theta)\prod_{i}p(y_{i}|f_{\theta}(x_{i}))\). To make predictions, we approximate the posterior predictive \(p(y_{\star}|x_{\star},\mathcal{D})=\mathbb{E}_{p(\theta|\mathcal{D})}[p(y_{ \star}|f_{\theta}(x_{\star}))]\).

Improving BNN priors has been a long-standing goal for the BNN community, primarily through improved human-designed priors. One approach is to improve the prior over the network's parameters (Louizos et al., 2017; Nalisnick, 2018). Others place priors directly over predictive functions (Flamm-Shepherd et al., 2017; Sun et al., 2019; Matsubara et al., 2021; Nalisnick et al., 2021; Raj et al., 2023). Both approaches, however, present challenges--the mapping between the network's parameters and predictive functions is complex, while directly specifying our beliefs over predictive functions is itself a highly challenging task. For these reasons, as well as computational convenience, isotropic Gaussian priors over network parameters remain the most common choice (Fortuin, 2022).

Figure 1: **Self-Supervised Bayesian Neural Networks**. (a) Pre-training in self-supervised BNNs corresponds to unsupervised prior learning. We learn a prior predictive distribution such that augmented images likely have the same label and distinct images likely have different labels. (b) Self-supervised BNN priors assign higher probabilities to semantically consistent image pairs having the same label compared to semantically inconsistent image pairs. Here, semantically consistent image pairs have the same ground-truth label, and semantically inconsistent image pairs have different ground-truth labels. The plot shows a kernel density estimate of the log-probability that same-class and different-class image pairs are assigned the same label under the prior. (c) Unlike self-supervised prior predictives, conventional BNN prior predictives assign similar probabilities to semantically consistent and semantically inconsistent image pairs having the same label.
Self-Supervised Bayesian Neural Networks

Conventional BNNs are unable to harness unlabelled data for improved uncertainty estimation and label efficiency. To overcome this limitation, we introduce _Self-Supervised Bayesian Neural Networks_. At a high level, they use unlabelled data and data augmentation to generate pseudo-labelled datasets, which are then conditioned on in a modified probabilistic model.

**Problem Specification** Suppose \(\mathcal{D}^{u}=\{x_{i}^{u}\}_{i=1}^{N}\) is an unlabelled dataset of examples \(x_{i}^{u}\in\mathbb{R}^{n}\) with \(x_{i}\sim P_{X}^{u}\); \(P_{X}^{u}\) is the underlying data distribution. Let \(\mathcal{D}^{t}=\{(x_{i}^{t},y_{i}^{t})\}_{i=1}^{T}\) be a labelled dataset corresponding to a supervised "downstream" task, where \(y_{i}^{t}\) is the target associated with \(x_{i}^{t}\). We have \(x_{i}\sim P_{X}^{t}\) and \(y_{i}|x_{i}\sim P_{Y|X}^{t}\). We want to use \(\mathcal{D}^{u}\) to improve our predictions on the downstream task.

### Incorporating Unlabelled Data into BNNs

The natural way to benefit from \(\mathcal{D}^{u}\) would be to use it to inform our beliefs about the model parameters, \(\theta\), through the posterior \(p(\theta|\mathcal{D}^{u},\mathcal{D}^{t})\propto p(\theta|\mathcal{D}^{u})\,p( \mathcal{D}^{t}|\mathcal{D}^{u},\theta)\). But if we are working with conventional BNNs, which are explicitly models for supervised prediction, \(p(\theta|\mathcal{D}^{u})=p(\theta)\). Further, as predictions depend only on the parameters, \(p(\mathcal{D}^{t}|\mathcal{D}^{u},\theta)=p(\mathcal{D}^{t}|\theta)\), which then means \(p(\theta|\mathcal{D}^{u},\mathcal{D}^{t})=p(\theta|\mathcal{D}^{t})\). Said differently, for a conventional BNN, we cannot incorporate \(\mathcal{D}^{u}\) by conditioning on it. One could proceed in different ways, for instance, by using \(\mathcal{D}^{u}\) to choose/learn the prior \(p(\theta)\), or using a generative model that defines a likelihood for the inputs \(p(x|\theta)\).

In this work, we instead note that if we were able to convert \(\mathcal{D}^{u}\) into labelled data, we could condition a discriminative model on it. If we then introduced a probabilistic link between generated labelled data, derived from \(\mathcal{D}^{u}\), and the predictions over the downstream task labels, we could harness \(\mathcal{D}^{u}\) through probabilistic conditioning. We now introduce each of these elements in turn.

First, to generate labelled data from \(\mathcal{D}^{u}\), i.e., to formulate _self-supervision_, we take inspiration from contrastive learning (Oord et al., 2019; Chen et al., 2020, 2020, 2020; Grill et al., 2020, 2020).1 We intuitively believe that distinct examples are unlikely to be semantically consistent, and thus are unlikely to have the same downstream label. Further, if we can specify data transformations that preserve semantic information, we also believe that augmented versions of the same example likely have the same downstream label. The correspondence between such beliefs and a prior distribution over predictive functions or network parameters is unclear, so we instead use \(\mathcal{D}^{u}\) and data transformations to generate pseudo-labelled data, and include this additional subjective information by conditioning on those datasets rather than a prior over parameters or functions.

Footnote 1: Our framework also encompasses other self-supervised tasks, e.g., next-token prediction and masked-token prediction for language models, though we do not consider them here.

Concretely, suppose we have a set of data augmentations \(\mathcal{A}=\{a:\mathbb{R}^{n}\rightarrow\mathbb{R}^{n}\}\) that preserve semantic content. We use \(\mathcal{A}\) and \(\mathcal{D}^{u}\) to generate a _contrastive dataset_\(\mathcal{D}^{c}\) that reflects our subjective beliefs by:

1. Drawing \(M\) examples from \(\mathcal{D}^{u}\) at random, \(\{\hat{x}_{i}\}_{i=1}^{M}\); \(i\) indexes the subset, not \(\mathcal{D}^{u}\)
2. For each \(x_{i}\), sampling \(a^{A},a^{B}\sim\mathcal{A}\) and augmenting, giving \(\tilde{x}_{i}^{A}=a^{A}(\hat{x}_{i})\) and \(\tilde{x}_{i}^{B}=a^{B}(\hat{x}_{i})\)
3. Forming \(\mathcal{D}^{c}\) by assigning \(\tilde{x}_{i}^{A}\) and \(\tilde{x}_{i}^{B}\) the same class label, which is the subset index \(i\)

We thus have \(\mathcal{D}^{c}=\{(x_{i}^{c},y_{i}^{c})\}_{i=1}^{2M}=\{(\tilde{x}_{i}^{A},i) \}_{i=1}^{M}\cup\{(\tilde{x}_{i}^{B},i)\}_{i=1}^{M}\), where the labels are between \(1\) and \(M\). The task is to predict the subset index corresponding to each augmented example. We can repeat this process \(L\) times and create a set of contrastive task datasets, \(\{\mathcal{D}_{j}^{c}\}_{j=1}^{L}\). Here, we consider the number of generated contrastive datasets \(L\) to be a fixed, finite hyper-parameter, but we discuss the implications of generating a potentially infinite number of datasets in Appendix A.

Second, to link the labelled data derived from \(\mathcal{D}^{u}\) and the downstream task predictions, we use parameter sharing (see Fig. 2). We introduce parameters \(\theta_{j}^{c}\) for each \(\mathcal{D}_{j}^{c}\), parameters \(\theta^{t}\) for \(\mathcal{D}^{t}\), and shared-parameters \(\theta^{s}\) that are used for both the downstream and the contrastive tasks. A contrastive dataset \(\mathcal{D}_{j}^{c}\) thus informs downstream predictions through \(\theta^{s}\). For example, \(\theta^{t}\) and \(\theta_{j}^{c}\) could be the parameters of the last layer of a neural network, while \(\theta^{s}\) could be parameters of earlier layers.

We now discuss different options for learning in this conceptual framework. Using the Bayesian approach, one would place priors over \(\theta^{s}\), \(\theta^{t}\), and each \(\theta_{j}^{c}\). This then defines a posterior distribution given the observed data \(\{\mathcal{D}_{j}^{c}\}_{j=1}^{L}\) and \(\mathcal{D}^{t}\). To make predictions on the downstream task, which depend on \(\theta^{s}\) and \(\theta^{t}\) only, we would then use the posterior predictive:
\[p(y^{t}_{\star}|x_{\star},\{\mathcal{D}^{c}_{j}\}_{j=1}^{L}, \mathcal{D}^{t}) \tag{1}\] \[\quad=\mathbb{E}_{p(\theta^{s}|\{\mathcal{D}^{c}_{j}\}_{j=1}^{L}, \mathcal{D}^{t})}[\mathbb{E}_{p(\theta^{s}|\theta^{s},\mathcal{D}^{t})}[p(y^{t }_{\star}|x_{\star},\theta^{s},\theta^{t})]],\]

where we have (i) noted that the downstream task parameters \(\theta^{t}\) are independent of \(\{\mathcal{D}^{c}_{j}\}_{j=1}^{L}\) given the shared parameters \(\theta^{s}\) and (ii) integrated over each \(\theta^{c}_{j}\) and \(\theta^{t}\) in the definition of \(p(\theta^{s}|\{\mathcal{D}^{c}_{j}\}_{j=1}^{L},\mathcal{D}^{t})\).

Alternatively, one can learn a point estimate for \(\theta^{s}\), e.g., with MAP estimation, and perform full posterior inference for \(\theta^{t}\) and \(\theta^{c}_{j}\) only. This would be a _partially stochastic_ network, which Sharma et al. (2022) showed often outperform fully stochastic networks while being more practical, and can also be justified as we can generate many contrastive tasks. This corresponds to model learning as used in deep kernels and variational autoencoders (Kingma and Welling, 2013; Wilson et al., 2015; Rezende and Mohamed, 2015; Wilson et al., 2016).

We now show these modifications incorporate information from \(\mathcal{D}^{u}\) and the augmentations \(\mathcal{A}\) into the downstream task prior predictive. Under the Bayesian approach, self-supervised BNNs push forward the posterior over \(\theta^{t}\) and \(p(\theta^{s}|\{\mathcal{D}^{c}_{j}\}_{j=1}^{L},\mathcal{D}^{t})\propto p( \theta^{s}|\{\mathcal{D}^{c}_{j}\}_{j=1}^{T})\,p(\mathcal{D}^{t}|\theta^{s})\) through the network to make predictions. In contrast, a standard BNN uses \(p(\theta|\mathcal{D}^{t})\propto p(\theta)\,p(\mathcal{D}^{t}|\theta)\). We see that self-supervised BNNs use \(\{\mathcal{D}^{c}_{j}\}_{j=1}^{L}\) to update the prior over the shared parameters \(p(\theta^{s})\). Alternatively, if we learnt a point estimate for \(\theta^{s}\) and used the Bayesian approach for \(\theta^{t}\), this would define the prior predictive:

\[p(y^{t}_{\star}|x_{\star},\{\mathcal{D}^{c}_{j}\}_{j=1}^{L})=\mathbb{E}_{p( \theta^{s})}[p(y^{t}_{\star}|x_{\star},\theta^{s}_{\star},\theta^{t})], \tag{2}\]

where \(\theta^{s}_{\star}\) is a point estimate for \(\theta^{s}\) learnt using \(\{\mathcal{D}^{c}_{j}\}_{j=1}^{L}\). Since the contrastive datasets were constructed to reflect our prior beliefs, both conditioning and training on them should be helpful.

### Practical Self-Supervised Bayesian Neural Networks

We now use our framework to propose a practical two-step algorithm for self-supervised BNNs.

**Preliminaries** We focus on image-classification problems. We use an encoder \(f_{\theta^{s}}(\cdot)\) that maps images to representations and is shared across the contrastive tasks and the downstream task. The shared parameters \(\theta^{s}\) thus are the base encoder's parameters. We also normalise the representations produced by this encoder. For the downstream dataset, we use a linear readout layer from the encoder representations, i.e., we have \(\theta^{t}=\{W^{t},b^{t}\}\) and \(y^{t}_{i}\sim\mathrm{softmax}(W^{t}f_{\theta^{s}}(x_{i})+b^{t})\). The rows of \(W^{t}_{j}\) are thus class template vectors. For the contrastive tasks, we use a linear layer without biases, i.e., \(\theta^{c}_{j}=W^{c}_{j}\), and \(j\) indexes contrastive tasks. We place Gaussian priors over \(\theta^{s}\), \(\theta^{t}\), and each \(\theta^{c}_{j}\).

Pre-training \(\theta^{s}\) (Step I)Here, we learn a point estimate for the base encoder parameters \(\theta^{s}\), which induces a prior predictive distribution over the downstream task labels (see Eq. 2). To learn \(\theta^{s}\), we want to optimise the (potentially penalised) log-likelihood \(\log p(\{\mathcal{D}^{c}_{j}\}_{j=1}^{L},\mathcal{D}^{t}|\theta^{s})\), but this would require integrating over \(\theta^{t}\) and each \(\theta^{c}\). Instead, we use the evidence lower bound (ELBO):

\[\widehat{\mathcal{L}}^{c}_{j}(\theta^{s})=\mathbb{E}_{q(\theta^{c}_{j})}[\log p (\mathcal{D}^{c}_{j}|\theta^{s},\theta^{c}_{j})]-D_{\text{KL}}(q(\theta^{c}_{j} )||p(\theta^{c}_{j}))\leq\log p(\mathcal{D}^{c}_{j}|\theta^{s}), \tag{3}\]

where \(q(\theta^{c}_{j})\) is a variational distribution over the contrastive task parameters. Rather than learning a different variational distribution for each contrastive task \(j\), we amortise the inference and exploit the structure of the contrastive task. The contrastive task is to predict the source image index from pairs of augmented images using a linear layer from an encoder that produces normalised representations. We define \(\omega_{i}=0.5(f_{\theta^{s}}(\tilde{x}^{A}_{i})+f_{\theta^{s}}(\tilde{x}^{B}_{ i}))\), i.e., \(\omega_{i}\) is the mean representation for each augmented pair of images, and because the rows of \(W^{c}_{j}\) correspond to class templates, we thus use:

\[q(W^{c}_{j};\tau,\sigma^{2})=\mathcal{N}(\mu^{c}_{j},\sigma^{2}I),\text{with}\ \mu^{c}_{j}=\begin{bmatrix}\omega^{T}_{1}&\dots&\omega^{T}_{M}\end{bmatrix}/\tau \tag{4}\]

Figure 2: **BNN Probabilistic Models** (a) Probabilistic model for conventional BNNs. (b) Probabilistic model for self-supervised BNNs. We share parameters between different tasks, which allows us to condition on generated self-supervised data when making predictions on a new task. \(j\) indexes self-supervised tasks, \(i\) indexes datapoints.
In words, the mean representation of each augmented image pair is the class template for each source image, which should solve the contrastive task well. The variational parameters \(\tau\) and \(\sigma^{2}\) determine the magnitude of the linear layer and the per-parameter variance, are shared across contrastive tasks \(j\), and are learnt by maximising Eq. (3) with reparameterisation gradients (Kingma and Welling, 2013).

Both the contrastive tasks and downstream task provide information about the base encoder parameters \(\theta^{s}\). One option would be to learn the base encoder parameters \(\theta^{s}\) using only data derived from \(\mathcal{D}^{u}\) (Eq. 3), which would correspond to a standard self-supervised learning setup. In this case, the learnt prior would be task-agnostic. An alternative approach is to use both \(\mathcal{D}^{t}\) and \(\mathcal{D}^{u}\) to learn \(\theta^{s}\), which corresponds to a semi-supervised setup. To do so, we can use the ELBO for the downstream data:

\[\tilde{\mathcal{L}}^{t}(\theta^{s})=\mathbb{E}_{q(\theta^{t})}[\log p( \mathcal{D}^{t}|\theta^{t},\theta^{s})]-D_{\text{KL}}[q(\theta^{t})||p(\theta ^{t})]\leq\log p(\mathcal{D}^{t}|\theta^{s}), \tag{5}\]

where \(q(\theta^{t})=\mathcal{N}(\theta^{t};\mu^{t},\Sigma^{t})\) is a variational distribution over the downstream task parameters; \(\Sigma^{t}\) is diagonal. We can then maximise \(\sum_{j}\tilde{\mathcal{L}}^{c}_{j}(\theta^{s})+\alpha\cdot\tilde{\mathcal{L} }^{t}(\theta^{s})\), where \(\alpha\) is a hyper-parameter that controls the weighting between the downstream task and contrastive task datasets. We consider both variants of our approach, using Self-Supervised BNNs to refer to the variant that pre-trains only with \(\{\mathcal{D}^{c}_{j}\}_{j=1}^{L}\) and \(\text{Self-Supervised BNNs}^{*}\) to refer to the variant that uses both \(\{\mathcal{D}^{c}_{j}\}_{j=1}^{L}\) and \(\mathcal{D}^{t}\).

Downstream Evaluation (Step II)Having learnt a point estimate for \(\theta^{s}\), we can use any approximate inference algorithm to infer \(\theta^{t}\). Here, we use a post-hoc Laplace approximation (Mackay, 1992; Daxberger et al., 2021).

Algorithm 1 summarises Self-Supervised BNNs, which learn \(\theta^{s}\) with \(\mathcal{D}^{u}\) only. We found tempering with the mean-per-parameter KL divergence, \(\bar{D}_{\text{KL}}\), improved performance, in line with other work (e.g., Krishnan et al., 2022). Moreover, we generate a new \(\mathcal{D}^{c}_{j}\) per gradient step so \(L\) corresponds to the number of gradient steps. We also regularise with \(\log p(\theta^{s})\), which here corresponds to weight decay. Finally, following best-practice for contrastive learning (Chen et al., 2020), we use a non-linear _projection head_\(g_{\psi}(\cdot)\)_only_ for the contrastive tasks. For further details, see Appendix A.

Understanding the Approximate Posterior To better understand \(q(W^{c}_{j};\tau,\sigma^{2})\) (Eq. 4), we evaluate the likelihood term of Eq. (3) at the approximate posterior's mean: \(\mu^{c}_{j}\). Define \(\tilde{z}^{A}_{i}=f_{\theta^{s}}(\tilde{x}^{A}_{i})\) and \(\tilde{z}^{B}_{i}=f_{\theta^{s}}(\tilde{x}^{B}_{i})\), and recall we have \(\omega_{i}=0.5(\tilde{z}^{A}_{i}+\tilde{z}^{B}_{i})\). Then:

\[\log p(\mathcal{D}^{c}_{j}|\theta^{s},\mu^{c}_{j})=\sum_{i=1}^{M}\Big{[}\log \frac{\exp\omega_{i}^{T}\tilde{z}^{A}_{i}/\tau}{\sum_{j=1}^{M}\exp\omega_{j}^ {T}\tilde{z}^{A}_{i}/\tau}+\log\frac{\exp\omega_{i}^{T}\tilde{z}^{B}_{i}/\tau }{\sum_{j=1}^{M}\exp\omega_{j}^{T}\tilde{z}^{B}_{i}/\tau}\Big{]}, \tag{6}\]

which we see takes the form of a (negative) normalised-temperature cross-entropy loss (NT-XENT), as used in contrastive learning algorithms like SimCLR (Chen et al., 2020). However, there are some differences between these objectives. In terms of learning the base-encoder parameters, our objective: (i) is a principled lower bound for the log-marginal likelihood, \(\log p(\mathcal{D}^{c}_{j}|\theta^{s})\); (ii) injects noise around \(\mu^{c}_{j}\), which may be a helpful regularisation (Srivastava et al., 2014); and (iii) adapts the temperature \(\tau\) throughout training, which may be beneficial (Huang et al., 2022).

Pre-training as Prior LearningSince our objective function is a principled lower bound on the log-marginal likelihood, it is similar to type-II maximum likelihood (ML), which is often used to learn parameters for deep kernels (Wilson et al., 2015) of Gaussian processes (Williams and Rasmussen, 2006), and recently also for BNNs (Immer et al., 2021). As such, similar to type-II ML, our approach can be understood as a form of prior learning. Although we learn only a point-estimate for \(\theta^{s}\), this fixed value induces a prior distribution over predictive functions through the task-specific prior \(p(\theta^{t})\)
However, while normal type-II ML learns this prior using the observed data itself, our approach maximises a marginal likelihood derived from unsupervised data. Because we maximise an ELBO during pre-training, our approach can similarly be understood as variational model learning.

## 4 How Good Are Self-Supervised BNN Prior Predictives?

We showed our approach incorporates unlabelled data into the downstream task prior predictive distribution (Eq. 2). We also argued that as the generated contrastive data reflect our subjective beliefs about the semantic similarity of different image pairs, incorporating the unlabelled data should improve the prior predictive. We now examine whether this is indeed the case.

However, although prior predictive checks are standard in the applied statistics community (Gelman et al., 1995), they are challenging to apply to BNNs due to the high dimensionality of the input space. To proceed, we note, intuitively, a suitable prior should reflect a belief that _the higher the semantic similarity between pairs of inputs, the more likely these inputs are to have the same label_. Therefore, rather than inspecting the prior predictive at single points in input space, we examine the _joint_ prior predictive of _pairs_ of inputs with known semantic relationships. Indeed, it is far easier to reason about the relationship between examples than to reason about distributions over high-dimensional functions.

Here, we focus on image classification and consider the following pairs of images: (i) an image from the validation set (the "base image") and an augmented version of the same image; (ii) a base image and another image of the same class; and (iii) a base image and an image of a different class. As these image pair groups have decreasing semantic similarity, we want the first group to be the most likely to have the same label, and the last group to be the least likely.

Experiment DetailsWe investigate the different priors on CIFAR10. For the BNN, we follow Izmailov et al. (2021) and use a ResNet-20-FRN with a \(\mathcal{N}(0,1/5)\) prior over the parameters. For the self-supervised BNN, we learn a base encoder of the same architecture with \(\mathcal{D}^{u}\) only and sample from the task prior predictive using Eq. (2). \(\theta^{t}\) are the parameters of the linear readout layer. We also pre-train a base encoder with SimCLR. For further details, see Appendix C.3.

Graphical Evaluation

First, we visualise the BNN and self-supervised BNN prior predictive (Fig. 1 and 3). The BNN prior predictive reflects a belief that all three image pair groups are similarly likely to have the same label, and thus does not capture semantic information well. In contrast, the self-supervised prior reflects a belief that image pairs with higher semantic similarity are more likely to have the same label. In particular, the self-supervised prior is able to distinguish between image pairs of the same class and of different classes, _even without access to any ground-truth labels_.

Quantitative EvaluationWe now quantify how well different prior predictives reflect data semantics. We define \(\rho(x,z)\) as the probability that inputs \(x,z\) have the same label under the prior predictive, i.e., \(\rho(x,z)=\mathbb{E}_{\theta}[p(y(x)=y(z)|\theta)]\) where \(y(x)\) is the label corresponding to input \(x\). We want a prior where image pairs with higher semantic similarity have higher values of \(\rho\). Therefore, to quantify the suitability of the prior predictive, we evaluate the frequency \(\alpha\) under the data distribution that the

\begin{table}
\begin{tabular}{l c}
**Prior Predictive** & **Prior Score \(\alpha\)** \\ \hline BNN — Gaussian & 0.261 \(\pm\) 0.024 \\ BNN — Laplace & 0.269 \(\pm\) 0.007 \\ SimCLR & **0.670 \(\pm\) 0.015** \\ Self-Supervised BNN & **0.680 \(\pm\) 0.063** \\ \end{tabular}
\end{table}
Table 1: **Prior Evaluation Scores**. Mean and standard deviation across three seeds shown. Self-supervised priors are better than standard BNN priors.

Figure 3: **BNN Prior Predictives. We investigate prior predictives by computing the probability \(\rho\) that particular image pairs have the same label under the prior, and examining the distribution of \(\rho\) across different sets of image pairs. We consider three sets of differing semantic similarity: (i) augmented images; (ii) images of the same class; and (iii) images of different classes. Left: Conventional BNN prior. Right: Self-supervised BNN learnt prior predictive. The self-supervised learnt prior reflects the semantic similarity of the different image pairs better than the BNN prior, which is reflected in the prior evaluation score, \(\alpha\).**
ranking of \(\rho\)s that correspond to the aforementioned three sets of image pairs indeed match the ranking of semantic similarities. Mathematically, we have \(\alpha=\mathbb{E}_{z,z_{13}}[\mathbb{I}[\rho(x,z_{1})>\rho(x,z_{2})>\rho(x,z_{3}))]\), where each of the \(\rho\)-terms above corresponds to a different image pair group: \(\rho(x,z_{1})\) corresponds to augmented image pairs, \(\rho(x,z_{2})\) to same-class image pairs, and \(\rho(x,z_{3})\) to different-class image pairs. \(\alpha\) is therefore a generalisation of the AUROC metric.

In Table 1, we see that conventional BNN priors reflect semantic similarity much less than self-supervised BNN priors. Further, SimCLR, a contrastive learning approach, also induces a prior predictive that well reflects semantic similarity well if the trained encoder is combined with a stochastic last-layer, which may underpin its success.

## 5 Experiments

We saw that self-supervised BNN prior predictives reflect semantic similarity better than conventional BNNs (SS4). We now investigate whether this translates to improved predictions in semi-supervised and active learning setups, which mimic scenarios where there is an abundance of unlabelled data, but labelling data points is expensive. We find that our approach combines the label efficiency of self-supervised methods and the uncertainty estimates of Bayesian approaches.

### Semi-Supervised Learning

**Experiment Details** We consider CIFAR10 and CIFAR100, reserving a validation set of 1000 examples from the test set and using the remaining 9000 examples as the test set. We evaluate the performance of different baselines when conditioning on 50, 500, 5000, and 50000 labels. All baselines use a ResNet-18 modified for CIFAR image size. For self-supervised BNNs, we use a non-linear projection head for pre-training, as standard, and use the full train set for pre-training with the data augmentations suggested by Chen et al. (2020). We consider two variants of self-supervised BNNs: Self-Supervised BNNs learn the base-encoder parameters using \(\mathcal{D}^{u}\) only, while Self-Supervised BNNs* also utilise \(\mathcal{D}^{t}\). For evaluation, we use a post-hoc Laplace approximation on the last-layer. For the conventional BNN baselines, we use MAP, SWAG (Maddox et al., 2019), a deep ensemble (Lakshminarayanan et al., 2017), and last-layer Laplace (Daxberger et al., 2021). We include standard data augmentation for the baselines, which were chosen because they support batch normalisation (Ioffe and Szegedy, 2015). We also compare to SimCLR, where we consider both the standard linear-evaluation protocol (i.e., maximum likelihood) as well as a post-hoc Laplace approximation on the last-layer. See Appendix C for further details.

Fig. 4 shows the performance of different BNNs in terms of test accuracy and expected calibration error. Self-supervised BNNs substantially outperform conventional BNNs at all but the largest data set size, a finding in line with our results that self-supervised BNN prior predictives are better than conventional BNN priors (SS4). We also find self-supervised BNNs provide well-calibrated uncertainty estimates at all dataset sizes, and that incorporating labelled data when learning the base encoder improves accuracy. Indeed, at the largest dataset size, Self-Supervised BNNs*, which leverage the labelled examples as well as unlabelled data to learn the base encoder parameters, are competitive in terms of accuracy and calibration with the deep ensemble, which is the strongest baseline but uses five trained networks. These results highlight the benefit of incorporating unlabelled data into BNNs.

Figure 4: **BNN Label Efficiency**. We compute the test accuracy and expected calibration error (ECE) when observing different numbers of labels. Self-Supervised BNNs pre-train the base-encoder parameters using \(\mathcal{D}^{u}\) only, while Self-Supervised BNNs* also use \(\mathcal{D}^{t}\). Mean and standard deviation across 3 seeds shown. Self-supervised BNNs substantially outperform conventional BNNs at most dataset sizes, highlighting the benefit of incorporating unlabelled data into the BNN framework.
Further, we evaluate the out-of-distribution generalisation performance of different BNNs. To do so, we condition BNNs on different numbers of labels of CIFAR10 as before, but evaluate the test accuracy and calibration on the CIFAR-10-C dataset (Hendrycks and Dietterich, 2019), following Ovadia et al. (2019). Again, self-supervised BNNs can benefit from unlabelled training examples, while conventional BNNs are unable to do so.

In Fig. 5, we see that self-supervised BNNs not only outperform conventional BNNs at most dataset sizes in terms of accuracy, but they consistently offer well-calibrated uncertainty estimates. On the largest dataset size, self-supervised BNNs are outperformed by SWAG and deep ensembles, which are strong baselines but use several trained networks. However, deep ensembles are poorly calibrated at small dataset sizes. These results provide further evidence for the benefits of incorporating unlabelled data into BNNs.

We now compare our approach with standard SimCLR on CIFAR10. Recall that, relative to SimCLR, our approach uses a modified pre-training objective that is a lower bound on a marginal likelihood and performs approximate inference over \(\theta^{t}\) when making predictions. We also consider combining SimCLR with approximate inference for the last-layer parameters to disentangle the effects of the pre-training objective and evaluation protocol.

In Fig. 6, Self-Supervised BNNs match SimCLR in terms of accuracy, while offering better-calibrated uncertainty estimates. All approaches have high accuracy at low data regimes, highlighting the benefit of leveraging unlabelled data. We find that incorporating \(\mathcal{D}^{t}\) when learning the base-encoder (Self-Supervised BNNs*) substantially improves accuracy, though surprisingly, it slightly hurts calibration. In contrast, the Laplace evaluation protocol improves calibration, and indeed, SimCLR combined with approximate inference is a strong baseline. This matches our previous observation that SimCLR pre-training and a prior over \(\theta^{t}\) induce a prior predictive that well captures image semantics (SS4).

### Active Learning

Experiment DetailsWe consider low-budget active learning, which simulates a scenario where labelling examples is extremely expensive. We use the CIFAR10 training set as the unlabelled pool set from which to label points. We assume an initial train set of 50 labelled points, randomly selected, and a validation set of the same size. We acquire 10 labels per acquisition round up to 500 labels and evaluate using the full test set. We compare self-supervised BNNs, SimCLR, and a deep ensemble. We use BALD (Houlsby et al., 2011) as the acquisition function for the deep ensemble and self-supervised BNN, because they provide epistemic uncertainty estimates, while we use predictive entropy for SimCLR. We also consider uniform selection. See Appendix C for more details.

Figure 5: **Out-of-Distribution Generalisation**. We train BNNs on the standard CIFAR10 dataset and evaluate on CIFAR-10-C with corruption intensity 5 (Hendrycks and Dietterich, 2019). Results averaged across corruptions, mean and std. across 3 seeds shown. Self-Supervised BNNs pre-train the base-encoder parameters using \(\mathcal{D}^{u}\) only, while Self-Supervised BNNs* also use \(\mathcal{D}^{t}\). Self-supervised BNNs outperform conventional BNNs at most dataset sizes.

Figure 6: **Comparison to SimCLR** on CIFAR10. We compare self-supervised BNNs with (i) standard SimCLR, and (ii) SimCLR with a post-hoc Laplace approximation over the last-layer. Self-Supervised BNNs pre-train the base-encoder parameters using \(\mathcal{D}^{u}\) only, while Self-Supervised BNNs* also utilise \(\mathcal{D}^{t}\) during pre-training. Approximate inference over the task-specific parameters improves calibration, while pre-training with both \(\mathcal{D}^{t}\) and \(\mathcal{D}^{u}\) improves accuracy.
In Fig. 7, we see that the methods that leverage unlabelled data perform the best. In particular, the self-supervised BNN with BALD selection achieves the highest accuracy across most numbers of labels, and is the only method for which active learning improves performance. Surprisingly, uniform selection is a strong baseline.

## 6 Related Work

Improving BNN PriorsWe demonstrated that BNNs have poor prior predictive distributions (SS4), a concern shared by others (e.g., Wenzel et al., 2020; Noci et al., 2021; Izmailov et al., 2021). The most common approaches to remedy this are through designing better priors, typically over network parameters Louizos et al. (2017); Nalisnick, 2018; Atanov et al., 2019; Fortuin et al., 2021) or predictive functions directly Sun et al. (2019); Tran et al. (2020); Matsubara et al. (2021), see Fortuin (2022) for an overview). In contrast, our approach incorporates vast stores of unlabelled data into the prior predictive distribution through variational model learning. Similarly, other work also _learns_ priors, but typically using labelled data e.g., by using meta-learning Garnelo et al. (2018); Rothfuss et al. (2021) or type-II maximum likelihood Wilson et al. (2015); Immer et al. (2021). Finally, Shwartz-Ziv et al. (2022) use generic transfer learning, potentially from an unsupervised task, for better BNN priors. While they consider fully stochastic networks, we use partially stochastic networks with a pre-training objective that is a principled lower bound on a log-marginal likelihood derived from unlabelled data. A further difference is that our work incorporates pre-training within the probabilistic modelling framework itself.

Understanding Contrastive LearningOur work shows that contrastive pre-training induces a prior predictive distribution that captures semantic similarity well (SS4), and further offers a Bayesian interpretation of contrastive learning (SS3). There has been much other work on understanding contrastive learning (e.g., Wang and Isola, 2020; Wang and Liu, 2021). Some work appeals to the InfoMax principle Becker and Hinton (1992), which maximises the mutual information between _representations_ of two views of an input datum, while our framework operates in predictive space. Zimmermann et al. (2022) argue that contrastive learning inverts the data-generating process, while Aitchison (2021) cast InfoNCE as the objective of a self-supervised variational auto-encoder.

Semi-Supervised Deep Generative ModelsDeep generative models (DGMs) are an alternative approach for label-efficient learning Kingma and Welling (2013); Kingma et al. (2014); Joy et al. (2020). DGMs generate supervision by reconstructing data through compact representations, i.e., they are _generative_. Our framework, however, is _discriminative_--supervision works through pseudo-labelled tasks. Ganev and Aitchison (2021) formulate several semi-supervised learning objectives as lower bounds of log-likelihoods in a probabilistic model of data curation. Finally, Sansone and Manhaeve (2022) unify self-supervised learning and generative modelling under one framework.

## 7 Conclusion

We introduced _Self-Supervised Bayesian Neural Networks_, which leverage unlabelled data for improved prior predictive distributions (SS3). We showed that self-supervised BNNs, as well as other self-supervised methods, learn prior predictives that reflect semantics of the data (SS4). In our experiments, self-supervised BNNs combine the label-efficiency of self-supervised methods and the uncertainty estimates of Bayesian approaches (SS5), and thus are a principled, practical, and performant alternative to conventional BNNs. We strongly encourage practitioners to incorporate unlabelled data into their Bayesian Neural Networks through unsupervised prior learning.

Figure 7: **Low-Budget Active Learning** on CIFAR10. We compare: (i) a self-supervised BNN; (ii) SimCLR; and (iii) a deep ensemble. For the self-supervised BNN and the ensemble, we acquire points with BALD; predictive entropy for SimCLR. We also use uniform sampling for all methods. Mean and std. shown (3 seeds). The methods that incorporate unlabelled data perform best, and self-supervised BNNs are the only method where active learning outperforms uniform sampling.
## Acknowledgements

MS was supported by the EPSRC Centre for Doctoral Training in Autonomous Intelligent Machines and Systems (EP/S024050/1), and thanks Rob Burbea for inspiration and support. VF was supported by a Postdoc Mobility Fellowship from the Swiss National Science Foundation, a Research Fellowship from St John's College Cambridge, and a Branco Weiss Fellowship.
