On the Importance of Embedding Norms in Self-Supervised Learning

Thank you for the helpful suggestions and the in-depth review! We include experiments and discussion in response to your questions on the SSL methods used in the paper, the training epochs and the applicability/novelty of our analysis. Responses to individual questions are below:

"while the paper makes a significant contribution to understanding the role of embedding norms in SSL, the study is limited to just two SSL methods"

We agree with the reviewer that adding more methods would be beneficial. To address this concern, we have run additional experiments on BYOL (resnet backbone; optimizes the cosine similarity), Mocov3 (ViT backbone; optimizes InfoNCE) and Dino (ViT backbone; not cos.sim.-based) on ImageNet-100. The results can be found in the response to Reviewer FWxF and are in accordance with the results in the paper. Please let us know if this has helped to address the concern.

"I am curious whether the same phenomenon would occur with embedding norm in clustering and whitening-based methods, such as SwAV and W-MSE."

Our key theoretical contribution (norms grow, gradients decrease by norm) holds generally for loss functions based on normalized embeddings. As a result, it applies to SwAV, as well as the similar ProtoNCE, and W-MSE. The latter two directly use the cosine similarity, so Prop. 3.1 and Thm. 3.4 apply directly. Additional methods which use normalized embeddings for which our results apply include NN-CLR, ProtoNCE, BYOL, CLIP, DCL, and BEITv2, to name a few. We will make this clearer when revising the paper and will include a formalization of the general statement (that our findings hold for loss functions which depend on normalized embeddings).

"Methods like Barlow Twins[1] and VICReg[2] do not adopt Embedding Norm; incorporating it leads to collapse in these methods."

We agree that this is an interesting direction for future work. As the reviewer points out, several SSL models perform worse when the embeddings are normalized. However, our analysis does not apply to these methods and we are not able to test why Barlow Twins and VICREG fail in the normalized setting in the limited rebuttal period. We will clarify this during the revision.

"theorem 3.4 is unrealistic in my opinion, as it only involves the cosine similarity of positive pairs and neglects the interactions of negative samples"; "the theoretical contributions, while useful, lack novel insights"

We maintain that theorem 3.4 is a novel theoretical contribution and note that it holds immediately for the non-contrastive models SimSiam and BYOL, both of which only optimize the cosine similarity between positive samples. While it does not directly apply to the InfoNCE loss, we believe that our extensive analysis in Section 6 describes how embedding norms and convergence interact in contrastive settings. We also refer to our rebuttal to reviewer Jmax, where we extended the simulations in Section 4.1 to further analyze how weight-decay affects convergence.

"The lack of original theoretical contributions (such as new theorems or proofs) limits the novelty of the paper’s theoretical aspect."

We recognize that our theoretical contributions build upon existing work in deep metric learning. However, our paper's novelty lies not in creating new mathematical frameworks, but in bringing theory and observations together and showing that they are causally related. The theory about gradient norms had not previously been experimentally verified and the experimental observations about SSL confidence did not have a theoretic underpinning. Our theory, simulations and experiments were designed to validate and explore these connections.

"the training duration used in the experiments (128–256 epochs) is relatively short"

To address this concern, we have run on Cifar10/100 for 1000 epochs and report the kNN accuracies below:

The performance improvements of cut-initialization and GradScale persist after 1000 epoch training and even grow in the case of Cifar100. We thank the reviewer for the suggestion.

The ImageNet experiments in Tables 2 and 3 were already run for the requested 500 epochs.

---

We hope we've addressed your concerns but please let us know if there is something else we can do to convince you further.

Thank you for the extensive analysis and suggestions for how to improve our work! Among other things, our rebuttal includes experiments on the Tiny-ImageNet dataset, on additional models, additional probes and on how the embedding's rank corresponds to the embedding norm. We respond to individual points below:

"I found the lack of analysis/experiments on ImageNet a concern."

Unfortunately, due to computational constraints, this is infeasible for us. However, to address this concern, we trained SimCLR for 500 epochs on the Tiny-Imagenet dataset, which contains 200 classes, each with 500 training samples. The results are consistent with our other experiments:

We will include these when updating the paper.

