Adaptive Multi-head Contrastive Learning
This work addresses a bottleneck in contrastive learning for computer vision by improving robustness to augmentation diversity, though it is incremental as it builds on existing methods.
The paper tackles the problem of enforcing high similarity for positive pairs and low similarity for negative pairs in contrastive learning, which can be unattainable due to diverse augmentations and varying sample similarities, by proposing Adaptive Multi-Head Contrastive Learning (AMCL) that uses multiple projection heads with adaptive temperatures, resulting in consistent improvements across methods like SimCLR, MoCo, and Barlow Twins, especially with multiple augmentations.
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.