Adaptive Multi-head Contrastive Learning

TL;DR

Adaptive Multi-Head Contrastive Learning (AMCL) introduces multiple projection heads with adaptive temperatures to better handle diverse augmentations and intra/inter-sample similarities, improving contrastive learning performance.

Contribution

It proposes a novel multi-head contrastive framework with adaptive temperature re-weighting, enhancing existing contrastive methods like SimCLR, MoCo, and Barlow Twins.

Findings

Consistent performance improvement across various backbones.

Enhanced results with multiple augmentation strategies.

Effective in improving contrastive learning outcomes.

Abstract

In contrastive learning, two views of an original image, generated by different augmentations, are considered a positive pair, and their similarity is required to be high. Similarly, two views of distinct images form a negative pair, with encouraged low similarity. Typically, a single similarity measure, provided by a lone projection head, evaluates positive and negative sample pairs. However, due to diverse augmentation strategies and varying intra-sample similarity, views from the same image may not always be similar. Additionally, owing to inter-sample similarity, views from different images may be more akin than those from the same image. Consequently, enforcing high similarity for positive pairs and low similarity for negative pairs may be unattainable, and in some cases, such enforcement could detrimentally impact performance. To address this challenge, we propose using multiple projection heads, each producing a distinct set of features. Our pre-training loss function emerges from a solution to the maximum likelihood estimation over head-wise posterior distributions of positive samples given observations. This loss incorporates the similarity measure over positive and negative pairs, each re-weighted by an individual adaptive temperature, regulated to prevent ill solutions. Our approach, Adaptive Multi-Head Contrastive Learning (AMCL), can be applied to and experimentally enhances several popular contrastive learning methods such as SimCLR, MoCo, and Barlow Twins. The improvement remains consistent across various backbones and linear probing epochs, and becomes more significant when employing multiple augmentation methods.