Contrastive Self-Supervised Learning with Hard Negative Pair Mining

Problem definition: Un/self-supervised learning tries to learn a good representation from a large scale unlabeled dataset $\mathcal{D}=\{I_{1},I_{2},\cdots,I_{N}\}$, where each $I$ represents an image. For an image sampled from the dataset $I_{i}\sim\mathcal{D}$, we can obtain two representational views of $I_{i}$ by constructing one student sub-network $\mathcal{S}\left(\cdot;{\theta}_{S}\right)$ and one teacher sub-network $\mathcal{T}(\cdot;{\theta}_{T})$ Shin (2020). To let the network learn various invariant, we employ advanced data augmentation $\mathcal{A}$, including color jittering, horizontal flipping, Gaussian blurring and random cropping, in the data generation process for the teacher sub-network. We then obtain two views of representations for image $I_{i}$ as

where $U_{i}$ is the representation from the teacher sub-network and ${U}^{\prime}_{i}$ is the representation from the student sub-network.

The self-supervised learning tries to build pretext tasks from these unlabeled data. The generated representation views $U_{i}$ and ${U}^{\prime}_{i}$ from the student and teacher sub-networks can be considered as a positive pair, which belong to the same cluster. The adversarial contrastive self-supervised learning tries to yield a compact representation for images of the same cluster by minimizing their normalized $L_{2}$ distance in the representational space. The intra-cluster distance can be defined

where images are randomly sampled from the dataset $I_{i}\sim\mathcal{D}$, $\|\cdot\|_{\infty}$ is the infinity norm, i.e., the maximum of the absolute value of elements in the vector.

It is not efficient to train a self-supervised network by solely using positive pairs of samples. Current self-supervised learning uses large batch size Chen et al. (2020a), memory bank Wu et al. (2018) or large dynamic dictionary He and others (2020) to achieve promising results. Adding negative image pairs can significantly improve the training efficiency of a self-supervised learning model.

We heuristically construct negative pairs in the self-supervised learning framework by mining hard negative pairs of images. For two different image $I_{i}$ and image $I_{j}$, we measure the dissimilarity of the two images by the normalized $L_{2}$ distance in the representation space

There exists large numbers of negative pairs of samples. Hard samples have been widely proved to improve the performance of a deep learning model Ren et al. (2015); Lin and others (2017). In the self-supervised learning framework, we define the hard negative pairs to be image pairs of small dissimilarity according to Equation 3. We try to maximize the normalized $L_{2}$ distance or dissimilarity of negative image pairs. The contrastive loss for negative pairs can be derived

where images are randomly sampled from the dataset $I_{i}\sim\mathcal{D}$, $\tilde{\mathcal{B}}_{i}$ is the hard negative sample set of the current batch $\mathcal{B}_{i}$ for image $I_{i}$. The hard negative sample set $\tilde{\mathcal{B}}_{i}$ can be constructed

We construct hard negative pairs on-the-fly in training, which can be used to efficiently train the self-supervised network.

To stabilize the training and avoid a collapsing solution in the self-supervised learning Chen and He (2020), we block the gradient for the student sub-network $\mathcal{S}(\cdot;\theta_{S})$. We employ the exponential moving average (EMA) to update the parameters $\theta_{S}$ in the student sub-network Tarvainen and Valpola (2017)

where $\tau$ is a smoothing coefficient to tune the update strength of the student sub-network.

In the back-propagation, we only use the gradient to update the parameters of the teacher sub-network. The overall loss function can be derived as

where $0<\alpha_{1}<1$ and $0<\alpha_{2}<1$ are the fixed coefficients to tune the trade-off between the intra-cluster loss and inter-cluster loss. During the back-propagation, we employ the gradient clipping to stabilize the training.

In our method, we employ hard negative pair mining (HNPM) to add negative image pairs in the training, and use a normalized $L_{2}$ distance in the loss function. We will demonstrate that minimizing the loss of our method is equivalent with minimizing the InfoNCE loss Oord et al. (2018). To simplify the analysis, we temporarily remove the hard negative pair mining mechanism in our method in the derivation of connection with InfoNCE.

The InfoNCE loss Oord et al. (2018) can be written as

where $U_{i}$, $U_{i}^{\prime}$ are calculated from the teacher sub-network and the student sub-network, $f_{k}(\cdot,\cdot)$ models the mutual information between the encoded representations in the InfoNCE and can use similarity loss as a surrogate loss to approximate the mutual information.

