Online Deep Clustering for Unsupervised Representation Learning

Introduction

Unsupervised representation learning aims at learning transferable image or video representations without manual annotations. Among them, clustering-based representation learning methods emerge as a promising direction in this area. Different from recovering-based approaches , clustering-based methods require little domain knowledge while achieving encouraging performances. Compared to contrastive representation learning that captures merely intra-image invariance, clustering-based methods are able to explore inter-image similarity. Unlike conventional clustering that is typically performed on fixed features , these works jointly optimize clustering and feature learning.

While evaluations of early works are mostly performed on small datasets, Deep Clustering (DC) proposed by Caron et al. is the first attempt to scale up clustering-based representation learning. DC alternates between deep feature clustering and CNN parameters update. In particular, at the start of each epoch, it performs off-line clustering algorithms on the entire dataset to obtain pseudo-labels as the supervision for the next epoch. Off-line clustering inevitably permutes the assigned labels in different epochs, i.e., even if some of the clusters do not change, their indices after clustering will be permuted randomly. As a result, parameters in the classifier cannot be inherited from the last epoch and they have to be randomly initialized before each epoch. The mechanism introduces training instability and exposes representations to a high risk of representation corruption. As shown in Figure 1 (a), network update in DC is interrupted by feature extraction and clustering in each epoch. This is in contrast to the conventional supervised classification that is performed in an uninterrupted manner using fixed labels, where an iteration consists of forward and backward propagations of the network.

In this work, we seek to devise a joint clustering and feature learning paradigm with high stability. To reduce the discrepancy of training mechanism between DC and supervised learning, we decompose the clustering process into mini-batch-wise label update, and integrate this update process into iterations of network update. Based on this intuition, we propose Online Deep Clustering (ODC) for joint clustering and feature learning. Specifically, an ODC iteration consists of forward and backward propagations, label re-assignment, and centroids update. For label update, ODC reuses the features in the forward propagation, thus avoiding additional feature extraction. To facilitate online label re-assignment and centroids update, we design and maintain two dynamic memory modules, i.e., samples memory to store samples’ labels and features, and centroids memory for centroids evolution. In this way, ODC is trained in an uninterrupted manner similar to supervised classification, while no manual annotation is required. During the training process, labels and network parameters evolve shoulder-to-shoulder, rather than alternatingly. Since labels are updated in each iteration continuously and instantly, the classifier in the CNN also evolves more steadily, resulting in a much more steady loss curve as shown in Figure 1 (b).

While ODC alone achieves compelling unsupervised representation learning performance on various benchmarks, it can be naturally used to fine-tune models that have been trained using other unsupervised learning approaches. Extensive experiments show that the steadiness of ODC helps it to perform superiorly over DC as an unsupervised fine-tuning tool. We conclude our contributions as follows: 1) we propose ODC that learns image representations in an unsupervised manner with high stability. 2) ODC also serves as a unified unsupervised fine-tuning scheme that further improves previous self-supervised representation learning approaches. 3) Promising performances are observed on different benchmarks, indicating the great potential of joint clustering and feature learning.

Related Work

Unsupervised Representation Learning. Many unsupervised visual representation learning algorithms are based on generative models, which usually use a latent representation bottleneck to reconstruct input images. Existing generation-based models include Auto-Encoders , Restricted Boltzman Machines , Variational Auto-Encoders and Generative Adversarial Networks , some of which have shown powerful ability in generating images or videos . By learning to generate examples, these models can learn meaningful latent representations that can be used for downstream tasks .

Another popular form of unsupervised representation learning is self-supervised learning, where a pretext task is designed to derive proxy labels from raw data. Representations are learned by encouraging a CNN to predict the proxy labels from the data. Various pretext tasks have been explored, e.g., predicting relative patch locations within an image , solving jigsaw puzzles , colorizing grayscale images , inpainting of missing pixels , cross-channel prediction , counting visual primitives , predicting image rotations , and multiview contrastive learning . For videos, self-derived supervision signals come from temporal continuity or motion consistency .

Joint Clustering and Feature Learning. Clustering-based unsupervised representation learning is of particular interest recently. Various methods are proposed to jointly optimize feature learning and image clustering. Notably, these methods have shown great potential in learning unsupervised features on small datasets . To scale up to large datasets like ImageNet , Caron et al. propose DeepCluster to cluster features and update CNN with subsequent assigned pseudo-labels for each epoch. In a subsequent study, Caron et al. propose DeeperCluster to leverage self-supervision and clustering, and validate the representation learning ability of their approaches on non-curated data. Although deep clustering methods are capable of learning good representations from large-scale unlabeled data, the alternating update of feature clustering and CNN parameters update leads to instability in training.

