Unsupervised Learning of Visual Features by Contrasting Cluster Assignments โ€” part 2/3

Cleaned-text code points [14620, 29877). Bibliography omitted here; source.md includes it.

In order to make our method online, we compute the codes using only the image features within a batch. Intuitively, as the prototypes ๐‚ \mathbf{C} are used across different batches, SwAV clusters multiple instances to the prototypes. We compute codes using the prototypes ๐‚ \mathbf{C} such that all the examples in a batch are equally partitioned by the prototypes. This equipartition constraint ensures that the codes for different images in a batch are distinct, thus preventing the trivial solution where every image has the same code. Given B B feature vectors ๐™ = [ ๐ณ 1 , โ€ฆ , ๐ณ B ] \mathbf{Z}=[\mathbf{z}_{1},\dots,\mathbf{z}_{B}] , we are interested in mapping them to the prototypes ๐‚ = [ ๐œ 1 , โ€ฆ , ๐œ K ] \mathbf{C}=[\mathbf{c}_{1},\dots,\mathbf{c}_{K}] . We denote this mapping or codes by ๐ = [ ๐ช 1 , โ€ฆ , ๐ช B ] \mathbf{Q}=[\mathbf{q}_{1},\dots,\mathbf{q}_{B}] , and optimize ๐ \mathbf{Q} to maximize the similarity between the features and the prototypes , i.e ., max ๐ โˆˆ ๐’ฌ โก Tr โ€‹ ( ๐ โŠค โ€‹ ๐‚ โŠค โ€‹ ๐™ ) + ฮต โ€‹ H โ€‹ ( ๐ ) , \max_{\mathbf{Q}\in\mathcal{Q}}\ \text{Tr}\left(\mathbf{Q}^{\top}\mathbf{C}^{\top}\mathbf{Z}\right)+\varepsilon H(\mathbf{Q}), (3) where H H is the entropy function, H ( ๐ ) = โˆ’ โˆ‘ i โ€‹ j ๐ i โ€‹ j log ๐ i โ€‹ j H(\mathbf{Q})=-\sum_{ij}\mathbf{Q}_{ij}\log\mathbf{Q}_{ij} and ฮต \varepsilon is a parameter that controls the smoothness of the mapping. We observe that a strong entropy regularization ( i.e . using a high ฮต \varepsilon ) generally leads to a trivial solution where all samples collapse into an unique representation and are all assigned uniformely to all prototypes. Hence, in practice we keep ฮต \varepsilon low. Asano et al . [ 2 ] enforce an equal partition by constraining the matrix ๐ \mathbf{Q} to belong to the transportation polytope. They work on the full dataset, and we propose to adapt their solution to work on minibatches by restricting the transportation polytope to the minibatch: ๐’ฌ = { ๐ โˆˆ โ„ + K ร— B | ๐๐Ÿ B = 1 K ๐Ÿ K , ๐ โŠค ๐Ÿ K = 1 B ๐Ÿ B } , \mathcal{Q}=\left\{\mathbf{Q}\in\mathbb{R}_{+}^{K\times B}~|~\mathbf{Q}\mathbf{1}_{B}=\frac{1}{K}\mathbf{1}_{K},\mathbf{Q}^{\top}\mathbf{1}_{K}=\frac{1}{B}\mathbf{1}_{B}\right\}, (4) where ๐Ÿ K \mathbf{1}_{K} denotes the vector of ones in dimension K K . These constraints enforce that on average each prototype is selected at least B K \frac{B}{K} times in the batch.