We define the similarity loss as the reciprocal of normalized $L_{2}$ distance of the encoded representations. The InfoNCE loss can then be defined as

The second part of derived loss in equation 9 is the same with our negative pair loss in the equation 4 if we temporarily neglect our hard negative sample pair mining for each batch. Minimizing the first part of equation 9 is equivalent with minimizing $\mathbb{E}_{I_{i}\sim\mathcal{D}}[(\frac{U_{i}}{\|U_{i}\|_{\infty}}-\frac{U_{i}^{\prime}}{\|U_{i}^{\prime}\|_{\infty}})^{2}]$, which is the positive pair loss in the equation 2. From the above derivation, we conclude, with the proper relaxation and assumption, minimizing the our loss is equivalent with minimizing the InfoNCE loss.

Next we demonstrate that the hard negative pair mining (HNPM) leads to stable training. Without the trade-off factors $\alpha_{1}$ and $\alpha_{2}$, the loss can be written as

Without loss of generality, we remove the normalization constraint and denote $\frac{U_{i}}{\|U_{i}\|_{\infty}}$ as $U_{i}$.

The hard negative pair mining (HNPM) always explores negative pairs with $L_{2}$ distance smaller than 1, which guarantees $(U_{j}-U_{i}^{\prime})^{2}$ is bounded to be smaller than 1. We use $M$ to denote the upper bound of negative pair loss.

Next we can further prove equation 12 can be optimized stably, and the first part of equation 12, i.e., the loss of positive pairs, can be decreased consecutively by escaping undesirable equilibria. If the model stacks into an undesirable equilibrium solution, the feature representation of teacher sub-network can be denoted as $\mathbb{E}[U_{i}^{\prime}|U_{i}]$ from the update rule in equation 6. The loss of positive pairs $\mathcal{L}_{P}$ can be derived as

Let $Z$ denote an additional variability induced by stochasticities in the training dynamics. We always have a solution leading to a lower loss during the training, which escapes the current equilibrium, because

From the above derivation, the learning is stable with the benefit of hard negative pair mining and student sub-network updating rule.

Because of our advanced learning strategy, we do not use any pretrained model as the backbone in our implementation. To generate multi-view representations, we employ data augmentation to model various variations in different views.

We use residual networks as the student sub-network $\mathcal{S}(\cdot,\theta_{S})$ and teacher sub-network $\mathcal{T}(\cdot,\theta_{T})$. The two coefficients of the loss in equation 7, $\alpha_{1}$ is set to 0.8 and $\alpha_{2}$ is set to 0.1. We employ the gradient clipping strategy in the back-propagation where we set the maximum norm of gradient clipping as $1.0$. The Adam optimizer is used to minimize the loss in equation 7. The batch size is 160. The learning rate is set as $0.1$ and we use a cosine annealing schedule for the learning rate with the maximum number of iterations as $100$. The smoothing coefficient $\tau$ in the update of student sub-network in equation 6 is set as $0.5$.

We employ data augmentation for teacher sub-network on-the-fly during training. We firstly apply color jittering with brightness of $0.8$, contrast of $0.8$, saturation of $0.8$, and hue of $0.2$ to random 80% training images in each batch. Then we convert random 20% images to gray scale, and horizontally flip 50% images. After that, we smooth random 10% images with a random Gaussian kernel of size $3\times 3$ and standard deviation of $1.5\times 1.5$. Finally, we crop each image with random crop size of scale range $[0.8,1.0]$. We use the mean of $[0.485,0.456,0.406]$ and the standard deviation of $[0.229,0.224,0.225]$ to normalize RGB channels.

The linear evaluation can be used to evaluate the accuracy of self-supervised learning (SSL) by freezing the SSL model and training a separate linear classifier after the SSL model Grill and others (2020); Kornblith et al. (2019); Zhang and others (2016). We compare our method with previous state-of-the-art approaches with the ResNet-50 encoder and other ResNet encoders on ImageNet in the Table 1 and Table 2, respectively. The top-1 and top-5 accuracy are listed. With the standard ResNet-50 encoder He et al. (2016), our method obtains 77.1% top-1 accuracy and 93.7% top-5 accuracy, which outperform previous state-of-the-art top-1 and top-5 results by 2.8% and 2.1%, respectively. Most surprisingly, our method achieves 0.6% better accuracy than the accuracy, 76.5%, of the supervised baseline from SimCLR Chen et al. (2020a).

