Understanding and Improving the Role of Projection Head in Self-Supervised Learning
This work addresses a fundamental design issue in self-supervised learning for computer vision researchers, offering an incremental improvement over existing methods.
The paper tackles the unclear necessity of using a learnable projection head in self-supervised contrastive learning, which is typically discarded after training, by proposing an alternative optimization scheme that treats the projection head as a parametric component for the InfoNCE objective, resulting in improved performance on multiple image classification datasets.
Self-supervised learning (SSL) aims to produce useful feature representations without access to any human-labeled data annotations. Due to the success of recent SSL methods based on contrastive learning, such as SimCLR, this problem has gained popularity. Most current contrastive learning approaches append a parametrized projection head to the end of some backbone network to optimize the InfoNCE objective and then discard the learned projection head after training. This raises a fundamental question: Why is a learnable projection head required if we are to discard it after training? In this work, we first perform a systematic study on the behavior of SSL training focusing on the role of the projection head layers. By formulating the projection head as a parametric component for the InfoNCE objective rather than a part of the network, we present an alternative optimization scheme for training contrastive learning based SSL frameworks. Our experimental study on multiple image classification datasets demonstrates the effectiveness of the proposed approach over alternatives in the SSL literature.