Preventing Dimensional Collapse in Self-Supervised Learning via Orthogonality Regularization

Summary

Driven by the dimensional collapse phenomenon observed in self-supervised learning, this research adopts a methodology designed to prevent the reduction in dimensionality of weight matrices. The empirical findings indicate a reduction in dimensional collapse in certain instances.

Strengths

The proposed method seamlessly integrates with existing SSL methods.

Weaknesses

The proposed method lacks novelty, essentially applying the loss formulation from Barlow Twins to the weight matrix. Moreover, the experimental results are insufficient to substantiate the efficacy of the method due to the following reasons: 1. Many of the SSL methodologies utilized to assess the OR are no longer considered state-of-the-art. Therefore, it is recommended to evaluate the OR using state-of-the-arts methods. 2. Table 1 reveals that the effectiveness of SO or SRIP appears to exhibit considerable variability, rendering the proposed methods unstable. 3. Tables 3 and 5 exclusively demonstrate the evaluation of OR for the BYOL method, omitting the efficacy assessments of other SSL methods. Additionally, in Table 4, the datasets are characterized as small-scale, which contradicts the claims made by the authors.

Questions

1. What are the advantages of regularizing weights over regularization of representation? 2. What is the performance of SRIP in Table 2 and 5?

Rating

5

Confidence

4

Soundness

2

Presentation

2

Contribution

1

Limitations

Please refer to the weakness.

Summary

In order to combat the problem of dimensionality collapse in self-supervised learning, the authors of this work introduce orthogonal regularization on encoders during pretraining. Taking inspiration from previous work on supervised orthogonality regularization, the authors study the impact of OR on eigenvalues of weights when considering self-supervised learning and find that OR learned weights/features contain more meaningful features and avoid dimensional collapse better. The recommended result seems to be effective for both contrastive and non-contrastive SSL methods.

Strengths

Overall I think the idea is simple and the paper is well written. Particularly the observations regarding the eigenvalues of SSL features and weights and the impact of OR are informative. In addition, I believe the OR improvements on transfer learning and downstream tasks are considerable and could become a staple part of SSL.

Weaknesses

1) I would like to point out the improvements of OR on SSL methods is not as major on the simple datasets such as Cifar100 set according to Table 1. However, the changes in larger datasets with BYOL are more noticeable but I would like to ask why results for other methods were not included for larger datasets. 2) Particularly with regards to Fig 3, to me it doesn’t look as if the feature eigenvalue changes much with the OR except for the representations. Perhaps a combination of other OR methods could help ? Or adding orthogonality regularization with the features. 3) With regards to classification accuracy on Table 5, higher batch size and longer training somewhat diminish the impact of OR improvement. Perhaps this result could be used to encourage OR when training with smaller batches or less epochs is required ?

Questions

1) For section 3.2 under what case would we have input > output ? 2) Could the authors provide some explanation for what entails whitening of the features ?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

None

Summary

This work proposes a method to reduce the dimensional collapse issue in self-supervised pretrained networks. Dimensional collapse occurs when a few large eigenvalues dominate the eigenspace. While previous research focused solely on the output representations of the pretrained networks, this work shows that dimensional collapse also occurs on the intermediate representations and weight matrices. To address this issue, the authors use orthogonal regularization (OR) objectives for the weight matrices in linear layers (dense or convolutional) within the network. The OR objective promotes orthogonality in the weight matrices by minimizing the distance between the Gram matrix of each weight matrix and the identity matrix. They demonstrate that using the OR regularization objective leads to a smoother decay of eigenvalues in the (intermediate) representations and weight matrices of a self-supervised pretrained ResNet18. The authors evaluate the OR regularization objective by pretraining various backbones (RN18, RN50, WideResNet28w2, and ViT-tiny/small/base) with multiple joint-embedding SSL methods (MoCov2, MoCov3, DINO, BYOL, NNBYOL) on CIFAR-100. They then train linear classification models on the frozen features, observing improved accuracy in nearly all cases with OR. Finally, they pretrain a RN50 backbone with BYOL on ImageNet-1k and evaluate it through linear probing on ImageNet and fine-tuning on Pascal object detection, showing improved results with OR.

