Preventing Dimensional Collapse in Self-Supervised Learning via Orthogonality Regularization

We appreciate your comments and feedback. In addition to the general response, we address your itemized concerns here.

The proposed method lacks novelty, essentially applying the loss formulation from Barlow Twins to the weight matrix.

Methods like Barlow Twins try to avoid dimensional collapse of intermediate features as well as weights by whitening the projector output.
Our OR, on the other hand, acts directly on the weight matrix to indirectly avoid the dimensional collapse of features.
The implementation and motivation of these two methods are different, and no one has previously explored the dimensional collapse of the weight matrix in SSL, nor has anyone attempted to constrain the weight matrix in SSL to obtain an improvement in the results (except for weight decay).

In a broader sense, OR is just one of the methods we use to constrain the weight matrix in SSL, as it is intuitive to avoid the collapse of the weight matrix and has deterministic effects on the features. We only employed the widely used form of OR, but we wanted to share it with the community given its plug-and-play and interpretability in SSL methods.

We complement the experiments based on feature whitening methods such as Barlow Twins and VICRGE, OR can further improve their performance. We also visualize and analyze BYOL w/ vic, VICREG, and VICREG w/ OR, which are not able to avoid dimensional collapse of their weight matrices by feature whitening alone. Simply adding feature whitening to the original SSL method will even degrade its performance, unlike our plug-and-play OR.
Please check the performance table and see the pdf material of the general response.

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.

Thank you for your advice. The SSL methods we have now include 6 methods implemented by solo-learn (MOCOv2plus, MOCOv3, DINO, NNBYOL, BYOL, VICREG) and 10 methods implemented by the LIGHTLY framework (BarlowTwins, BYOL, DCL, DCLW, DINO, Moco, NNCLR, SimCLR, SimSiam, SwaV). Removing duplicates, there are in total 13 different SSL methods.
Please check the performance table in the general response.

We are currently experimenting as many SSL methods as possible with that have been built as benchmarks for fair comparison (LIGHTLY, solo-learn). If there are SOTA SSL methods that we still need to include, please let us know and we will try to add OR to it.

Table 1 reveals that the effectiveness of SO or SRIP appears to exhibit considerable variability, rendering the proposed methods unstable.

Our ORs, both SO and SRIP, work on all kinds of SSLs and all kinds of backbones. Only when MOCOv2 plus replaces the backbone with ResNet50, OR makes the Top-1 go down, but its Top-5 all goes up.
Please note that we did not adjust any hyperparameters when changing the backbone or adding OR.

Tables 3 and 5 exclusively demonstrate the evaluation of OR for the BYOL method, omitting the efficacy assessments of other SSL methods.

This is because training on large-scale data takes too much time (3 days for 100 epochs), and then BYOL presents the most stable and best performance on cifar100, so we only tested BYOL for now.
Due to our limited time and computational resources, the results of DINO and DINO (SO) on ImageNet with larger batchsize (2048) and epoch (300) are provided.
Also, we report BarlowTwins and BYOL results under Imagenet100 (ResNet 18, epoch 200, batch size 256).
Please check the performance table in the general response.

Additionally, in Table 4, the datasets are characterized as small-scale, which contradicts the claims made by the authors.

Table 4 refers to the results of SSL models pre-trained on large-scale datasets to perform linear probing on small datasets, which is consistent with numerous SSL studies. We will add more descriptions to the table later to avoid ambiguity.

What are the advantages of regularizing weights over regularization of representation?

We supplement the experiments based on feature whitening methods such as Barlow Twins and VICRGE, OR can further improve their performance. We also visualize and analyze BYOL w/ vic, VICREG, and VICREG w/ OR, which are not able to avoid dimensional collapse of their weight matrices by feature whitening alone. Simply adding feature whitening to the original SSL method will even degrade its performance, unlike our plug-and-play OR.
Please check the performance table and see the pdf material of the general response.

What is the performance of SRIP in Table 2 and 5?

The VITs experiment in Table 2 and the BYOLl 200 epoch experiment in Table 5 were too time-consuming, so we only tested the effect of SO. The results of SO are very similar to those of SRIP, with shorter time and less memory usage. In general, we advocate the use of SO for large models. However, we will still add the results of SRIP with the same settings in the later version.

