CVPR 2020 (Oral)

Momentum Contrast for Unsupervised Visual Representation Learning

Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, Ross Girshick • Facebook AI Research (FAIR)

"We present Momentum Contrast (MoCo) for unsupervised visual representation learning. By framing contrastive learning as a dynamic dictionary look-up, we build a large and consistent dictionary using a queue and a momentum-updated encoder."

Dynamic Queue Dictionary Momentum Key Encoder Shuffling Batch Norm Outperforms Supervised Pre-training

Why This Matters

While unsupervised representation learning (like GPT and BERT) revolutionized Natural Language Processing, computer vision long remained dominated by supervised pre-training on ImageNet.

The fundamental challenge lies in the difference between their signal spaces. Language has a discrete, structured vocabulary to build tokenized dictionaries. Visual signals, by contrast, are continuous, high-dimensional, and lack structured tokens.

MoCo addresses this by framing unsupervised visual learning as building a large and consistent dynamic dictionary. This simple paradigm shift allows models to learn rich visual concepts without requiring human-annotated labels, closing the performance gap between supervised and unsupervised learning on downstream vision tasks.

The Big Idea

Contrastive learning trains an encoder to perform a dictionary look-up: an encoded query $q$ should match its positive key $k_+$ and be highly dissimilar to negative keys $k_-$. MoCo maintains this dictionary as a Queue of data samples (decoupling dictionary size from mini-batch size) and updates the key encoder using a Momentum-based moving average to ensure representation consistency.

How It Works

1. Contrastive Loss as Dictionary Look-up

The query $q$ and keys $k$ are generated by encoders. The model is trained using the InfoNCE loss, which behaves like a multi-class classifier trying to identify the positive key $k_+$ among a dictionary of size $K+1$:

$$\mathcal{L}_q = -\log \frac{\exp(q \cdot k_+ / \tau)}{\sum_{i=0}^{K} \exp(q \cdot k_i / \tau)}$$

Where $\tau$ is a temperature hyper-parameter controlling the scale of the distribution.

Interactive Widget: InfoNCE Loss Calculator

Adjust the similarity scores and temperature to see how the InfoNCE loss and positive prediction probability change dynamically.

Positive Pair Similarity ($q \cdot k_+$) 0.80

Higher means the query is closer to its augmented view.

Avg Negative Similarity ($q \cdot k_i$) 0.10

Average similarity to random negative samples.

Negatives ($K$) 4096

Temp ($\tau$) 0.07

Calculated InfoNCE Loss

0.24

Positive Key Softmax Prob: 78.5%

2. Queue-based Dictionary

To achieve high representation quality, we need a large dictionary. Standard end-to-end setups couple dictionary size with mini-batch size, limited by GPU memory. MoCo decouples this by using a Queue.

Interactive Widget: Contrastive Loss Architectures

Click the tabs to compare MoCo's queue mechanism with traditional end-to-end and memory bank approaches.

3. Momentum Update

Using a queue introduces a massive challenge: how do we update the key encoder? Backpropagating through a queue of size $K=65,536$ is mathematically impossible due to memory limits.

MoCo solves this by updating the key encoder $\theta_k$ slowly using a momentum-based moving average of the query encoder $\theta_q$:

$$\theta_k \leftarrow m\theta_k + (1-m)\theta_q$$

Where $m \in [0, 1)$ is the momentum coefficient. A large momentum (e.g., $m=0.999$) ensures that the keys in the queue are encoded by highly consistent encoders.

Interactive Widget: Momentum Tracking Simulator

Simulate how the Key Encoder parameters ($\theta_k$) track the Query Encoder ($\theta_q$) over training steps. Notice how high momentum values smooth out sudden updates, maintaining consistency.

Momentum Coefficient ($m$)

0.90

Note: MoCo uses a default $m = 0.999$. Lower values (like $0.9$) cause the encoder to evolve too rapidly, reducing representation consistency and leading to training instability.

4. Shuffling Batch Normalization

During development, the authors noticed that standard Batch Normalization (BN) allowed the model to "cheat" the pretext task. Because BN aggregates statistics across the mini-batch on a single GPU, information leaks from other samples in the batch. The network easily exploits this intra-batch communication to lower the contrastive loss without learning generalizable features.

The Cheating Problem

With standard BN, the model uses batch statistics as a "signature" to easily identify which key matches which query, instead of learning actual visual structures.

The Solution: Shuffling BN. MoCo trains with multiple GPUs. For the key encoder, the sample order in the mini-batch is shuffled across GPUs before encoding, and shuffled back afterwards. This ensures that the query and its positive key are processed with batch statistics computed from completely different subsets of samples.

Results & Benchmarks

MoCo achieves state-of-the-art visual representation learning. Crucially, it closes the gap with supervised learning when transferred to downstream tasks like object detection and segmentation.

ImageNet Linear Classification (Top-1 Accuracy)

MoCo (ResNet-50) 60.6%

InstDisc (Wu et al.) 54.0%

LocalAgg (Zhuang et al.) 58.8%

MoCo (R50w4x - Wider ResNet) 68.6%

PASCAL VOC Object Detection (AP50)

MoCo Pre-trained (IG-1B) 82.2%

MoCo Pre-trained (ImageNet-1M) 81.5%

Supervised Pre-training (ImageNet-1M) 81.3%

Random Initialization (Scratch) 60.2%

Limitations & Open Questions

Pretext Task Simplicity: MoCo uses a basic instance-discrimination task (views of the same image match, others mismatch). While effective, it doesn't leverage richer pretext tasks like masked auto-encoding or spatial jigsaw puzzles.
Scale Saturation: Moving from ImageNet-1M to Instagram-1B (1 billion public images) yields relatively small improvements. This suggests that simply scaling raw data without better pretext tasks or filtering has diminishing returns.
Batch Normalization Dependency: The requirement for shuffling BN makes implementation complex and hardware-dependent, as it relies on specific multi-GPU synchronization.

Glossary

Contrastive Learning

A machine learning technique where the model learns representations by contrasting positive pairs (similar items) against negative pairs (dissimilar items).

Momentum Encoder

An encoder whose parameters are updated slowly using an exponential moving average of another encoder's weights, rather than directly via gradient descent.

Pretext Task

A self-supervised learning task used to pre-train a model, where the objective is formulated automatically from the data without requiring human annotations.

Citation

@inproceedings{he2020momentum,
  title={Momentum contrast for unsupervised visual representation learning},
  author={He, Kaiming and Fan, Haoqi and Wu, Yuxin and Xie, Saining and Girshick, Ross},
  booktitle={Proceedings of the IEEE/CVF conference on computer vision and pattern recognition},
  pages={9729--9738},
  year={2020}
}