Deciphering the Projection Head: Representation Evaluation Self-supervised Learning

1 Introduction

Extracting meaningful representations from a large amount of unlabeled data is a crucial task in self-supervised learning. With the rapid progress in the SSL [Chen et al., 2020a; He et al., 2020; Wang et al., 2021; Caron et al., 2020; Grill et al., 2020; Chen and He, 2021; Caron et al., 2021; He et al., 2022; Chen et al., ], simple classifiers learned from pre-trained representations can achieve comparable performance to those from supervised learning. Despite its empirical success, the underlying mechanisms of SSL still require further exploration. Many efforts have been devoted to studying the loss function [Sohn, 2016; Oord et al., 2018; Wang and Isola, 2020; Li et al., 2020; Khosla et al., 2020; Wang and Liu, 2021; Hu et al., 2022], and the construction of positive pairs [Arora et al., 2019; Wen and Li, 2021; Wang and Isola, 2020; Tian et al., 2021; Wang and Liu, 2021; Wang et al., 2022a], while less are paid on the investigation of the model architectural components. Typically, the SSL architecture includes two components: an encoder and a projection head. The encoder is usually a discriminative model like ResNet [He et al., 2016]

or ViT [Dosovitskiy et al., 2020] that aims to extract semantic features; the projection head is a multi-layer perceptron used in pre-training. After the pre-training, the projection head is discarded, and the encoder outputs (representation vectors) are used for the downstream tasks. The inclusion of projection head during the pre-training significantly improves the SSL method in downstream performance [Chen et al., 2020a], leading to its widespread adoption in various SSL architectures [Chen et al., 2020a; He et al., 2020; Chen et al., 2020b; Caron et al., 2020; Grill et al., 2020; Zbontar et al., 2021; Chen and He, 2021; He et al., 2022]. However, the mechanism behind the contribution of the projection head during pre-training has not been fully understood.

Through an in-depth analysis of the pre-training objective values across the intermediate layers of SSL, we discover that the projection head mainly promotes the uniformity objective, i.e., mapping the samples closer to the uniform distribution. In contrast, the encoder mainly serves to enhance alignment, i.e., minimizing the distance between similar samples, such as a training data point and its augmentations. We demonstrate that this phenomenon widely exists in SSL methods, including contrastive methods such as SimCLR and MoCo-V2, and noncontrastive methods such as SimSiam. While augmentation invariance (alignment) may relate more to extracting semantic features [Purushwalkam and Gupta, 2020; Von Kügelgen et al., 2021; Wang et al., 2022a; Wang et al., 2022c], the projection head can effectively prevent the SSL model maps the inputs to the identical constant output, known as trivial constant solution [Jing et al., 2021; Li et al., 2022].

Building on our understanding of the projection head, we propose a novel approach called Representation Evaluation Design (RED) for SSL methods. RED involves creating a shortcut connection between the representation and projection vectors, allowing the benefits (lower entropy, general augmentation, and downstream guidance) of the representation vector to bypass the projection head and guide the training directly. As a result, RED-SSL strengthens the augmentation invariance of representation (extracting semantic information [Von Kügelgen et al., 2021; Purushwalkam and Gupta, 2020; Wang et al., 2022a]) without impairing the uniformity of the projection head. Through comprehensive comparison experiments between the baseline SSL methods (SimCLR, MoCoV2, SimSiam) and the RED-version (RED-SimCLR, REDMoCo-V2, and RED-SimSiam) in various datasets (CIFAR-

10, CIFAR-100 [Krizhevsky, 2009], ImageNet [Deng et al., 2009]), we observe a consistent performance improvement in the downstream classification tasks using k-nearest neighbor (kNN) [Wu et al., 2018] and linear classifier. Furthermore, the RED-SSL learned representations exhibit stronger robustness to augmentation in the downstream evaluation.

Our main contributions are summarized as follows.

• We uncover that the projection head mainly serves as a uniform projector, regardless of whether the uniformity appears explicitly in the loss function of SSL. Thus, the projection head enables the encoder to focus on boosting alignment without degenerating to a collapsed constant. It explains the combination of the encoder and projection head outperforms the individual encoder.
• We reveal that the encoder outputs (representation vectors) exhibit superiority in terms of augmentation robustness, lower entropy, and better downstream task performance than the outputs of the projection head (projection vectors). It explains why the representation vectors are used in the downstream tasks.
• We propose a representation evaluation design (RED) that bridges the representation information and the SSL objective functions. Extensive experiments on different SSL methods, various datasets, and different classifiers demonstrate that our proposed design can consistently improve the downstream task performance of the baseline models and is more robust to unseen augmentations and out-of-distribution data.