We appreciate your comments and feedback. In addition to the general response, we address your itemized concerns here.

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.

This is because training on large-scale data takes too much time (3 days for 100 epochs), and then BYOL presents the most stable and best performance on cifar100, so we only tested BYOL for now.
Due to our limited time and computational resources, the results of DINO and DINO (SO) on ImageNet with larger batchsize (2048) and epoch (300) are provided.
Also, we report BarlowTwins and BYOL results under Imagenet100 (ResNet 18, epoch 200, batch size 256).
Please check the performance table in the general response.

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.

Interestingly, the effects of OR on intermediate features do not appear to be as large as that of representations.
We visualize intermediate features of VICREG (based on feature whitening), and VICREG_OR. Although OR can improve the performance of VICREG, we can find that the change in their intermediate features and even the representation is not obvious, but the weight matrix has a significant change.
Perhaps we need to train a deeper network to observe the effect of OR on the deeper features, we leave this part for future work.
Please check the performance table and see the pdf material of the general response.

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 ?

Thank you for your suggestions! Our current results on Imagenet show that SO is able to improve more with small batchsize and small epochs. To test the effects of OR with large batchsize and long epochs, the results of DINO and DINO (SO) on ImageNet with larger batchsize (2048) and epoch (300) are provided.
In the follow-up work, we will definitely make up more epoch experiments to fully test the effect of OR.
In addition, we added the results of various SSL methods on LIGHTLY framework cifar10, (resnet18), 200 epochs, and 400 epochs, the OR can get the same improvement.
Please check the performance table in the general response.

For section 3.2 under what case would we have input > output ?

For convolutional layers, input = convolution kernel area * input feature dimension, output = output feature dimension, e.g. for the first convolution in ResNet18, input = 7*7*3, output = 64, and input>output.
For linear layers, input is greater than output as long as the input feature dimension is greater than the output feature dimension.

Could the authors provide some explanation for what entails whitening of the features?

OR causes features to whiten because the orthogonalized weight matrices, which make the feature spectra more consistent across depth layers of features, prevent deeper features (e.g., representations) from collapsing and showing a whitened nature $1,2$ .
We can observe that the distribution of features in each layer is much closer after adding the OR constraints (Table 3 (d) brown and yellow lines in the paper).

$1$ Huang, Lei, et al. "Orthogonal weight normalization: Solution to optimization over multiple dependent stiefel manifolds in deep neural networks." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 32. No. 1. 2018.

$2$ Wang, Jiayun, et al. "Orthogonal convolutional neural networks." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2020.

We appreciate your comments and feedback. In addition to the general response, we address your itemized concerns here.

The soundness of some experiments

Thanks for pointing out this. Solo-learn only provides the official config file (yaml) for resnet18 on cifar100, on which we replaced the backbone (VITs, ResNet50,WideResnet) and added the OR. To ensure a fair comparison, we did not modify any other hyperparameters. Replacing a backbone may require tuning these parameters to achieve the best results. For now, we just want to show that OR can be plug-and-play under all kinds of backbone. More importantly, we believe our OR method could boost the performance under different hyperparameter settings. Please check the performance table in the general response.

Additionally, the results using Dino

In fact, VIT-small outperforms all the ResNet backbone results in DINO under the cifar100 experiment, please refer to Tables 1 and 2 in the paper. Results of DINO in the ResNet series are poor, and we believe this is due to the parameters in the config files of its repository. For example, in DINO, its default config uses the lars optimizer, warmup_teacher_temperature_epochs: 0, warmup_start_lr: 3e-05. However, the optimizer used for checkpoint results is SGD, lars=Ture, warmup_teacher_temperature_epochs: 50, warmup_start_lr: 3e-03. These differences in hyperparameters cause our results to be different from their reported results. Our baseline and OR methods run entirely under their officially provided hyperparameter config. To ensure fair comparisons, we did not modify the parameters in the original configs.

The performance of BYOL with OR

