Investigating the Benefits of Projection Head for Representation Learning
Authors: Yihao Xue
Eric Gan
Jiayi Ni
Siddharth Joshi
Baharan Mirzasoleiman
Venue: ICLR 2024 poster (2024)
Keywords: representation learning
representation learning theory
contrastive learning
robustness
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Abstract: An effective technique for obtaining high-quality representations is adding a projection head on top of the encoder during training, then discarding it and using the pre-projection representations. Despite its proven practical effectiveness, the reason behind the success of this technique is poorly understood. The pre-projection representations are not directly optimized by the loss function, raising the question: what makes them better? In this work, we provide a rigorous theoretical answer to this question. We start by examining linear models trained with self-supervised contrastive loss. We reveal that the implicit bias of  training algorithms leads to layer-wise progressive feature weighting, where features become increasingly unequal as we go deeper into the layers. Consequently, lower layers tend to have more normalized and less specialized representations. We theoretically characterize scenarios where such representations are more beneficial, highlighting the intricate interplay between data augmentation and input features. Additionally, we demonstrate that introducing non-linearity into the network allows lower layers to learn features that are completely absent in higher layers. Finally, we show how this mechanism improves the robustness in supervised contrastive learning and supervised learning. We empirically validate our results through various experiments on CIFAR-10/100, UrbanCars and shifted versions of ImageNet. We also introduce a potential alternative to projection head, which offers a more interpretable and controllable design.

Published as a conference paper at ICLR 2024
INVESTIGATING THE BENEFITS OF PROJECTION HEAD FOR REPRESENTATION LEARNING
Yihao Xue, Eric Gan, Jiayi Ni, Siddharth Joshi, Baharan Mirzasoleiman Department of Computer Science, University of California, Los Angeles y yihacxuelg ucla. udl,, egan8eg.use .ecla.edu,, www.aayllll111211 www.ulllaea.. s joshii 8044cs cs. ucla edu, baharan www.llllcee. edu
ABSTRACT
An effective technique for obtaining high-quality representations is adding pro- jection head on top of the encoder during training, then discarding it and using the pre-projection representations. Despite its proven practical effectiveness, the rea- son behind the success of this technique poorly understood. The pre-projection representations are not directly optimized by the loss function, raising the question: what makes them better? In this work, we provide rigorous theoretical answer to this question. We start by examining linear models trained with self- sulfrsuuer contrastive loss. We reveal that the implicit bias of training algorithms leads to layer- wise progressive feature weighting, where features become increasingly unequal we go deeper into the layers. Consequently, lower layers tend to have more normal- ized and less specialized representations. We theoretically characterize scenarios where such representations are more beneficial, highlighting the intricate interplay between data augmentation and input features. Additionally, we demonstrate that introducing non con-liity into the network allows lower layers to learn features that are completely absent higher layers. Finally, we show how this mechanism improves the robustness in supervised contrastive learning and supervised learning. We empirically validate our results through various experiments on (TFAR-1O/100 UrbanCars and shifted versions of ImageNet. We also introduce potential alterna- tive to projection head, which offers a more interpretable and controllable design.
INTRODUCTION
Representation learning has been the subject of extensive study in the last decade (Chen et al. 2020 Zbontar et al. 2021 Khosla et al. 2020 Ermolov et al. 202I). Despite the great progress, learning representations that generalize well across various domains has remained challenging. Among the existing techniques, contrastive self-supervised learning has gained lot attention, due to its ability to learn robust representations that generalize better across various domains. Contrastive learning (CL) learns representations by maximizing the agreement between different augmented views of the same examples and minimizing that of different examples (Chen al. 2020). Pre-training with CL is often essential before transfering models to new or shifted domains Hendrycks et al. 2019).
key factor that enables contrastive learning to learn robust representations is projection head, shallow MLP that is used during pretraining and is discarded afterwards (Chen et al. (020), which has become standard for learning high-quality representations (He et al.) 2020 Yeh al. 2022
Bardes et al., 2022 Chuang et al. 2020 Grill et al. 20200 Garrido et [2.2] Zbontar et al. 2021
particularly advantageous when there misalignment between training and downstream objectives (Bordes et al. (003). The benefit of projection head extends beyond self- self-supervised learning to other representation learning methods, including supervised contrastive learning (Khosla al. 2020). However, the mechanism by which the projection head improves the generalizability and robustness of representations is poorly understood.
Ermolov et al. 202 Joshi & Mirzasoleiman 2023).
Theoretically analyzing the effect of projection head for CL is challenging, as one needs to understand feature learning both before and after the projection head and compare them. A few recent studies (Tian et al.) 2021 Wang et al. 2021 Wen & Li (022) have focused non-contrastive self-supervised methods and demonstrated that the projection head can mitigate the dimensional collapse problem, where all inputs are mapped to the same representation. However, such results

---

Published as a conference paper at ICLR 2024
do not extend to contrastive methods due to its inherently different learning mechanism. Besides, they do not provide an understanding of the projection head's 's role in enhancing robustness under misalignment between pretraining and downstream objectives.
In this work, we analyze the effect of projection head on the quality and robustness of representations learned by CL., and extend results supervised CL (SCL) and supervised learning (SL.. First, theoretically show that linear models progressively assign weights to features as they operate through the layers. Thus, in deeper layers the features are represented more unequally and the representations are more specialized toward the pretraining objective. Moreover, non-linear activations exacerbate this effect allowing lower layers to learn features that are entirely absent in the outputs of the projection head. We demonstrate that projection head provably improves the robustness and generalizability of the representations, when data augmentation harms useful features of the pretraining data, or when features relevant to the downstream task are too weak or too strong in the pretraining data. Finally, we extend our results beyond CL to SCL and SI.. In this setting, we reveal that lower layers can learn subclass-level features that are not represented in the final layer, demonstrating how representations before the final representation layer can significantly reduce class/neural collapse, a problem previously observed in the final representations.
We conduct extensive experiments to confirm our theoretical analysis. First, we construct a semi- synthetic dataset by adding MNIST (LeCun 1998) digits to CIFARIO (Krizhevsky et al., 2009) images, and confirm that using projection head yields superior representations when data augmenta- tion harms the rowestream-relevant features, when these features either very strong very weak during pretraining. Then, we study supervised learning and demonstrate that using projection head results in better course-to-fine transferability on CIFAR 100, superior performance of few-shot adaption to distribution shift on UrbanCars (Li et al. 2023), and better robustness against natural distribution shifts in ImageNet (Xiao al.) 2020; Hendrycks et al. 2021a|b) b). We also demonstrate how fixed reweighting head can achieve performance comparable to the projection head, providing further evidence for our theoretical conclusions and potentially inspiring future design.
2 ADDITIONAL RELATED WORK
Projection head. While projection head is widely used, the underlying reasons for its benefit have not been fully understood. The empirical study by Bordes et al. (222)) suggests that the benefit of the projection head is influenced by optimization, data, and downstream task, and is especially significant when there is misalignment between the training and downstream tasks. On the theoretical side, Jing et al. (202T) suggests that projection head can alleviate dimensional collapse, but did not provide insights into the significant role of the projection head in cases of misalignment between pretraining and downstream tasks. Recently, Gui et al. (2023) analyzed the training of linear projection head using CL. However, their analysis is performed in the case where representations are fixed, and the results only show how training linear head on these representations can lead to worse post-proge poojprooecttin representations. This does not align with practical scenarios where encoder and projection are trained simultaneously and does not reveal why adding a projection head is needed in the first place.
Generalizability and transferability of representations. A longstanding goal in machine learning is to acquire representations that generalize and transfer across various tasks. Two key challenges emerge in this pursuit. The first arises from shifts in labels between training and downstream tasks. Notably, the issue of class or neural collapse (Papyan et al. 2020 Han et al. 2021 Zhu et al. 2021
Zhou et al. 2022b a Lu & Steinerberger [022] [ang et al.) 22021; et al. 2022] Chen et al.
2022 Graf al.) 202 Xue et al. (023). affecting both SCL and SL. where representations within
the same class become indistinguishable at a subclass level, rendering them unfit for fine-grained labeling. We theoretically demonstrate that such issue can be alleviated by taking the pre-projection head representations. The second challenge arises from shifts in input distributions, where neural networks rely heavily on patterns specific to the training distribution that do not generalize (Zhu et al. 2016, Geirhos et al.] [2018] IIyas et al.] 2019 Barbu et al.) 2019 Recht et al.) 2019. Sagawa et al. 2019: Xiao et 20.0 Taori et al. 2020 Koh et al., 2021: Shankar et al. (021). We explore
the benefits of projection head in the above two scenarios, but noting that our primary focus lies in
elucidating the broader concepts at play rather than addressing each specific problem.
3 EXPLORING THE ROLE OF PROJECTION HEAD IN SELF-SUPERVISED CL
In this section we consider self-supervised contrastive learning, a scenario where the concept of the projection head is frequently employed and has a crucial role in achieving optimal performance.
2

---