"I have seen linear probe... favored in SSL literature"

We focused on kNN as it is known to lower-bound the other probes [1] and is a good indicator of model performance [2]. However, to address this concern, we have also run linear probes on several datasets:

Tiny-ImageNet linear probe results are presented in the table above.

CIFAR-100 linear probe results can be found below:

In both cases, the linear probe results are in line with the kNN ones. We will include linear probe evaluations in the revision.

"using popular SSL methods would make the paper interesting to readers as well as the authors"

We agree that adding more methods would be beneficial. To this end, we have run additional experiments on BYOL (resnet backbone; optimizes the cosine similarity), Mocov3 (ViT backbone; optimizes InfoNCE) and Dinov2 (ViT backbone; not cos.sim.-based) on ImageNet-100. The results can be found in the response to Reviewer FWxF and are consistent with our other experiments.

"the use of embedding norm makes it harder to interpret (why is 0.74 considered a good indicator of "OOD" in Figure 3)"

The absolute embedding norms will differ between models and would be difficult to use directly as a confidence metric. Instead, we propose using the relative embedding norms. Put simply, if an embedding norm is smaller than most norms seen on the training set, then the sample is likely OOD. Thus, in Figure 3, we have normalized the values by the training set's mean embedding norm, implying that the expected in-distribution norm is 1. The closer to 0 the value is, the more likely the sample is to be OOD. We will make this more clear in the revision.

"lack of discussion on the uses of embedding rank and relationship to downstream performance"

We thank the reviewer for these references, we will include them in the related work section. Although the references address a different aspect of SSL embeddings than we do (they discuss the rank of all the embeddings whereas we discuss the norm of individual embeddings), we agree that these topics are related. Specifically, our work shows that the embedding norms grow in regions of the latent space which have high density (Section 4.2). This implies that these regions may induce large eigenvalues on the covariance matrix or be clusterable as in CLID.

As a preliminary test, we include below a table which shows the rank of the Cifar10 embedding space over training epochs and the corresponding mean embedding norm. We calculate rank as the number of principal components required to capture 99 percent of the latent space's variance. Interestingly, we find that the rank starts growing at roughly the same epoch where the norms stop growing, implying a potential correlation. We will include this in the discussion section at the end of the paper.

We also note that this is likely related to the study on dimensional collapse in SSL methods, such as [3].

"the new interventions need testing at a larger scale"

We emphasize that these interventions were included to study the behavior of the embedding-norm effect in practice and were not intended to obtain SOTA results. However, we agree that further testing would be helpful. To this end, we hope the analysis on additional models (BYOL, MoCov3 and Dino on ImageNet-100) and the additional experiments on Tiny-Imagenet have addressed this concern.

Figure 3

We will amend the figure.

References

[1]: Oquab, Maxime, et al. "Dinov2: Learning robust visual features without supervision." TMLR 2023.

[2]: Marks, Markus, et al. "A closer look at benchmarking self-supervised pre-training with image classification." arXiv preprint 2024.

[3]: Tian, Yuandong, et al. "Understanding self-supervised learning dynamics without contrastive pairs." ICML 2021.

---

Please let us know if these have addressed your concerns.

Thank you for your thorough review. Regarding Theorem 3.4's applicability, we appreciate your insights about potential mitigations - weight decay, scaling gradients by embedding norms, and adaptive optimizers. We note that our paper systematically analyzes how the first two affect SSL training in practice. For completeness, we include adaptive optimizer experiments below.

We now respond to individual comments:

"The extension to the broader InfoNCE objective is trivial"

We agree that the extension from cosine similarity to the InfoNCE objective is straightforward. However, we believe that this is a strength of our paper.

The InfoNCE objective function is used in countless settings (see reply for reviewer FWxF) and our work shows that, in each of these, it has a dependence on the embedding norms. Although a few of these ideas were known in the deep metric learning literature, they had not been experimentally validated or extended to standard contrastive losses. In short, it is relevant for the SSL community to be aware of how the embedding norms interact with the InfoNCE loss function.

"the core analysis reaffirms known results about the cosine-distance gradient"