We checked the parameters in its checkpoint and BYOL trained on ImageNet with a batch size of 128 instead of 64 in its yaml config file. We think this is the reason for the different results. Our current results on ImageNet1k show that SO can perform better with small batchsize and small epochs. Although the benefit of OR on ImageNet1k is smaller after increasing the epoch on ImageNet1k, it still improves transfer learning substantially.

As for the results for larger epochs and larger batch sizes, the experiments of DINO and DINO (SO) on ImageNet with larger batchsize (2048) and epoch (300) are provided.In the follow-up work, we will conduct more experiments to fully test the effect of OR. Please check the performance table in the general response.

Minor issue:

Sorry for our mistake. The SSL methods we have now include 6 methods implemented by solo-learn (MOCOv2plus, MOCOv3, DINO, NNBYOL, BYOL, VICREG) and 10 methods implemented by the LIGHTLY framework (BarlowTwins, BYOL, DCL, DCLW, DINO, Moco, NNCLR, SimCLR, SimSiam, SwaV). Removing duplicates, there are in total 13 different SSL methods.

Additionally, since SSL

Your suggestions are very constructive. We tried to add OR to VICREG and Barlow Twins (resnet18, cifar10). Both gained a boost from OR, suggesting that simply requiring the output of the projector to be whitened is not enough to avoid weight matrix collapse. We additionally visualized the weights of the former, please see the pdf material of the general response.

Using one learning rate (LR)

Thanks for the suggestion, which is very reasonable. The representations produced by different methods may require different learning rates for the downstream tasks. We are only using the solo-learn settings for now. Considering the time constraints during the rebuttal period, we will try to adjust the learning rate later to get the best results.
For fair comparisons, we added the results of 10 SSL methods using KNN classification under the LIGHTLY framework to visualize the quality of the representations better.

Memory Overhead of OR: How much memory overhead does orthogonal regularization (OR) introduce?

For a specific backbone, OR traverses the parameters (linear and convolutional layers) in it.
Let's assume that the largest W is a square matrix with shape n*n, then the memory overhead needed is to compute W.T @ W which is also an n*n matrix, and to introduce an identical matrix I, which is also an n*n matrix, and storing the distance between the two requires an n*n matrix. So in total, the memory overhead of the forward propagation is
4 *n*n*(Floating Point Precision)/4. Using BarlowTwins in LIGHTLY with resnet18 as an example, the maximum memory usage without OR is 4.87GB, with SO it is 4.89GB. Therefore, the memory overhead of OR is quite low.

Explanation for Lower Gamma

You are asking a very good question to promote understanding of OR loss.
Because our OR is a penalized summation of all parameters, the larger the number of backbone parameters, the OR loss will increase rapidly, in order to be balanced with the SSL loss, we have only appropriately scaled down the Gmma by a multiplicity of 0.1. We checked the SO loss of the randomly initialized ResNet18 and ResNet50 backbones (gamma= 1), which is 5873.4062 for the former and 217609.7031 for the latter, and the increase in the number of model parameters brings about a nearly 40-fold increase in OR loss.

Balanced Eigenvalues and Performance Discrepancy

Your notice is well worth further discussion. The whiteness (balanced eigenvalues) of representations and their performance do not grow in tandem, see Table 4 in $1$ .
We visualized the original VICREG, and we can see that the original VICREG is much better than the version with the projector removed, but the whiteness degree of its representations is not as high as the one with the projector removed. Please see the pdf material of the general response.

$1$ He, Bobby, and Mete Ozay. "Exploring the gap between collapsed & whitened features in self-supervised learning." International Conference on Machine Learning. PMLR, 2022.

Thanks to the authors for the thorough response and the additional experiments.

(A) Regarding experiments. It is not a fair comparison if one uses suboptimal hyperparameters for a baseline and then introduces a new hyperparameter that can be tuned. It does not provide a lot of insights of the usefulness of the method to a well tuned baseline. It only shows that the proposed regularization has some kind of effect on the training. The experiments with WideResnet and Dino in its current form are simply a distraction. Either the dataset is not large enough to apply a WideResNet or the hyperparameters do not work at all. For Dino, it is true that Dino ViT ourperforms Dino ResNet. But all other methods outperform the best Dino ViT result even with ResNet18. Therefore, the used hyperparameters or training schedule of the Dino setting do not work at all.