Improvements to Self-supervised Learning. Some works aim at improving previous self-supervised learning approaches from different perspectives. For instance, Larsson et al. give a first in-depth analysis on colorization as a pretext task and provide some insights on improving its effectiveness. Mundhenk et al. explore a set of methods to avoid some trivial shortcuts like chromatic aberration on context-based self-supervised learning. Noroozi et al. improve the performance of self-supervised models using a clustering-based knowledge transfer method that allows a deeper network during pre-training. Wang et al. and Doersch et al. exploit multiple cues contained in different pretext tasks to improve self-supervised models. Recently, some works have studied extensively the architectures and scaling ability on existing self-supervised approaches. Complementary to these works, ODC serves as a flexible and unified unsupervised fine-tuning scheme to boost general self-supervised learning methods although it can be used alone to perform unsupervised representation learning from scratch.

Methodology

In the following sub-sections, we first discuss the differences between the proposed ODC to the conventional DC in Sec. 3.1. We then recommend some useful strategies to maintain stable cluster size while using ODC in Sec. 3.2. We finally explain how one can use ODC for unsupervised fine-tuning (Sec. 3.3) and the implementation details of ODC (Sec. 3.4).

We first discuss the basic idea of DC and then detail the proposed ODC. To learn representations, DC alternates between off-line feature clustering and network back-propagation with pseudo-labels. The off-line clustering process requires deep feature extraction on the entire training set, followed by a global clustering algorithm, e.g., K-Means clustering. The global clustering permutes the pseudo labels vastly, requiring the network to adapt to new labels rapidly in the subsequent epoch.

Framework Overview. Different from DC, ODC does not require an extra feature extraction process. Besides, labels evolve alongside the network parameters update smoothly. This is made possible by the newly introduced samples and centroids memories. As shown in Fig. 2, the samples memory stores features and pseudo-labels of the entire dataset; while the centroids memory stores the features of class centroids, i.e., the mean feature of all samples in a class. A “class” here represents a temporary cluster that evolves continuously during training. Labels and network parameters are updated simultaneously during uninterrupted iterations of ODC. Specific techniques including loss re-weighting and dealing with small clusters are introduced to avoid ODC from getting stuck into trivial solutions.

An ODC Iteration. Assuming that we are given with a randomly initialized network $f_{\theta}\left(*\right)$ along with a linear classifier $g_{w}\left(*\right)$, the goal is to train the backbone parameters $\theta$ to produce highly discriminative representations. To prepare for ODC, the samples and centroids memories are initialized via a global clustering process, e.g., K-Means. Next, one can perform uninterrupted ODC iteratively.

An ODC iteration contains four steps. First, given a batch of input images $\left\{x\right\}$, the network maps the images into compact feature vectors $F=f_{\theta}\left(x\right)$. Second, we read pseudo-labels for this batch from the samples memory. With the pseudo-labels, we update the network with stochastic gradient descent to solve the following problem:

where $y_{n}$ is the current pseudo label from the samples memory, $B$ denotes the size of each mini-batch. Third, $f_{\theta}\left(x\right)$ after L2 normalization is reused to update the samples memory:

where $F_{m}\left(x\right)$ is the feature of $x$ in the samples memory, $m\in\left(0,1\right]$ is a momentum coefficient. Simultaneously, each involved sample is assigned with a new label by finding the nearest centroid following:

where $C_{y}$ denotes the centroid feature of class $y$. Finally, the involved centroids, including those in which new members join, and those from which old members leave, are recorded. They are updated every $k$-th iterations by averaging the features of all samples belonging to their corresponding centroid.

2 Handling Clustering Distribution in ODC

Loss Re-weighting. To avoid the training from collapsing into a few huge clusters, DC adopts uniform sampling before each epoch. However, for ODC, the number of samples over the clusters changes in each iteration. Using uniform sampling requires one to re-sample the entire dataset in each iteration, a process that is deemed redundant and costly. We propose an alternative approach, i.e., re-weighting the loss according to the number of samples in each class. To verify their equivalence, we implement a DC model with loss re-weighting and empirically find that the performance remains unchanged when the weight follows $w_{c}\propto\frac{1}{\sqrt{N_{c}}}$, where $N_{c}$ denotes the number of samples in class $c$. Hence, we adopt the same loss re-weighting formulation for ODC. With loss re-weighting, samples in smaller clusters contribute more towards backpropagation, thus pushing the decision boundary farther to accept more potential samples.

