Adaptive Multi-head Contrastive Learning
Lei Wang Piotr Koniusz Tom Gedeon Liang Zheng
In contrastive learning, two views of an original image generated by
different augmentations are considered as a positive pair whose similarity is
required to be high. Moreover, two views of two different images are considered
as a negative pair, and their similarity is encouraged to be low. Normally, a
single similarity measure given by a single projection head is used to evaluate
positive and negative sample pairs, respectively. However, due to the various
augmentation strategies and varying intra-sample similarity, augmented views
from the same image are often not similar. Moreover, due to inter-sample
similarity, augmented views of two different images may be more similar than
augmented views from the same image. As such, enforcing a high similarity for
positive pairs and a low similarity for negative pairs may not always be
achievable, and in the case of some pairs, forcing so may be detrimental to the
performance. To address this issue, we propose to use multiple projection
heads, each producing a separate set of features. Our loss function for
pre-training emerges from a solution to the maximum likelihood estimation over
head-wise posterior distributions of positive samples given observations. The
loss contains the similarity measure over positive and negative pairs, each
re-weighted by an individual adaptive temperature that is regularized to
prevent ill solutions. Our adaptive multi-head contrastive learning (AMCL) can
be applied to and experimentally improves several popular contrastive learning
methods such as SimCLR, MoCo and Barlow Twins. Such improvement is consistent
under various backbones and linear probing epoches and is more significant when
multiple augmentation methods are used.