U NCERTAINTY -A WARE S ELF -S UPERVISED L EARNING WITH I NDEPENDENT S UB -N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
Self-supervised learning methods are state-of-the- art across a wide range of tasks in computer vision natural language processing
and multimodal analysis
Esti- mating the epistemic – or model – uncertainty of self-supervised model predic tions is critical for building trustworthy machine learning systems in crucial applications, such as medical diagnosis and autonomous driving.
Ensembling a neural network is an effective approach to estimating uncertainty.
However, deep ensembles induce high computational costs and memory demand.
This is even more challenging in the case of self-supervised deep learning, where even a single network is computationally demanding.
Moreover, most existing model uncertainty techniques are built for supervised deep learning.
Motivated by this, we propose a novel approach to making self-supervised learning probabilistic
We introduce an uncertainty-aware training regime for self-supervised models with an ensemble of independent sub-networks and a novel loss function for encouraging diversity
Our method builds a sub-model ensemble with high diversity
– and consequently well calibrated estimates of model uncertainty
– at low computational overhead over a single model, while performing on par with deep self-supervised ensembles.
Extensive experiments across different tasks, such as in distribution
gen
- eralization, out
-of-distribution detection, dataset corruption and semi-supervised settings, demonstrate that our approach increases prediction reliability.
We show that our method achieves both excellent accuracy and calibration, improving over existing ensemble methods in a wide range of self-supervised architectures for computer vision, natural language processing and genomics data
I NTRODUCTION
Estimating uncertainty is an important component for developing reliable and trustworthy machine learning algorithms used to make predictions for important decisions, such as medical diagnosis (Azizi et al., 2021), drug discovery (Jim ´ enez-Luna et al., 2020), and autonomous driving (Kaya et al., 2022;
Can et al., 2022).
In recent years, self-supervised learning methods have achieved cuttingedge performance across a wide range of tasks in natural language processing (NLP; (
Devlin et al., 2018; Brown et al., 2020), computer vision (Chen et al., 2020; Bardes et al., 2021;
Grill et al., 2020), multimodal learning (Radford et al., 2021; Li et al., 2022; Shi et al., 2022), and bioinformatics (G¨und¨uz et al., 2021).
In contrast to supervised techniques, these models learn representations of the data without relying on costly human annotation.
Despite remarkable progress in recent years, self-supervised models do not allow practitioners to inspect the model’s confidence.
This problem is non-trivial, given the degree to which critical applications rely on self-supervised methods.
Therefore, quantifying the predictive uncertainty of self-supervised models is key to more reliable downstream tasks.
Here, we follow the definition of reliability as described by Plex (Tran et al., 2022), in which the ability of a model to work consistently across many tasks is assessed.
In particular, Tran et al.
(2022) introduces three general desiderata of reliable machine learning systems: a model should generalize robustly to new tasks , as well as new datasets , and represent the associated uncertainty in a faithful manner.
Predictive uncertainty is classified as epistemic (model) and aleatoric (data) uncertainty.
Aleatoric uncertainty is inherent to the data and arises from natural stochasticity in observations, for instance, due to class overlap, label noise, or multi-label disagreement.
On the other hand, epistemic uncer- tainty concerns the learning hypothesis parameterized through model weights.
It is caused by limited data or knowledge and is therefore reducible.
Addressing this source of uncertainty is crucial in order to build methods that generalize robustly to unseen data and avoid wrong, overconfident predictions.
Importantly, evaluating the quality of predictive uncertainties is challenging, as the ‘ground-truth’ uncertainty is usually not available.
Bayesian neural networks (BNNs; Neal (2012)) and neural ensemble networks (Hansen & Salamon, 1990) are the common approaches for capturing distributions over parameters of neural networks.
While the Bayesian paradigm, with its principled approach to inference, is the canonical answer to the problem of uncertainty quantification, it is not well-suited to self-supervised methods.
Bayesian deep learning, which usually requires samples from the posterior distribution, scales poorly to the large architectures that prevail in this field and relies fundamentally on the availability of true labels y for maximum-a-posterior (MAP) estimation.
During pretext-task learning (pretraining) in self-supervised and unsupervised problems, we do not have access to such label information.
This prevents us from taking advantage of BNNs during the pretraining of self-supervised models.
On the other hand, deep ensembles (Lakshminarayanan et al., 2017) have been proposed as a heuristic alternative to Bayesian learning.
By adding additional parameter instances through training multiple networks in parallel, deep ensembles make learning probabilistic.
Despite their lack of theoretical foundation, deep ensembles have been successful in practice and become gold standard in estimating uncertainty of neural networks.
The high performance of deep ensembles, however, comes at the expense of drastically increased costs.
Subsequent work, such as BatchEnsemble (Wen et al., 2020), Masksemble (Durasov et al., 2021), MIMO (Havasi et al., 2021), and FiLMEnsemble (Ozgur Turkoglu et al., 2022) has attempted to address this problem by creating ensembles more efficiently.
While these models succeed at speeding up training and/or inference, they exhibit a substantial performance gap compared to the deep ensemble, lacking both prediction quality and calibration.
In this work, we introduce a simple and scalable framework for incorporating predictive uncertainty estimates into the self-supervised learning approach while preserving performance with a negligible increase in computational cost.
Our contributions can be summarized as follows:
• We propose a novel way to extend self-supervised learning to be probabilistic, which enables self-supervised models to provide uncertainty information for predictions.
• We develop a complementary uncertainty-aware loss function to enforce diversity among the independent sub-networks.
• We perform extensive empirical analyses to highlight the benefits of learning probabilistic representations.
We demonstrate that this inexpensive probabilistic modification achieves very competitive (in most cases, better) predictive performance: 1) on in-distribution (IND) and out-of-distribution (OOD) tasks; 2) in semi-supervised settings; 3) learns a better predictive performance-uncertainty trade-off than compared baselines (i
e., exhibits high predictive performance and low uncertainty on IND datasets as well as high predictive performance and high uncertainty on OOD datasets).
R ELATED WORK Self-supervised learning
For most large-scale modeling problems, learning under full supervision is severely inhibited by the scarcity of annotated samples.
Self-supervised learning techniques, which solve pretext tasks (Devlin et al., 2018) to generate labels from (typically abundant) unlabeled data, have proven to be a powerful remedy to this bottleneck.
The learned feature maps can serve as a starting point for downstream supervised tasks, such as classification, object detection or sentiment analysis, with a substantially reduced need for labeled examples (Jaiswal et al., 2020).
Alternatively, the downstream application may directly use the extracted representation for problems such as anomaly OOD detection.
Self-supervision has been shown to enhance robustness against numerous sources of corruption and to even improve on purely supervised approaches in certain tasks (Hendrycks & Dietterich, 2019).
However, its success depends on the quality of pretraining, for which (Tran et al., 2022) recently proposed a number of desiderata.
While there have been attempts to make pretraining more robust by preventing embedding collapse (Bardes et al., 2021) or boosting performance in OOD detection (Winkens et al., 2020; Sehwag et al., 2021), the aspect τ 𝑓 θ 𝐶𝑜𝑛𝑐𝑎𝑡(𝐸 𝑍 ,𝐸 𝑍 ′ )
ℓ 𝑆𝑆𝐿 + λ ℓ 𝑑𝑖𝑣 τ′
Self-supervised Network Contrastive/Self-distillation/BERT
Ensemble of Independent Sub-Networks 𝑔𝑓 𝜃 𝑧 1,..,𝑀 𝑔𝑥 𝑥′
𝐼
Embedding 𝑔 𝑀 𝑔 𝑀 𝐶𝑜𝑛𝑐𝑎𝑡(𝑆𝑡𝑑 𝑍 ,𝑆𝑡𝑑 𝑍 ′ ) 𝑧′ 1,..,𝑀
Representation 𝑋 of uncertainty-awareness has been studied to a lesser extent in the self-supervised context.
Motivated by this, we present a simple way to make self-supervised learning uncertainty-aware during pretext-task learning.
Model uncertainty estimation and uncertainty-awareness Uncertainty can be viewed as a state of missing information in which it is impossible to exactly describe an outcome of interest (Ghahramani, 2015).
In predictive modeling, awareness of such uncertainty is obviously desirable.
Uncertainty estimation helps practitioners assess the trustworthiness of model outcomes, thereby permitting suitable interventions when the model is likely to be wrong.
Bayesian learning handles uncertainty in a natural manner but is not applicable to self-supervised learning in the absence of labeled data.
For example, SNGP (Liu et al., 2020) creates uncertainty-awareness via the preservation of input-space distances.
Subsequent studies, including modifications using Mahalanobis instead of
Euclidean distances (Vazhentsev et al., 2022), suggest that SNGP works well in representation learning for computer vision and NLP tasks.
Crucially, however, the underlying single-model approach is inherently unable to capture the epistemic part of uncertainty.
The stochastic embeddings learned in (Park et al., 2022) resolve this but depend on the use of contrastive learning, as does loss-temperature scaling (Zhang et al., 2021), and a number of Gaussianity assumptions.
(Hu & Khan, 2021) promote a more holistic approach rooted in evidential learning yet rely on additional datasets from outside the training manifold.
Sidestepping such shortcomings, we propose to employ the well-established ensembling approach to enforce uncertainty-awareness in a simple and effective manner.
Ensemble learning Deep Ensembles (Lakshminarayanan et al., 2017) comprise a set of M neural networks that independently train on the same data using random initialization.
Deep ensembles often outperform other approaches in terms of calibration and predictive accuracy (Ovadia et al., 2019; Gustafsson et al., 2020; Ashukha et al., 2020), but their naive application incurs high computational complexity, as training, memory, and inference cost multiplies with the number of base learners.
BatchEnsemble (Wen et al., 2020) introduces multiple low-rank matrices with little training and storage demand, whose Hadamard products with a shared global weight matrix mimic an ensemble of models.
FiLM-Ensemble (Ozgur Turkoglu et al., 2022) uses feature-wise linear modulation to construct an ensemble for uncertainty estimation.
Masksensemble (Durasov et al., 2021) builds upon Monte Carlo dropout (Gal & Ghahramani, 2016) and proposes a learnable (rather than a random) selection of masks used to drop certain network neurons.
MIMO (Havasi et al., 2021) uses ensembles of sub-networks diverging only at the beginning and end of the parent architecture – thus sharing the vast majority of weights – in order to obtain multiple predictions with a single forward pass.
At test time, several copies of each sample are fed to the enlarged input layer, and the multi-head last layer returns an according number of predictions.
Although these methods reduce the inference time and computational resources required at training, the benefits are limited to larger pretraining model that is used in self-supervised learning.
M ETHOD
We propose a simple principle to: 1) estimate uncertainty during pretraining of self-supervised deep learning, and 2) improve predictive uncertainty.
As depicted in Figure.
1, our proposed method can be readily applied to the most recent trends in self-supervised learning (Caron et al., 2021; Grill et al., 2020; Chen et al., 2020; Devlin et al., 2019; G¨und¨uz et al., 2021; Klein & Nabi, 2022) and is based on a joint embedding architecture.
In the following sections, we first describe our proposed uncertainty-aware self-supervised model, followed by the diversity loss, and then discuss ensemble diversity and computational cost.
Uncertainty-aware self-supervised learning via independent sub-networks
Given a randomly sampled mini-batch of data X = { x k } Nk
=1 , x ∈ X ⊆ R p , the transformer function derives two augmented views ˜ x
= τ ( x ) , ˜ x ′ = τ ′ ( x ) for each sample in X .
The augmented views are obtained by sampling τ, τ ′ from a distribution over suitable data augmentations, such as masking parts of sequences (Baevski et al., 2022; Devlin et al., 2019), partially masking image patches (He et al., 2022), or applying image augmentation techniques (Chen et al., 2020).
The two augmented views ˜ x and ˜ x ′ are then fed to an encoder network f θ with trainable parameters θ ⊆ R d .
The encoder (e.g., ResNet-50 (He et al., 2016), ViT (Dosovitskiy et al., 2020)) maps distorted samples to a set of corresponding features.
We call the output of the encoder the representation .
Afterwards, the representation features are transformed by M independent sub-networks { g ϕ m } Mm =1 with trainable parameters ϕ m .
The ensemble constructs from the representation M different q -dimensional embedding vectors { z m } Mm =1 , { z ′ m } Mm =1 , respectively, for ˜ x and ˜ x ′ .
We modify the conventional self-supervised loss and replace the usual z m , z ′ m with the mean values ¯ z , ¯ z ′ over m = 1 , . . .
, M .
Averaging over the embeddings generated by the M sub-networks improves predictive performance.
For example, in the case of contrastive learning (Chen et al., 2020), the self-supervised loss ℓ ssl with temperature t > 0 and cosine similarity sim( · , · ) is computed as:
[Equation]
Diversity is a key component of successful model ensembles (Fort et al., 2019).
To this end, we design a new loss function for encouraging diversity during the training of the sub-networks.
We define the diversity regularization term ℓ div as a hinge loss over the difference of the standard deviation across the embedding vectors { z k,m } Mm =1 , { z ′ k,m } Mm =1 to a minimum diversity of α > 0 .
The standard deviation is the square root of the element-wise variance { σ 2 k,o } qo =1 :
[Equation]
where ϵ > 0 is a small scalar preventing numerical instabilities.
The diversity regularization function is then given by:
[Equation]
[Equation]
The objective of the diversity loss is to encourage disagreement among sub-networks by enforcing the element-wise standard deviations to be close to α > 0 and to thus prevent the embeddings from collapsing to the same vector.
Figure 2a motivates the importance of the diversity loss on the total sum of standard deviation between different sub-networks, which increases by adding the diversity loss.
The total loss is calculated by combining the self-supervised loss (Eq. 1) and the diversity loss (Eq. 3), where the degree of regularization is controlled by a tunable hyperparameter λ ≥ 0 :
[Equation]
[Equation]
Finally, the total loss is aggregated over all the pairs in minibatch X :
[Equation]
[Figure]
t
[Figure]
i
[Figure]
Diversity analysis
Diversity is an important feature for powerful model ensembles and reflects the degree of independence among its members (Zhang & Ma, 2012).
We follow (Fort et al., 2019) to quantify the diversities among the ensemble of sub-networks.
Specifically, we report the diversities in terms of disagreement score between the members’ predictive distributions.
Diversity disagreement is defined as the fraction of distance disagreement divided by 1 - accuracy , where the distance disagreement between two classification models h i and h j is calculated as
[Equation]
, with N denoting the number of samples.
Figure 2b compares the diversity disagreement between our method, a deep ensemble
[Equation]
and the single-network baseline.
The results clearly indicate that our proposed method achieves comparable results with deep selfsupervised ensembles in terms of both accuracy and diversity disagreement.
Computational cost We analyze the efficiency of our proposed method in terms of the resources needed for training.
Table.
1 compares the relative cost of our proposed method with the selfsupervised baseline and ensembles in terms of the number of parameters, memory demand, and computation time.
A more detailed description of relative cost is provided in Figure.
6 Appendix A
[Equation]
[Table]
E XPERIMENTAL SETUP
We perform several experiments with a variety of self-supervised methods to improve reliability during both pretext-task learning and downstream tasks (fine-tuning).
The source code and our models are provided in https://anonymous.4open.science/r/Uncertainty_aware_
SSL-95D7/README.md
Deep self-supervised network architecture Our proposed approach builds on two recent popular self-supervised models in computer vision: i) SimCLR (Chen et al., 2020) is a contrastive learning framework that learns representations by maximizing agreement on two different augmentations of the same image, employing a contrastive loss in the latent embedding space of a convolutional network architecture (e.g., ResNet-50 (He et al., 2016)), and ii) DINO (Caron et al., 2021) is a self-distillation framework in which a student vision transformer (ViT; Dosovitskiy et al.
(2021)) learns to predict global features from local image patches supervised by the cross-entropy loss from a momentum teacher ViT’s embeddings.
Furthermore, we study the impact of our approach in NLP and modify SCD (Klein & Nabi, 2022), which applies the bidirectional training of transformers to language modeling.
Here, the objective is self-supervised contrastive divergence loss.
Lastly, we implement a probabilistic version of Self-GenomeNet (G¨und¨uz et al., 2021), a contrastive self- supervised learning algorithm for learning representations of genome sequences.
More detailed descriptions of the employed configurations are provided in Appendix A.3.
Deep independent sub-networks We implement M independent sub-networks on top of the encoder, for which many possible architectures are conceivable.
For our experiments on computer vision datasets, we consider a convex architecture where each network includes a multi-layer perceptron (MLP) with two layers of 2048 and 128 neurons, respectively, with ReLU as a non-linearity and followed by batch normalization (Ioffe, 2017).
Each sub-network has its own independent set of weights.
For the NLP dataset, the projector MLP contains three layers of 4096 neurons each, also using ReLU activation’s as well as batch normalization.
For the genomics dataset, our ensemble of sub-networks includes one fully connected layer with an embedding size of 256.
Optimization For all experiments on image datasets based on DINO and SimCLR, we follow the suggested hyperparameters and configurations by the paper (Caron et al., 2021; Chen et al., 2020).
Implementation details for pretraining with DINO on the 1000-classes ImagetNet dataset without labels are as follows: coefficients ϵ , α , and λ are respectively set to 0 .
0001 , 0 . 15 , and 2 in Eq.2, 3, and 4.
We provide more details in ablation studies (Section 6) on the number of sub-networks and the coefficients λ and α used in the loss function.
The encoder network f θ is either a ResNet50 (He et al., 2016) with 2048 output units when the baseline is SimCLR (Chen et al., 2020) or ViTs (Dosovitskiy et al., 2020) with 384 output units when the baseline is DINO (Caron et al., 2021).
The best prediction and calibration performance is achieved when the number of sub-networks is 5.
The training protocol follows the suggested settings by Caron et al.
Datasets We use the following datasets in our experiments: CIFAR-10/100 (Krizhevsky, 2009) are subsets of the tiny images dataset.
Both datasets include 50,000 images for training and 10, validation images of size 32 × 32 with 10 and 100 classes, respectively.
SVH (Netzer et al., 2011) is a digit classification benchmark dataset that contains 600,000 32 × 32 RGB images of printed digits (from 0 to 9) cropped from pictures of house number plates.
ImageNet (Deng et al., 2009), also known as ILSVRC 2012, contains 1,000 classes, with 1.28 million training images and 50, validation images.
For the NLP task, we train on a dataset of 1 million randomly sampled sentences from Wikipedia articles (Huggingface, 2021) and evaluate our models on 7 different semantic textual similarity datasets from the SentEval benchmark suite (Conneau & Kiela, 2018): MR (movie reviews), CR (product reviews), SUBJ (subjectivity status), MPQA (opinion-polarity), SST-2 (sentiment analysis), TREC (question-type classification), and MRPC (paraphrase detection).
The
T6SS effector protein dataset is a public real-world bacteria dataset (SecReT6, (Li et al., 2015)) with actual label scarcity.
The sequence length of the genome sample is 1000nt in all experiments.
Tasks We examine and benchmark a model’s reliability on different tasks considering evaluation protocols by self-supervised learning (Chen et al., 2020) and Plex’s benchmarking tasks (Tran et al., 2022).
Specifically, we evaluate our model on the basis of uncertainty-aware IND generalization ,
OOD detection , semi-supervised learning , corrupted dataset evaluation (see Section 5), and transfer learning to other datasets and tasks (see Section B.2) .
Evaluation metrics
We report prediction/calibration performance with the following metrics, where upward arrows indicate that higher values are desirable, et vice versa .
Top accuracy
: share of test observations for which the correct class is predicted.
AUROC ↑ : area under the ROC curve arising from different combinations of false-positive and false-negative rates (here: with positive and negative classes referring to being in and out of distribution, respectively) for a gradually increasing classification threshold.
Negative log-likelihood (NLL) ↓ : negative log-likelihood of test observations under the estimated parameters.
Expected calibration error (ECE) ;(Naeini et al., 2015) ↓ : mean absolute difference between accuracy and confidence (highest posterior probability among predicted classes) across equally-spaced confidence bins, weighted by relative number of samples per bin.
Thresholded adaptive calibration error (TACE) ; (Nixon et al., 2019)) ↓ : modified ECE with bins of equal sample size, rather than equal interval width, and omitting predictions with posterior probabilities falling below a certain threshold (here: 0.01) that often dominate the calibration in tasks with many classes.
Compared methods
We compare our method to the following contenders.
Baseline: selfsupervised architectures (i.e., SimCLR, DINO, SCD, or Self-GenomeNet, depending on the task) without probabilistic modification.
Ensemble-SSL : deep ensemble comprising multiple of the aforementioned baseline networks.
Monte Carlo (MC) dropout : (Gal & Ghahramani, 2016) baseline networks with dropout regularization applied during pretraining of baseline encoder.
BatchEnsemble : baseline encoder with BatchEnsemble applied during pretraining.
R ESULTS AND DISCUSSION
In-distribution generalization IND generalization (or prediction calibration ) quantifies how well model confidence aligns with model accuracy.
We perform several experiments on small and large image datasets as well as the genomics dataset to evaluate and compare the predictive performance of our proposed model in IND generalization.
Here, the base encoder f θ is frozen after unsupervised pretraining, and the model is trained on a supervised linear classifier.
The linear classifier is a fully connected layer followed by softmax, which is placed on top of f θ after removing the ensemble of sub-networks.
High predictive scores and low uncertainty scores are desired.
Figure 3 illustrates the predictive probability of correctness for our model on CIFAR-10, CIFAR100, ImageNet, and T6SS dataset in terms of Top-1 accuracy, ECE, and NLL, respectively.
Based on Figure 3, our method achieves better calibration (ECE and NLL) than the deep ensemble of selfsupervised models, MC-dropout, and BatchEnsemble, with substantial margins at large ensemble sizes.
More detailed descriptions are provided in Appendix B.1 (Tables 5, 6, 7, and 8).
[Figure]
[Table]
Here, we compare our method with BatchEnsemble (BatchEns), MC-dropout (MCDO), ensemble of deep self-supervised networks (SSL-Ens), as well as the baseline.
Detailed descriptions of IND generalization for each dataset are presented in Appendix B.1 (Tables 5, 6, 7, and 8).
Out-of-distribution detection OOD detection shows how well a model can recognize test samples from the classes that are unseen during training (Geng et al., 2020).
We perform several experiments to compare the model generalization from IND to OOD datasets
[Table]
and to predict the uncertainty of the models on OOD datasets.
Evaluation is performed directly after unsupervised pretraining without a fine-tuning step.
Table 2 shows the AUROC on different OOD sets for our model, baseline and deep self-supervised ensemble.
Our approach improves over all compared methods
[Table]
[Table]
Semi-supervised evaluation We explore and compare the performance of our proposed method in the low-data regime.
Again, the encoder f θ is frozen after self-supervised pretraining
[Table]
and the model is trained on a supervised linear classifier using 1% and 10% of the dataset.
The linear classifier is a fully connected layer followed by softmax.
Table 3 shows the result in terms of top-1 accuracy,
ECE
[Table]
and NLL.
The results indicate that our method considerably outperforms other methods in the low-data regime – both in terms of calibration and accuracy.
[Table]
Corrupted dataset evaluation
Another important component of model robustness is its ability to make accurate predictions when the test data distribution changes.
Here, we evaluate model robustness under covariate shift.
We employ a configuration similar to the one found in Tran et al.
Figure 4 summarizes the improved performance across metrics of interest.
The results confirm that our method outperforms the baseline and achieves comparable predictive performance as a deep self-supervised ensemble – both in terms of calibration (TACE) and AUROC.
[Figure]
[Table]
[Figure]
[Table]
[Figure]
A BLATION STUDY
In order to build intuition around the behavior and the observed performance of the proposed method, we further investigate the following aspects of our approach in multiple ablation studies exploring: (1) the number M of sub-networks, (2) the role of each component of the proposed loss, and (3) the impact of our approach at the time of pretraining vs. at the finetuning step.
Number of sub-networks We train M individual deep neural networks on top of the representation layer.
The networks receive the same inputs but are parameterized with different weights and biases.
Here, we provide more details regarding our experiments on IND generalization by considering varying M .
Fig.
5a compares the performance in terms of top-1 accuracy, ECE, and NLL for
CIFAR-10 and CIFAR-100.
Based on the quantitative results depicted in Fig.
5a, the predictive performance improves in both datasets when increasing the number of sub-networks ( M ) until a certain point.
For example, in the case of CIFAR-10, when M = 3 , our performance is 91 .
9% ; increasing M to 10 levels top-1 accuracy up to 92 .
6% , while the ECE and NLL decrease from 0 .and 0 .
249 to 0 .
023 and 0 .
222 , respectively.
These findings underline that training our sub-networks with a suitable number of heads can lead to a better representation of the data and better calibration.
Analysis of loss
The total loss (Eq. 4) is calculated by the combination of self-supervised loss (Eq. 1) and diversity loss (Eq. 3), where the mean value of the embeddings across the ensemble of
[Figure]
sub-networks is fed to the self-supervised loss, and the corresponding standard deviation is used for the diversity loss.
First, we note that the use of our diversity regularizer indeed improves uncertainty awareness.
The results in Fig.
3 show the impact of our loss function in relation to the baseline.
By comparing the first and fifth rows of Table
[Figure]
it can be inferred that our proposed loss function results in a much lower ECE ( 0 . 016 ) than the network trained by SimCLR (baseline) with 0 .
039 on
CIFAR-10 dataset.
Similarly, the first and third rows of Table 7 compare the predictive probability of correctness of DINO (baseline) and our model on ImageNet.
Second, we explore different hyperparameter configurations to find the optimal values for α and λ in Fig.
5b,5c.
Note that, in practice, α and λ must be optimized jointly.
The best top-1 accuracy in our case is achieved when α and λ are set to 0.08 and 1.5, respectively, on the CIFAR-10 dataset.
Efficient ensemble of sub-networks at pretraining vs. finetuning We performed additional experiments to study the impact of efficient ensemble and uncertainty-aware loss i) during pretraining, ii) during finetuning, and iii) during both pretraining and finetuning.
As shown in Table 4, pretraining with an ensemble of sub-networks is beneficial, and additional fine-tuning with multiple head can further improve performance.
[Table]
C ONCLUSION
In this paper, we presented a novel uncertainty-aware self-supervised framework.
We achieved high predictive performance and good calibration using a simple yet effective idea – an ensemble of independent sub-networks that can estimate model uncertainty during pretraining of a self-supervised framework.
We introduced a new loss function to encourage diversity among different sub-networks.
It is straightforward to add our method to many existing self-supervised learning frameworks during pretraining.
Our extensive experimental results show that our proposed method outperforms, or is on par with, an ensemble of self-supervised methods in many different experimental settings.
A I MPLEMENTATION DETAILS
[Figure]
Figure 6 illustrates relative computation cost – as compared to the baseline – in terms of the number of parameters, computation time, and memory required between our model and a deep selfsupervised ensemble.
[Figure]
D ATA AUGMENTATION FOR COMPUTER VISION DATASETS
We define a random transformation function T that applies a combination of crop, horizontal flip, color jitter, and grayscale.
Similar to Chen et al.
, we perform crops with a random size from 0 .
2 to 1 .
0 of the original area and a random aspect ratio from 3 / 4 to 4 / 3 of the original aspect ratio.
We also apply horizontal mirroring with a probability of 0 . 5 .
Then, we apply grayscale with a probability of 0 .
2 as well as color jittering with a probability of 0 .
8 and a configuration of (0 . 4 , 0 . 4 , 0 . 4 , 0 . 1) .
However, for ImageNet, we define augmentation based on the original DINO from their official repository.
In all experiments, at the testing phase, we apply only resize and center crop.
[Figure]
H YPERPARAMETERS FOR SELF - SUPERVISED NETWORK ARCHITECTURES
SimCLR (Chen et al., 2020): we use ResNet-50 as a backbone, a loss temperature of 0.07, batch size 512, and a cosine-annealing learning rate scheduler.
The embedding size is 2048, and we train for 800 epochs during pretraining.
DINO (Caron et al., 2021): we use ViT-small as a backbone, patch size 16, batch size 1024, and a cosine-annealing learning rate scheduler.
The embedding size is 384/1536, and we train for 100 epochs during pretraining.
B A DDITIONAL RESULTS B. U NCERTAINTY - AWARENESS IND GENERALIZATION
Tables 5, 6, and 7 present results for the predictive performance and calibration of our model on
CIFAR-10, CIFAR-100, and ImageNet, respectively.
Based on Table 5, our method achieves better calibration than the deep ensemble of self-supervised networks, MC-dropout, and BatchEnsemble, with substantial margins at large ensemble sizes.
[Table]
[Table] [Table]
[Table]
We also performed experiments on a dataset of 1-dimensional genomic sequences – the T6SS identification of effector proteins– to demonstrate that uncertainty-aware subnetworks can also be readily combined with existing models for 1-dimensional datasets and models.
Based on Table 8, our method improves the accuracy and the calibration compared to the baseline
[Table]
B. T RANSFER TO OTHER TASKS AND DATASET
We further assess the generalization capacity of the learned representation on learning a new task in NLP.
We train our model without any labels on a dataset of sentences from Wikipedia (Huggingface, 2021) and fine-tune the pretrained representation on seven different semantic textual similarity datasets from the SentEval benchmark suite (Conneau & Kiela, 2018): MR (movie reviews), CR (product reviews), SUBJ (subjectivity status), MPQA (opinion-polarity), SST-2 (sentiment analysis), TREC (question-type classification)
[Figure]
and MRPC (paraphrase detection).
Then, we evaluate on the test set of each dataset.
Figure 7 provides a comparison of transfer learning performance of our self-supervised approach for different tasks.
Our results in Figure 7 indicate that our approach performs comparably to or better than the baseline method
[Figure]
We test the performance of the trained model on ImageNet (Deng et al., 2009) on CIFAR10 Krizhevsky (2009) dataset where the model is trained for 100 epochs.
[Figure] [Table]