Projection Head is Secretly an Information Bottleneck

TL;DR

This work analyzes the projection head in contrastive learning from an information-theoretic perspective, showing that an effective projector should act as an information bottleneck between encoder features $Z_1$ and the self-supervised target $R$ to maximize downstream utility $I(Y;Z_1)$ while minimizing $I(Z_1;Z_2)$. It derives lower and upper bounds on $I(Y;Z_1)$ in terms of $I(Z_1;R)$, $I(Z_1;Z_2)$, and $I(R;Y)$, providing a principled design rule: the projector should filter information irrelevant to the contrastive objective. Guided by this principle, the paper introduces training and structural regularizations, including a matrix MI surrogate bottleneck term, discretized projection, and sparse autoencoder approaches, and validates them across CIFAR-10, CIFAR-100, and ImageNet-100 with SimCLR and Barlow Twins, achieving consistent downstream improvements. The results bridge theory and practice in projector design and offer actionable techniques for principled enhancements in contrastive representation learning.

Abstract

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.

Figures (10)

• Figure 1: The general construction of contrastive learning can be displayed as the information flow model above, where $Y$ denote the ground-truth labels in downstream tasks, $X$ are the input samples, $Z_1, Z_2$ denote the encoder and projector features, and $R$ represent the self-supervised targets.
• Figure 2: The change processes of estimated guarantees and practical downstream performance of encoder features trained with SimCLR on CIFAR-100 for 200 epochs. The trend indicates that the bounds provide a fairly accurate estimation of the variations in downstream performance.
• Figure 3: Correlation between downstream task accuracy and the estimated theoretical bounds on CIFAR-10 and CIFAR-100. Different points represent the encoder features learned by SimCLR with different projectors.
• Figure 4: Empirical understandings of proposed methods. (a), (b) show the correlation between estimated theoretical guarantees and downstream performance. The proposed methods improve the downstream performance by filtering out the irrelevant information, leading to a better downstream performance guarantee. (c), (d), (e) demonstrate the influence of different regularization parameters. With stronger regularizations, the downstream performance increases first and then decreases.The experiments are conducted on the models trained by SimCLR on CIFAR-100.
• Figure 5: Tendency during training of SimCLR on CIFAR-10 and ImageNet-100.

Theorems & Definitions (14)

• Theorem 3.1
• proof : Proof Sketch
• Theorem 3.2
• proof : Proof Sketch
• Definition 3.3: Matrix-based $\alpha$-order (Rényi) entropy
• Definition 3.4: Matrix-based mutual information
• Theorem 4.1
• Lemma A.1: Information cutoff theory
• proof
• Lemma A.2: Information processing theory