(B) Memory overhead. Thank you, the fact that memory overhead is feasible and does not scale with the used batchsize is great.

(C) Usefulness of the proposed method. One the one hand, the introduction of a new hyperparameter and the increased effort to tune it should lead to a significant increase in a well tuned baseline that is trained for sufficiently long time. The current work does not convince me that this is the case. On the other hand, the improvements with short training schedules and the improved transfer capabilities are noteworthy. To me this is very similar to the strengths and weaknesses of Barlow Twins (sensitive to the number of features in the projector space). The inductive bias of the independence regularization speeds up learning and creates a representation that is robust to shifts, but can also be detrimental to fit a function in domain, given enough samples and training time. A study how to the gamma scales with model size and training schedule that might reduce the necessary tuning of the introduced hyperparameter would be of interest. Futhermore a direct comparison to BarlowTwins/VicReg to specifically measure the benefits of improved efficiency regarding training time or training data as well as robustness to transfer learning could improve the work in my opinion.

I will keep my score, but I am not opposed to the acceptance of the work given the predominantly positive scores by other reviewers. EDIT: In expectation that the authors incorporate the feedback into their final version, I will raise my score.

My point was that consistent results do not directly render an experiment useful, as there is too much room for improvement when the tuning is not done correctly.

I want to clarify that I did not want to critique the new experiments, which are small scale, but look solid and cover a lot of methods.

We appreciate your comments and feedback. In addition to the general response, we address your itemized concerns here.

(W1)