Quick clarification: we believe that our core analysis lies in studying the relationship between SSL training and the embedding norms with (i) new theory (Thm 3.2 and 3.4), (ii) new experiments and (iii) new mitigation strategies. Such a comprehensive analysis was previously absent from the literature.

"Could we not... use adaptive learning rate rules to prevent the slow down"

We believe the reviewer is implying that the results in Theorem 3.4 do not apply under, for example, the Adam optimizer. If so, we disagree. The gradient under Adam optimization still inversely depends on the embedding norm and, consequently, the embedding norms slow down the training process. To test this, we have trained SimCLR and SimSiam on the Cifar datasets using the Adam optimizer for 100 epochs and find consistent results:

We also note that the Adam optimizer is not common in cos.sim.-based SSL models.

"Could we not choose the learning rate proportional to the embedding norm?"

We agree. In fact, this is precisely what our GradScale does: it multiplies the gradient on each sample by its embedding's norm. We find this improves training but can lead to instability due to a positive feedback loop which is introduced: the embedding norm grows with the magnitude of the gradient and the gradient is now multiplied by the embedding norm. We include the embedding norms from the SimCLR w/ Adam Cifar10 training runs as an example of how the norms are affected by our mitigation strategies:

"How do regularization terms in the overall loss (such as weight decay) impact the conclusions drawn?"

We politely note that we extensively analyzed how weight-decay improves SSL convergence in Section 6 and the appendix, particularly Table 1 and Figures 5, S3.

Additionally, weight decay would not change the correctness of Theorem 3.4: each individual step still depends on the square of the embedding norm. Nonetheless, we see the reviewer's point that weight decay will interact with the bound in Theorem 3.4, possibly leading to faster convergence. Practically, however, it is difficult to find a regularization strength large enough to mitigate the slow-down in Theorem 3.4 without training diverging.

To evaluate this, we have extended the simulation from Section 4.1 by augmenting the objective with a regularization term on the embedding norm to simulate how weight-decay affects convergence. Thus, we now optimize $\mathcal{L} = -\hat{z}_i^\top \hat{z}_j + \gamma \cdot ||z_i||_2$. The results are listed in the tables below. We find that while weight-decay can speed up convergence, the number of required steps still depends quadratically on the initial embedding norm. However, if the weight decay is made too large, convergence never occurs. This is in line with our experiments in Section 6.

For $\gamma = 0.5$:

For $\gamma = 1$:

For $\gamma = 10$, the simulation diverges.

"Would it be possible to constrain embeddings to the sphere?"

This is something we considered: it would mean forcing the outputs onto $\mathcal{S}^{d-1}$ and may work well with the loss from Koishekenov et al. [1]. We chose to implement the GradScale layer as an alternative as it accomplishes a similar effect in principle.

References

[1]: Koishekenov, et al. "Geometric contrastive learning." VIPriors Workshop, CVPR 2023.

---

Please let us know if we have addressed your concerns.

Thank you for the kind words regarding our paper. We respond to your questions below.

"The weakness would be the proposed methods, especially the cut-initialization, are relatively brutal and less innovative"

Regarding the cut-initialization, we agree with this point and it is something we deliberated on for quite some time. However, we could not come up with an alternative method for ensuring that the embedding norms are small at initialization. An alternative option is to apply cut-initialization at only the final layer, but this performs roughly equivalently. We also note that cut-initialization seems to accomplish the desired task very well, as evidenced by Tables 1-3 in the paper.

If the reviewer has suggestions as to what we may do as an alternative to cut-initialization, we would be happy to try them out.

"Besides k-NN accuracy, have you considered other measures of latent space quality?"

We recognize that there are several probes that can be used. We focused on kNN as it is known to lower-bound the other probes [1] and is a good indicator of model performance [2]. To address the reviewer's question, we have run the linear probe on various models at the 500 epochs mark and see the same performance improvements:

We will include linear probe evaluations in the revision of the paper.

[1]: Oquab, Maxime, et al. "Dinov2: Learning robust visual features without supervision." TMLR 2023.

[2]: Marks, Markus, et al. "A closer look at benchmarking self-supervised pre-training with image classification." arXiv preprint 2024.

