Revisiting Pretraining for Semi-Supervised Learning in the Low-Label Regime

The success of many computer vision tasks depends on training deep neural networks with large amount of labeled data which is expensive to acquire. Semi-supervised learning (SSL) aims to address learning with small labeled data and large unlabeled data [1, 2, 22, 20, 23, 3].

Transfer learning aims to reuse the data or model from source task/domain to target task/domain [27]. One of the key challenges in transfer learning is the domain shift between source and target datasets and domain adaptation aims to tackle this issue. When both source and target data are available during model training, UDA [28] introduced a discriminator to align the distribution between source and target domains. Nevertheless, the assumption of access to target domain data during model training restricts the application and recently source-free domain adaptation, a.k.a. test time training, emerges as a solution which no longer requires the access to labeled data during adaptation [29, 30, 31] at the cost of modifying the learning objective during training. To avoid altering training objective, which could be very expensive for large-scale pretraining tasks, Tent [32] adopted a self-training like approach to adapt only the affine transformation parameters for batchnorm layers by minimizing an entropy loss. Despite requiring no changes to training objective, Tent requires the prediction task to be consistent between source and target domains. In contrast to source-free domain adaptation and test time training, our target pretrain does not modify training objective in the source domain nor assume the prediction task to be identical between the source and target domains. A recent related work addressed finetuning on target domain in an unsupervised manner [33]. To avoid large deviations from pretrained model parameters they simultaneously finetune on both target and source domain data. Therefore, the adaptation still requires access to source domain data.

Constrative learning [4, 10, 11, 12] is particularly effective in pretraining backbone network on large unlabeled dataset by enforcing the image-wise features subject to two augmentation to be consistent. As such the network is able to attend to the discriminative areas for feature extraction. To focus on dense prediction downstream tasks, e.g. segmentation, detection, etc., contrasting pixel-wise features was proposed [34, 35] to train the backbone network to generate local features invariant to augmentations. Contrastive learning has been demonstrated to be friendly to data-efficient learning [12] as the feature representation becomes more amenable for training classifiers after pretraining. However, using contrastive learning to adapt pretrained weights to heterogeneous target dataset is yet explored and in this work we employ contrastive learning as a target pretraining approach and demonstrate effectiveness when label-data on target dataset is sparse.

We use $\mathbf{Z}=f(\mathbf{X};\Phi)$ to denote a backbone network, parameterized by $\Phi$, which takes input image $\mathbf{X}$ and produces a feature embedding $\mathbf{Z}$. We further denote a projection head, parameterized by $\Psi$, as $\hat{\mathbf{Y}}=g(\mathbf{Z};\Psi)$ which outputs a lower dimension embedding. Contrastive learning is formulated as minimizing a contrastive loss defined over two augmentations of the same input image. Formally, the InfoNCE loss [37] is adopted as contrastive loss as in Eq. (1).

In this formulation, the positive pair $(i,j)$ is two different augmentations of the same input image, and negative pairs $(i,k)$ are constructed as different images and $s_{i,j}$ defines a cosine similarity between feature encodings of pair $(i,j)$. It is identified in [4] that a large batchsize is essential to contrastive learning performance which is computationally expensive. Instead of maintaining both positive and negative pairs, BYOL [10] relies on similarity of features between a target network and an online network with target network being a parameter moving average of online network; the output of target network is used as a target for the online network to match against. Because of the slowly updating target network, trivial solution is avoided. BYOL optimizes the cosine distance, as in Eq. 2, between the representation subject to two augmentations as below, where $\hat{\mathbf{Y}}_{1}=g(f(t_{1}(\mathbf{X})))$ and $\hat{\mathbf{Y}}_{2}=g(\hat{f}(t_{2}(\mathbf{X})))$ are respectively the features after projection head with two augmentations $t_{1}(\cdot)$ and $t_{2}(\cdot)$ and $q(\cdot)$ is a predictor.

The consistency loss $\mathcal{L}_{con}$ can be instantiated as off-the-shelf semi-supervised learners, e.g. FixMatch [3].

One of the most critical ingredients in contrastive training is data augmentation. An investigation into effective data augmentation was carried out in [4] and many contrastive approaches follow the augmentation policies [10, 34]. For a dense prediction task, [34] developed a contrastive learning paradigm by comparing pixel-wise features and it turns out to be effective for detection and segmentation tasks. Therefore, we follow the dense contrastive learning strategy when target pretraining is applied to segmentation tasks. However, we believe the data augmentation strategies that work well for classification task may not be the best for segmentation.

CNNs are known to translation equivariant, therefore cropping and weak resizing as augmentation is not stronger enough to force the network learn useful information on unlabeled data. Under such weak geometric transformations, the network will simply output feature maps subject to the same transformation. Therefore, we propose to include a differentiable affine transformation [36] to achieve more diverse geometric poses in augmented data.

An affine transformation involves a combination of scaling, rotation, shearing and translation. Each time one affine transformation is constructed by randomly chosing each component with $50\%$ chance and drawing parameters from a predefined distribution. For input image, we apply the affine transformation after all non-differentiable data augmentation operations. As with STN [36], we use bilinear interpolation to estimate the pixel intensity after transformation. This transformation is differentiable w.r.t. input image and so does its inversed one.

We denote the transformation applied to image $\mathbf{X}$ as $t(\mathbf{X})$. To enable computing pixel-wise contrastive loss, we further apply the inverse transformation to the output of both online and target networks as $t^{-1}(f(t(\mathbf{X})))$. With the inverse transformation, predictions on two arbitrarily augmented images are directly comparable at pixel-level.