Once a continuous solution ๐ โˆ— \mathbf{Q}^{*} to Prob. ( 3 ) is found, a discrete code can be obtained by using a rounding procedure [ 2 ] . Empirically, we found that discrete codes work well when computing codes in an offline manner on the full dataset as in Asano et al . [ 2 ] . However, in the online setting where we use only minibatches, using the discrete codes performs worse than using the continuous codes. An explanation is that the rounding needed to obtain discrete codes is a more aggressive optimization step than gradient updates. While it makes the model converge rapidly, it leads to a worse solution. We thus preserve the soft code ๐ โˆ— \mathbf{Q}^{*} instead of rounding it. These soft codes ๐ โˆ— \mathbf{Q}^{*} are the solution of Prob. ( 3 ) over the set ๐’ฌ \mathcal{Q} and takes the form of a normalized exponential matrix [ 13 ] : ๐ โˆ— = Diag โ€‹ ( ๐ฎ ) โ€‹ exp โก ( ๐‚ โŠค โ€‹ ๐™ ฮต ) โ€‹ Diag โ€‹ ( ๐ฏ ) , \mathbf{Q}^{*}=\text{Diag}(\mathbf{u})\exp\left(\frac{\mathbf{C}^{\top}\mathbf{Z}}{\varepsilon}\right)\text{Diag}(\mathbf{v}), (5) where ๐ฎ \mathbf{u} and ๐ฏ \mathbf{v} are renormalization vectors in โ„ K \mathbb{R}^{K} and โ„ B \mathbb{R}^{B} respectively. The renormalization vectors are computed using a small number of matrix multiplications using the iterative Sinkhorn-Knopp algorithm [ 13 ] . In practice, we observe that using only 3 3 iterations is fast and sufficient to obtain good performance. Indeed, this algorithm can be efficiently implemented on GPU, and the alignment of 4 4 K features to 3 3 K codes takes 35 35 ms in our experiments, see ยง 4 .

##### Working with small batches.

When the number B B of batch features is too small compared to the number of prototypes K K , it is impossible to equally partition the batch into the K K prototypes. Therefore, when working with small batches, we use features from the previous batches to augment the size of ๐™ \mathbf{Z} in Prob. ( 3 ). Then, we only use the codes of the batch features in our training loss. In practice, we store around 3 3 K features, i.e ., in the same range as the number of code vectors. This means that we only keep features from the last 15 15 batches with a batch size of 256 256 , while contrastive methods typically need to store the last 65 65 K instances obtained from the last 250 250 batches [ 24 ] .

### 3.2 Multi-crop: Augmenting views with smaller images

As noted in prior works [ 10 , 44 ] , comparing random crops of an image plays a central role by capturing information in terms of relations between parts of a scene or an object. Unfortunately, increasing the number of crops or โ€œviewsโ€ quadratically increases the memory and compute requirements. We propose a multi-crop strategy where we use two standard resolution crops and sample V V additional low resolution crops that cover only small parts of the image. Using low resolution images ensures only a small increase in the compute cost. Specifically, we generalize the loss of Eq ( 1 ): L โก ( ๐ณ t 1 , ๐ณ t 2 , โ€ฆ , ๐ณ t V + 2 ) = โˆ‘ i โˆˆ { 1 , 2 } โˆ‘ v = 1 V + 2 ๐Ÿ v โ‰  i โ€‹ โ„“ โ€‹ ( ๐ณ t v , ๐ช t i ) . L(\mathbf{z}_{t_{1}},\mathbf{z}_{t_{2}},\dots,\mathbf{z}_{t_{V+2}})=\sum_{i\in\{1,2\}}\sum_{v=1}^{V+2}\mathbf{1}_{v\neq i}\ell(\mathbf{z}_{t_{v}},\mathbf{q}_{t_{i}}). (6) Note that we compute codes using only the full resolution crops. Indeed, computing codes for all crops increases the computational time and we observe in practice that it also alters the transfer performance of the resulting network. An explanation is that using only partial information (small crops cover only small area of images) degrades the assignment quality. Figure 3 shows that multi-crop improves the performance of several self-supervised methods and is a promising augmentation strategy.

## 4 Main Results

We analyze the features learned by SwAV by transfer learning on multiple datasets. We implement in SwAV the improvements used in SimCLR, i.e ., LARS [ 64 ] , cosine learning rate [ 40 , 44 ] and the MLP projection head [ 10 ] . We provide the full details and hyperparameters for pretraining and transfer learning in the Appendix.

### 4.1 Evaluating the unsupervised features on ImageNet

