Continual Contrastive Learning for Image Classification

The core idea of self-supervised learning is to train a neural network on the large-scale unlabeled data by designing specific pretext tasks, so that the network can learn a general and transferable feature representation. At present, instance discrimination [7], or more specifically, contrastive learning [8][9][10][11], has made breakthroughs and achieves new state-of-the-art performance on feature representation learning. In instance discrimination methods, one image falls into one category by distinguishing from other images, and the learning objective is to maximize the mutual information of two views from the same image generated by different data augmentation operations.

When training a self-supervised learning network, a large-scale data set is indispensable. However, directly training a whole large data set may be unaffordable in many scenarios. Moreover, new data is usually generated continuously. Collecting old and new data together and training the network again is a time and resources consuming solution. Thus, a continuous style of self-supervised learning method is imperative when it comes to application.

Class-Incremental Learning (CIL) aims to learn a classification model with the number of classes increasing step by step. Existing works on CIL often adopt rehearsal methods and knowledge distillation to tackle the catastrophic forgetting problem. Rehearsal-based methods [1][12] try to select a representative set of samples from the old data. Specifically, these works use data labels or use a classifier trained with labeled data, to select the samples. However, in self-supervised continual learning, there is no access to the data labels. Thus, in our method, we propose a novel sampling strategy using feature variance measurement, which does not rely on data labels. Distillation-based methods [13][14] use knowledge distillation loss as the regularization term to preserve previous knowledge when learning new data. Many distillation-based methods use supervised knowledge distillation, i.e., the predicted label logits of the new model are enforced to be close to those of the old model. However, directly applying supervised knowledge distillation loss to self-supervised continual learning will cause some problems since it lacks the classification predictor. Thus, we adopt self-supervised knowledge distillation loss [15] in our method, which does not require the classification head. Moreover, unlike traditional knowledge distillation methods which often fix the teacher network during distillation, we propose the momentum teacher design.

First, we introduce MoCoV2 [6] briefly for better understanding. MoCoV2 [6] contains two encoders, $f_{q}$ and $f_{k}$, and a memory bank. Given an input image $x$, we first transform it to two views $x_{q}$ and $x_{k}$ with two different augmentations. Then, we get query vector $q$ and its positive key sample $k_{+}$ by feeding $x_{q}$, $x_{k}$ into $f_{q}$ and $f_{k}$ respectively. The memory bank stores negative key samples $Q=\{{k_{1}},{k_{2}},{k_{3}},...,{k_{n}}\}$ of $q$. Finally, MoCoV2 [6] adopts a contrastive loss to pull $q$ to $k_{+}$ and push away $q$ from $Q=\{{k_{1}},{k_{2}},{k_{3}},...,{k_{n}}\}$.

Specifically, the contrastive loss is defined as:

where $\tau$ is the hyper-parameter of temperature and $n$ is the number of negative samples.

Moreover, two encoders $f_{q}$ and $f_{k}$ share the same architecture, i.e., the backbone of the network followed by an extra MLP, but with different weights $\theta_{q}$ and $\theta_{k}$. The parameters $\theta_{q}$ are updated by back-propagation and the parameters $\theta_{k}$ are updated by momentum strategy, that is,

Rehearsal is widely used in supervised continual learning. By replaying the stored old samples from previous tasks, we find rehearsal methods can also help alleviate catastrophe forgetting in self-supervised continual learning. Besides, inspired by many supervised continual learning methods, we utilize self-supervised knowledge distillation to transfer the contrastive information when the network learns a new task.

In continual learning, given a set of datasets $\{D_{1},D_{2},...,D_{t},...,D_{N}\}$, the network will be trained continuously on the dataset $D_{t}$ at time t. In rehearsal methods, we store a small fraction of data from $D_{t}$ after the training of $D_{t}$ is finished. Instead of randomly storing samples from old data, we propose a novel sampling strategy with the feature variance measurement. Specifically, when the training on $D_{t}$ is finished, we input all images of $D_{t}$ into the encoder $f_{q}$ to extract feature vectors. Then we group the vectors by the K-Means algorithm and divide them into $C$ classes. In our experiments, we set the $C$ equal to the number of classes in $D_{t}$. Then, for images $x$ in class $c$, we have different views $\{\mathcal{T}_{1}(x),\mathcal{T}_{2}(x),...,\mathcal{T}_{l}(x)\}$ of $x$ after applying different data augmentations $\mathcal{T}$ mentioned in the MoCoV2 [6]. After that, we calculate the variance of the feature vectors of these views $\{f_{q}(\mathcal{T}_{1}(x)),f_{q}(\mathcal{T}_{2}(x)),$ $...,f_{q}(\mathcal{T}_{l}(x))\}$. We found $l=6$ is enough for measuring the variance. Finally, we store the first $n$ images with the smallest variance for each class. Essentially, the images with the smallest variance are ones well learned by the network, thus the contrastive information of these images can be easily transferred to the new task through self-supervised knowledge distillation.