Table 2 reports the accuracy of self-supervised learning methods using deeper and wider ResNet encoders based on linear evaluation. Our method with ResNet-200 (2$\times$) obtains 81.9% top-1 and 96.4% top-5 accuracy which increase previous best top-1 and top-5 accuracy by 2.3% and 1.6%, respectively. With ResNet-50 (2$\times$) and ResNet-50 (4$\times$) encoders, our method also achieves better accuracy than those of CMC Tian et al. (2019), SimCLRv2 Chen et al. (2020b) and BYOL Grill and others (2020) with the same encoder.

Semi-supervised learning can also be used to evaluate the accuracy of self-supervised learning (SSL) by fine-tuning representation with a small subset of the training set Grill and others (2020). In this experiment, we use the fixed data splits of 1% and 10% of the training set in ImageNet, which are the same as Grill and others (2020). We also use the top-1 and top-5 accuracy as the evaluation metric for the semi-supervised learning. The comparison using the ResNet-50 encoder and deeper and wider ResNet encoders are listed in Table 3 and Table 4, respectively. Our method achieves 80.2% top-5 accuracy based on a ResNet-50 encoder which improves previous best result by 1.8% using only 1% training labels in the Table 3. Using 10% training labels, our method achieves 73.4% and 92.5% for the top-1 and top-5 accuracy, which improve previous best top-1 and top-5 accuracy by 4.6% and 3.5%.

The result with ResNet of various depths, widths, selective kernel convolution Li and others (2019) configurations are listed in Table 4. Our method achieves the best top-1 and top-5 accuracy for all the experimental configurations. Specifically, based on ResNet-50 (2$\times$) encoder, our method achieves 65.7% and 78.6% top-1 accuracy using 1% training labels and 10% training labels, which improves previous best top-1 accuracy by 3.5% and 5.1%. Based on ResNet-200 (2$\times$), our method obtains 76.5% and 80.7% top-1 accuracy using 1% training labels and 10% training labels, which improves the accuracy of BYOL Grill and others (2020) by 5.3% and 3.0%.

Transfer learning is another widely used task to evaluate the accuracy of self-supervised learning (SSL) methods. Transfer learning can be used to evaluate the generalization ability of the learned SSL model. Practically, both linear evaluation, i.e., only training the last classification layer, and fine-tuning the whole network based on the target dataset can be employed for the evaluation of transfer learning. The comparison of transfer learning with linear evaluation and fine-tuning are listed in the Table 5 and Table 6.

For the linear evaluation of transfer learning task, our method achieves better accuracy than previous state-of-the-art approaches on six out of seven widely used transfer learning datasets in Table 5. We provide the accuracy improvement in Table 5, and our method improves 2.3%, 3.5%, 4.5%, 2.7%, 4.9% and 0.9% on Flood101, SUN397, Cars, Pets, VOC 2007 and Flowers datasets, respectively. On average, the transfer learning accuracy of our method is 2.5% higher than previous best results based on linear evaluation. For the transfer learning with a fine-tuning task, our method achieves the best accuracy on four out of seven tasks in Table 6.

Coefficient $\tau$ in the update of student sub-network We investigate the accuracy of our method with linear evaluation based on the ResNet-50 encoder with respective to the smoothing coefficient $\tau$ of the exponential moving average (EMA) in Table 7. The bigger the $\tau$ is, the smaller update the student sub-network performs. When the $\tau$ is 0, it means that we copy the weights of the teacher sub-network to update the student sub-network completely in each step. When the $\tau$ is 1, it means that the student sub-network is never updated. We find that at moving average coefficient value of 0.5, we obtain the best top-1 accuracy, 77.1%, based on linear evaluation. Neither the moving average coefficient $\tau$ of 0 nor 1 generates good performance.

Hard negative pair mining (HNPM) We conduct ablation study on hard negative pair mining (HNPM) based on liner evaluation task using the ResNet-200 (2$\times$) encoder on the ImageNet dataset. Training with all negative pairs, i.e., without HNPM, is denoted as “w/o HNPM + block student gradient”, and our method is trained with HNPM, which is denoted as “w/ HNPM + block student gradient”. The loss and accuracy comparison $w.r.t.$ the training epochs for the two methods are in Fig. 2. With hard negative pair mining, the training of our method is much more stable and it achieves lower loss and higher accuracy than that without hard negative pair mining.

Blocking gradient of student sub-network We also conduct ablation study on blocking gradient of student sub-network in Fig. 2. Training without blocking the gradient of student sub-network is denoted as “w/ HNPM + student gradient”. Our method achieves lower loss and higher accuracy than that with gradient updating of student sub-network.