---

Please let us know if there is anything else which we can address.

Thank you for your thoughtful and in-depth review! Below, we address your concerns by adding experiments on Tiny-ImageNet and testing additional models (BYOL, MoCo v3, Dino), which confirm the generality of our theoretical and experimental results. Furthermore, we clarify the novelty of our analysis: while some of our methods and ideas have precedent, we are the first to consolidate them and show unambiguously how the embedding norms interact with self-supervised learning. We believe that knowledge of these interactions is valuable to the ICML community.

Detailed responses are below:

"[The paper] could benefit from experiments on larger dataset like Imagenet"

We agree that it would be nice to run on ImageNet. However, due to computational constraints, this is infeasible for us. We therefore refer to our response to Reviewer wGR1, where we include experiments on Tiny-Imagenet. The results are consistent with what we see on other datasets.

"I wonder if other methods also exhibit similar behaviour."

To address both this and the broader question of our results' generalizability, we have trained BYOL (which optimizes the cos. sim.), MoCo v3 (which optimizes InfoNCE) and Dino (which is not cos.sim.-based) on ImageNet-100 for 100 epochs with and without cut-init. We find that cut-init improves BYOL's accuracy at 100 epochs from 26% to 54% and MoCo v3's accuracy from 54% to 58%. These are in accordance with the results in our paper (in particular, Figure S4 from the paper). Our theory only applies to losses with normalized embeddings, thus not to Dino, VicReg or Barlow Twins. Indeed, we find that cut-init drops Dino's accuracy from 44% to 29%. This implies that trying to resolve the embedding norm effect in settings where it does not occur hurts performance. Nevertheless, many other prominent SSL methods do rely on normalized embeddings (e.g., W-MSE, NNCLR, ProtoNCE, SWAV, CLIP, DCL, and BEiTv2), highlighting how widely applicable our analysis is. We will discuss this in the revision.

"it is unfair to say there has been no explanation for network confidence via norm"

While we agree Kirchhof et al. (2022) have discussed the embedding norm as a concentration parameter for the vMF distribution, they only stated that it behaves this way and gave brief empirical evidence of this. In fact, Kirchhof et al. reference Scott et al. (2021) who explicitly raise why this holds as an open question: ``Why does the COSINE embedding convey a confidence signal in the norm?'' Our work therefore gives a clear explanation for this phenomenon for the first time. Furthermore, Scott et al. was in the supervised setting.

"...weight decay and cut-init are already known methods..."

While weight-decay is a known method, its standard discussion in the literature is entirely different from ours and focuses on overfitting. We only know of one reference, Wang et al. (2017), which suggests that weight decay may help with the embedding norm effect but this was not tested.

Similarly, cut-init only exists in transformer architectures and was proposed to manage the attention. We are not aware of it in any convolutional architectures or for managing embedding norms. Thus, the idea of dividing the weights at initialization has not been discussed with regards to the convergence rate and in this sense is a novel contribution of our paper.

Both of these are complementary to the regularization term for the variance from Zhang et al. (2020).

"It is not clear whether GradScale is a novel contribution of this paper or not."

GradScale is a novel contribution, but our goal with it was not to propose a new SOTA method. Instead, we use it as a way to analyze how embedding norms affect SSL training. That is, by using GradScale we can show that removing the gradient's relationship to the embedding norm can improve generalization, especially on imbalanced datasets.

"When describing GradScale, p=0 and p=1 are switched."

The $p$ parameter in GradScale is the power of the embedding norm that gets additionally multiplied to the original gradient. Thus, for $p=0$, the gradient is multiplied by $\|z\|^0 = 1$, leaving the gradient untouched.

"the authors argue that negative contribution is averaged across many samples and should be smaller than the positive contribution. I'm not sure this is true"

We agree with this point and will remove this from the paper. However, we note that this has no effect on the correctness of Prop 3.2. The main point is that all gradients of the InfoNCE loss are orthogonal to the points they act on, thus increasing the embedding points' norm, independent of the sizes of the attractive and repulsive contributions.

---

We hope we have addressed your concerns. Please let us know if there is anything we can do to convince you further.