We demonstrate the effectiveness of proposed contrastive target pretraining strategy on four datasets. For classification tasks, we evaluate on CIFAR-10 [5], CIFAR-100 [5] and SVHN [38]. We further evaluate on additional semantic segmentation datasets, Pascal VOC 2012 [39], CityScapes [40] and ISIC2017 [41]. CIFAR-10 is a image classification datasets with 50,000 training images and 10,000 testing images evenly categorized into 10 classes. For semi-supervised learning setting, we evaluate on 40, 250 and 1000 labeled data. CIFAR-100 extended from CIFAR-10 with 10 times more categories and the same amount of training/testing samples. Therefore, the categories are more fine-grained than CIFAR-10. For semi-supervised learning setting, we evaluate on 250 and 1000 labeled data. SVHN is a street view house numbers recognition dataset with all numbers cropped and registered. There are 73257 digits for training, 26032 digits for testing respectively. For semi-supervised learning setting, we evaluate on 40 labeled data. For CIFAR-10, CIFAR-100 and SVHN, we follow the standard data splits for semi-supervised evaluation as in [3]. Pascal VOC 2012 has been widely adopted for semantic segmentation tasks. The dataset consists of 1464 training images and 1449 validation images. Following the practice in [7, 34] we augment the training set with additional training images, resulting in 10582 training images in total. For semi-supervised learning, we evaluate on $1\%$, $2\%$ and $5\%$ of the labeled data. CityScapes was collected from cars diving in urban environment with aim to segment objects commonly seen on the street. It has 2975 training and 500 validation samples, with 19 semantic classes. We evaluate on 100 and 372 labeled data. ISIC2017 is a medical image segmentation dataset detecting skin lesions. It has 2000 training images and features a significant difference from the source ImageNet dataset. We evaluate SSL at 10, 20 and 50 labeled data.

For segmentation tasks, we adopt CutMix [7] as the semi-supervised leaner. We adopt the following hyperparameters, the initial learning rate $1e-5$, batchsize $20$/$8$ for PascalVOC and CityScapes respectively, iterations per epoch $1000$ and the number of total training epochs $40$. The Adam optimizer is adopted for both datasets. We randomly crop out a $321\times 321$ and $256\times 512$ region for training on PascalVOC and CityScapes respectively. For CityScapes, we downsample the image to $512\times 512$ to allow more efficient training.

In this section, we carry out ablation study to verify the importance of weight regularization and affine transformation as augmentation. We carry out ablation on PascalVOC segmentation task. As shown in Tab. IV, we observe both weight regularization and affine transformation as augmentation improve the generalization of further pretrained model weights.

We perform additional evaluations on alternative contrastive learning and semi-supervised learning methods on CIFAR-10 with 40 labeled samples. We perform target pretraining with MoCo v2 [44] in addition to BYOL on CIFAR-10 with ImageNet pretrained weights and SSL finetuning with UDA [28] and MeanTeacher in addition to FixMatch. From Tab. V, we observe consistent improvements with alternative contrastive learning methods and we observe the advantage of target pretraining with additional SSL methods as well.

We discuss in this part the hypotheses of why this simple target pretraining is so effective. Some early works managed to explain deep neural network pretraining from providing a regularization effect on model parameters [45]. For example, pretraining will initialize the parameters in the basin of attraction of a good local optimum. Therefore, the following supervised finetuning will be easier to generalize. Another recent empirical study [9] argued that when there is enough labeled data, training from scratch is comparable to finetuning from pretrained weights. However, a breakdown point was observed in the low labelling regime, indicating good initial weights is necessary when labeled data is low.

In contrast to the previous attempts to explain, we would like to analyze from a covariate shift point of view. W.l.o.g, we describe a general supervised pretraining and finetuning example. The conclusion also applies to contrastive pretrain and semi-supervised finetuning. For simplicity, we denote an additional unlabelled pretraining dataset as $\mathcal{D}_{pre}$. The objective of supervised pretraining can be seen as minimizing the following negative log likelihood, where we can decompose the conditional probability into the backbone $p(\mathbf{Z}|\mathbf{X};\Phi)$ and classifier $p(\mathbf{Y}|\mathbf{Z};\Theta)$.

The objective of finetuning can be written as below.

The data distribution between pretraining and finetuning stages are often not identical thus subject to a covariate shift, i.e. $p_{pre}(\mathbf{X})\neq p_{ft}(\mathbf{X})$. It is reasonable to believe the optimal backbone networks subject to pretraining and finetuning objectives respectively are not identical $p(\mathbf{Z}|\mathbf{X};\Phi_{ft}^{*})\neq p(\mathbf{Z}|\mathbf{X};\Phi_{pre}^{*})$ either. As a result, directly re-using $p(\mathbf{Z}|\mathbf{X};\Phi_{pre})$ for the supervised finetuning task would be sub-optimal. Further pretrain the backbone on the whole target dataset and the backbone parameters are initialized by the pretrained parameters, i.e. $\Phi_{ft}^{0}=\Phi_{pre}^{*}$, as below.

Compared with direct finetuning from $\Phi_{pre}$, target pretrained weights $\Phi_{ft}$ is first updated on the whole target dataset $\mathcal{D}_{L}\cup\mathcal{D}_{U}$. With covariate shift being eliminated, $\Phi_{ft}$ will move towards the optimal weights for target task. Eventually, $\Phi^{*}_{ft}$ as initial weights for supervised finetuning can ease the difficulty in optimization for finetuning. To validate the above hypotheses we carry out the following empirical studies.