P ROTOTYPICAL C ONTRASTIVE L EARNING OF U NSUPERVISED R EPRESENTATIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
This paper presents Prototypical Contrastive Learning (PCL), an unsupervised representation learning method that bridges contrastive learning with clustering.
PCL not only learns low-level features for the task of instance discrimination, but more importantly, it encodes semantic structures discovered by clustering into the learned embedding space.
Speciﬁcally, we introduce prototypes as latent variables to help ﬁnd the maximum-likelihood estimation of the network parameters in an
Expectation-Maximization framework.
We iteratively perform E-step as ﬁnding the distribution of prototypes via clustering and M-step as optimizing the network via contrastive learning.
We propose ProtoNCE loss, a generalized version of the
InfoNCE loss for contrastive learning, which encourages representations to be closer to their assigned prototypes.
PCL outperforms state-of-the-art instance-wise contrastive learning methods on multiple benchmarks with substantial improvement in low-resource transfer learning.
I NTRODUCTION
Unsupervised visual representation learning aims to learn image representations from pixels themselves without relying on semantic annotations, and recent advances are largely driven by instance discrimination tasks (Wu et al., 2018;
Ye et al., 2019;
He et al., 2020; Misra & van der Maaten, 2020; Hjelm et al., 2019; Oord et al., 2018; Tian et al., 2019).
These methods usually consist of two key components: image transformation and contrastive loss.
Image transformation aims to generate multiple embeddings that represent the same image, by data augmentation (Ye et al., 2019;
Bachman et al., 2019; Chen et al., 2020a), patch perturbation (Misra & van der Maaten, 2020), or using momentum features (He et al., 2020).
The contrastive loss, in the form of a noise contrastive estimator (Gutmann & Hyvärinen, 2010), aims to bring closer samples from the same instance and separate samples from different instances.
Essentially, instance-wise contrastive learning leads to an embedding space where all instances are well-separated, and each instance is locally smooth ( i .
e .
input with perturbations have similar representations).
Despite their improved performance, instance discrimination methods share a common weakness: the representation is not encouraged to encode the semantic structure of data.
This problem arises because instance-wise contrastive learning treats two samples as a negative pair as long as they are from different instances, regardless of their semantic similarity.
This is magniﬁed by the fact that thousands of negative samples are generated to form the contrastive loss, leading to many negative pairs that share similar semantics but are undesirably pushed apart in the embedding space.
In this paper, we propose prototypical contrastive learning (PCL), a new framework for unsupervised representation learning that implicitly encodes the semantic structure of data into the embedding space.
Figure 1 shows an illustration of PCL.
A prototype is deﬁned as “a representative embedding for a group of semantically similar instances”.
We assign several prototypes of different granularity to each instance, and construct a contrastive loss which enforces the embedding of a sample to be more similar to its corresponding prototypes compared to other prototypes.
In practice, we can ﬁnd prototypes by performing clustering on the embeddings.
We formulate prototypical contrastive learning as an Expectation-Maximization (EM) algorithm, where the goal is to ﬁnd the parameters of a Deep Neural Network (DNN) that best describes the data distribution, by iteratively approximating and maximizing the log-likelihood function.
Speciﬁcally, we introduce prototypes as additional latent variables, and estimate their probability in the E-step by
[Figure]
performing k -means clustering.
In the M-step, we update the network parameters by minimizing our proposed contrastive loss, namely ProtoNCE .
We show that minimizing ProtoNCE is equivalent to maximizing the estimated log-likelihood, under the assumption that the data distribution around each prototype is isotropic Gaussian.
Under the EM framework, the widely used instance discrimination task can be explained as a special case of prototypical contrastive learning, where the prototype for each instance is its augmented feature, and the Gaussian distribution around each prototype has the same ﬁxed variance.
The contributions of this paper can be summarized as follows:
• We propose prototypical contrastive learning, a novel framework for unsupervised representation learning that bridges contrastive learning and clustering.
The learned representation is encouraged to capture the hierarchical semantic structure of the dataset.
• We give a theoretical framework that formulates PCL as an Expectation-Maximization (EM) based algorithm.
The iterative steps of clustering and representation learning can be interpreted as approximating and maximizing the log-likelihood function.
The previous methods based on instance discrimination form a special case in the proposed EM framework.
• We propose ProtoNCE, a new contrastive loss which improves the widely used InfoNCE by dynamically estimating the concentration for the feature distribution around each prototype.
ProtoNCE also includes an InfoNCE term in which the instance embeddings can be interpreted as instancebased prototypes.
We provide explanations for PCL from an information theory perspective, by showing that the learned prototypes contain more information about the image classes.
• PCL outperforms instance-wise contrastive learning on multiple benchmarks with substantial improvements in low-resource transfer learning.
PCL also leads to better clustering results.
R ELATED WORK
Our work is closely related to two main branches of studies in unsupervised/self-supervised learning: instance-wise contrastive learning and deep unsupervised clustering.
Instance-wise contrastive learning (Wu et al., 2018;
Ye et al., 2019;
He et al., 2020; Misra & van der Maaten, 2020; Zhuang et al., 2019; Hjelm et al., 2019; Oord et al., 2018; Tian et al., 2019;
Chen et al., 2020a) aims to learn an embedding space where samples ( e .
g .
crops) from the same instance ( e .
g .
an image) are pulled closer and samples from different instances are pushed apart.
To construct the contrastive loss, positive instance features and negative instance features are generated for each sample.
Different contrastive learning methods vary in their strategy to generate instance features.
The memory bank approach (Wu et al., 2018) stores the features of all samples calculated in the previous step.
The end-to-end approach (Ye et al., 2019; Tian et al., 2019; Chen et al., 2020a) generates instance features using all samples within the current mini-batch.
The momentum encoder approach (He et al., 2020) encodes samples on-the-ﬂy by a momentum-updated encoder, and maintains a queue of instance features.
Despite their improved performance, the existing methods based on instance-wise contrastive learning have the following two major limitations, which can be addressed by the proposed PCL framework.
• The task of instance discrimination could be solved by exploiting low-level image differences, thus the learned embeddings do not necessarily capture high-level semantics.
This is supported by the fact that the accuracy of instance classiﬁcation often rapidly rises to a high level (>90% within 10 epochs) and further training gives limited informative signals.
A recent study also shows that better performance of instance discrimination could worsen the performance on downstream tasks (Tschannen et al., 2020).
• A sufﬁciently large number of negative instances need to be sampled, which inevitably yields negative pairs that share similar semantic meaning and should be closer in the embedding space.
However, they are undesirably pushed apart by the contrastive loss.
Such problem is deﬁned as class collision in (Saunshi et al., 2019) and is shown to hurt representation learning.
Essentially, instance discrimination learns an embedding space that only preserves the local smoothness around each instance but largely ignores the global semantic structure of the dataset.
Deep unsupervised clustering.
Clustering based methods have been proposed for deep unsupervised learning.
Xie et al.
; Yang et al.
; Liao et al.
; Yang et al.
; Chang et al.
; Ji et al.
) jointly learn image embeddings and cluster assignments, but they have not shown the ability to learn transferable representations from a large scale of images.
Closer to our work, DeepCluster (Caron et al., 2018) performs iterative clustering and unsupervised representation learning, which is further improved by Zhan et al.
) with online clustering.
However, our method is conceptually different from DeepCluster.
In DeepCluster, the cluster assignments are considered as pseudo-labels and a classiﬁcation objective is optimized, which results in two weaknesses: (1) the high-dimensional features from the penultimate layer of a ConvNet are not optimal for clustering and need to be PCA-reduced; (2) an additional linear classiﬁcation layer is frequently re-initialized which interferes with representation learning.
In our method, representation learning happens directly in a low-dimensional embedding space, by optimizing a contrastive loss on the prototypes (cluster centroids).
Concurrent to our work, SwAV (Caron et al., 2020) also brings together a clustering objective with contrastive learning.
However, their clustering is performed on a subset of the data.
Self-supervised pretext tasks.
Another line of self-supervised learning methods focus on training
DNNs to solve pretext tasks, which usually involve hiding certain information about the input and training the network to recover those missing information.
Examples include image inpainting (Pathak et al., 2016), colorization (Zhang et al., 2016; 2017), prediction of patch orderings (Doersch et al., 2015; Noroozi & Favaro, 2016) and image transformations (Dosovitskiy et al., 2014; Gidaris et al., 2018; Caron et al., 2019; Zhang et al., 2019).
Compared to heuristic pretext task designs, the proposed
PCL is a more general learning framework with better theoretical justiﬁcation.
P ROTOTYPICAL C ONTRASTIVE L EARNING 3. P RELIMINARIES
Given a training set X = { x 1 , x 2 , ..., x n } of n images, unsupervised visual representation learning aims to learn an embedding function f ✓ (realized via a DNN) that maps X to V
= { v 1 , v 2 , ..., v n } with v i =
f ✓ ( x i ) , such that v i best describes x i .
Instance-wise contrastive learning achieves this objective by optimizing a contrastive loss function, such as InfoNCE (Oord et al., 2018;
He et al., 2020), deﬁned as:
[Equation]
where v i {+and v 
 i are+} positive embeddings for instance i , and v 0 j includes one positive embedding and r negative embeddings for other instances, and ⌧ is a temperature hyper-parameter.