We evaluate the features of a ResNet-50 [ 27 ] trained with SwAV on ImageNet by two experiments: linear classification on frozen features and semi-supervised learning by finetuning with few labels. When using frozen features ( fig. 2 left), SwAV outperforms the state of the art by + 4.2 % +4.2\% top-1 accuracy and is only 1.2 % 1.2\% below the performance of a fully supervised model. Note that we train SwAV during 800 800 epochs with large batches ( 4096 4096 ). We refer to fig. 3 for results with shorter trainings and to table 3 for experiments with small batches. On semi-supervised learning ( table 1 ), SwAV outperforms other self-supervised methods and is on par with state-of-the-art semi-supervised models [ 51 ] , despite the fact that SwAV is not specifically designed for semi-supervised learning.

Variants of ResNet-50. Figure 2 (right) shows the performance of multiple variants of ResNet-50 with different widths [ 35 ] . The performance of our model increases with the width of the model, and follows a similar trend to the one obtained with supervised learning. When compared with concurrent work like SimCLR, we see that SwAV reduces the difference with supervised models even further. Indeed, for large architectures, our method shrinks the gap with supervised training to 0.6 % 0.6\% .

### 4.2 Transferring unsupervised features to downstream tasks

We test the generalization of ResNet-50 features trained with SwAV on ImageNet (without labels) by transferring to several downstream vision tasks. In table 2 , we compare the performance of SwAV features with ImageNet supervised pretraining. First, we report the linear classification performance on the Places205 [ 67 ] , VOC07 [ 17 ] , and iNaturalist2018 [ 54 ] datasets. Our method outperforms supervised features on all three datasets. Note that SwAV is the first self-supervised method to surpass ImageNet supervised features on these datasets. Second, we report network finetuning on object detection on VOC07+12 using Faster R-CNN [ 50 ] with a R50-C4 backbone and on COCO [ 38 ] using Mask R-CNN [ 26 ] with a R50-FPN backbone and finally using DETR detector [ 6 ] . DETR is a recent object detection framework that reaches competitive performance with Faster R-CNN while being conceptually simpler and trainable end-to-end. Interestingly, unlike Faster R-CNN [ 50 ] , using a pretrained backbone for DETR is crucial to obtain good results compared to training from scratch [ 6 ] . In table 2 , we show that SwAV outperforms the supervised pretrained model on both VOC07+12 and COCO datasets. Note that this is line with previous works that also show that self-supervision can outperform supervised pretraining on object detection [ 44 , 24 , 19 ] . We report more detection evaluation metrics and results from other self-supervised methods in the Appendix. Overall, our SwAV ResNet-50 model surpasses supervised ImageNet pretraining on all the considered transfer tasks and datasets. We have released this model so other researchers might also benefit by replacing the ImageNet supervised network with our model.

### 4.3 Training with small batches

We train SwAV with small batches of 256 256 images on 4 4 GPUs and compare with MoCov2 and SimCLR trained in the same setup. In table 3 , we see that SwAV maintains state-of-the-art performance even when trained in the small batch setting. Note that SwAV only stores a queue of 3,840 3,840 features. In comparison, to obtain good performance, MoCov2 needs to store 65,536 65,536 features while keeping an additional momentum encoder network. When SwAV is trained using ร— 160 + ร— 96 2\!\times\!160+4\!\times\!96 crops, SwAV has a running time 1.2 ร— 1.2\times higher than SimCLR with ร— 224 2\!\times\!224 crops and is around 1.4 ร— 1.4\times slower than MoCov2 due to the additional back-propagation [ 11 ] . Hence, one epoch of MoCov2 or SimCLR is faster in wall clock time than one of SwAV, but these methods need more epochs for good downstream performance. Indeed, as shown in table 3 , SwAV learns much faster and reaches higher performance in 4 ร— 4\times fewer epochs: 72 % 72\% after 200 200 epochs ( 102 102 hours) while MoCov2 needs 800 800 epochs to achieve 71.1 % 71.1\% . Increasing the resolution and the number of epochs, SwAV reaches 74.3 % 74.3\% with a small batch size, a small number of stored features and no momentum encoder. Finally, note that SwAV could be combined with a momentum mechanism and a large queue [ 24 ] ; we leave these explorations to future work.