Dealing with Small Clusters. Loss re-weighting helps to prevent the formation of huge clusters. Nevertheless, we still face the risk of having some small clusters collapsing into empty clusters. To overcome this problem, we propose to process and eliminate extremely small clusters in advance before they collapse. Denoting normal clusters as $\mathcal{C}_{n}$ whose sizes are larger than a threshold, and small clusters as $\mathcal{C}_{s}$ whose sizes are not, for $c\in\mathcal{C}_{s}$, we first assign samples in $c$ to the nearest centroids in $C_{n}$ to make $c$ empty. Next, we split the largest cluster $c_{max}\in\mathcal{C}_{n}$ into two sub-clusters by K-Means and randomly choose one of the sub-clusters as the new $c$. We repeat the process until all clusters belong to $\mathcal{C}_{n}$. Though this process alters some clusters abruptly, it only affects a small portion of samples which are involved in this process.

Dimensionality Reduction. Some of the backbone networks map an image to a high-dimensional vector, e.g., AlexNet produces 4,096-dimensional features and ResNet-50 yields 2,048-dimensional features, leading to high space and time complexities in subsequent clustering. DC performed PCA on features across the entire dataset to reduce dimension. However, for ODC, the features of different samples have varying timestamps, leading to incompatible statistics among samples. Hence, PCA is not applicable anymore. It is also costly to perform PCA in each iteration. We therefore add a non-linear head layer of {fc-bn-relu-dropout-fc-relu} to reduce high dimensional features into 256 dimensions. It is jointly tuned during ODC iterations. The head layer is removed for downstream tasks.

3 ODC for Unsupervised Fine-tuning

Compared with self-supervised learning approaches that tend to capture intra-image semantics, clustering-based methods focus more on inter-image information. Hence, DC and ODC are naturally complementary to previous self-supervised learning approaches. As DC and ODC are not restricted to a specifically designed objective, like rotation angle or color prediction, they readily serve as an unsupervised fine-tuning scheme to boost the performance of existing self-supervised approaches. In this paper, we study the effectiveness of DC and ODC as a fine-tuning process with initialization from different self-supervised learning methods.

4 Implementation Details

Data Pre-processing. We use ImageNet that contains 1.28M images without labels for training. Images are first randomly cropped to have a resolution of 224x224 with augmentation including random flipping and rotation ($\pm$$$). DC adopts a Sobel filter on the images to avoid exploiting color as the shortcut. Such a pre-processing step requires the downstream tasks to include the Sobel layer as well, which potentially limit its application. We find that strong color jittering shows the same effect as a Sobel filter in avoiding shortcuts, while it allows normal RGB images as inputs. Specifically, we adopt PyTorch style color jitter transform with brightness factor$\left(0.6,1.4\right)$, contrast factor$\left(0.6,1.4\right)$, saturation factor$\left(0,2\right)$, and hue factor$\left(-0.5,0.5\right)$. Besides, we randomly convert images to grayscale with a probability of$0.2$. The random color jittering and grayscale applied on training samples randomize the similarity measured in color. This discourages the network from exploiting trivial information from color.

Training of ODC. We use ResNet-50 as our backbone. Considering that most early works use AlexNet, we also perform experiments on AlexNet for comparison. Following , we use AlexNet architecture without Local Response Normalization and add batch normalization layers. The ODC models for AlexNet and ResNet-50 are trained from scratch. The batch size is 512 allocated to 8 GPUs. The learning rate is constantly $0.04$ for AlexNet and $0.06$ for ResNet-50 for 400 epochs, and decayed by $0.1$ for further 40 epochs. Following DC, the number of clusters is set as 10,000, which is 10 times larger than the annotated number of classes of ImageNet. The momentum coefficient $m$ is set as $0.5$. The threshold to identify small clusters is set as 20. Varying this threshold does not affect the results significantly, provided that it does not exceed the average number of samples in a cluster. The centroids memory is updated in every 10 iterations. The centroids update frequency constitutes a trade-off between learning efficacy and efficiency. In our experiments, we observe that as long as the frequency is restricted to a reasonable range, the performance of ODC is not sensitive to it.

Experiments

After pre-training the ODC model, we evaluate the quality of unsupervised features on standard downstream tasks including ImageNet classification, Places205 classification, VOC2007 SVM classification, and VOC2007 Low-shot classification. We provide the details of each benchmark and show our competing results as follows.

