[Paper Review] Prototypical Contrastive Learning of Unsupervised Representations

Junnan Li, Pan Zhou+2 more|arXiv (Cornell University)|May 11, 2020

Domain Adaptation and Few-Shot Learning55 references483 citations

TL;DR

Prototypical Contrastive Learning (PCL) combines contrastive learning with clustering by introducing prototypes and an EM-based training loop, yielding representations that encode semantic structure and improve transfer in unsupervised settings.

ABSTRACT

This paper presents Prototypical Contrastive Learning (PCL), an unsupervised representation learning method that addresses the fundamental limitations of instance-wise contrastive learning. PCL not only learns low-level features for the task of instance discrimination, but more importantly, it implicitly encodes semantic structures of the data into the learned embedding space. Specifically, we introduce prototypes as latent variables to help find the maximum-likelihood estimation of the network parameters in an Expectation-Maximization framework. We iteratively perform E-step as finding 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. Code and pretrained models are available at https://github.com/salesforce/PCL.

Motivation & Objective

• Bridge contrastive learning and clustering to capture semantic structure in unsupervised visual representations.
• Introduce prototypes as latent variables to enable EM-style optimization of network parameters.
• Develop ProtoNCE, a generalized contrastive loss that uses prototype-based concentration estimates.
• Demonstrate improved transfer learning and clustering quality across benchmarks.

Proposed method

• Formulate PCL as an EM algorithm with prototypes as latent variables.
• E-step: assign each sample to clusters by performing k-means on momentum-encoded features to obtain prototypes.
• M-step: update network parameters by minimizing ProtoNCE, a loss combining instance-level InfoNCE with prototype-based terms.
• Introduce per-prototype concentration estimation phi to adapt the softness of assignments.
• Cluster samples with multiple granularities by using several K values and aggregating losses.
• ProtoNCE reduces to a form similar to InfoNCE, while enabling semantic structure learning via prototypes.

Experimental results

Research questions

• RQ1 Can integrating clustering prototypes into contrastive learning yield embeddings that reflect semantic structure beyond instance discrimination?
• RQ2 Does the EM-based prototypical framework improve transfer learning and clustering quality compared to pure instance-wise contrastive methods?

Key findings

• PCL outperforms state-of-the-art instance-wise contrastive methods on multiple benchmarks, especially in low-resource transfer settings.
• ProtoNCE incorporates dynamic concentration estimates that balance cluster sizes and prevent trivial collapse.
• Prototypes exhibit higher mutual information with class labels than instance features, indicating richer semantic encoding.
• PCL yields better clustering metrics (AMI) on ImageNet and improves object detection performance when used as a frozen backbone.
• PCL v2 and related implementations further enhance linear classification and semi-supervised transfer results.