## 5 Ablation Study

### 5.1 Clustering-based self-supervised learning

##### Improving prior clustering-based approaches.

In this section, we re-implement and improve previously published clustering-based models in order to assess if they can compete with recent contrastive methods such as SimCLR. In particular, we consider two clustering-based models: DeepCluster-v2 and SeLa-v2, which are obtained by applying various training improvements introduced in other self-supervised learning papers to DeepCluster [ 7 ] and SeLa [ 2 ] . Among these improvements are the use of stronger data augmentation [ 10 ] , MLP projection head [ 10 ] , cosine learning rate schedule [ 44 ] , temperature parameter [ 58 ] , memory bank [ 58 ] , multi-clustering [ 2 ] , etc. Full implementation details can be found in the Appendix. Besides, we also improve DeepCluster model by introducing explicit comparisons to k-means centroids, which increase stability and performance. Indeed, a main issue in DeepCluster is that there is no correspondance between two consecutive cluster assignments. Hence, the final classification layer learned for an assignment becomes irrelevant for the following one and thus needs to be re-initialized from scratch at each epoch. This considerably disrupts the convnet training. In DeepCluster-v2, instead of learning a classification layer predicting the cluster assignments, we perform explicit comparison between features and centroids.

##### Comparing clustering with contrastive instance learning.

In fig. 3 (left), we make a best effort fair comparison between clustering-based and contrastive instance (SimCLR) methods by implementating these methods with the same data augmentation, number of epochs, batch-sizes, etc. In this setting, we observe that SwAV and DeepCluster-v2 outperform SimCLR by 2 % 2\% without multi-crop and by 3.5 % 3.5\% with multi-crop . This suggests the learning potential of clustering-based methods over instance classification.

##### Advantage of SwAV compared to DeepCluster-v2.

In fig. 3 (left), we observe that SwAV performs on par with DeepCluster-v2. In addition, we train DeepCluster-v2 in SwAV best setting (800 epochs - 8 crops) and obtain 75.2 % 75.2\% top-1 accuracy on ImageNet (versus 75.3 % 75.3\% for SwAV). However, unlike SwAV, DeepCluster-v2 is not online which makes it impractical for extremely large datasets ( ยง 5.4 ). For billion scale trainings for example, a single pass on the dataset is usually performed [ 24 ] . DeepCluster-v2 cannot be trained for only one epoch since it works by performing several passes on the dataset to regularly update centroids and cluster assignments for each image.

As a matter of fact, DeepCluster-v2 can be interpreted as a special case of our proposed swapping mechanism: swapping is done across epochs rather than within a batch. Given a crop of an image DeepCluster-v2 predicts the assignment of another crop, which was obtained at the previous epoch. SwAV swaps assignments directly at the batch level and can thus work online.

### 5.2 Applying the multi-crop strategy to different methods

In fig. 3 (left), we report the impact of applying our multi-crop strategy on the performance of a selection of other methods. Details of how we apply multi-crop to SimCLR loss can be found in the Appendix. We see that the multi-crop strategy consistently improves the performance for all the considered methods by a significant margin of โˆ’ 4 % 2\!-\!4\% top-1 accuracy. Interestingly, multi-crop seems to benefit more clustering-based methods than contrastive methods. We note that multi-crop does not improve the supervised model.

### 5.3 Impact of longer training

In fig. 3 (right), we show the impact of the number of training epochs on performance for SwAV with multi-crop . We train separate models for 100 100 , 200 200 , 400 400 and 800 800 epochs and report the top-1 accuracy on ImageNet using the linear classification evaluation. We train each ResNet-50 on 64 64 V100 16GB GPUs and a batch size of 4096 4096 . While SwAV benefits from longer training, it already achieves strong performance after 100 100 epochs, i.e ., 72.1 % 72.1\% in 6h15.

### 5.4 Unsupervised pretraining on a large uncurated dataset