Thank you for your very constructive suggestions! Following your suggestions, we added more experiments of BYOL with vic loss (feature whitening their predictor's output). Simply adding feature whitening to the original SSL method will even degrade its performance, in contrast, our plug-and-play OR method can improve the performance. We also visualize weights and features of BYOL w/ vic, VICREG, and VICREG w/ OR (please see the pdf material), which are not able to avoid dimensional collapse of their weight matrices by feature whitening alone. Please check the performance table and see the pdf material of the general response.

(W2)

Thanks for the comment, in which we respectfully disagree. Although OR has been explored under supervised and semi-supervised learning, this paper is the first time to explore its effect on dimensional collapse in SSL and attempt to turn it into a plug-and-play method for enhancing SSL. It is also the first time that OR has been applied to a transformer architecture.
We only employed the widely used form of OR, but we wanted to share it with the community given its plug-and-play and interpretability in SSL methods.

To further highlight our contribution in a more systematic study of dimensional collapse in SSL, we have added relevant experiments and visualizations based on your suggestions:
BYOL, BYOL with OR, and BYOL with whitening in terms of their effects on the eigenvalues of intermediate representations and weight matrices.

The contribution of this paper is rephrased and summarized as follows: From our observations, although existing methods such as VICREG deal with the collapse of features via whitening features, they can not prevent the dimensional collapse of the weight matrices. We are not only the first work to systematically study dimensional collapse at all levels in SSL but also provide new solutions to address the dimensional collapse by constraining the weight matrix.

OR is just one of the methods we use to constrain the weight matrix in SSL, as it is intuitive to avoid the collapse of the weight matrix and has deterministic effects on the features.
We only employed the widely used form of OR, but we wanted to share it with the community given its plug-and-play and interpretability in SSL methods.

(W3)

Thanks for pointing this out. We followed the latest Solo-learn config file exactly and ran the results directly. We carefully compared the config in its GitHub repository with the config provided by checkpoints and found that there are differences in its hyperparameters, which is most likely the reason why our reproduced results are not the same as its report results.

For example, in DINO, its default config uses the lars optimizer, warmup_teacher_temperature_epochs: 0, warmup_start_lr: 3e-05. However, the optimizer used for reporting results is SGD, lars=Ture, warmup_teacher_temperature_epochs: 50, warmup_start_lr: 3e-03. These differences in hyperparameters result in our results being different from the results of their checkpoint reporting.

Our baseline and OR methods run entirely under their officially provided hyperparameter configs. It is difficult to make direct use of the parameters in its checkpoints. solo-learn has been updated and its hyperparameter form has been changed.

Fortunately, our OR methods are plug-and-play enough to compare the impact of SO with 10 SSL methods in another open-source SSL framework Lightly, which includes BarlowTwins , BYOL ,DCL, DCLW , DINO, Moco, NNCLR, SimCLR, SimSiam, SwaV. In our environment, its original baseline results are higher, and OR can further enhance these SSL methods.
Please check the performance table and see the pdf material of the general response.

(W4)

Thank you for your suggestions! Our current results on Imagenet show that SO is able to improve more with small batchsize and small epochs. Although the benefit of OR on ImageNet1k is smaller after increasing the epoch on ImageNet1k, it still improves transfer learning substantially. As for the results for larger epochs and larger batchsize, due to our limited time and computational resources, the results of DINO and DINO (SO) on ImageNet with larger batchsize (2048) and epoch (300) are provided. In the follow-up work, we will definitely make up more epoch experiments to fully test the effect of OR.
In addition, we added the results of various SSL methods on LIGHTLY framework cifar10, resnet18, 200 epochs, and 400 epochs, the OR can get the same improvement.
Also, we report BarlowTwins and BYOL results under Imagenet100 (ResNet 18, epoch 200, batch size 256).
Please check the performance table in the general response.

(W5)

By visualizing weight matrices and features, we can see that OR avoids the dimensional collapse of weights as well as features (Figure 3 in the paper), and it even likewise improves on feature whitening methods such as VICREG and BarlowTwins. Please check the performance table in the general response.

Without OR, the current SSL methods will not completely collapse as they introduced techniques such as using negative samples, self-distillation, clustering, and feature whitening. In this paper, we visualize the decaying eigenspace of weight matrices and features to observe the dimensional collapse, and we demonstrate the OR can effectively mitigate, if not prevent, the dimensional collapse.

Thank you for your kind response and suggestions.

(A)

We are very happy that we have solved this part of your concerns and it has greatly improved the quality of our work.

(B)

Indeed, as suggested by Reviewer KCGP, we intend in the future to make changes to the form of the OR to avoid the adjustment of the OR hyperparameters by the increase in model size.

(C)

Thanks again for your recognition!

(D)

DINO on Solo-learn results are really puzzling. For this reason, we performed DINO ResNet50 300 epoch experiments on the LIGHTLY framework. However, the results are still far from those in the original DINO article, and we found that this is due to the fact that LIGHTLY does not perform multi-crop, which has a fatal impact on the performance. Without multi-crop, as shown in the following table (results from the paper of DINO), we believe that LIGHTLY's results are consistent with the original DINO article.

In the future, to be more convincing, we need to add multi-crop to LIGHTLY or conduct experiments in the DINO repository.

We appreciate your comments and feedback. In addition to the general response, we address your itemized concerns here.

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.

Thanks for the advice, it helps a lot for those who intend to use OR!
Yes, both are almost the same in terms of effectiveness. SRIP is more time-consuming and takes up more memory than SO.
So we chose SO for very time-consuming experiments (ImageNet and ImageNet100). Additionally, SRIP has theoretically weaker (softer) regularization constraints than SO because it uses the spectral norm instead of the L2 norm.

Based on our experience with large-scale datasets and OR's results for the feature whitening method VICREG, we recommend SO. We add 10 SSL methods in the LIGHTLY framework, all of which use SO, and all of which obtain improved results. Please check the performance table in the general response.

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.

Thanks for raising this question. Indeed, dropout is also able to prevent feature co-adaptation. In SSL, there is no dropout in ResNet, but dropout is enabled in VITs. Our OR approach significantly improves the DINO with VITs (Table 2 in the paper), suggesting that single dropout cannot adequately prevent dimensional collapse at the three levels of the weight matrix, intermediate features, and representations.
We can visualize Dropout's impact on all three in future work. Exploring the role of Dropout and OR in SSL is very interesting for further research.

In fact, because of the widespread use of Dropout in the Transformer architecture, there has been some exploration of the impact of dropout on the features of SSL using the Transformer architecture. Dropout encourages the use of global features rather than local features in the Transformer architectures $1$ .

$1$ Luo, Jian, et al. "Dropout regularization for self-supervised learning of transformer encoder speech representation." arXiv preprint arXiv:2107.04227 (2021).