In MoCo (He et al., 2020), these embeddings are obtained by feeding x i to a momentum encoder parametrized by ✓ 0 , v 0 i =
f ✓ 0 ( x i ) , where ✓ 0 is a moving average of ✓ .
In prototypical contrastive learning, we use prototypes c instead of v 0 , and replace the ﬁxed temperature ⌧ with a per-prototype concentration estimation  .
An overview of our training framework is shown in Figure 2, where clustering and representation learning are performed iteratively at each epoch.
Next, we will delineate the theoretical framework of PCL based on EM.
A pseudo-code of our algorithm is given in appendix B.
[Figure]
3. PCL AS E XPECTATION -M AXIMIZATION
Our objective is to ﬁnd the network parameters ✓ that maximizes the log-likelihood function of the observed n samples:
[Equation]
We assume that the observed data { x i } ni =1 are related to latent variable C
= { c i } ki =1 which denotes the prototypes of the data.
In this way, we can re-write the log-likelihood function as:
[Equation]
It is hard to optimize this function directly, so we use a surrogate function to lower-bound it: n
[Equation]
i
[Equation]
where Q ( c i ) denotes some distribution over c ’s (
P c i 2 C Q ( c i ) = 1 ), and the last step of derivation uses Jensen’s inequality.
To make the inequality hold with equality, we require p ( x i ,c
i ; ✓ ) Q ( c i ) to be a constant.
Therefore, we have:
[Equation]
By ignoring the constant
P
[Equation]
=
[Equation]
Q
[Equation]
c i ) log Q ( c i ) in eqn.(4), we should maximize:
[Equation]
E-
step.
In this step, we aim to estimate p ( c i ; x i , ✓ ) .
To this end, we perform k -means on the features v 0 i = f ✓ 0 ( x i ) given by the momentum encoder to obtain k clusters.
We deﬁne prototype c i as the centroid for the i -th cluster.
Then, we compute p ( c i ; x i , ✓ ) =
( x i 2 c i ) , where ( x i 2 c i )
=
if x i belongs to the cluster represented by c i ; otherwise ( x i 2 c i ) = 0 .
Similar to MoCo, we found features from the momentum encoder yield more consistent clusters.
M-step.
Based on the E-step, we are ready to maximize the lower-bound in eqn.(6).
[Equation]
Under the assumption of a uniform prior over cluster centroids, we have:
[Equation] [Equation]
where we set the prior probability p ( c i ; ✓ ) for each c i as 1 /k since we are not provided any samples.
We assume that the distribution around each prototype is an isotropic Gaussian, which leads to:
[Equation]
where v i = f ✓ ( x i ) and x i 2 c s .
If we apply ` 2 -normalization to both v and c , then ( v  c ) 2 = 2  2 v · c .
Combining this with eqn.(3, 4, 6, 7, 8, 9), we can write maximum log-likelihood estimation as
[Equation]
where  /  2 denotes the concentration level of the feature distribution around a prototype and will be introduced later.
Note that eqn.(10) has a similar form as the InfoNCE loss in eqn.(1).
Therefore,
InfoNCE can be interpreted as a special case of the maximum log-likelihood estimation, where the prototype for a feature v i is the augmented feature v 0 i from the same instance ( i .
e .
c = v 0 ), and the concentration of the feature distribution around each instance is ﬁxed ( i .
e .
= ⌧ ).
In practice, we take the same approach as NCE and sample r negative prototypes to calculate the normalization term.
We also cluster the samples M times with different number of clusters
K = { k m } Mm =1 , which enjoys a more robust probability estimation of prototypes that encode the hierarchical structure.
Furthermore, we add the InfoNCE loss to retain the property of local smoothness and help bootstrap clustering.
Our overall objective, namely ProtoNCE , is deﬁned as
[Equation]
3. C ONCENTRATION ESTIMATION
The distribution of embeddings around each prototype has different level of concentration.
We use  to denote the concentration estimation, where a smaller  indicates larger concentration.
Here we calculate  using the momentum features { v 0 z } Zz =1 that are within the same cluster as a prototype c .
The desired  should be small (high concentration) if (1) the average distance between v 0 z and c is small, and (2) the cluster contains more feature points ( i .
e .
Z is large).
Therefore, we deﬁne  as:
[Equation]
where ↵ is a smooth parameter to ensure that small clusters do not have an overly-large  .
We normalize  for each set of prototypes C m such that they have a mean of ⌧ .
In the ProtoNCE loss (eqn.(11)),  m s acts as a scaling factor on the similarity between an embedding v i and its prototype c ms .
With the proposed  , the similarity in a loose cluster (larger  ) are down-scaled, pulling embeddings closer to the prototype.
On the contrary, embeddings in a tight cluster (smaller  ) have an up-scaled similarity, thus less encouraged to approach the prototype.
Therefore, learning with ProtoNCE yields more balanced clusters with similar concentration, as shown in Figure 3(a).
It prevents a trivial solution where most embeddings collapse to a single cluster, a problem that could only be heuristically addressed by data-resampling in DeepCluster (Caron et al., 2018).
3. M UTUAL INFORMATION ANALYSIS
It has been shown that minimizing InfoNCE is maximizing a lower bound on the mutual information (MI) between representations V and V
(Oord et al., 2018).
Similarly, minimizing the proposed
ProtoNCE can be considered as simultaneously maximizing the mutual information between V and all the prototypes { V 0 , C 1 , ..., C M } .
This leads to better representation learning, for two reasons.
First, the encoder would learn the shared information among prototypes, and ignore the individual noise that exists in each prototype.
The shared information is more likely to capture higher-level semantic knowledge.
Second, we show that compared to instance features, prototypes have a larger mutual information with the class labels .
We estimate the mutual information between the
[Figure]
instance features (or their assigned prototypes) and the ground-truth class labels for all images in
ImageNet (Deng et al., 2009) training set, following the method in (Ross, 2014).
We compare the obtained MI of our method (ProtoNCE) and that of MoCo (InfoNCE).
As shown in Figure 3(b), compared to instance features, the prototypes have a larger MI with the class labels due to the effect of clustering.
Furthermore, compared to InfoNCE, training on ProtoNCE can increase the MI of prototypes as training proceeds, indicating that better representations are learned to form more semantically-meaningful clusters.
3. P ROTOTYPES AS LINEAR CLASSIFIER
Another interpretation of PCL can provide more insights into the nature of the learned prototypes.
The optimization in eqn.(10) is similar to optimizing the cluster-assignment probability p ( s ; x i , ✓ ) using the cross-entropy loss, where the prototypes c represent weights for a linear classiﬁer.
With k -means clustering, the linear classiﬁer has a ﬁxed set of weights as the mean vectors for the representations in each cluster, c = 1 Z
P
Zz =1 v 0 z .
A similar idea has been used for few-shot learning (Snell et al., 2017), where a non-parametric prototypical classiﬁer performs better than a parametric linear classiﬁer.
E XPERIMENTS
We evaluate PCL on transfer learning tasks, based on the principle that good representations should transfer with limited supervision and ﬁne-tuning.
We follow the settings in MoCo, therefore direct comparisons with MoCo could demonstrate the improvement from the prototypical contrastive loss.
4. I MPLEMENTATION DETAILS
To enable a fair comparison, we follow the same setting as MoCo.
We perform training on the
ImageNet-1M dataset.
A ResNet-
(He et al., 2016) is adopted as the encoder, whose last fullyconnected layer outputs a 128-D and L2-normalized feature.
We follow previous works (He et al., 2020; Wu et al., 2018) and perform data augmentation with random crop, random color jittering, random horizontal ﬂip, and random grayscale conversion.
We use SGD as our optimizer, with a weight decay of 0.0001, a momentum of 0.9, and a batch size of 256.
We train for 200 epochs, where we warm-up the network in the ﬁrst 20 epochs by only using the InfoNCE loss.
The initial learning rate is 0.03, and is multiplied by 0.1 at 120 and 160 epochs.
In terms of the hyper-parameters, we set ⌧ = 0 . 1 , ↵ = 10 , r = 16000 , and number of clusters K = { 25000 , 50000 , 100000 } . We also experiment with PCL v2 using improvements introduced by Chen et al.
(2020a;b), which includes a
MLP projection layer, stronger data augmentation with additional Gaussian blur, and temperature ⌧ = 0 .
We adopt faiss (Johnson et al., 2017) for efﬁcient k -means clustering.
The clustering is performed per-epoch on center-cropped images. We ﬁnd over-clustering to be beneﬁcial. Recent advances in self-supervised learning have been propelled by huge compute which is inaccessible to many researchers. We instead target a more commonly accessible training resource for PCL with
NVIDIA-V100 GPUs and approximately 5 days [Table] [Table] [Table]
Low-shot classiﬁcation.
We evaluate the learned representation on image classiﬁcation tasks with few training samples per-category.
We follow the setup in Goyal et al.
and train linear SVMs using ﬁxed representations on two datasets: Places
(Zhou et al., 2014) for scene recognition and PASCAL VOC
(Everingham et al., 2010) for object classiﬁcation.
We vary the number k of samples per-class and report the average result across 5 independent runs (standard deviation is reported in appendix C).
Table 1 shows the results, in which our method substantially outperforms both MoCo and SimCLR.
Semi-supervised image classiﬁcation.
We perform semi-supervised learning experiments to evaluate whether the learned representation can provide a good basis for ﬁne-tuning.
Following the setup from Wu et al.
; Misra & van der Maaten (2020), we randomly select a subset (1% or 10%) of
ImageNet training data (with labels), and ﬁne-tune the self-supervised trained model on these subsets.
Table 2 reports the top-5 accuracy on ImageNet validation set.
Our method sets a new state-of-the-art under 200 training epochs, outperforming both self-supervised learning methods and semi-supervised learning methods.
The standard deviation across 5 runs is low ( < 0 .
6 for 1% labels).
[Table]
4. I MAGE CLASSIFICATION BENCHMARKS
Linear classiﬁers.
Next, we train linear classiﬁers on ﬁxed image representations using the entire labeled training data.
We evaluate the performance of such linear classiﬁers on three datasets:
ImageNet, VOC
[Table]
and Places205.
Table 3 reports the results.
PCL outperforms MoCo under direct comparison, which demonstrate the advantage of the proposed prototypical contrastive loss [Table]
KNN classiﬁers.
We perform k-nearest neighbor (kNN) classiﬁcation on ImageNet.
For a query image with feature v , we take its top k nearest neighbors from the momentum features, and perform weighted-combination of their labels where the weights are calculated by exp( v · v 0 i / ⌧ ) .
Tablereports the accuracy.
Our method substantially outperforms previous methods
[Table]
4. C LUSTERING EVALUATION
In Table 5, we evaluate the k -means clustering performance on ImageNet using representations learned by different methods.
PCL leads to substantially higher adjusted mutual information (AMI) score.
Details are given in appendix F.
[Table]
4. O BJECT DETECTION
We assess the representation on object detection.
Following Goyal et al.
(2019), we train a Faster
R-CNN (Ren et al., 2015) on VOC07 or VOC07+12, and evaluate on the test set of VOC07.
We keep the pretrained backbone frozen to better evaluate the learned representation, and use the same schedule for all methods.
Table 6 reports the average mAP across three runs.
Our method substantially closes the gap between self-supervised methods and supervised training.
In appendix D, we show the results for ﬁne-tuning the pretrained model for object detection and instance segmentation on
COCO (Lin et al., 2014), where PCL outperforms both MoCo and supervised training.
[Table]
V ISUALIZATION OF LEARNED REPRESENTATION
In Figure 4, we visualize the unsupervised learned representation of ImageNet training images using t-SNE (Maaten & Hinton, 2008).
Compared to the representation learned by MoCo, the representation learned by the proposed PCL forms more separated clusters, which also suggests representation of lower entropy.
C ONCLUSION
This paper proposes Prototypical Contrastive Learning, a generic unsupervised representation learning framework that ﬁnds network parameters to maximize the log-likelihood of the observed data.
We introduce prototypes as latent variables, and perform iterative clustering and representation learning in an EM-based framework.
PCL learns an embedding space which encodes the semantic structure of data, by training on the proposed ProtoNCE loss.
Our extensive experiments on multiple benchmarks demonstrate the advantage of PCL for unsupervised representation learning.