Strengths

- The used method of reducing dimensional collapse of SSL methods is simple and can be easily integrated on various SSL methods and backbones. - The strongest aspect of this work is that it shows that OR consistently improves various joint-embedding SSL methods and backbones. However, the improvements (in linear accuracy) are small in the case of well-tuned baselines.

Weaknesses

- **(W1)** The introduction's claimed contribution, "We **systematically study** the phenomenon of dimensional collapse in SSL, including how whitening representations and network depth affect the dimensional collapse of weight matrices and hidden features," seems to me to be an exaggeration. In the main paper, there is only one experiment pretraining ResNet18 with BYOL w/ and w/o OR on CIFAR-10, and another experiment in the supplementary material using VICReg, which are not directly comparable. Given the emphasis on this contribution, a more appropriate approach would be to include an analysis of whitening's impact on dimensional collapse in the main paper (e.g., adding another column with whitening-based results in Figure 3). This analysis should be conducted under the same conditions as the experiments in Figure 3, for example, using BYOL with a projector and whitening. This setup would enable a straightforward comparison of vanilla BYOL, BYOL with OR, and BYOL with whitening in terms of their effects on the eigenvalues of intermediate representations and weight matrices. - **(W2)**: **Limited technical contribution**: as discussed in the related work sections of this work, orthogonality regularization of the weight matrices within neural network has been extensively studied before (related work subsection 2.3). In particular, the OR objectives studied here are directly taken from Bansal et al. 2018, and straight-forwardly applied to the examined self-supervised approaches, without any modification at all. So, the main novelty of this work is to apply in a plug-and-play way the OR regularization objectives of Bansal et al. 2018 on self-supervised methods. - **(W3)** About the CIFAR100 results: in the results provided in solo-learn (https://github.com/vturrisi/solo-learn), which the authors also use for their experiments, some baselines (i.e., SSL methods that do not use OR) have higher results than what is reported in this work and often higher even to the results reported in this work with the OR regularization objective. In particular, this is the case for MoCov2+ (69.89% in solo-learn vs 69.51% here), MoCov3 (68.83% in solo-learn vs 67.13% here), DINO (66.76% in solo-learn vs 59.13% in here). For BYOL however, is the other way around (70.46% in solo-learn vs 71.15% in here). What is the reason for the differences in the results? - **(W4)** The improvement that OR offers over a well-tuned BYOL baseline on ImageNet (200 pretraining epochs and 256 batch-size) is marginal: from 69.76 to 70.16. It seems that it offers significant performance improvements in mostly un-tuned baselines, like the results when pretraining BYOL for 100 and for 64 batch-size. In reality, even 200 epochs are not adequate for pretraining BYOL that, its performance reaches 74.3 for 1000 epochs (see BYOL paper). So, it is unclear whether OR could offer any improvement for such long pretraining settings that are typically used in SSL. - **(W5)** Not my primary concern, but although OR is shown to provide a small improvement on the linear accuracy of some SSL methods (e.g., going from 69.76 to 70.16 for BYOL pretrained with 200 epochs on ImageNet-1k), no evidence is shown that OR actually avoids dimension collapse in settings where otherwise (i.e., without OR) the SSL method will completely collapse and fail to learn any meaningful representation.

Questions

I would like the authors to address my concerns (W1), (W2), (W3) and (W4), as described in the weaknesses section. Some more specific questions: - About W1: can the authors provide a comparison of vanilla BYOL, BYOL with OR, and BYOL with whitening in terms of their effects on the eigenvalues of intermediate representations and weight matrices. - About W3: What is the reason for the differences in the results reported in this work and in in solo-learn (https://github.com/vturrisi/solo-learn)? - About W4: Can the authors demonstrate that OR improves even well-tuned baselines like BYOL with 1000 epochs (74.3 in BYOL paper) or MoCo-v2 with 800 epochs (71.1% in https://arxiv.org/pdf/2003.04297), or MoCo-v3 (74.6 for RN50 see: https://github.com/facebookresearch/moco-v3). Any of them would be good.

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

Yes, they are discussed.

Summary

This paper aims to improve self-supervised learning (SSL) by introducing orthogonality regularization (OR) during pretraining to prevent dimensional collapse. OR improses orthogonality constraint on weight matrices, not only preventing dimensional collapse of representation but also regularizing the weights from collapsing. The experimental results show that the benefit of the proposed method is consistent and widely applicable across datasets and architecture choices.

Strengths

* This paper effectively communicates both the problem it addresses and the solution it presents. The motivation and solutions are compelling and presented in an easily understandable manner. * The benefit of the proposed method appears to be consistent and significant.

Weaknesses

* More discussion on the difference between the proposed orthogonal regularization methods would have been helpful. On the surface, the two methods are almost equally good.

Questions

* In practice, a practitioner has to choose between SO and SRIP. What would be the authors' recommendation? * Do you expect that applying dropout on the encoder's hidden neurons would result in a similar effect, as dropout is also meant to prevent feature co-adaptation? It is quite likely that the existing encoders have already employed dropout, so I wonder if dropout is insufficient to prevent the collapse.

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

NA

Summary

Dimensional collapse, characterized by the dominance of a few large eigenvalues, is a phenomenon in self-supervised learning (SSL) that can lead to redundant features and weight matrices. To address this, the authors propose orthogonal regularization (OR) during pretraining to promote orthogonality in weight matrices, rather than using common approaches that regularize features via the objective (e.g., Barlow Twins, VicReg). The authors conduct experiments with five SSL methods from the SoloLearn framework to demonstrate the benefit of OR.

Strengths

The paper is well-written and provides a good introduction to the topic of representation collapse. The motivation for the approach is clearly stated, and the experimental setting is well documented. With small architectures like ResNet18, the effect of orthogonal regularization (OR) appears impressive at first glance. Although OR introduces some computational overhead, it is relatively small and may diminish further with larger batch sizes.

Weaknesses

The soundness of some experiments seems doubtful. Results for WideResNet28x2 are consistently below those for ResNet18. Additionally, the results using Dino are generally poor, leading to all ViT experiments being below any ResNet baseline. The performance of BYOL with OR, when trained for only 100 epochs with a batch size of 64, is even below the reported accuracy in SoloLearn. Considering the optimal training duration for BYOL is approximately 1000 epochs with a batch size of 4096, this comparison appears to be in a severely undertrained regime. Slightly longer training, 200 epochs with a batch size of 256, shows that the benefit of the OR approach diminishes rapidly. Minor issue: Line 44 states, "We evaluate the effect of adding OR in seven modern SSL methods during pretraining," but results are shown for only five methods. Additionally, since SSL approaches like Barlow Twins and VicReg aim to achieve orthogonality as well, there should be an experiment evaluating the effect of OR on these methods. Using one learning rate (LR) for all methods in linear probing might not be optimal. It is cheap to attach multiple linear probing heads and find the most suitable LR for each method to ensure fair comparisons.

Questions

Memory Overhead of OR: How much memory overhead does orthogonal regularization (OR) introduce? Explanation for Lower Gamma in Larger Models: The paper uses a lower gamma for larger models, but why? One might think that larger models would require more regularization. Results for Barlow Twins and VicReg: Why were results for Barlow Twins or VicReg not reported? Is the OR approach redundant when combined with these methods, which also aim to achieve orthogonality? Balanced Eigenvalues and Performance Discrepancy: The eigenvalues of VicReg without a predictor (Appendix Fig 4) appear more balanced than those of BYOL OR (Fig 3), yet VicReg's performance is far worse.

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

3

Limitations

The authors adequately addressed the limitations.

Reply for rebuttal

Thank you for your response. Your answers clarified my questions. Other reviewers seem to have concerns regarding empirical results, of which I don't have a strong opinion, since I do not have much hands-on experience with SSL. In general, I believe the paper is well-structured and easy to understand. Hope the concerns regarding the experiments are resolved well during the discussion.