Knowledge distillation is another way to address catastrophe forgetting problems. Hence, we introduce it to our method to further enhance the ability for handling catastrophe forgetting. Following some self-supervised knowledge distillation methods [15][16], we utilize the contrastive similarity between images from the teacher network as the distillation target for the student network. Specifically, given the sampled images $x$ from the old dataset, we have $x_{q}$ after one data augmentation. Then, $x$ and $x_{q}$ are mapped and normalized into feature vector representations $\textbf{z}^{T}=f_{q}^{T}(x)$, $\textbf{z}^{T}_{q}=f_{q}^{T}(x_{q})$, $\textbf{z}^{S}=f_{q}^{S}(x)$ and $\textbf{z}^{S}_{q}=f_{q}^{S}(x_{q})\in\mathbb{R}^{B\times D}$, where $B$ is the number of images, $D$ is the feature dimension, and $f_{q}^{T}$ and $f_{q}^{S}$ denote the query encoder of teacher and student respectively. We compute the similarity matrix $P^{T}(\textbf{z}^{T},\textbf{z}^{T}_{q})\in\mathbb{R}^{B\times B}$ between the feature vectors $\textbf{z}^{T}$ and $\textbf{z}^{T}_{q}$ extracted by teacher. After that, the softmax function with the temperature scaling factor $\tau$ is applied to $P^{T}$ for normalization. We obtain $P^{S}$ with a similar operation from the student. Finally, we adopt KL-divergence loss to $P^{T}$ and $P^{S}$:

Besides, compared to supervised continual learning, which often fixes the teacher network during distillation, we newly find that a momentum teacher brings better performance. Thus, for each training epoch, we update the teacher network with a momentum style:

where $m_{t}=0.996$ in our experiments.

As mentioned in the section of the introduction, learning feature representations of a new task will cause the mixture of feature space among the old data and the new data. As illustrated in Fig. 3, when the network learns the feature representation (orange circulars) of new data, it may interfere with the feature representation (blue triangles) of old data. Thus, we propose an extra sample queue (red triangles) to push away the feature of new data from the feature of old data.

Specifically, the extra sample queue only contains the negative samples of the sampled images from old data. For a data mini-batch $B=B_{S}\cup B_{D}$, where $B_{S}$ is the sampled images from old data and $B_{D}$ is the images from new data, we sent them to two encoders of MoCoV2 [6] and then obtain feature representations $\textbf{z}^{S}_{q}=q_{S}\cup q_{D}$ and $\textbf{z}^{S}_{k}=k_{S}\cup k_{D}$. The network discriminates data of the new dataset from data of the old datasets by computing the contrastive loss $\mathcal{L}_{ESQ}$ between $\textbf{z}^{S}_{q}$, $\textbf{z}^{S}_{k}$, and extra sample queue. In every iteration, we will pick out the feature $k_{S}$, and update the extra sample queue with $k_{S}$.

To sum up, the total loss is given by:

where $\lambda_{1}$, $\lambda_{2}$, and $\lambda_{3}$ are balancing weights. In our experiments, we set $\lambda_{1}=0.9$, $\lambda_{2}=0.1$ and $\lambda_{3}=0.1$.

Besides our method, we implemented and tested two naive continual contrastive learning methods, Finetuning and Simple Rehearsal. Finetuning learns new tasks one by one without taking any technique to prevent catastrophic forgetting. Simple Rehearsal stores a constant number of images randomly after training on the dataset $D_{t-1}$. For the subsequent task $t$, it adds the stored images into the new training set $D_{t}$ to train the model. It’s worth noting that, Simple Rehearsal does not use techniques like knowledge distillation or the extra sample queue.

Table 1 shows the performance of different continual contrastive learning methods on CIFAR-100 and ImageNet. One can see that our method outperforms the others from each column of Table 1. More specifically, the top-1 accuracy is improved from $48.40\%$ to $50.10\%$ on CIFAR-100, from $52.62\%$ to $55.48\%$ on ImageNet-Sub, and from $61.50\%$ to $62.79\%$ on ImageNet-Full under 10 incremental settings. These results show that our method works well on both small and large datasets, and demonstrate the effectiveness of our proposed continual contrastive learning method.

Besides MoCoV2 [6], we also apply our method to other contrastive learning methods, including SimCLR [9] and InsDisc [7]. More detail can be found in Appendix. The results on CIFAR-100 are shown in Fig. 4. We can see that our method also improves the linear classification accuracy of other contrastive learning methods, which demonstrates the generalization of our method.

To test the effectiveness of each component in our approach, we conduct ablation studies on ImageNet-Sub under 10 incremental steps setting.

In the first experiment, we evaluate the main components of our method: sampling strategy with feature variance measurement, self-supervised knowledge distillation, and the extra sample queue. As described in Section 3.2 and 3.3, data sampling strategy and knowledge distillation are exploited to distill the contrastive information from the previous contrast process, and the extra sample queue is exploited to separate the feature space of the new data and old data. As shown in Table 2, our novel sampling strategy, knowledge distillation, and extra sample queue boost the performance consistently and gains $2.86\%$ accuracy improvement when all of them are applied.

Besides, since it is impossible to decide the category number of the unlabeled data, the hyper-parameter K in the K-Means algorithm in our experiment is the prior knowledge. Fig. 5 shows that the top-1 accuracy on ImageNet-Sub varies according to the choice of the hyper-parameter K. We find when K varies within a certain range, from 5 to 15, the performance of our method has little change. However, when the gap between K and the real number of classes is too large, the accuracy decreases rapidly.

We further explore the effect of the size of the extra sample queue and conduct the experiment on CIFAR-100. Table 3 shows that when the size of the extra sample queue increases, the performance is not improved consistently. A possible reason is that, when the size is large, the update rate of the extra sample queue is slow. Thus, many feature vectors in the extra sample queue are from previous iterations, which are detrimental to the update of the model.

In this paper, we try to solve the problem of continual contrastive learning. The experiment results show that our method works well on both small and large datasets. However, our work still has many limitations. First, our method is only available for the contrastive learning method. When applying our method to other self-supervised learning methods, like GAN [18], BYOL [11], and MAE [19], some additional revisions are required. Second, though our method narrows the gap between continual contrastive learning and its upper bound, there is still a large gap between them. We hope our work can give researchers some insights into the catastrophic forgetting problems in contrastive learning, and inspire more research in this direction.