Published as a conference paper at ICLR 2024
Pretraining data distribution. We define the following data distribution and data augmentation used for pretraining with contrastive learning. We choose this setting for clarity in presenting our results, but we note that our results can be generalized to more complex cases.
Definition 3.1 (Input distribution of pretraining data). The input data used for pretraining follows distribution D, where, for an input E IRS, its i-th element is randomly drawn from {-δi, 6z).
Definition 3.2 (Data augmentation). Given an input x, its augmentation follows the distribution A(x). The augmentation operates on an input x in two steps. First, for each i, with probability ag, it alters i-th element in x by randomizing its sign, resulting in modified input denoted by x'. Then, random noise & satisfying Eig = O and ElŞt [ = is added to x', yielding augmentation x' + f. Each coordinate of the input is an independent feature. At high level, Φ; represents the magnitude of a feature in the input data, while a; quantifies the level of disruption introduced by the augmentation. If a, = 1, it means that this feature is completely 'destroyed' during augmentation, such that there is no correlation between positive pair (two augmented versions of the same input) at this coordinate. This data augmentation is designed to mimic practical scenarios where data augmentation techniques intentionally modify different features to varying degrees. For example, some features, like color and texture, may be altered more significantly than core features of the main object. Additionally, the noise & with variance x² at each coordinate is included because data augmentation used in practice unintentionally introduces certain level of noise.
Model. We consider 2-layer linear and non-linear models for the sake of clarity, as it suffices to demonstrate our main findings, although our results generalize to multiple layers. Given an input x E Rd, the output after the 1-th layer denoted as fi(r), with f2(x)) h(f (z)). We will consider linear and non-linear functions for fa (-) and h(-) in our analysis. The second layer h(-) serves as the projection head and the first layer f.()) serves as the encoder.
Contrastive loss. We consider the following spectral loss, which has been widely used in previous
theoretical and empirical studies (HaoChen et al.] 2021 Xue et al., 2022 Saunshi et al. 2022]
HaoChen & Ma 2022
Garrido et al. 2022 Xue et al. 2023). Given model representing function
f(.). loss is [fL(f) = 2wE =
f(z) )1 +
f(A(zz))) ((
After training a model to minimize this loss, we utilize it in downstream tasks where we feed different inputs, which may or may not follow the same distribution as the pretraining data, into the model. We can leverage the representations provided by the model for various purposes, most commonly to train a linear model on these representations and predict labels for the downstream data. When we use the two-layer model, i.e., when LCL (f2) is minimized, we have the option to choose either the post-projectionn representations generated by fa()) or the pre-projection representations from (-) for the downstream task.
Clarification on what we compare. To gain a comprehensive understanding, it's necessary to compare these three cases: (1) pre-projection, where we minimize Lcr((2) first and discard h(-), using only Si (-) for the downstream task; (2) post-projection, where we minimize Les (f2) and use fe()) for the downstream task: (3) no-projection, where we minimize (CL(fi) and using f-) for the downstream task. The goal is to determine when pre-projection outperforms both of the other two and understand the reasons behind However, in practical deep neural networks, there is typically no significant difference between pr-pooeectoon and post-projection as both scenarios use the network's final output for the downstream task and the difference is mainly an additional layer in post- projjcctin..o. Given that the networks are sufficiently large and expressive, this one-layer difference does not significantly change the representations achieved at the output. Therefore, in theory, we consider settings where fi (-) and fz(.) have the same expressiveness and yield equivalent representations when used to minimize the loss, and then solely comparing pre projeecoooinn and post-projeee projjecrooetton is sufficient.
3.1 LAYER- WISE PROGRESSIVE FEATURE WEIGHTING IN LINEAR MODELS
Linear network. We consider linear model in which f2(Φ) (2) W iz, with W, E Rexd and W2 E RPXP representing the weights of the first and second layers, respectively. Both the hidden and output dimensions are p, to ensure consistent dimensionality between pre- projection and post- projeprooection representations for fair comparison.
3.1.1 STRUCTURE OF LAYER WEIGHTS
We begin by examining the weights at different layers within a model. Our investigations reveals the relationship between these layer weights, which generally hold regardless of the data distribution.
3

---

Published as a conference paper at ICLR 2024
Firstly, consider the weights of the minimum norm minimizer of the CL loss. This is pertinent because gradient- gradient-based algorithms are shown to prefer minimizers with small norms (Neyshabur et al.
2014
Gumasekar et al. (017). Furthermore, many theoretical studies on CL (Ji et al. 2021 Liu et al.
2021
Nakada et al. 2023) have considered regularization in the form of W W F. which promotes
small norm, and Xue al. (2023) have shown that the minimum norm provides explanation for many intriguing phenomena in CL.
Theorem 3.3 (Weights of the minimum norm minimizer), The global minimizer of the CL loss LCL with the smallest norm, defined as [[[ W,Wl + [ W2ll?. satisfies W W.W W W2
In addition, establish that similar conclusion holds for models trained using gradient flow, which is a continuous version of gradient descent widely adopted theoretical analysis. In gradient flow, at any time l, the weight updates are given by 2
(
8 ow! www..WWW...... (W(1), where i = 1,2.
Theorem 3.4 (Weights of the model trained with gradient flow, proved (Arora et al. 2018 Suppose the initialization satisfies ( (") (0 W _ Using gradient flow, at any time t,
we have w((T w (
(t)
3.1.2 LAYER- WISE PROGRESSIVE FEATURE WEIGHTING
Given that W W W have the same singular values, and the left singular vectors of Wa match the right singular vectors of W2. Consequently, the singular values of the joint weight matrix W2W1 are the squares of those in the first layer, Wi 1- As a result, the differences in weights assigned to the features are smaller when the input is passed through the first layer than when the input is passed through the whole network. To illustrate this concept, we analyze the model trained on the data distribution given by Definition 3.1 and analyze the resulting representations.
What insights can we gain from Theorems 3.3 and 3.4
W 2. both layers
The following analysis holds for both models obtained from either the minimum norm minimizer of the loss (as in Theorem 3.3) or the model trained using gradient flow under the assumption that the model converges to a global minimum (as in Theorem 3.4). as they equivalent.
In our input data, we refer to the d independent coordinates as input features. Our interest lies in under- standing the weight assigned to each feature in the pre- and post-projection representations. To achieve this, we examine the quantities [e,), i = 1,..., d, and 1 = 1,2, with e; denoting i-th standard basis. These quantities represent the scale of the representation of unit feature at each coordinate. The following theorem (see proof in Appendix (A2)) shows the weights of the features in each layer:
Theorem 3.5. Define 3: == (1 (1 and
%i = (1-o,)0 Let II = j2......... Id) be 0,-0"
of+o"
permutation of indices (1,2,... d} such that 3n 2 ... Bir Then after pretraining,
[] = if i (i,.... Smintaphh else 0.
What does the model do? According this theorem, the model follows two key steps: (1) Feature selection: The model selects the top features with the highest & values, which experience a low level of disruption from augmentation (a²) and/or have large feature magniude (o+). (2) Feature weighting: These selected features are scaled by % at each layer, with zero weight assigned to the remaining features. The rescaling serves a dual purpose: (a) moderating the features by assigning small weights to either overly strong or overly weak features in terms of their magnitude, as indicated by % + when di approaches either 0 or +oo; (b) giving larger weights to features that are less disrupted by augmentation.
The difference between fi(-) and f2(-). (1) Both learn the same features but assign different weights to them. (2) fi(-) treats features more equally, exhibiting smaller gap between feature weights.
3.1.3 WIIY AND WHEN CAN MORE NORMALIZED FEATURES BENEFIT A DOWNSTREAM TASK?
To address this question, we analyze the representations of downstream data drawn from the following distribution: each input æ E Id is a vector, and its i-th element is independently drawn from the set
(-ô,, ).). Note that G's may differ from o, 's as in woal-world scenarios, the input data in the down- stream task may follow different distribution than the pretraining inputs. Additionally, each input
4

---

Published as a conference paper at ICLR 2024
is labeled as sign(e,. where e,. is the j" standard basis vector. In simpler terms, the label is determined by the sign of the j"-th coordinate of the input. that is the downstreean-eleevnn feature.
To evaluate the informativeness of the learned representations for the downstream task, imput the data (without any data augmentation, as is typically the case in practice) into the model that has been pretrained with the CL objective. We then evaluate the quality of these representations each layer by analyzing the sample complexity of the hard SVM trained with labels on these representations, which can be equivalently viewed as training a linear model with logistic losses using gradient descent Soudry et al. (2018). This aligns with the standard linear evaluation protocol in practice (Ye et al. 2019 Oord et 2018 Bachman et al. 2019] Kolesnikov et al. 2019 Chen et al. (020). data distribution that is separable with a (7, P)- -argin (see details in Appendix A.3 it well-known that the sample complexity only grows with r = )/ Bartlett & Shawe- Taylor 1999). Hence, we refer to as the sample complexity indicator and compare its values for pre-projection and post-projection representations. The following theorem shows the conditions under which one has higher sample complexity indicator than the other. Note that smaller sample complexity is preferable.
Theorem 3.6. Let Ti and r2 be the sample complexity indicators for pre-projection and post- projection representations, respectively. Define A = Eisismin(d.)) and jett 33.
If A depends on the strength of the downstream-relevar feature and the weights of features. While determining A's A's value may seem complex, in general, the key factor is whether the model assigns sufficient weight to the downstream-reelevan feature. If this feature is underweighted by the model, using the first layer is beneficial. To better understand when this occurs, we provide the following interpretable examples
<. O then 11 < r2, and if À > then ri > r2-
Corollary 3.7. In each of the following examples < O, i.e., the pre-projection representations are preferred. (1) Data augmentation disrupts the feature too much. Example: all features in both the pretraining and downstream data have a magnitude of 1, and 0:z. is the p-th smallest among all {a] indicating that the data augmentation disrupts useful feature the most among the p features that will be learned. (2) The downstrean-relevamt feature is too weak in pretraining. Example: all 04,'s 's are equal, p 2 2, Vi. Φi S o, and 6z- is the p-th largest among all {ss Yo (3) Multiple features are selected by the model, with the downstrean- pretraining. Example: all ar, 's are equal, 2 2, vi + j Φz- > max of o/oij.
being too strong in
It might seem surprising include scenario 3 above, as having a strong in pretraining may be expected to benefit the downstream task. However, as we discussed after Theorem 3.5 the model moderates the features and assigns small weights to overly strong ones. This can potentially explain why the use of a projection head remains beneficial in cases where the pretraining task seems to be good match for the downstream task. Furthermore, we will validate each of the above three observations in experiments in Section 5.1 We also provide discussion on multi-layer models in Appendix
maw.strrram--elevaan feature
3.2 LOWER LAYERS CAN LEARN MORE FEATURES THAN HIGHER LAYERS VIA NON-LINEARITY
In the previous section on linear models, it's worth noting that both layers select the same features, albeit with different weightings. Now, let's consider a scenario where augmentation disrupts the useful feature excessively, for example, when p,. = 1 such that the useful feature is assigned a weight of Y' = 0. In such cases, neither layer would learn this feature. However, in this section, we'll explore an interesting aspect of non-linear models. Specifically, we'll demonstrate that pre-projection head representations can learn features that are weighted zero in post-projection head representations.
For clarity, we will present our result in the simplest case, although it holds in broader scenarios. For the pretraining data, we let d 2 2, and assume é1 62 = I, O. Additionally, we set p2 1 and pi = O, meaning that the augmentation completely 'destroys' feature 2 while fully preserving feature Consequently, during CL, the pretraining objective discourages the learning of feature 2.
Non- diagonal network. We consider diagonal non-lr network with fi (I) = a (w) @ xF bi), f2(x) = h(S (x)) = (w, @ fi(),, b2).
(1)
where w1 W2. b, b2 E Rd are the trainable weights and trainable biases, and o(: represents the symmetrized ReI. activation function, defined as (a, b) ReL b) Relau( -a - b). applied
5

---

Published as a conference paper at ICLR 2024
element- wise. In this model, each coordinate in the input is processed independently without any cross-coordinate connections. This design not only simplifies our analysis but also aligns with our definition where the features at all coordinates no motivation for even a fully connected model to combine the features. Moreover, this model is sufficient for characterizing the feature selection and weighting processes described in the previous section, enabling us to understand the key aspects clearly. Our results derived with this model also extend to fully connected ReLU networks, as we will empirically demonstrate in Section S.I
independent. Because of this definition, there is
We train the model using gradient flow to minimize the contrastive loss Lor. (f2). Interestingly, in the following theorem, where we compare the pre-projection and post-proce projjeroentoonn representations, we will observe that feature 2, which is discouraged from being leared during the pretraining process, has zero weight post-projection but a non-zero weight pre-projection.
Theorem 3.8. If at initialization w = [wil ( I.. .
I..,.,.. (0) = [w w
(0) (0)
...]T. = bb))) =
[b(0) ((0) [T. with | S vbu and |u22 1( 12 bo) 2 by then as t - co, (l + 0. (e2) vbo.
(0) (0)
The theorem indicates that pre-projection representations are more transferable. The 'destroyed' feature, feature 2, is weighted zero after the projection head but non-zero before the projection head. Therefore, using the first layer representations allows us to successfully learn downstream tasks where feature 2 is the downstream relevant feature, whereas we can't do so with the second layer representations. In practice, augmentations are imperfect and may inadvertently distort important features. Additionally, there is no -illssa-li augmentation suitable for all downstream tasks. With non-linear models, the projection head can save us from losing valuable information.
We note that the above advantage of the pre-projection representation diminishes when we use weight decay. Indeed, we can show that (e2)] - 0 in this case as weight decay tends to shrink the weights considered unnecessary for minimizing the loss. However, we note that in practical scenarios, weight decay values are often small, c.g., 10 and training epochs are finite, e.g., 100, as the default setting inCChen et al. (022)). Therefore, we can estimate the effect as approximately (1-10-5 100 ~ 1-10-1 for weight that receives zero gradient from other sources, indicating limited effect. Consequently, the benefit of pre-projection representations remains evident with reasonable weight decay. However, using excessive weight decay will essentially remove this benefit, as shown empirically in Section 5.1
4 HOW THE IDEA OF PROJECTION HEAD CAN BENEFIT SUPERVISED LEARNING
The idea of projection head has been adopted by supervised contrastive learning Khoslaa et al. 2020). yet its effects have not been systematically studied nor theoretically explored. Here, we provide the- oretical insights into its benefits in both supervised CL (SCL) and standard supervised learning (SL).
4.1 THE GENERAL INSIGHTS FROM LINEAR MODELS
The insights gained from Section 3. can be carried over to understand the case of supervised learning as well. The conclusion in Theorem 3.4 actually holds in more generality. In fact, when W to minimize any loss function via gradient flow with initialization (0) (( W [+1 (0) one can derive conclusion in the form of W (t) ( Wi+ (1)
employing a multi-layer linear model represented as f(z) = W WW
the proof of this relationship can be found in Arora et al. 2018). The singular values of the matrix represented by the model grow exponentially as we go deeper, making the model more specialized toward the pretraining task. Therefore, when considering cutting off at lower layers, it leads to more generalizable and less specialized representations. This is particularly beneficial when the target matrix W*, determined by the training data to be represented by the model using features that are important for downstream tasks.
W1, doesn't ' adequately weight
4.2 ALLEVIATING CLASS COLLAPSE AND NEURAL COLLAPSE NON-LINEAR MODELS
When considering representations' transferabilitty a crucial aspect is their utility for finer-grained downstream tasks (c.g., distinguishing dog breeds) compared to the pretraining task (e.g..
distinguishing dogs from cats). In both SCL and standard supervised learning, common challenge arises, where representations within each class become indistinguishable at a finer grained level.
This is known as class collapse in SCL (Chen et al. 2022; Graf et al., 2021; Xue et al. 2023) and
neural collapse in standard SL Papyan et al. (2020); Han et al. (202 1b: Zhu et al. (2021): Zhou et al.
(2022h a): [u & Steinerberger (2022); Fang et al. (2021 Hui et al. (202)).

---

Published as a conference paper at ICLR 2024
weightings of features
weighangs of features (nurmalizon (elmalzzed)
stings of featuures treemnalooe
LONE
pre
pe post
(nomallzzed))
post
post
14/M
0. 0.13
(disruption augmentation)
4: (strength of features]
Figure 1: Weights of features in two-layer fully connected ReL.U network trained with CL. Left: With all features having equal strength, those that are more disrupted by augmentation have smaller/zero weights. Right: With augmentation treating all features equally. features with the largest/smallest strength weighted less compared to those with intermediate strength. In right and left, weights are more equal pre- pre-projection.
Figure 2: Left:: Weights of features in two-layer fully connected ReLU network trained using SCL.. subclass feature is not represented post-projection but is represented pre-projection. Right: As result, the four subclasses are only separable in pre-projection representations.
Soop the cigst one
En-evval on minist
625
Drop digit
33.0
60 pe-ordlction
675
13.0
000
pot projection pre-2rgecion
6
825
pripetion
02
10-2
10s1
pmbaiility digts
strength of mnist digits
weight decay
An illustration
(b) Effect of data aug-
(c) Effect of strength
Effect of wd
(a) The data augmentation keeps the digit for one image in the positive pair and randomly drops the digit for the other with probability Purop- Pre-projection is more beneficial (b) with more inappropriate augmentation (large Paro)) during pretraining, (c) when digits are very weak/strong during pretraining, and (d) when weight decay is smaller.
Figure 3: Results MNIST-on-CIFARE
we show that lower layers can learn subclass-level features not represented in the final layer in supervised learning, suggesting that using and then discarding the projection head can mitigate the issue of class/neural collapse. To theoretically characterize this, we consider the following data distribution.
Interestingly,
Definition 4.1. We have four subclasses represented by (y. yeub), where {1, [1,-1} is the class label, and Youb E 11. For each input vector RR in subclass (y, yout), I [y Youb .... with the other coordinates being independent and symmetrically distributed.
We consider training the non-linear diagonal network defined in Equation on the above data distribution. For supervised CL, we consider the following spectral loss, which natural generalization from that for self sulf-supe (Col) (HaoChen al. [021 Xue al. 2022
Saunshi
et al. 2022 HaoChen & Ma 2022 Garrido et al. 2022 Xue et al. 2023).
(sc = - 22
class [f2(x)] fa(x f2(x*) + indepenennyy drawnl (f2(x) fz(x ()(xx)) ")
For standard supervised learning, consider the Mean Squared Error (MSH) loss. Since the network's output (2) is d-dimensional, we map it to a scalar by adding linear layer with all weights set 1. Let 1 Id be vector with all elements equal to then the loss can be written as:
CSL (f2) (Z) (f2(z) 1-
According Definition 4.1 we can use (e:) and (e2)) determine the model's 's weighting of the class feature and the subclass feature, respectively, in the i-th layer representations. Next theorem demonstrates that the features indicating subclasses are learned pre-projection but not post-projection.
Theorem 4.2. Under the same assumption about initialization as in Theorem 3.8 and w W22 > 0, for either the model trained minimize Lsor(f2) or Lst(J2) with gradient flow, as I oo.
[l)z(e=)) + 0. (ex)ll 2 Vbn.
EXPERIMENTS
5.1 CONTROLLED EXPERIMENTS ON SELF-SUPERVISED CL
The empirical benefits of the projection head are already well-established, and evidence can be found in numerous large-scale experiments. Instead, in this subsection, focus controlled
7

---

Published as a conference paper at ICLR 2024
proof-of-concept experiments, allowing to the underlying mechanisms that are not shown in existing works and validate our theoretical results from Section
Synthetic data. We train two-layer ReLU network on data drawn from the distribution defined in Definition 3.1 with the second ReLU layer serving as the projection head. The network is randomly initialized. More details are in Appendix C. We conduct experiments in two settings. (1) In setting 1, we let all five features have equal strength and let the augmentation disrupts features differently, with feature I perfectly preserved and feature 5 completely randomized. Figure left shows the weights assigned to features at pre- and post-projection. Consistent with the conclusion from Theorem 3.5 features disrupted more by data augmentation are assigned smaller weights. However, these features are weighted more equally pre- pre-projection Notably, feature 5 has zero weight post- projjprooettion but non-7ero weights pre-projection, aligning with the finding in Theorem 3.8 (2) In setting 2. we aim to demonstrate the moderating effect reflected in Theorem 3.5 We let o 's be equal, controlling the factor of data augmentation, then plot the weights of features with different strength in Figure right. Features being either too strong or too weak are weighted less than intermediate strength Teatures, but in pre-projection, they are less underweighted compared to post-projection.
We design the following dataset for pretraining, with the downstream task on the original MNIST. The pretraining dataset is based on CIFAR- 10, where each image combines CIFAR-10 10 image with MNIST digit. Specifically, the pixel values within the digit area calculated as (1 - 8) x CIFAR-10 10 image + x MNIST digit, while the values outside this area remain unchanged as in the CIFAR-10 image. Here s controls the strength of the wwwnssreem-relevane feature (MNIST digits) during pretraining. We use the following data augmentation during pretraining, For each positive pair, we keep one image's digit, while dropping the other's 's digit with probability Pdrop- A larger Pdrop means that the augmentation disrupts the downstream--elleann feature more. See Figure for for an illustration. The downstream task is classification on the original MNIST dataset. We feed MNIST digits to the pretrained model and train a linear classifier on representations. We pretrain ResNet- 18s with one-hidden-layer MLP as the projection head using the popular SimCLR loss
WAIST-On-CIAARE 10.
CCen et al. 2020) with varying 8 and pdrop and report the evaluation results in Figure
Figure 3b shows the downstream accuracy against Parop with s 1. Pre-projection representations outperform post-projection representations, and the gap widens as Pdrop increases. This confirms that inappropriate augmentations can significantly amplify the superiority of pre-projection representa- tions, aligning with Corollary 3.7 Figure Be shows the downstream accuracy against the strength of digits, with Purop = O. While the pre pre-projection accuracy barely changes, the post- post-projection accuracy increases and decreases, as s increases, making the benefit of using pre-projection more pronounced when either the digit is too weak or too strong, aligning with Corollary (37 Figure 3d with Parop 1 demonstrates the impact of weight decay. Larger weight decay diminishes the benefits of pre- prejprojection representations, aligning with our discussion in Section 3.2 and can even turn them into a detriment. However, pre-projection remains superior with reasonable weight decay.
In Appendix [.I] we present two additional experiments involving minimally modified CIFAR- 10 im- ages to further support our conclusions. We also discuss the effect of early stopping in Appendix D.2
5.2 PROJECTION HEAD IN SUPERVISED LEARNING
We present the following experiments showing broader implications of projection head in SCL and SL.
Coarse-to-fine transferability on synthetic data. We consider the data distribution in Definition 4.1 two-layer ReLU network is trained from random initialization on such data. Figure 2 visualizes the weight assigned to each input component by the model (left) and the representations (right) at each layer. Consistent with the conclusion from Theorem 4.2 we observe that the subclass feature is assigned weight of zero post-projection but has non-zero weight pre- pro-erojctiion. As a result, pre- pre-proeection representations do not suffer from class collapse, making them more transfcrable.
Coarse-to-fine transferability on CIFAR-100. We pretrain a ResNet-18s on CIF with 20 coarse-grained labels and conducted linear evaluation on representations using both these 20 labels and 100 fine-grained labels, separately (details in Appendix C.T). The results
Table 1: collapse, allowing for better fine-grained classification.
The pre-projection reps. suffer less from class/neural
coarse
fine
pre
post
pre
post
SCL
55.240.4
53.6
36.0+o.1
25.7
SL
53. 1
53.2+0.7
33.7
29.7
8

---

Published as a conference paper at ICLR 2024
Table 2: The fixed reweighting head can yield improvements comparable those of the trainable projection head. For the distribution shift scenario, we report the average test accuracies for 2, 4, 8, 16, 32, 64, and 128 shots adaption. More detaied results are in Appendix C.2
Scenario
Dataset
Alg.
Performance Measure
Performance
vanilla
proj
reweight
synthetic coarse-to-fine coarse-to-fing distribution shift
M-on-C CIFARI00 CIFARI00
SSCL
digit clf. acc. fine-grained elf. acc.
77.0
97.3
97.3
SCL
21.8
36.0
30.2
SL
fine-grained clf. acc.
31.44
33.7
32.2
UrbanCars
SL
few-shot adaption acc.
82.2
86.1
87.0
are compared in Table I Pre-projection representations yield higher accuracy in general, but with the benefit much more significant in the fine-grained downstream task. This emphasizes that pre- projecroetoonn representations are less susceptible to class/neural collapse, resulting in more distinguishable representations within each pretraining class.
Few-shot adaption to distribution shift on UrbanCars. Subpopulation shift typically in- volves scenarios where some data groups are underrepresented in the source domain but well-
represented in the target domain. Recent works have demonstrated the effectiveness of retraining the final linear layer using target data to adapt to distribution shifts Rosen- feld et al. (2022): Kirichenko et al. (2022); Mehta et al. (2022). However, has been shown that this approach leads suboptimal results when the target data for training the linear layer is scarce Chen et al. (2023). This could be attributed to the suboptimal quality of penultimate representations. Here, in Figure 4 we demonstrate that by applying the projection head technique and performing adaption on pre-projection representations can lead to better performance. The details are provided in the Appendix C. We also present experiments on shifted ImageNets in Appendix D.3
Targut That Accuncy
90-
s0
proj. pre proj. vanilla
70-
16 32 larget
64 128
Figure 4: Performance of few-shot adaption.
6 REPLACING THE PROJECTION HEAD WITH A FIXED REWEIGHTING HEAD
Based on the observations above, we now explore an alternative design that employs a fixed re- weighting head instead of a trainable projection head. Let = [r1,.... ................................ Tpp be the representation output by the encoder. The reweighting head h(m()) acts as (T) [r], T2....., p pp] where K 1 a hyperparameter. This mimics the role of the projection head described earlier. During pretraining, the reweighting head unequally weight features by assigning different weights to the dimensions of representations, with x controlling how unequal they are. This allows the encoder before the head to weight input features more equally. Note that with a slightly large k. dimensions at the end will be almost "turned off' after the head, which can result in representations after the head representing fewer features than before. This mimics the effect discussed for non-linear models in Sections 3 and We evaluate this approach in the scenarios from Sec. 5.2 and demonstrate that the fixed reweighting head can achieve improvements comparable to those of the projection head, as shown in Table This has two implications: (1) it serves as further evidence for our theoretical indicating that the reweighting effect accounts for all or most of the improvements achieved by the projection head. (2) it opens up possibilities for designing more straightforward and interpretable alternatives to the projection head. It would be interesting for future work to explore the effects of making trainable rather than fixed.
conclusions in Sections [3] and [
CONCLUSION
We provided the first theoretically rigorous explanation for the intriguing empirical success of using the projection head for self-super contrastive learning, supervised contrastive learning, and supervised learning. We demonstrated that lower layers represent features more evenly in linear networks and can represent more features non-linear networks. This enhances the generalizability and transferability of the representations, especially when downstream-relevan features are weak in the pretraining data or heavily distorted by data augmentations. Interestingly, we also show the benefits when the downstream-relevant features are prominent in the pretraining data. We validated our theoretical findings through extensive experiments on both synthetic and real-world data. Finally, we demonstrate how a fixed reweighting head can achieve performance comparable to the projection head, providing further evidence to support our thcoretical conclusions. We also hope that this will offer valuable guidance for future design choices.

---

Published as a conference paper at ICLR 2024
Acknowledgements This research is partially supported by the National Science Foundation CA- REER Award 2146492.
REFERENCES
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit accelcration by overparameterization. In International Conference on Machine Learning, pp. 244-253. PMLR, 2018.
Philip Bachman, Devon Hjelm, and William Buchwalter. Learning representations by maximizing mutual information across views. Advances in neural information processing systems, 32, 2019.
Andrei Barbu, David Mayo, Julian Alverio, William Luo, Christopher Wang, Dan Gutfreund, Josh Tenenbaum, and Boris Katz. Objectnet: A large-scale consroondolledd dataset for pushing the limits of object recognition models. Advances in neural information processing systems, 32, 2019.
Adrien Bardes, Jean Ponce, and Yann LeCun, Variance-imarrinee--vvvaenne regularization for self-supervised learning. ICLR, Vicreg, 1:2, 2022.
Peter L. Bartlett and John Shawe-Taylor. Generalization performacce of support vector machines and other pattern classifiers. 1999. URL htt.:/ppp..pppp...aaaaaca.o.aoaa
61853586
Florian Bordes, Randall Balestriero, Quentin Garrido, Adrien Bardes, and Pascal Vincent. Guillotine regularization: Why removing layers is needed improve generalization in self-supervised learning. Transactions of Machine Learning Research (TMLR), 2023.
Annie S Chen, Yoonho Lee, Amrith Setlur, Sergey Levine, and Chelsea Finn. Project and probe: adaptation by interpolating orthogonal features. In The Twelfth International Conference Learning Representations, 2023.
Mayee Chen, Daniel Y Fu, Avanika Narayan, Michael Zhang, Zhao Song, Kayvon Fatahalian, and Christopher Ré. Perfectly balanced: Improving transfer and robustness of supervised contrastive learning. International Conference on Machine Learning, pp. 3090-3122, PMLR, 2022.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020.
Ching-Yao Yao Chuang, Joshua Robinson, Yen-Chen Lin, Antonio Torralba, and Stefanie Jegelka. De- biased contrastive learning. Advances in neural information processing systems, 33:8765-8775, 2020.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 248-255. leee, 2009.
Aleksandr Ermolov, Aliaksandr Siarohin, Enver Sangineto, and Nicu Sebe. Whitening for self- supervised representation learning. In International Conference on Machine Learning, pp. 3015- 3024. PMLR. 2021.
Cong Fang, Hangfeng He, Qi Long, and Weijie J Su. Exploring deep neural networks via layer- peer-peeeeed model: Minority collapse in imbalanced training. Proceedings of the National Academy of Sciences, 118(4))e210309111 18, 2021.
Quentin Garrido, Yubei Chen, Adrien Bardes, Laurent Najman, and Yann Lecun. On the duality between contrastive and non-contrastive self-supervised learning. arXiv preprint arXiv:2206.02577 025044 2022.
Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A Wichmann, and Wieland Brendel. Imagenet-traime trained cnns are biased towards texture; increasing shape bias improves accuracy and robustness. arXiv preprint arXiv: /8/ /2231, 2018.
10

---

Published as a conference paper at ICLR 2024
Florian Graf, Christoph Hofer, Mare Niethammer, and Roland Kwitt. Dissecting supervised con- strastive learning. In International Conference on Machine Learning, pp. 3821-3830. PMLR, 2021.
Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems, 33:21271-21284 2020.
Yu Gui, Cong Ma, and Yiqiao Zhong. Unraveling projection heads in contrastive learning: Insights from expansion and shrinkage. arXiv preprint arXiv:2306.00333, 2023.
Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. Advances in neural information processing systems, 30, 2017.
XY Han, Vardan Papyan, and David Donoho. Neural collapse under mse loss: Proximity and dynamics on the central path. arXiv preprint S.kiv:2106.02077. 2021.
Jeff Z HaoChen and Tengyu Ma. A theoretical study of inductive biases in contrastive learning. arXiv preprint ar XrXi::2211. 14699, 2022.
Jeff Z HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma. Provable guarantees for self-supervised deep learning with spectral contrastive loss. Advances in Neural Information Processing Systems, 34:5000-5011, 2021.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 9729-9738, 2020.
Dan Hendrycks, Kimin Lee, and Mantas Mazeika. Using pre-training can improve model robustness and uncertainty. In International conference on machine learning, pp. 2712-2721. PMLR, 2019.
Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang. Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, et al. The many faces of robustness: A critical analysis of -uttt--disti generalization. In Proceedings of the IEEE/CVF International Conference Computer Vision, pp. 8340-8349, 2021a.
Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. Natural adversarial examples, 2021b.
Like Hui, Mikhail Belkin, and Preetum Nakkiran. Limitations of neural collapse for understanding generalization in deep learning. arXiv preprint arXiv:2202.08384, 2022.
Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Logan Engstrom, Brandon Tran, and Aleksander Madry. Adversarial examples are not bugs, they are features. Advances in neural information processing systems, 32, 2019.
Wenlong Ji, Zhun Deng, Ryumei Nakada, James Zou, and Linjun Zhang. The power of contrast for feature learning: A theoretical analysis. arXiv preprint 10. 2021.
Li Jing, Pascal Vincent, Yann LeCun, and Yuandong Tian. Understanding dimensional collapse in contrastive self-supervised learning. arXiv preprint arXiv:21 2021.
Siddharth Joshi and Baharan Mirzasoleiman. Data-efficient contrastive self-supervised learning: Most beneficial examples for supervised learning contribute the least. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning. volume 202 of Pro- ceedings of Machine Learning Research, pp. 15350. PMLR, 23-29 Jul 2023. URL fittps:/prrceeddii
m mlr. www.:mmrr...rs.eemm wotps///22/////oo.3oooo html
Siddharth Joshi, Yu Yang. Yihao Xue, Wenhan Yang, and Baharan Mirzasoleiman. Towards mitigating spurious correlations in the wild: A benchmark and more realistic dataset, 2023.

---

Published as a conference paper at ICLR 2024
Prannay Khosla, Piotr Teterwak, Chen Wang. Aaron Sama,, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishman. Supervised contrastive learning. Advances neural information processing systems, 33:18661-18677 2020.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training sufficient for robustness to spurious correlations. arXiv preprint arXiv:2204.029333 2022.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Bal- subramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al. Wilds: A benchmark of in-the-wild distribution shifts. In International Conference Machine Learning. pp. 5637-5664. PMLR, 2021.
Alexander Kolesnikov, Xiaohua Zhai, and Lucas Beyer. Revisiting self-supervised visual repre- sentation learning. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 1920-1929, 2019.
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. 2009.
Yann LeCun. The mnist database of handwritten digits. http://yann. lecun. com/exdb/mmist/, 1998.
Zhiheng Li, Ivan Evtimov, Albert Gordo, Caner Hazirbas, Tal Hassner, Cristian Canton Ferrer, Chenliang Xu, and Mark Ibrahim. A whac-a-mole dilemma: Shortcuts come multiples where mitigating one amplifies others. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 20071-20082 2023.
Hong Liu, Jeff Z HaoChen, Adrien Gaidon, and Tengyu Ma. Self-supervised learning is more robust to dataset imbalance. arXiv preprint arXiv:21 2021.
Jianfeng Lu and Stefan Steinerberger. Neural collapse under cross-entropy loss. Applied and Computational Harmonic Analysis, 59:224-241, 2022.
Raghav Mchta, Vitor Albiero, Li Chen, Ivan Evtimov, Tamar Glaser, Zhiheng Li, and Tal Hass- ner. You only need good embeddings extractor to fix spurious correlations. arXiv preprint arXiv:22/2.062544 2022.
Ryumei Nakada, Halil Ibrahim Gulluk, Zhun Deng, Wenlong Ji, James Zou, and Linjun Zhang. Understanding multimodal contrastive learning and incorporating unpaired data. In International Conference Artificial Intelligence and Statistics, pp. 4348-4380 PMLR, 2023.
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint arXiv: pete:411.66661 2014.
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint / 1807003748 2018.
Vardan Papyan, XY Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40): 24652-244 24663, 2020.
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? In International conference on machine learning, pp. 5389-5400, PMLR, 2019.
Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. Domain-adjusted regression or: Erm may already learn features sufficient for out-of-distribution generalization. arXiv preprint arXiv:2202.06886. 2022.
Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint arXiv: /9/ /. 108731, 2019.
12

---

Published as a conference paper at ICLR 2024
Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang. An investigation of why overparameterization exacerbates spurious correlations. In International Conference on Machine PMLR, 2020.
Learning, pp.
Nikunj Saunshi, Jordan Ash, Surbhi Goel, Dipendra Misra, Cyril Zhang, Sanjeev Arora, Sham Kakade, and Akshay Krishnamurthy. Understanding contrastive learning requires incorporating inductive biases. arXiv preprint arXiv:2202. 14037, 2022.
Vaishaal Shankar, Achal Dave, Rebecca Roelofs, Deva Ramanan, Benjamin Recht, and Ludwig Schmidt. Do image classifiers generalize across time? In Proceedings of the IEEE/CVF Interna- tional Conference on Computer Vision, pp. 9661-9669, 2021.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research, 19(1): 1822-2876 2018.
Rohan Taori, Achal Dave, Vaishaal Shankar, Nicholas Carlini, Benjamin Recht, and Ludwig Schmidt. Measuring robustmess to natural distribution shifts in image classification. Advances in Neural Information Processing Systems, 33: 18583-88311 18583-18599,, 2020.
Yuandong Tian, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning dynamics without contrastive pairs. In International Conference on Machine Learning, pp. 10268-10278 PMLR, 2021.
Xiang Wang, Xinlei Chen, Simon S Du, and Yuandong Tian. Towards demystifying representation learning with non-contrastive self-suuprvision. arXiv preprint arXiv:2110.049447 2021.
Zixin Wen and Yuanzhi Li. The mechanism of prediction head in non-contrastive sulf-supe learning. Advances in Neural Information Processing Systems, 35:24794-24809, 2022.
Kai Xiao, Logan Engstrom, Andrew Ilyas, and Aleksander Madry, Noise signal: The role of image backgrounds in object recognition. arXiv preprint 2Xiv:2006.09994, 2020.
Yihao Xue, Kyle Whitecross, and Baharan Mirzasoleiman. Investigating why contrastive learning benefits robustness against label noise. In International Conference on Machine Learning, pp. 24851-24871. PMLR, 2022.
Yihao Xue, Siddharth Joshi, Eric Gan, Pin- Yu Chen, and Baharan Mirzasoleiman. Which features are learnt by contrastive learning? on the role of simplicity bias in class collapse and feature suppression. arXiv preprint arXiv:2305. /6536, 2023.
Mang Ye, Xu Zhang, Pong C Yuen, and Shih-Fu Chang. Unsupervised embedding learning via invariant and spreading instance feature. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 6210-6219, 2019.
Chun-Hsiao Yeh, Cheng-Yas Yao Hong, Yen-Chi Hsu, Tyng-Luh Liu, Yubei Chen, and Yann LeCun. Decoupled contrastive learning. In European Conference on Computer Vision, pp. Springer, 2022.
Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self-supervised learning via redundancy reduction. In International Conference on Machine Learning, pp. 12310- 12320. PMLR. 2021.
Jinxin Zhou, Xiao Li, Tianyu Ding, Chong You, Qing Qu, and Zhihui Zhu. On the optimization landscape of neural collapse under mse loss: Global optimality with unconstrained features. In International Conference Machine Learning, pp. 27179-27202. PMLR, 2022a.
Jinxin Zhou, Chong You, Xiao Li, Kangning Liu, Sheng Liu, Qing Qu, and Zhihui Zhu. Are all losses created equal: A neural collapse perspective. arXiv preprini arXiv:22 192, 2022b.
Zhihui Zhu, Tianyu Ding, Jinxin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu. A geometric analysis of neural collapse with unconstrained features. Advances in Neural Information Processing Systems, 34:29820-29833 2021.
Zhuotun Zhu, Lingxi Xie, and Alan L Yuille. Object recognition with and without objects. arXiv preprint arXiv: 1. 1,5596, 2016.
13

---

Published as a conference paper at ICLR 2024
A THEORETICAL ANALYSIS
A.1 STRUCTURE OF WEIGHTS
Theorem A.1 (Weights of the minimum norm minimizer). The global minimizer of the CL loss LCL with the smallest norm, defined as [ W 1 + | W2ll satisfies
W = W W
Proof. Observe that
W WT | + W. W2lls [[,W 11 + W, W2 1
-Tr(WW WI)² 2 + (W, W2)²)
-Tr(WW W W2) W, W W. W.W.W. WWW,W,Wi)
Tr(( - W2)²) - TT W w2W1)
Consider any value of W = WW,W, that is global minimizer of the CL loss. This determines the value TW I W. W W.,). Hence the norm is minimized when Tr((W W W. W 2)² =
But (W W W W 2) is positive definite this occurs if and only if W W W W 2-
On the other hand, this minimum is achievable, since if W has SVD W = UVV T, then W
V W₂ + is solution satisfying W W W W2-
> Theorem A.2 (Weights of the model trained with gradient flow, proved in Arora et al. (2018)
Suppose the initialization satisfies
W ( I (0) W (0))
Using gradient flow, at any time t, we have
W (OW (1) W (1).
Proof. We reiterate the proof. Let Z f(X) be the matrix of model outputs. Calculate that the gradients are
aC
W. UL X
OW; 0Z
OC OC
X W
OW, 2 Oz
In addition, the chain rule and gradient flow gives
d
Oc
OC OW
W WT) -
W
di
OW
(W W2) = Oc
W OC OW
W2
di
Ow
Substituting, we see that di
w.wT) = di (W, W2). The conclusion follows.
A.2 WEIGHTS OF FRATURES
The following analysis is performed for self-supervised CL loss, under the condition that W, W,W WT
W W
Define the covariance of augmented examples
M = ((A())(((())) Diag( (0) 02 22..0 o?1),
(2)
and the covariance of augmentation centers
M = E.A[A(())))() Diag([(1 22n)) ... 2
14

---

Published as a conference paper at ICLR 2024
Let W 1- By invoking Lemma B.2 in Xue et al. (022)) about the minimizer of the loss, given the fact that M, M are full rank, have
W WM M M)p
(4)
where notation [A], represents the matrix composed of the first p eigenvalues and eigenvectors of positive semidefinite A (if 2 rankA then [A]p = A). Therefore, substituting yields
20 1-201) ... Diag( o o²
W W [Dia((
of + a² =
(5)
Note that [Diag(
... MM matrix that only keeps the p largest diagonal entries of ....).
(1 1- P4
of
Diag( (1-20:) 8
01+a"
To avoid cluttered notations, below we consider the case where p I d. However, we note that similar analysis holds for p > d. Now we can obtain W UvD, where U E RPxP is a matrix with orthonormal rows. Then
Ilf)(e)) |[eel | : D Di. a]l vD
(6)
which yields the conclusion about (l.f( in Theorem 3.5
Now let's ss examine [ (e:)ll Let U V and V be W, and W2'sSVD, respectively. Given
that W, W = W, W2, we have
SSS U V2
(7)
implying that the non-zero singular values W and We matches, and those singular values corresponding columns in U, match those in 22. Given that W2W, = W = UVD, we have
W =U2S.SVV = UVD because singular vectors match
S UVD
because U2 RPxp is unitary.
S.V. UVDD
because singular values match.
(8)
Given equation 8 we have
(e) =[[Wee
- ,s V
-|SS V eall
-IU DD:
D(i, i].
which completes the proof of Theorem 3.5
A.3 COMPARING SAMPLE COMPLEXITY
Definition A.3. We say that a data distribution P separable with a (7 p) margin if where exists w", 6" s.f. | and
(x, 8) || S A y(w x + b*) 7} 1).
is well-known that the sample complexity Bartlett & Shawe- Shawe-Taylor (1999) ). Therefore, refer to r as the sample complexity indicator. From the analysis Section A.2] we know that
ward-nargin SVM only grows with r = (p/n)² (e.g..
(e,) UUU VDe VDli (U (U,
Since U,U, is unitary and U has orthornormal rows, UU also has orthornormal rows. Thus, (e,))'ss 's are orthonormal. The same conclusion holds for fi(e,)'s 's as well.
f((e,-) [Tlley-) IT and 6* = 0, we can obtain 7 = llf Ol
Now, given Definition A.3 by letting w* =
for the i-th layer's representations. We also have v Σ Ii Thus, the sample
complexity indicator r; = E3-,, [11a for = 1,2. Substituting the values of (e, into the comparison between and 72, with some algebraic manipulation yields Theorem 3.6
[I,, I[le,
15

---

Published as a conference paper at ICLR 2024
A.4 ANALYSIS FOR NON-LINEAR MODELS
We introduce the following two lemmas which allow us to analyze coordinates of the model separately, Lemma AA. Suppose the model can be decomposed coordinate-wise fi...., fy and each fi is odd. Also suppose the dataset 'D follows a coordinate wise symmetric distribution, namely the pdf satisfies
pla1....., -ii- Jii 1.. ;(x ;2i-,,, Dit,
(9)
for any i. Then the contrastive loss can be decomposed coordinate-wise
-EI5((=))(((2]]]++++++ (f(()))((777) ((KK))))((()))+({({(){(()))))())
(10)
Proof.
f= -2) f(x)+1 B[(((x)
f(x)-)
(11)
(
= -E
+
(12)
27
-E
Σ ((()))((attt))
+E
Bos
(13)
27
2E
Σ M(())((()))))
+E
Bo.-Uni(-L.))
Σ (1f((z))(((T))
(14)
2TE.
Σ {()))(()))
+E
(((()))((07)))
(15)
P
(x(ff + (((())) (HH
(16)
Note that this hold for both supervised and unsupervised contrastive loss, since they only differ in how positive pairs defined.
Lemma A.5. Under the same assumptions as A.4 MSE loss as defined in i can be decomposed coordinate-wiz
Proof.
= E(g,y)
[(f(x) T 10--)) y]
(17)
27
= Eig,y)
y +
f((x;))
(18)
P
p
E((,)))
y²
yfs (m:) + E. -Ro,-Uii((-L1)
Σ
(19)
1=1
27
P
E(a.y)
3² Σ
(x.) + -Em-Unif(-Lp)
Σ
a1f((s))
(20)
P
2
Σ
uf((ar)) [1(()))
(21)
(1 11-)),,,
+
E ( (fi((z)) -2)
(22)
16

---

Published as a conference paper at ICLR 2024
is easy to check that the setting with the diagonal network and given data distribution satisfies the conditions of the lemma.
By the above lemmas, we can consider optimizing over each coordinate separately.
Theorem A.6 (Contrastive Loss). Assume |222 I vb(0) and [w/2 (0) + oo, [2() - and |fi (e2) 2 Vb(O).
(0)
(( ((w12) b(0) b(0), then as
Proof. Let f2 (-1,.... 2n) be the embeddings outputted by the second layer. Then the coordinate-win wise decomposition of the contrastive loss takes the form
- E 1)2
(23)
From here it is clear that the an optimal solution maps the second coordinate of every embedding to zero. Writing down the gradients for the weights and threshold,
aL
Fa>D(1200 (w22 11242, (12), b22) (W12 b12)]
Ow22
0L
E P222042200 (w22cr (12222 b12), b22 W220 ( (w12222)22
Owi2
0L
Ea>Dl faw --2220 (W2207 (w 222 b12). b22)l
0b22
aL
0b12 = Faa (w220 (2020(01222 (#12422 b12), d22 W22
Observe that [w12| and |222| are both decreasing and b12, b22 are hoth increasing throughout training, so that as + oc, the second coordinate of all embeddings goes to zero. Namely, this means that
U222
1112
b22
(24)
But since weights are decreasing and thresholds are increasing, |u:2 (0) that (lw (12)
v Vb(0)
< V
b22
It follows
V 61
V1(0),
Rearranging gives
((0) b
as desired.
Theorem A.7 (MSE Loss), Assume |w22 I Vb(o and |122 K[wiz ((0) ((0), and Waz and
(0)
(0)
(0)
W12 have the same sign, then as t - oc, [f((e2) + 0 and Ifi (ez) 2 / /(0)
Proof. Let f2(x) (zI,..., zn) be the embeddings outputted by the second layer. Then the coordinate-wise decomposition of the loss takes the form
= [<> 1)2
21 +
(25)
where C independent of the weights.
From here it is clear that the an optimal solution maps the second coordinate of every embedding to zero. W.L.O.G., assume w w/2 0. Writing down the gradients for the weights and threshold,
0L.
Ex- 22200' (wan (12222 b12), b22) (w (1222 b12)]
Ow22
0L
Ez-D 22220'
( 12 b12), b22 o220 (w 1222))))))))
Ow
aT.
Ez- 122200 ( (12222 b12). byr)l
0b
0L
Bu=pl-2ze (W22a (1022o(w1222 (222 b12). b22 [222
0b12
17

---

Published as a conference paper at ICLR 2024
10
sample complexity indicator
0.45
0.5
0.55
0.6
10a
6 8
10
12
layer index
Figure 5: The value of the sample complexity indicator at different layers. Legends show the weight assigned to the downstream relevant weight by the pretrained model. We observe that the optimal layer shifts lower when the pretrained model assigns less weight to the downstream relevant features.
Observe that W12 and W22 are both decreasing and b12, 622 are both increasing throughout training, so that as + oo, the second coordinate all embeddings goes to zero. Namely, this means that
22
W12
22
(26)
But since weights are decreasing and thresholds are increasing, w S Vbl0) S <V
It follows
that (w b(1) z 12 2 Vb(0), Rearranging gives
Vo(0) + -11
W12
as desired.
B DISCUSSION ON MULTI-LAYER LINEAR MODEL
As mentioned in Section Theorem 3.4 can be extended to multi-layer models, which gives us (ttt) = W (t)W (t) Based on this, going through similar process in Appendix
A.2land Appendix A.3 we will obtain the following expression of the sample complexity indicator for each layer d:
E 2/L 7i
Cj $3
2/12 L E.' 35-
where
(1 (1-ag)os
if j E f-i.....-miiaad.. Imin(d,3)
Cj
0.
else
The definitions of j1. ... (m.n(d.p)) remain the same as in Theorem 3.5, and other quantities are consistent with the definitions provided in Section 3. Here, cj represents the weight the full model would allocate the j-th feature. Although the expression appears complex, some intuition can be gleaned from extreme scenarios: If cy- is the largest, indicating that the full model assigns the most weight to the downstream relevant feature, Ti decreases with l. This means one should just use the final-layer representations. 2. Conversely, if cy. is the smallest among non-zero Cj' ss indicating that the full model assigns the least weight to the downstream relevant feature, rt increases with this case, using the lowest layer would be preferable. Applying these observations in conjunction with the relationship between and ay, éj. similar conclusions to those in Corollary 3.7 regarding the impact of augmentation and feature strength for multi-layer models can be drawn.
18

---

Published as a conference paper at ICLR 2024
The greater the mismatch between the pretraining and downstream tasks, the lower the optimal layer tends to be. What about situations that are more intricate, occurring between the above extreme cases? To explore this, consider the following setting for simulation. We let
L = 12, 6, = 1, 5. 9, 6 = 0.1, Vj 2 10,
j j== =,
C = 0.4, S8
Cj 0.6, vj. S 10.
Then, we vary cy-, the weight assigned to the downstream relevant feature from 0.4 to 0.6, and plot T vs / (i.c., sample complexity indicator vs depth) under each value of oj- in Figure 5 We observe that, for cj. 0.45, 0.5, 0.55, the best sample complexity is achieved by some intermediate layer. Additionally, the optimal layer (corresponding to the bottom of the U-shaped curve) becomes lower as the weight assigned to the relevant feature decreases, which indicates larger mismatch between the pretraining and downstream tasks.
Challenges in locating the optimal layer. Even in this simplified scenario, observe that the depth of the optimal layer is influenced by various factors, including the position of the downstream- relevant feature, the strength of features in the downstream task, and the weights assigned features during pretraining. The last one is further a function of features, noise and augmentations for pretraining data. In practical settings, we acknowledge that more factors may come into play, such as the model architecture, which varies across layers. Therefore, determining the exact optimal layer for downstream tasks is very challenging and represents an intriguing and valuable avenue for future exploration. We believe the analytical framework established in this paper, capable of expressing downstream sample complexity closed form through various elements in pretraining and downstream tasks, and explaining several observed phenomena (e.g., those depicted in Figure can significantly aid advancing research in this direction.
C EXPERIMENTAL DETAILS
C.I SETUPS
Synthetic Experiments, CL. We train two-layer (symmetrized) ReLu Networks with momentum
SGD, with momentum set to 0.9. We use the spectral CL loss. (1) In setting 1, we set d = 5, p = 20
and Φ1 = = 6s = 1, indicating that all five features have equal strength. We set a1 02, dra, 044, Os to 0, 0.25, 0.5, 0.75, 1, respectively, meaning that features are disrupted by the augmentation to different extents. We let o = 0.01 and set learning rate to 0.05. (2) In setting 2, we set d = 9, = 20, 6 3:2/2', Vi, and 0.1. We use learning rate 0.01.
Synthetic Experiments, SCL. The data distribution is the same as described in Definition 4.1 We let other coordinates be randomly drawn from {-0.001, 0.0013. We let d 5.p 20, learning rate = 0.01.
0.1,
MNIST-on-CIFAR-III The main data generation process is outlined in Section 5.1 To provide further details, after processing the digits in each image, we apply standard data augmentations, RandomRessizedCrop and ColourDistortion. Additionally,
resize the digits set their height to 16 pixels. We train ResNet-18 models. By default, we use a temperature of 0.5 and minimize the SimCLR loss with the Adam optimizer. Our training batch size is 512. with learning rate of 0.001 and weight decay set to 1 x 10-6 6 layer MLP with an output dimension of 128, and the hidden dimension is set to match the output dimension of the ResNet- 18 encoder.
We train for 400 epochs. The projection head is a one-hidden
Coarse-to-fine transfer on CIFAR-100. On CIFAR-100, 100, we refer to the 10 super-classes 'coarse' and the 100 classes as 'fine". (1) SCL. The pretraining is conducted with the 10 coarse-grained labels, using the SCL loss in Khosla et al. (2020) with temperature 0.5. We use train a ResNet- 18 with momentum SGD, using learning rate O.I. momentum 0.9 and weight decay 1c-6. We train with batch size set to 512 for 400 epochs. The projection head is a one-hidden layer MLP with an output dimension of 1 128, and the hidden dimension is set match the output dimension of the ResNet- 18 encoder. (2) SL. We use momentum SGD, with learning rate = 0.1 and momentum = 0.9. We train for 200 epochs with batchsize 128 and weight decay 5e-4.
Few-shot adaption on UrbanCars. UrbanCars is constructed by Li et al. (2023), ). features multiple spurious correlations. The task is classifying images as either urban cars or country cars. Each image
19

---

Published as a conference paper at ICLR 2024
Target Test Accuracy
90
85
post proj. pre proj.
80
post reweight
75
pre reweight
70
vanilla
2
8
16
32
64
128
Target Data
Size
Figure 6: Target test accuracy with respect the number of target examples for adaptation, comparing different representations.
has background (BG) and co-occurring object (CoObj). The BG is selected from either urban or country backgrounds, and the CoObj selected from either urban or country objects. the source distribution, for each class, images with common BG and CoObj constitute 90.25%, images with uncommon BG and common CoObj, or with common BG and uncommon CoObj. constitute 4.75%, and images with both uncommon BG and CoObj constitute 0.25%. This dataset presents challenge due to multiple spurious correlations/shortee the subpopulations that are most underrepresented in the source distribution, i.e., images with both uncommon BG and CoObj. We train the model on the source data using the SGD optimizer with batch size of 128, a learning rate of 0.01, and weight decay of 0.000001 for 50 epochs. The linear layer is then trained on top of representations using few (2 to 128) data from the target distribution. Following Chen et al. 2023). for the training of the linear layer, we employ the Adam optimizer Kingma & Ba (2014) with batch size of 64 and train for 100 epochs. We tune both the learning rate and weight decay in the range of 0.1, 0.01, 0.001, and report the configuration that yields the best result. Each experiment is repeated 10 times, and we report the average. We use the implementation from Joshi et al. (2023) for these experiments.
We let the target distribution contains only
Experiments with reweighting heads. All setups remain the same before. We set the values of to 1.05, 1.5, 1.2, and 1.01 for the four experiments in Table 2 from top to bottom, respectively.
C.2 DETAILS RESULTS FOR THE LAST ROW OF TABLE
presents the target test accuracy with respect the number of target examples for adaptation, comparing different representations. We observe that for both the projection head and the reweighting head, the representations before the head are better than those after, and also better than the vanilla approach, which does not add any head during pretraining- Furthermore, pre-reweighting-hees representations outperform pre-projection-head representations, underscoring the potential of the reweighting head.
D ADDITIONAL EXPERIMENTS
D. TWO EXPERIMENTS ABOUT FEATURE STRENGTH NATURAL IMAGES
Dataset and Downstream Task. The experiment is based on images in CIFAR-10, 10, aiming to validate the second and third points outlined in Corollary 3.7 using real data. Given the subjective nature of defining features in natural images, less controversial definition. This choice enables a controlled experiment with minimal modification to the original natural images. Our downstream task involves predicting whether given image is categorized as 'red'' "green,' or 'blue' based on the channel with the largest mean value. To vary the the strength of the color feature in the pretraining data, consider the following two approaches.
select the feature 'color,' which offers straightforward and
20

---

Published as a conference paper at ICLR 2024
Oattre Stargh Uaser Ere A65
Lewerbo Ere AecM Comment Fraage
Jaber Emplieet Colrr During dfaci
Fal
113
115
74
Cerbrat-teen
(a) Acc v.s. Contrast.
(b) Gap v.s. Contrast.
(c) Acc v.s. color dist.
(d) Gap v.s. color dist.
Figure 7: Accuracy and gap when the level of color distinguishability in the pretraining data is varied. (a)(c) are for approach (b)(d) are for approach 2.
We note that the second one involves no change to the images themselves, making sure all of them are natural.
1. Varying Color Distinguishability by Processing the Image. During the training, we process each image as follows: Let R, G, and B denote the average pixel value for the three channels. We first find the largest two values among R. G, and B and calculates u as their mean. For the matrix at each channel, denoted MR. Mc. and MB for the red, green, and blue channels respectively, we update each matrix as follows:
RR u) O' +
Mn MR
GG U) a + u
MG - Mc x
(B u) a + u
MB + MR x
B
Then, the values are capped between 0 and 1. a 2 strength. A larger increases the values in the dominant channel while decreasing the values in other channels. Intuitively, larger a means larger gap between the dominant channel and the other two channels, making the information about the dominant channel more obvious. We varied the in the range of [0, 0.2, 0.5, 1, 2, 3, 4, 5).
parameter that we denoted contrast
2. Varying Color Distinguishability by Selecting a Subset. For each image, we choose the channel that has the greatest mean pixel value as the dominant channel, and compute the color distinguishability D (dominant channel's mean value / sum of the mean values for the three channels). We sort images in each class of CIFAR- 10 based on D from highest to lowest. Then we select 1000 images from each class starting from the K-th image to form the training set. Intuitively, a larger K leads to images with higher average color distinguishability D values being selected the pretraining data, and higher color distinguishability means that the information about the dominant channel is more obvious. We varied K in the range of [O, 500, 1000, 1500, 2000, 2500, 3000, 3500, 4000], which corresponds the average color distinguishability of [0.4254, 0.3950, 0.3815, 0.3717, 0.3640, 0.3575, 0.3517, 0.3463, 0.3406].
Training Details. We train ResNet18 with a one hidden layer projection head with hidden dimension 2048 using SimCLR (Chen et al. 2020) on the selected trainset images for 400 epochs. We use Adam optimizer with learning rate 0.001, weight decay 10-6, and batch size 512. We used temperature 0.5.
Results. Figure F presents the results for the above two experiments. In both cases, we see that the gap between pre-projection and post-projection roughly show an secreasing-increaste trend. In other words, using pre-projection is more beneficial when either the color distringuishabiliyy is very high or very low, confirming our theoretical results in Corollary 3.7
D.2 EFFECT OF EARLY STOPPING
We examine how the linear evaluation accuracy changes during training MNIST-on -IFAR-- -00. The setting is consistent with the setting described Section 5.1. Intuitively, early stopping should result in model that is less specialized towards the training objective, potentially benefiting the
21

---

Published as a conference paper at ICLR 2024
acc of li eval on mnist
acc of lin eval mnist
998
96
post projection on pre-projection
past projection pre pre-prajection
epoch
epoch
(a) Pitrop 0
(b) Pdroy 0.2
Figure 8: Linear evaluation accuracy during training. We see that early stopping reduces the gap
between pre-projection and post-projection when Pdrop 0, but not when Pdrop 0.2.
Original
Dfi--ff-iitiitttioon Accuracy
No Projection
Pre Projection
0.685
Post Projection
0.675
0.850
In-Distribution Accuracy
Figure Pre-projection representations exhibit better robustness. OOD accuracy evaluated across four shifted ImageNets.
Original
Out- Dr-Dsisiiibttton Accuracy
No Projection
Projection
Post Projection
0.935
3.500
0.945
5,050
0.955
3,900
0.655
In-Distributtin Accuracy
Figure 10: OOD-ID accuracy relation for models trained on Mixed-Rand.
downstream task when there is misalignment. However, perhaps counter intuitively, the results in Figure reveal that early stopping improves post-projection with good augmentation (Ddrap = 0), but with bad augmentation (Párop 0.2). A deeper analysis of training dynamics during the course of training is required to fully grasp the effect of early stopping.
D.3 ROBUSTNESS TO DISTRIBUTION SHIFT ON IMAGENET
It is well-known that deep learning models trained on ImageNet Deng et al. (2009) tend to rely heavily on backgrounds rather than the main objects for making predictions (Xiao et al. 2020 Zhu al. 2016), resulting in poor performance test sets where backgrounds are random absent.
Recently, Kirichenko et al. (2022 demonstrate that despite this poor performance, the representations before the linear classifier actually contain information about the objects. Consequently, training a new classifier on these representations using data that does not exhibit a correlation between backgrounds and classes can lead improved out-of-distributtion (OOD) accuracy. We demonstrate
22

---

Published as a conference paper at ICLR 2024
that the quality of representations be further enhanced by leveraging non-linear projection head. We take the ImageNet pretrained ResNet-50 model, and fine-tune the model with an dditional projection head on ImageNet for 50 epochs. We find that the pre-projection representations result in an improved in-distribution vs. OOD accuracy relationship (which is a standard measurement of robustness (Taori al. 2020)), as shown in Figure Figure(wwhrre where compare it with post-projection representations, representations provided by the original pretrained model and representations obtained by fine-tuning the model for 50 epochs without the projection head.
Experimental details. Following Kirichenko et al. (002)), we use the datasets in Backgrounds Challenge based on the ImageNet-9 dataset (Xiao al. 2020) along with the ImageNet-R (Hendrycks) et al. (021a) dataset. We use ImageNet-A (Hendrycks et al.) (021b) instead of Paintings-BG (Kirichenko et al., 2022) due to limited dataset availability. We consider Original, Mix-Rand, and ResNet-50
FG-only datasets for Backgrounds Challenge. We finetune the ImageNet-pretrainng encoder along with an additional randomly initialized projection head and a randomly initialized linear classifier. The projection head is MLP with one hidden layer of size 2048. We finetune the model for 50 epochs with batch size 256, momentum 0.9, weight decay 10 10-4, learning rate 0,01, and learning rate scheduler with step size 30 and 7 0.1. We also finetune the original ImageNet- pretrained ResNet-50 model with the same hyperparameters. During the training, we follow the same procedure as Kirichenko et al. (2022), using Mixed-Rand and Original. The sizes of Original and Mixed-Rand are the same. We use the same data preprocessing step as Kirichenno et al. (2022). We train DFR on varied-sized random subsets of the training data with Mixed-Rand data sizes in the range of (5000, 10000, 20000, 45405) ) for 1000 epochs using SGD with full batch, learning rate I, and weight decay equal 100 / size of the datasets, Rand, FG-Only, ImageNet-RR R,
train sets: Mixed-Rand and Combination of
data. For evaluation, we use four ImageNet-A, and one in-distribution dataset, Original. We average the accuracy of the four in- distribution datasets and present the disttii-ittiitoonnnnn accuracy vs. in-distribution accuracy plot for the four settings: pre-projection, post-projection, original (representing original mmageNet-pretraiine ResNet-50 model) and no no-projection (representing finetuned ResNet-50 model without projection head). The results for Combination of Mixed-Ram and Original and Mixed -andd are shown in
Figures and 10 respectively.
23


Reviews:
[{"review_id": "XIzBecsrmQ", "reviewer": "Reviewer_Kpzq", "summary": "This paper analyzed an important technique in contrastive learning: the projection head. The authors theoretically demonstrated the benefits of the projection layer via a simplified model. The theoretical analysis showed that lower layers represent features more evenly in linear networks and can represent more features in non-linear networks, which implies better generalization performance in downstream tasks. Empirically, they verify the theoretical findings on synthetic and real-world datasets.", "questions": "see my comments above.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The projection layer is one of the most important techniques in contrastive learning and the mechanism behind it is still under-explored. Consequently, this paper addresses an important problem.\n2. The theoretical analysis in this paper looks solid and insightful. And the empirical results verify the theoretical findings.", "weaknesses": "1. The theoretical analysis in this paper is based on a two-layer model. Is it possible to extend the results to multiple-layer networks? For example, should we discard other layers except for the projection head in downstream tasks in the deep networks? It would be better to provide more discussions about that.\n2. This paper demonstrates the benefits of the projection head. However, we can observe that the designs of the projector (e.g., the layers and the dimensions) also have a significant influence on the downstream performance. Is it possible to provide some insights about the design of the projector based on the theoretical analysis in this paper? \n3. As stated in this paper, the pre-projection representations are preferred in three different scenarios and the findings are verified on the synthetic datasets. However, the authors do not show similar results (e.g., the influence of data augmentations) on the real-world datasets. It would be better to provide more empirical findings on real-world datasets.\n4. I note that the abstract on the OpenReview website is different from that on the pdf file, which should be corrected.\n5. The forms of references are inconsistent. For example, some of the conferences are full titles while others are abbreviations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FYGj1JN1Vx", "reviewer": "Reviewer_EBwM", "summary": "The paper delves into a nuanced aspect of neural network architecture design, specifically the use of a projection head during the training phase. This technique has garnered attention due to its empirical success in enhancing representation quality. The core methodology involves appending a projection head atop the encoder during training, which is subsequently discarded, favoring the pre-projection layer representations for inference tasks.\n\nDespite its proven practical effectiveness, a comprehensive theoretical understanding of why the projection head enhances representation learning remains underdeveloped. The paper aims to bridge this gap by dissecting the mechanics of the projection head and elucidating its impact on the learning dynamics of neural networks. This investigation is critical as it addresses a disconnect between empirical practices and their theoretical foundations in the field of deep learning.\n\nThe projection head's primary role is hypothesized to act as a regularization mechanism, potentially aiding in learning more generalizable and robust features. By expanding the representational capacity during training, the projection head could encourage the encoder to learn a broader set of features, some of which may be discarded during the projection phase but still contribute to a richer feature space in the pre-projection layer. Moreover, the projection head could serve to disentangle the feature space, making it easier for the network to differentiate between relevant and irrelevant features. This disentanglement might facilitate better generalization to new, unseen data by reducing overfitting to the idiosyncrasies present in the training dataset.", "questions": "1. **How does the architecture of the projection head influence the learning dynamics and final representation quality?**\n   - The design choices within the projection head (e.g., the number of layers, types of activations, dropout rates) likely have a profound impact on its efficacy as a regularizer and feature disentangler. What are the optimal architectural configurations for different types of data and tasks? Investigating this could provide more nuanced guidelines for practitioners and lead to a deeper theoretical understanding of the projection head's role.\n\n2. **What is the impact of the projection head on the interpretability of the learned representations?**\n   - While the projection head might aid in learning more generalizable features, its impact on the interpretability of these features is unclear. Do the representations learned with a projection head offer better clarity in terms of feature importance or contribution to the final decision? Understanding this could bridge the gap between performance and explainability in neural networks.\n\n3. **Can the benefits of the projection head be replicated or enhanced by alternative or complementary techniques?**\n   - Are there other methods or architectural innovations that could either replicate the benefits of the projection head or enhance its effects? For instance, could certain types of normalization, attention mechanisms, or even different training paradigms offer similar or greater benefits regarding feature representation and generalization? Exploring this could lead to a broader set of tools for improving neural network training beyond the projection head.\n\nDelving into these questions could significantly enhance the paper's contribution, offering both a deeper theoretical understanding and more practical guidelines for employing projection heads in neural network training.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "**Empirical and Theoretical Integration**:\nThe paper bridges the gap between empirical success and theoretical understanding by critically investigating the role of the projection head in representation learning. By scrutinizing a technique that has demonstrated practical effectiveness without a solid theoretical foundation, the paper contributes to a more profound understanding of neural network architectures, potentially guiding future designs with a better-informed rationale.\n\n**Regularization and Feature Representation**:\nIt hypothesizes that the projection head is a regularization mechanism, allowing the encoder to explore a wider feature space during training. This could lead to more robust and generalizable representations, as the encoder is encouraged to capture a broader and more nuanced feature landscape. The paper's exploration of this aspect could elucidate how neural networks can be trained more effectively to learn generalizable features.\n\n**Feature Disentanglement**:\nThe projection head's potential to disentangle the feature space is a significant strength of the paper's hypothesis. By facilitating a clearer separation of relevant and irrelevant features, the projection head might aid in reducing overfitting and improving the model's ability to generalize to unseen data. This aspect of the paper could contribute valuable insights into how neural networks can be made more interpretable and reliable.", "weaknesses": "**Potential Overfitting Risks**:\nThe introduction of a projection head could potentially lead to overfitting, especially if not properly regularized or if used in conjunction with datasets that have a high degree of noise or variability. The paper should address these risks and propose strategies to mitigate them.\n\n**Generalizability and Applicability**:\nThe projection head's effectiveness might vary across different architectures, tasks, and data modalities. The paper could benefit from a more detailed exploration of these variations to understand where the projection head is most beneficial and where it might be detrimental. This would enhance the paper's technical depth and practical applicability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ips6r4KPLm", "reviewer": "Reviewer_B9NE", "summary": "This paper investigates the effectiveness of the projection head in self supervised contrastive learning, supervised contrastive learning and supervised learning. It provides theoretical analysis of the quality and robustness of the learned representations and their generalizability in simple linear and nonlinear models. Theoretical results are supported with experimental evaluation on several image datasets.", "questions": "- In Definition 3, could you add details about $\\phi_i$?\n- Theorem 3.6 essentially tells us when it is beneficial to use pre and post projection representations, with concrete guidelines given in Corollary 3.7. Do these results hold for non linear models as well? As a stretch question, I was wondering how could one infer the downstream-relevant features and the weights of features in practice? In other words, is there a way to use your results to identify what features will be relevant for the downstream task? I am aware this might be out of scope of this work. \n- What is $\\alpha_i$ in setting 2 in Sec 5.1? Could you comment on the results in Fig 1 right, where $\\phi_i$ = 0.2 and 0.4, in particular, why do we see the spikes in weights in pre-features for $\\phi_i$ = 0.2 and post-features for $\\phi_i$ = 0.4?\n- I do not understand the experiment in Fig 3b. In the text your say that “Figure 3b shows the downstream accuracy against $p_{drop}$ and s=1”. However, if s=1, it means that you only use MNIST images as input to the augmentation, and then you additionally drop the digit with probability $p_{drop}$. Doesn’t that mean your augmented image is completely black if $p_{drop}=1$?\n- In Table 2, WaterBird SL results are almost the same, especially for no projection and pre-projection. Why is that? \n- Could you explain better what exactly is plotted in Figure 4?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Thorough theoretical analysis and interesting insights.\n- Extensive related work.", "weaknesses": "- Experimental evaluation is done on very simple networks and small datasets. Although the results nicely support the theoretical results, it might be beneficial to include one more complex experiment.\n- The paper could benefit from a discussion on limitations and assumptions of the analysis.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1t4N5XRVu0", "reviewer": "Reviewer_p748", "summary": "This paper studied a very interesting question: what makes pre-projection representations better if they are not directly optimized? Based on theoretical analysis on some toy models, they proposed that the implicit bias of training algorithms makes deeper features more unequal, and hence lower layers tend to have more normalized and less specialized representations. Then they showed that lower layers are better in the following cases: (1) data augmentation disrupts useful feature; (2) downstream-relevant features are too weak/strong in pre-training. They also showed how this mechanism makes lower representations better for supervised contrastive learning and supervised learning. Finally, they conducted some experiments to verify their theoretical analyses.", "questions": "1. I want to know whether their data model (Def 3.1) is commonly used in literature? Even though the authors' theoretical analysis are insightful, I am worried that this data model is too simply to be applied to practical scenarios.\n2. Most of the cases where lower representations are better arises from the inappropriate data augmentation (like Thm 4.2), namely the pre-training signal does not align with the downstream problem. In such cases, overfiting the pre-training data (what post-projection layers intend to do) may lead poor downstream performance. In addition to using pre-projection layers, can other regularization methods (weight decay, dropout, early stopping, etc.) also lead to satisfied performance even without projection head? I noticed that the author mentioned in the end of Sec 3 that the advantage of pre-projection representations diminishes when using weight decay.\n3. The author mainly focused on two-layer neural networks. When there are multiple layers, how do the depth of representations affect the result? Is there any trade-off, like deeper layers have more representation power while less diversity/robustness? How to choose the appropriate depth of projection heads?\n\nOverall, I think this is an insightful work and I am glad to raise my score if the authors can clarify their over-simple model and give a slightly deeper answers to the above questions.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. This paper studied a very interesting question: what makes pre-projection representations better if they are not directly optimized?\n2. Based on theoretical analysis on some toy models, they proposed that the implicit bias of training algorithms makes deeper features more unequal, and hence lower layers tend to have more normalized and less specialized representations.\n3. Then they showed that lower layers are better in the following cases: (1) data augmentation disrupts useful feature; (2) downstream-relevant features are too weak/strong in pre-training.\n4. They also showed how this mechanism makes lower representations better for supervised contrastive learning and supervised learning.", "weaknesses": "My main concern is that their data and model are too simple. But if similar models are commonly used, it may be okay.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)