Re-implementation of Deep Clustering. Since the original paper of DC does not include ResNet-50, we implement a DC model with ResNet-50. The DC model adopts the same data augmentations as ODC, except that DC applies a Sobel filter on images. For fair comparisons, the training hyper-parameters of DC are the same as ODC except that we empirically find $lr=0.1$ is more suitable for DC.

ImageNet Classification. Following the setup in Zhang et al. , we keep the backbone including all convolution and batch normalization layers frozen, and train a 1000-way linear classifier on features from different depths of convolutional layers. The features are mapped to around 9000 dimensions via average pooling. We train all models for 100 epochs in total, using SGD with a momentum of 0.9 and batch size of 256. The learning rate is initialized as 0.01, decayed by a factor of 10 after every 30 epochs. Other hyper-parameters are set following Goyal et al. . We report top-1 center-crop accuracy on the official validation split of ImageNet.

For AlexNet, as shown in Table 1, ODC has a consistent improvement over DC in all conv layers, with the largest improvement (6.7%) observed in conv1 layer. The performance in conv1 layer surpasses the ImageNet pre-trained model. With regard to the best-performing layer, ODC achieves 41.4% on conv4 layer, outperforming the latest LA , ranking only second to Rot-Decoupling . Though ODC does not outperform Rot-Decoupling in its best performing layer, it provides a complementary perspective to rotation based methods.

ODC also scales well with deeper architectures. For ResNet-50, as shown in Table 2, ODC achieves 57.6% center-crop accuracy in the conv5 layer, which is 5.4% higher than the best performing layer of the re-implemented DC. Compared with the concurrent state-of-the-art method LA , our method produces competing results. Though the result of conv5 is slightly lower than LA, ODC outperforms LA from conv1 to conv4 layers by large margins. We observe a consistent performance increase from shallower layers to deeper layers, indicating that ODC makes full use of all residual layers.

Places205 Classification. Following Zhang et al. , to test the generalization ability on other domains, we also transfer the learned models to Places205 dataset that contains 2.45M images of 205 scene categories. Similar to the experiments on ImageNet, we train a 205-way linear classifier on top of each frozen convolutional layer on the train split of Places205, and report top-1 center-crop accuracy on the standard validation split. The evaluation setting and hyper-parameters are the same as those in the ImageNet classification task.

The results in Table 1 show that ODC with AlexNet as the backbone outperforms DC in all layers as well. ODC surpasses all previous works on conv1, conv3 and conv4 layers. Similar to the observation in the ImageNet classification task, ODC scales well on deeper architectures when it is transferred to Places205 with ResNet-50. As shown in Table 2, in all layers, ODC surpasses all previous works, with the largest margin (3.1%) to the runner-up observed in conv2 layer. With regard to the best performing layer, ODC reaches 49.3% center-crop accuracy in the conv5 layer, surpassing the re-implemented DC by 3.2% in the respective best layer. We observe the superiority of ODC in conv1 and conv2 layers over the supervised model using either Places labels or ImageNet labels. The transfer performance of our method in the Places205 classification task indicates that representations learned by ODC can generalize well to different domains from ImageNet.

VOC2007 SVM Classification. To further evaluate the generalization of learned features, we perform experiments on the VOC2007 transfer learning task that resembles real applications with smaller datasets. Following , we train linear SVMs on features extracted from the frozen backbone on the “trainval” split of VOC2007 and evaluate on the test split. We follow the same test setting and hyper-parameters used in , and report the best performing layers of different methods for ResNet-50. The results in Table 3 show that ODC surpasses previous approaches by a significant margin on the VOC2007 SVM classification task. With ODC, we achieve 78.2% mAP performance, which is 9.1% higher than DC. However, We also note that there is still a significant 9.8% performance gap between our ODC and the supervised model pre-trained with ImageNet labels, leaving room for further exploration.

Low-shot VOC2007 SVM Classification. Following , we also transfer our learned representations to a low-shot setting of VOC2007 SVM classification to test the quality of features when there are few training examples per category. We vary the number of positive samples in each class and train linear SVMs on the frozen ResNet-50 backbone using the same setting from VOC2007 SVM classification. We use the standard “trainval” split of VOC2007 in training and the test split in testing. We report the mean average precision (mAP) across five independent samples for various low-shot values in Figure 3. The final mAP results shown in Table 3 are observed as the averages of all low-shot values and all independent runs. The per-shot results are shown in Figure 3. ODC has a consistent improvement over DC for each shot, with the performance gap further increasing when more positive examples per class are allowed. We also observe that the performance gap between ODC and the supervised model pre-trained with ImageNet labels is gradually narrowed down with the increase of training shot values. Table 3 shows that ODC achieves 57.1% mAP performance in low-shot SVM classification on VOC2007, 10.2% higher than our counterpart DC. The low-shot results of ODC in this benchmark suggest that the learned features through ODC generalize well to low-shot classification.