2 RelatedWorks

The projection head in SSL. The projection head design was initially introduced in SimCLR [Chen et al., 2020a], which differentiates the pre-training and the downstream phase as separate objectives for the projection vector z and representation vector r, respectively. This projection head design is widely adopted by the later proposed methods [He et al., 2020; Chen et al., 2020b; Grill et al., 2020; Caron et al., 2020; Zbontar et al., 2021; Chen and He, 2021], where some architectures are displayed in Figure 1. [Chen et al., 2020a] indicates that the downstream classification does not increase monotonically as the number of layers of projection head increases. [Wang et al., 2022b] indicated that the projection head is the key to enhancing the transferability by measuring the downstream supervised accuracy between different network layers. However, for unlabeled SSL, assessing it solely based on supervised performance may raise doubts. In contrast, we examine the SSL objective values among the layers, which illustrates the benefits provided by the projection head more directly. [Gupta et al., 2022] regarded the projection head as a low-rank mapping such that the trained vectors can be more style-invariant and generalize better. We think that this does not capture the complete picture. For example, in the ablation study conducted on SimCLR [Chen et al., 2020a], altering the dimensionality of the non-linear projection head does not significantly impact the downstream performance. Additionally, in SimSiam [Chen and He, 2021], reducing the projection dimension adversely affects the downstream performance, indicating that low-rank projection substantially deteriorates the performance of SimSiam. In contrast, our interpretation of the projection head applies to various SSL architectures (including SimCLR, MoCo-V2, and SimSiam) and inspires designs that universally improve SSL performance.

Analysis of SSL. The theoretical understanding of SSL has been the focus of many works [Arora et al., 2019; Tosh et al., 2021; HaoChen et al., 2021; Wen and Li, 2021; Ji et al., 2021; Tian, 2022]. The analysis of the SSL loss and the architecture is relatively fewer. [Wang and Isola, 2020] split the InfoNCE loss as alignment and uniformity, and demonstrate that uniformity stands for uniformly distributing on the hypersphere based on Gaussian potential kernel. Compared with [Wang and Isola, 2020], our study integrates alignment and uniformity concepts into the analysis of SSL architectures. This integration enables a deeper exploration of SSL network design. Importantly, our analysis goes beyond contrastive SSL approaches and also includes non-contrastive SSL methods, providing a more comprehensive understanding. [Hu et al., 2022] relates contrastive SSL to the stochastic neighbor embedding(SNE). [Wang et al., 2022a] considers augmentation encourages the different intra-class samples to be overlapped, and thus positive alignment could attract the intra-class samples together. [Wang and Liu, 2021] states that uniformity guides the contrastive model to learn separable features, and proper temperature gives tolerance to semantically similar samples. [Wen and Li, 2022] reveals that the identityinitialized prediction head prevents BYOL from the training collapse. [Saunshi et al., 2022] point out that the inductive biases within the contrastive function class contribute to the downstream success.

3 Projection Head Is a Uniformity Projector

This section investigates the interplay between the encoder and the projection head during the pre-training process. Specifically, we investigate the alignment & uniformity values within the SSL architectures. We reveal that the projection head prioritizes the uniformity objective, and the encoder focuses on the alignment objective in the pre-training.

The InfoNCE loss [Sohn, 2016; Oord et al., 2018; Chen et al., 2020a; Wang and Isola, 2020; Yeh et al., 2022] is

We can define

where n is the batch size, zi is the projection vector of a sample i. In (1), z(1) i ↪ z(2) i are referred to as positive pairs, since they are from the same sample under different augmentations, and z(k) i ↪ z(l) j are regarded as negative pairs. In the InfoNCE loss Formula1, the first term corresponds to the positive pair alignment loss: Pni=1 (-z(1) i z(2) i ◁⌧), which is minimized if the projection z is invariant to the training augmentation. The second term represents uniformity loss: Pni=1 log Pnj=1 exp(z(k) i z(l) j ◁⌧), which is minimized if the projection z is distributed uniformly on the hypersphere [Wang and Isola, 2020]. Inspired by the alignment and uniformity formula in 1, we define the alignment and uniformity values of given vector x in (2), respectively. In our subsequent analysis, we compute the alignment and uniformity values for intermediate layer output vectors in various SSL networks.

