Unsupervised Learning of Visual Features by Contrasting Cluster Assignments — part 1/3
Cleaned-text code points [0, 14620). Bibliography omitted here; source.md includes it.
# Unsupervised Learning of Visual Features by Contrasting Cluster Assignments
###### Abstract
Unsupervised image representations have significantly reduced the gap with supervised pretraining, notably with the recent achievements of contrastive learning methods. These contrastive methods typically work online and rely on a large number of explicit pairwise feature comparisons, which is computationally challenging. In this paper, we propose an online algorithm, SwAV, that takes advantage of contrastive methods without requiring to compute pairwise comparisons. Specifically, our method simultaneously clusters the data while enforcing consistency between cluster assignments produced for different augmentations (or “views”) of the same image, instead of comparing features directly as in contrastive learning. Simply put, we use a “swapped” prediction mechanism where we predict the code of a view from the representation of another view. Our method can be trained with large and small batches and can scale to unlimited amounts of data. Compared to previous contrastive methods, our method is more memory efficient since it does not require a large memory bank or a special momentum network. In addition, we also propose a new data augmentation strategy, multi-crop , that uses a mix of views with different resolutions in place of two full-resolution views, without increasing the memory or compute requirements. We validate our findings by achieving 75.3 % 75.3\% top-1 accuracy on ImageNet with ResNet-50, as well as surpassing supervised pretraining on all the considered transfer tasks.
## 1 Introduction
Unsupervised visual representation learning, or self-supervised learning, aims at obtaining features without using manual annotations and is rapidly closing the performance gap with supervised pretraining in computer vision [ 10 , 24 , 44 ] . Many recent state-of-the-art methods build upon the instance discrimination task that considers each image of the dataset (or “instance”) and its transformations as a separate class [ 16 ] . This task yields representations that are able to discriminate between different images, while achieving some invariance to image transformations. Recent self-supervised methods that use instance discrimination rely on a combination of two elements: (i) a contrastive loss [ 23 ] and (ii) a set of image transformations. The contrastive loss removes the notion of instance classes by directly comparing image features while the image transformations define the invariances encoded in the features. Both elements are essential to the quality of the resulting networks [ 44 , 10 ] and our work improves upon both the objective function and the transformations.
The contrastive loss explicitly compares pairs of image representations to push away representations from different images while pulling together those from transformations, or views, of the same image. Since computing all the pairwise comparisons on a large dataset is not practical, most implementations approximate the loss by reducing the number of comparisons to random subsets of images during training [ 10 , 24 , 58 ] . An alternative to approximate the loss is to approximate the task—that is to relax the instance discrimination problem. For example, clustering-based methods discriminate between groups of images with similar features instead of individual images [ 7 ] . The objective in clustering is tractable, but it does not scale well with the dataset as it requires a pass over the entire dataset to form image “codes” ( i.e ., cluster assignments) that are used as targets during training. In this work, we use a different paradigm and propose to compute the codes online while enforcing consistency between codes obtained from views of the same image. Comparing cluster assignments allows to contrast different image views while not relying on explicit pairwise feature comparisons. Specifically, we propose a simple “swapped” prediction problem where we predict the code of a view from the representation of another view. We learn features by Sw apping A ssignments between multiple V iews of the same image ( SwAV ). The features and the codes are learned online, allowing our method to scale to potentially unlimited amounts of data. In addition, SwAV works with small and large batch sizes and does not need a large memory bank [ 58 ] or a momentum encoder [ 24 ] .
Besides our online clustering-based method, we also propose an improvement to the image transformations. Most contrastive methods compare one pair of transformations per image, even though there is evidence that comparing more views during training improves the resulting model [ 44 ] . In this work, we propose multi-crop that uses smaller-sized images to increase the number of views while not increasing the memory or computational requirements during training. We also observe that mapping small parts of a scene to more global views significantly boosts the performance. Directly working with downsized images introduces a bias in the features [ 53 ] , which can be avoided by using a mix of different sizes. Our strategy is simple, yet effective, and can be applied to many self-supervised methods with consistent gain in performance.
We validate our contributions by evaluating our method on several standard self-supervised benchmarks. In particular, on the ImageNet linear evaluation protocol, we reach 75.3 % 75.3\% top- 1 1 accuracy with a standard ResNet-50, and 78.5 % 78.5\% with a wider model. We also show that our multi-crop strategy is general, and improves the performance of different self-supervised methods, namely SimCLR [ 10 ] , DeepCluster [ 7 ] , and SeLa [ 2 ] , between 2 % 2\% and 4 % 4\% top-1 accuracy on ImageNet. Overall, we make the following contributions: • We propose a scalable online clustering loss that improves performance by + 2 % +2\% on ImageNet and works in both large and small batch settings without a large memory bank or a momentum encoder.
• We introduce the multi-crop strategy to increase the number of views of an image with no computational or memory overhead. We observe a consistent improvement of between 2 % 2\% and 4 % 4\% on ImageNet with this strategy on several self-supervised methods.
• Combining both technical contributions into a single model, we improve the performance of self-supervised by + 4.2 % +4.2\% on ImageNet with a standard ResNet and outperforms supervised ImageNet pretraining on multiple downstream tasks. This is the first method to do so without finetuning the features, i.e ., only with a linear classifier on top of frozen features.
## 2 Related Work
##### Instance and contrastive learning.
Instance-level classification considers each image in a dataset as its own class [ 5 , 16 , 58 ] . Dosovitskiy et al . [ 16 ] assign a class explicitly to each image and learn a linear classifier with as many classes as images in the dataset. As this approach becomes quickly intractable, Wu et al . [ 58 ] mitigate this issue by replacing the classifier with a memory bank that stores previously-computed representations. They rely on noise contrastive estimation [ 22 ] to compare instances, which is a special form of contrastive learning [ 29 , 47 ] . He et al . [ 24 ] improve the training of contrastive methods by storing representations from a momentum encoder instead of the trained network. More recently, Chen et al . [ 10 ] show that the memory bank can be entirely replaced with the elements from the same batch if the batch is large enough. In contrast to this line of works, we avoid comparing every pair of images by mapping the image features to a set of trainable prototype vectors.
##### Clustering for deep representation learning.
Our work is also related to clustering-based methods [ 2 , 4 , 7 , 8 , 30 , 59 , 63 , 68 , 19 , 62 ] . Caron et al . [ 7 ] show that k k -means assignments can be used as pseudo-labels to learn visual representations. This method scales to large uncurated dataset and can be used for pre-training of supervised networks [ 8 ] . However, their formulation is not principled and recently, Asano et al . [ 2 ] show how to cast the pseudo-label assignment problem as an instance of the optimal transport problem. We consider a similar formulation to map representations to prototype vectors, but unlike [ 2 ] we keep the soft assignment produced by the Sinkhorn-Knopp algorithm [ 13 ] instead of approximating it into a hard assignment. Besides, unlike Caron et al . [ 7 , 8 ] and Asano et al . [ 2 ] , we obtain online assignments which allows our method to scale gracefully to any dataset size.
##### Handcrafted pretext tasks.
Many self-supervised methods manipulate the input data to extract a supervised signal in the form of a pretext task [ 14 , 1 , 31 , 34 , 36 , 42 , 45 , 48 , 49 , 55 , 56 , 66 ] . We refer the reader to Jing et al . [ 32 ] for an exhaustive and detailed review of this literature. Of particular interest, Misra and van der Maaten [ 44 ] propose to encode the jigsaw puzzle task [ 46 ] as an invariant for contrastive learning. Jigsaw tiles are non-overlapping crops with small resolution that cover only part ( ∼ 20 % {\sim}20\% ) of the entire image area. In contrast, our multi-crop strategy consists in simply sampling multiple random crops with two different sizes: a standard size and a smaller one.
## 3 Method
Our goal is to learn visual features in an online fashion without supervision. To that effect, we propose an online clustering-based self-supervised method. Typical clustering-based methods [ 2 , 7 ] are offline in the sense that they alternate between a cluster assignment step where image features of the entire dataset are clustered, and a training step where the cluster assignments, i.e ., “codes” are predicted for different image views. Unfortunately, these methods are not suitable for online learning as they require multiple passes over the dataset to compute the image features necessary for clustering. In this section, we describe an alternative where we enforce consistency between codes from different augmentations of the same image. This solution is inspired by contrastive instance learning [ 58 ] as we do not consider the codes as a target, but only enforce consistent mapping between views of the same image. Our method can be interpreted as a way of contrasting between multiple image views by comparing their cluster assignments instead of their features.
More precisely, we compute a code from an augmented version of the image and predict this code from other augmented versions of the same image. Given two image features 𝐳 t \mathbf{z}_{t} and 𝐳 s \mathbf{z}_{s} from two different augmentations of the same image, we compute their codes 𝐪 t \mathbf{q}_{t} and 𝐪 s \mathbf{q}_{s} by matching these features to a set of K K prototypes { 𝐜 1 , … , 𝐜 K } \{\mathbf{c}_{1},\dots,\mathbf{c}_{K}\} . We then setup a “swapped” prediction problem with the following loss function: L ( 𝐳 t , 𝐳 s ) \displaystyle L(\mathbf{z}_{t},\mathbf{z}_{s}) = \displaystyle= ℓ ( 𝐳 t , 𝐪 s ) + ℓ ( 𝐳 s , 𝐪 t ) , \displaystyle\ell(\mathbf{z}_{t},\mathbf{q}_{s})+\ell(\mathbf{z}_{s},\mathbf{q}_{t}), (1) where the function ℓ ( 𝐳 , 𝐪 ) \ell({\mathbf{z}},{\mathbf{q}}) measures the fit between features 𝐳 {\mathbf{z}} and a code 𝐪 {\mathbf{q}} , as detailed later. Intuitively, our method compares the features 𝐳 t \mathbf{z}_{t} and 𝐳 s \mathbf{z}_{s} using the intermediate codes 𝐪 t \mathbf{q}_{t} and 𝐪 s \mathbf{q}_{s} . If these two features capture the same information, it should be possible to predict the code from the other feature. A similar comparison appears in contrastive learning where features are compared directly [ 58 ] . In fig. 1 , we illustrate the relation between contrastive learning and our method.
### 3.1 Online clustering
Each image 𝐱 n \mathbf{x}_{n} is transformed into an augmented view 𝐱 n t \mathbf{x}_{nt} by applying a transformation t t sampled from the set 𝒯 \mathcal{T} of image transformations. The augmented view is mapped to a vector representation by applying a non-linear mapping f θ f_{\theta} to 𝐱 n t \mathbf{x}_{nt} . The feature is then projected to the unit sphere, i.e ., 𝐳 n t = f θ ( 𝐱 n t ) / ‖ f θ ( 𝐱 n t ) ‖ 2 \mathbf{z}_{nt}=f_{\theta}(\mathbf{x}_{nt})/\|f_{\theta}(\mathbf{x}_{nt})\|_{2} . We then compute a code 𝐪 n t \mathbf{q}_{nt} from this feature by mapping 𝐳 n t \mathbf{z}_{nt} to a set of K K trainable prototypes vectors, { 𝐜 1 , … , 𝐜 K } \{\mathbf{c}_{1},\dots,\mathbf{c}_{K}\} . We denote by 𝐂 \mathbf{C} the matrix whose columns are the 𝐜 1 , … , 𝐜 k \mathbf{c}_{1},\dots,\mathbf{c}_{k} . We now describe how to compute these codes and update the prototypes online.
##### Swapped prediction problem.
The loss function in Eq. ( 1 ) has two terms that setup the “swapped” prediction problem of predicting the code 𝐪 t \mathbf{q}_{t} from the feature 𝐳 s \mathbf{z}_{s} , and 𝐪 s \mathbf{q}_{s} from 𝐳 t \mathbf{z}_{t} . Each term represents the cross entropy loss between the code and the probability obtained by taking a softmax of the dot products of 𝐳 i \mathbf{z}_{i} and all prototypes in 𝐂 \mathbf{C} , i.e ., ℓ ( 𝐳 t , 𝐪 s ) = − ∑ k 𝐪 s ( k ) log 𝐩 t ( k ) , where 𝐩 t ( k ) = exp ( 1 τ 𝐳 t ⊤ 𝐜 k ) ∑ k ′ exp ( 1 τ 𝐳 t ⊤ 𝐜 k ′ ) . \ell(\mathbf{z}_{t},\mathbf{q}_{s})=-\sum_{k}\mathbf{q}_{s}^{(k)}\log\mathbf{p}_{t}^{(k)},\quad\text{where}\quad\mathbf{p}_{t}^{(k)}=\frac{\exp\left(\frac{1}{\tau}\mathbf{z}_{t}^{\top}\mathbf{c}_{k}\right)}{\sum_{k^{\prime}}\exp\left(\frac{1}{\tau}\mathbf{z}_{t}^{\top}\mathbf{c}_{k^{\prime}}\right)}. (2) where τ \tau is a temperature parameter [ 58 ] . Taking this loss over all the images and pairs of data augmentations leads to the following loss function for the swapped prediction problem: − 1 N ∑ n = 1 N ∑ s , t ∼ 𝒯 [ 1 τ 𝐳 n t ⊤ 𝐂𝐪 n s + 1 τ 𝐳 n s ⊤ 𝐂𝐪 n t − log ∑ k = 1 K exp ( 𝐳 n t ⊤ 𝐜 k τ ) − log ∑ k = 1 K exp ( 𝐳 n s ⊤ 𝐜 k τ ) ] . -\frac{1}{N}\sum_{n=1}^{N}\sum_{s,t\sim\mathcal{T}}\left[\frac{1}{\tau}\mathbf{z}_{nt}^{\top}\mathbf{C}\mathbf{q}_{ns}+\frac{1}{\tau}\mathbf{z}_{ns}^{\top}\mathbf{C}\mathbf{q}_{nt}-\log\sum_{k=1}^{K}\exp\left(\frac{\mathbf{z}_{nt}^{\top}\mathbf{c}_{k}}{\tau}\right)-\log\sum_{k=1}^{K}\exp\left(\frac{\mathbf{z}_{ns}^{\top}\mathbf{c}_{k}}{\tau}\right)\right]. This loss function is jointly minimized with respect to the prototypes 𝐂 \mathbf{C} and the parameters θ \theta of the image encoder f θ f_{\theta} used to produce the features ( 𝐳 n t ) n , t (\mathbf{z}_{nt})_{n,t} .
##### Computing codes online.