2 Further Analysis

In this section, we further analyze our ODC model from different perspectives.

ODC as a Fine-tuning Scheme. The high efficiency enables ODC to easily serve as a rapid unsupervised fine-tuning scheme. To assess the fine-tuning ability of ODC, we also use our reimplemented DC to fine-tune other self-supervised models. The improvements over different self-supervised approaches are shown in Table 4. Compared with DC, we observe that ODC boosts the performance of each self-supervised approach by a significant margin. With ODC fine-tuning, we achieve 16.7% improvements for Col., 9.9% for Jig., 7.1% for Rot., and 7.9% for DC, respectively, on the VOC2007 SVM classification benchmark. By contrast, DC also yields fine-tuning improvements but lags far behind ODC.

Influence of the Hyper-parameters. The hyper-parameters of ODC include the frequency of updating the centroids memory, and the minimal size of clusters. To study the influence of the aforementioned two hyper-parameters, we train models with 90K images that are uniformly sampled from the original 1.28M ImageNet dataset, and evaluate the performance on VOC2007 SVM classification benchmark. Figure 4 shows the influence of the update frequency of centroids memory. We observe no significant decrease in the performance of ODC when the update frequency becomes lower, indicating that our method is insensitive to this hyper-parameter provided that it is within a reasonable range. The influence of the minimal size of small clusters is shown in Figure 4. The results show that a large threshold ( i.e. 160) on clusters size would lead to a performance drop. The result is not surprising. A cluster whose size is smaller than the minimal size is identified as a “small cluster”. An overly frequent processing of such small clusters (see Sec. 3.2) introduces instability in feature learning. The large threshold would also group images that should not have belonged to the same class. It is noteworthy that ODC does not experience a significant change in performance within a reasonable range of minimal cluster sizes.

Stability and Convergence. Figure 1 already demonstrates the superior stability of ODC over DC from the aspect of the loss curve. In Figure 5, we show the training stability and convergence of ODC throughout the full training iterations. To measure the stability of our models, we record the ratio of samples whose labels are changed in a batch. Intuitively, fewer label switchings suggest a higher stability. We report the ratio when different backbones are trained from scratch with ODC. The curves begin with the highest label-switching ratio, i.e., nearly 100% of samples in a batch experience a switch in their labels. Gradually, the label-switching ratio declines and converges to a relatively low value. Though there is always a small portion of samples altering their labels at last, ODC reaches a stable state.

Training on Long-Tailed Data. In all previous experimens, we train our models on the class-balanced ImageNet dataset. To evaluate the learning efficacy of ODC on long-tailed data, we perform experiments on downsampled long-tail ImageNet following . Specifically, we randomly downsample 300 classes with 100K images from the original ImageNet dataset to make different levels of long-tail ImageNet datasets, where the ratio of the largest class to the smallest class ranges from 1 (the non-long-tail level) to 64 (the highest long-tail level). Figure 6 shows the performance of ODC trained on different levels of long-tail ImageNet. We observe no significant performance drop even in the conditions with large long-tail degrees, suggesting the robustness of our method on long-tailed data.

Visualization of Clusters. We visualize some selected clusters as shown in Figure 7. Since the number of clusters is much larger than that of the original annotations, there will certainly be some clusters that represent new semantics beyond the annotated classes. We find new classes, e.g., “hand” and “feet”, and new relations, e.g., “animal in cage”, “person holds dog” and “person leads dog with a rope”, that are discovered by ODC. The phenomenon reveals the potential of unsupervised learning to capture new semantics beyond manual annotations.

Conclusion

We have proposed an effective joint clustering and feature learning paradigm for unsupervised representation learning. The proposed approach, Online Deep Clustering (ODC), attains effective and stable unsupervised training of deep neural networks, via decomposing feature clustering and integrating the process into iterations of network update. ODC performs compellingly as an unsupervised representation learning scheme alone. It can also be used to fine-tune and substantially improve previous self-supervised learning methods.

Acknowledgements. This work is supported by the SenseTime-NTU Collaboration Project, Singapore MOE AcRF Tier 1 (2018-T1-002-056), NTU SUG, NTU NAP, the Max Planck-NTU Joint Lab for Artificial Senses and Data Science and Artificial Intelligence Research Lab. We thank Yue Zhao for his participation in discussing the idea.

References