Figures 2 depict the alignment and uniformity values using Formula (2) for the intermediate layers of the SSL models: (a) MoCo-V2; (b) SimSiam. (SimCLR analysis is placed in Appendix B.1) All are trained for 200 epochs in CIFAR-10, and the encoder is ResNet-18. In the encoder part of Figures 2, the uniformity values fluctuate without showing a consistent decrease, while the alignment values decrease across diverse augmentation types (except for Gaussian augmentation, which lacks semantic meaningfulness). Moving into the projection head, there is a noticeable decrease in uniformity values, accompanied by an increase in alignment values. This phenomenon in the projection head suggests that the samples are mapped closer to a uniform distribution on the hyper-sphere, albeit at the expense of losing some augmentation invariance.

Table 1 presents the explicit changes in alignment and uniformity values (computed using Formula (2)) within the encoder and projection head. In the encoder part, alignment values decrease while uniformity values increase. Within the projection head, the decrease in uniformity values coincides with an increase in alignment values, indicating its role as a uniformity projector. In fact, with the alignment and uniformity values dropping within the encoder and projection head separately, the encoder output (representation layer) owns the lowest alignment value and possesses semantically meaningful features. This SSL inner investigation helps explain why introducing the projection head can improve the downstream performance of the encoder output (representation).

In SimSiam pre-training, uniformity terms are not included in the training objective. Nevertheless, similar observations found in contrastive SSL methods like SimCLR and MoCov2 can still be made. The significant decrease in uniformity values within the projection head, as shown in the upper subfigure of SimSiam (b) in Figure 2 and the SimSiam row in Table 1, indicates that the projection head implicitly promotes uniformity and prevents the collapse of SimSiam pre-training.

After uncovering the projection head is a uniformity projector through SSL architectural analysis, we aim to understand the effect on the encoder output layer (representation vectors) in the following analysis.

4 Representation Vector Analysis

Having revealed the projection head's role as a uniformity projector through SSL architectural analysis, we conduct a detailed analysis of the encoder output layer (representation vectors) in this section. Specifically, we demonstrate that the representation vectors exhibit superiority in augmentation robustness, lower entropy, and better downstream task performance compared to the projection vectors.

4.1 Robustness to Unseen Augmentations

Augmentation invariance plays a crucial role in guiding the extraction of semantic features in SSL [Von Kügelgen et al., 2021; Purushwalkam and Gupta, 2020; Wang et al., 2022a] and contributes to the achievements of noncontrastive models [Grill et al., 2020; Chen and He, 2021; He et al., 2022]. However, in the pre-training process, only a limited number of augmentation types are typically applied, while the invariance to unseen augmentation types can also be advantageous in extracting additional semantic information.

To study the robustness to unseen augmentations of the representation and projection vectors, we compare the cosine similarity of the positive pairs based on the representation and the projection vectors. Table 2 records the cosine similarity under different augmentations, calculated with representation and projection vectors, respectively. The "train augmentation" row in Table 2 refers to the augmentation types used during the training (random combination of RandomResizedCrop, HorizontalFlip, ColorJitter, and RandomGrayscale, specified in [Chen et al., 2020a]). The projection vectors exhibit a greater cosine similarity for the training augmentation, but for all other unseen augmentation types, the cosine similarity of the representation vectors is significantly higher than that of the projection vectors. The exception to this is Gaussian blur, which is not a semantically meaningful augmentation. This superior augmentation robustness indicates that the representation vectors are capable of extracting meaningful semantic information beyond the predetermined augmentation types.

4.2 Entropy Analysis

Uniformity is crucial in preventing a model from collapsing into a trivial constant state [Wang and Liu, 2021]. However, promoting uniformity can lead to increased entropy, as shown in Proposition 1. In this section, we demonstrate that the entropy of representation vectors (encoder output) is lower than that of projection vectors (projection head output). This is because the projection head prioritizes the uniformity objective, while the encoder focuses on the alignment objective.

Proposition 1. Encouraging uniformity is equivalent to reducing the KL divergence towards the uniform distribution, and approaching closer to the uniform distribution results in an entropy increase.

The proof of Proposition 1 is based on the derivation in (A.2) and (A.3) of Appendix A. Intuitively, since the uniform distribution is the distribution that has the largest entropy for bounded variables, and the projection head encourages uniformity, the output distribution of the projection head has a higher entropy than its original distribution. To provide a quantitative comparison, we use the discrete entropy estimator [Beirlant et al., 1997; Lombardi and Pant, 2016] defined as the following Formula, b refers to the m nonoverlapping sub-regions that constitute all samples. Defined as: ˆH = 1m Pmb=1 Pklabel=1 -ˆ p¶b↪label♦ log ˆ p¶b↪label♦, where: ˆ pb↪{label = i} = (labeled i sample counts in region b) ◁ (sample counts in region b).

