Projection Head is Secretly an Information Bottleneck
This provides a theoretical foundation for a widely used but poorly understood component in contrastive learning, potentially inspiring more principled designs in the field.
The paper tackles the lack of theoretical understanding of the projection head in contrastive learning by analyzing it from an information-theoretic perspective, revealing it acts as an information bottleneck to filter irrelevant information, and introduces modifications that improve downstream performance on datasets like CIFAR-10, CIFAR-100, and ImageNet-100.
Recently, contrastive learning has risen to be a promising paradigm for extracting meaningful data representations. Among various special designs, adding a projection head on top of the encoder during training and removing it for downstream tasks has proven to significantly enhance the performance of contrastive learning. However, despite its empirical success, the underlying mechanism of the projection head remains under-explored. In this paper, we develop an in-depth theoretical understanding of the projection head from the information-theoretic perspective. By establishing the theoretical guarantees on the downstream performance of the features before the projector, we reveal that an effective projector should act as an information bottleneck, filtering out the information irrelevant to the contrastive objective. Based on theoretical insights, we introduce modifications to projectors with training and structural regularizations. Empirically, our methods exhibit consistent improvement in the downstream performance across various real-world datasets, including CIFAR-10, CIFAR-100, and ImageNet-100. We believe our theoretical understanding on the role of the projection head will inspire more principled and advanced designs in this field. Code is available at https://github.com/PKU-ML/Projector_Theory.