Table 3 compares the estimated entropy calculated using representation and projection vectors from SimCLR, MoCoV2, and SimSiam, all trained for 200 epochs on the CIFAR-10 dataset using ResNet-18 as the encoder. The results verify that the estimated entropy of representation vectors is consistently lower than that of projection vectors, indicating that the encoder implicitly transfers the burden of uniformity to the projection head. Consequently, the representation vectors enjoy a smaller entropy and contain a wealth of information useful for downstream tasks. Our quantitative analysis of entropy provides a formal and comprehensive illustration, offering a contrast to the visually-oriented T-SNE scatter plots presented in [Wang et al., 2022b]. Figure B.3 in Appendix B.2 presents a visual evaluation of the entropy.

4.3 Relationship to DownstreamTask

With the projection head targeting the uniformity objective, SSL models allow the representation vectors to have higher alignment and therefore, contain more semantic information and are more closely related to downstream performance than projection vectors, as demonstrated in this subsection.

We calculate the mean of the pair-wise cosine similarity of representation r and projection vectors z in Formula 3

where ri and zi refer to representation and projection vectors of sample i. Define yi as the downstream misclassification of sample i. Specifically, yi = 1 represents misclassification while yi = 0 represents correct classification in the downstream task. The correlation between the sample average of pairwise cosine similarity s(r) ↪ s(z) and misclassification result y is denoted as ⇢(s(r) ↪ y) and ⇢(s(z) ↪ y). Intuitively, a sample with higher similarity to others is more difficult to differentiate and is thus more likely to be misclassified.

Table 4 displays the correlations ⇢(s(r) ↪ y), ⇢(s(z) ↪ y) for the representation and projection vectors, respectively. The results indicate that representation vectors are more closely related to downstream task performance, as they exhibit a higher correlation. To further verify that the quality of the representation vectors can influence the downstream task performance more directly, we compute the error rate split for the representation and projection vectors. Specifically, we split the samples into two groups according to the pair-wise cosine similarity of representation vectors r, as follows:

The error rate for each group is defined as

The split and the error rates of the projection vectors z: G (z) ↪ e(z) , can be obtained similarly. Intuitively, the error rate of G (r) 1 (or G (z) 1 ) is larger than that of G (r) 2 (or G (z) 2 ), since the samples in the former group are harder to be differentiated. The gap between the error rates of the two groups should be large if the cosine similarity accurately represents the difficulty of the downstream task. Table 4 exhibits that the gap in error rates for representation vectors is much larger than that for projection vectors, suggesting that the sample pair-wise similarity based on representation vectors is more closely related to downstream task performance.

5 Representation Evaluation Design in SSL

In Section 4, we demonstrate that representation exhibits superior robustness to universal augmentations, lower entropy, and stronger relevance to downstream tasks compared to projection vectors. These merits motivate us to propose Representation Evaluation Design (RED) in the SSL pre-training, which is applicable for both the contrastive and non-contrastive SSL.

Our motivation is to enhance the performance of SSL by targeting samples that are difficult to differentiate from the others. Building on our discovery that representation vectors r offer benefits including better augmentation invariance, lower entropy, and stronger downstream relevance than the projection vectors, we propose the representation evaluation weights w in (5) for usage in SSL pre-training loss.

In (5), r is the representation vector, ⌘ is the representation temperature parameter, and k is the percentile parameter. The weight w measures the reciprocal of the percentile of cosine similarity of each sample to others within a training batch. With re-weighting wi included in SSL (RED-SSL), the objective functions of RED-Contrastive and RED-Non-Contrastive are the following (where ⌧ is the temperature hyperparameter):

Incorporating w into the SSL loss in (6) and (7) assigns a higher loss to the sample with greater similarities to others in the batch, since this sample tends to be more challenging for the model in downstream tasks. Additionally, as representation vectors are more related to the downstream performance, the representation-based w actively guides the training to boost downstream accuracy. Furthermore, we leverage the superior augmentation robustness of representation vectors to adjust the alignment term of projection vectors, thereby equipping features with stronger augmentation robustness in RED-SSL.

With the logarithm expansion of RED-SSL losses in Formula (6) and (7), RED-SSL essentially establishes a shortcut connection between the representation based w and projectionbased SSL loss, illustrated in Figures 3. During the batch stochastic gradient descent, the gradients of logwi can bypass the projection head and reach the representation vector directly. This approach also implicitly helps to avoid vanishing gradients of the representation vector.

Remark 1. batch optimization is important in RED-SSL. In (6) and (7), RED-SSL essentially adds w term with the SSL loss and does not reweigh samples within the batch. Instead, the w of RED updates the model training in a batch-wise manner in batch optimization.

It should be noted that wi in (5) is defined by the k% percentile of the exponential pairwise product: percentilej (exp(rirj◁⌘)↪ k%) instead of the average: 1n Pj exp(rirj◁⌘). The percentile approach allows us to focus on more similar samples without considering other samples that are already well separated, resulting in a better assessment of sample-wise similarity within a training batch. In contrast, the average operation considers all samples equally in the calculation. Table 5 compares the downstream classification accuracy of MoCo-V2 trained with the baseline contrastive loss (InfoNCE) versus the percentile, and averagebased RED-contrastive loss. The model was trained for 200 epochs on the CIFAR-10/100 using ResNet-18 as the encoder.

6 Experiments

In this section, we compare the existing SSL models: SimCLR, MoCo-V2, and SimSiam, with our representation evaluation design SSL (RED-SSL): RED-SimCLR, RED-MoCo-V2, and RED-SimSiam, respectively. Our results demonstrate that RED consistently improves downstream performance across diverse architectures, datasets, and out-of-distribution (OOD) settings. Similar to the superior robustness of shortcut connections to perturbations [Cazenavette et al., 2021; Reshniak and Webster, 2020], we show that RED-SSL models also gain more robustness against unseen augmentations. We conduct an ablation study on the two hyperparameters within the re-weighting term w in RED, namely the percentile parameter k% and the representation temperature ⌘. The ablation results exhibit that RED can steadily boost downstream performance. However, due to space limitations, we include the ablation study in Appendix B.6.

6.1 DownstreamTask Performance

We evaluated the classification accuracy of SimCLR, MoCoV2, SimSiam, and their corresponding representation evaluation design models, RED-SimCLR, RED-MoCo-V2, and RED-SimSiam, on downstream classification tasks. Table 6 presents the classification accuracy of different SSL models on CIFAR-10 and CIFAR-100 data using k-nearest neighbors (kNN) and linear classifiers as downstream classifiers [Wu et al., 2018]. Additionally, we compared the downstream classification performance on ImageNet1000 in Table 7 and mixed-Gaussian simulated data in Appendix C. We list the training settings in AppendixB.3. Our results demonstrate that RED consistently improves the downstream task performance of popular SSL models, regardless of whether they are contrastive (e.g., SimCLR, MoCo-V2) or not (e.g., SimSiam).

6.2 Robustness to (Unseen) Augmentations

In Section 4.1, we discussed how representation vectors are more robust to unseen augmentations. By introducing reweights wi that establish a shortcut connection between the representation and projection vectors, we expect RED-SSL models to be more robust to diverse (unseen) augmentations and verify it in this subsection.

We used MoCo-V2 as the base model and compared its performance with RED-MoCo-V2. After training both models, we applied (unseen) augmentations to the test samples and evaluated the kNN classification accuracy of the augmented test data. The results are reported in Table 8 and training settings in Appendix B.4. Compared to the original MoCoV2, RED-MoCo-V2 consistently exhibits stronger robustness to diverse augmentations, whether they are specified during training (train augmentation) or unseen. These results demonstrate that RED-SSL captures more semantic information that is invariant to augmentations.

6.3 Other Experiments

SSL methods could suffer from out-of-distribution shift [Hu et al., 2022]. The results in Table 9 indicates that RED can improve the out-of-distribution generalization and training details in Appendix B.5. The ablation study in Appendix B.6 verifies that RED can steadily boost downstream performance.

7 Discussion

This paper presents a comprehensive analysis of the projection head design in SSL, revealing that it promotes uniformity and allows the encoder to focus on alignment. This explains why combining the encoder and projection head outperforms using the encoder alone. The projection head also ensures non-collapsed training in non-contrastive models like SimSiam. The encoder & projection head combination enables the encoder to spotlight boosting alignment without worrying about uniformity; thus, the representation vectors enjoy more robustness to augmentation, lower entropy, and better downstream task performance than the projection vectors. Drawing on these insights, we introduce the Representation Evaluation Design (RED), an adaptable approach to diverse SSL models. Our experiments demonstrate that RED-SSL outperforms baseline models in downstream tasks and exhibits greater robustness to augmentation and out-of-distribution data. Our research sheds light on SSL model structure and can inspire further research in this area.

We think our work can be extended in several ways. First, it will be meaningful to explore ways to further simplify the RED while preserving its effectiveness. Second, exploring the performance of RED on different architectures, such as NLP: SimCSE [2021] or multi-modalities architecture: CLIP [Radford et al., 2021], is an interesting topic for future exploration. We hope our work can inspire more effective and generalizable SSL methods.