Harnessing small projectors and multiple views for efficient vision pretraining

Thank you for your thorough review and valuable suggestions for improving our paper. We appreciate the time and effort you've invested in providing this feedback. We acknowledge the shortcomings in our initial presentation, particularly the absence of a comprehensive overview of existing self-supervised learning (SSL) methods such as SimCLR, BarlowTwins, and VICReg. We are committed to rectifying these issues to ensure our paper is informative and comprehensible to a broad audience. Below, we address each of the reviewer's specific concerns in detail:

1. Lack of Defined Key Terms: We apologize for the oversight in not defining some key terms within the text. To rectify this, we will incorporate a comprehensive preliminary section that outlines the foundational SSL methods referenced in our study. Additionally, we will append a glossary of formal terms, including $𝑘_{𝐷𝐴𝐵}$, in the Appendix and expand Section A to include a clear definition and intuitive explanation of Mercer features and Mercer’s Theorem.

2. Formal Restatement of Theorem 3.2: We appreciate the reviewer's attention to the formal restatement of Theorem 3.2. We acknowledge the confusion caused by not linking Theorem 3.2 to its formal counterpart in the Appendix. We will amend the main text to explicitly state that Theorem 3.2 is formally restated and proven as Theorem B.4 in the Appendix. Furthermore, we will explain the theorem's implications regarding the selection of redundant directions by SGD in linear networks.

Theorem B.4: (Formal) Let $\Gamma = V\Lambda V^T$ represent the eigendecomposition of $\Gamma$, and define $z$ as the projection of the weight vectors in $W$ onto singular vectors of $\Gamma$, V. Formally, $z = WV$. Assuming small initialization (as in Simon et al. (2023), i.e. $\| z_{pi}(0) \| << 1$ for all $p,i$, we can derive the following conclusions:

1. $sign(\frac{\Delta z_{pi}(t)}{z_{pi}(t)}) = sign(\lambda_i)$
2. For all $\lambda_i, \lambda_j > 0$, $\frac{z_{pi}(t)}{z_{pi}(0)} = (\frac{z_{pj}(t)}{z_{pj}(0)})^{\frac{\lambda_i}{\lambda_j}}$ where $\lambda_i$ denotes the $i^{th}$ singular value, i.e. $i^{th}$ element of diagonal matrix $\Lambda$.

As noted, the proof of the above theorem is presented in Appendix B (Pg 16-17). The insights derived from this theorem are as follows:

1. If $\lambda_i = 0$, $z_{pi}(t) = z_{pi}(0)$ and if $\lambda_i < 0$, then $lim_{t \to \infty} z_{pi}(t) = 0$.
2. The alignment between a weight vector, say $W_p$ and a singular vector of $\Gamma$, say $V_i$, i.e. $z_{pi}(t)$ depends on the corresponding singular values, i.e. $\lambda_i$ and $z_{pi}(0)$.

Theorem B.4 presents the implicit bias of gradient descent as an optimizer. We show that those eigenfunctions of $\Gamma$ that are aligned with the feature space of the network at initialization and correspond to high singular values are more likely to be learned. Therefore, under weak orthogonalization constraints, the learned representation space will over "represent" the strong singular vectors, thereby losing orthogonality. We will further discuss these insights and implications in Appendix B (Pg 18).

3. Typographical Errors: We thank the reviewer for identifying these errors and will diligently correct them to enhance the manuscript's clarity.

4. Impact of the Proposed Loss Objective: The reviewer is correct that the new loss formulation should agree with the less computationally efficient older formulations of the loss in the limit of long pretraining. The performance gap between the new loss formulation and original formulation, as identified by the reviewer, arises from the faster convergence of the efficient formulation in the limit of finite training budget (100 epochs in most of our experiments). To confirm this, we conducted extended experiments (400 epochs) to confirm that both formulations converge to similar levels of accuracy. This finding is detailed in the rebuttal document and Appendix E, Figure 18. We will update the main text to emphasize this equivalence.

We would like to thank the reviewer again for their thoughtful suggestions and excellent questions. We will carefully consider your points as we revise the presentation of our results and make the proposed changes. Please let us know if you have any other feedback or concerns.

Best regards,
The Authors

We thank the reviewer for their thorough and insightful comments. We are pleased that the reviewer found our analysis of the SSL loss formulation and the training dynamics exciting and relevant to the community. Below, we address the questions raised by the reviewer with enhanced clarity and additional insights:

1. Understanding Downstream Performance Claims: The reviewer's question is important and touches on the core of representation learning in SSL. Our framework focuses on accelerating convergence of the SSL loss itself, which inherently influences downstream performance. By optimizing the SSL loss more efficiently, we ensure that semantically similar images are mapped closely in the representation space, facilitating easier classification through methods like k-nearest neighbors or linear decoding. This is why reducing the SSL loss more rapidly leads to more rapid improvements in downstream accuracy. Our empirical results confirm this (e.g. see Figure 3).We will emphasize this critical link in the revised manuscript to ensure clarity and robustness in our claims.

2. Results with Barlow Twins and VICReg on ImageNet-100: We acknowledge the oversight in not including results from ImageNet-100. In the rebuttal, we have added plots for ResNet-18 pretrained with BarlowTwins on ImageNet-100. These results reinforce our recommendation in Section 4.1 that smaller projector dimensions can yield effective representations even in larger-scale pretraining scenarios. This addition aims to fully address the reviewer's concern regarding the scalability and applicability of our findings.

3. Connections to Maximum Manifold Capacity Representations (MMCR): We thank the reviewer for bringing the MMCR paper by Yerxa et al. to our attention, which is highly relevant. Yerxa et al. introduce a non-contrastive SSL objective that leverages multiple augmentations to estimate the manifold of object centroids. By maximizing the nuclear norm of the covariance matrix of these centroid representations, they effectively enhance the representation space's suitability for linear decoding of semantic categories. Although this approach is rooted in a statistical mechanical perspective, Schaeffer et al. [3] have convincingly linked it to the information-theoretic viewpoints prevalent in the SSL literature. Consequently, the MMCR loss formulation can be viewed as a functionally equivalent form to the loss formulations discussed in our work.
   While Yerxa et al. provide compelling empirical evidence, intuitive justifications, and posthoc analyses to underscore the benefits of multiple augmentations in improving semantic categorization, our work offers a complementary, theoretically grounded perspective. We elucidate how these augmentations not only contribute to a more efficient optimization process but also facilitate the early learning of superior features. This dual benefit—accelerated convergence and enhanced feature quality—underscores the efficacy of the MMCR framework in learning robust representations.

4. Contrasting results in Rankme and LiDAR: This is an excellent point to clarify, thank you for highlighting the connection. Our focus differs from RankMe and LiDAR, which analyze the dimensionality of the learned representation space in SSL algorithms. On the other hand, our study investigates the design space configurations underlying these algorithms, precisely the number of units in the projector head and the orthogonalization constraint hyperparameter in the loss function. We believe there’s no inherent contradiction between our findings and RankMe/LiDAR, and we complement the previous studies by providing a theoretical basis for how these design choices impact SSL representation learning.
   To elaborate on the relationship between the design configurations and the dimensionality of the features, Agrawal et al. [4], demonstrate that the same number of units in the projector can yield different representation space dimensionalities depending on the orthogonalization constraint. One can reduce the number of neurons in the projector head while maintaining high dimensionality in the representation space as long as one does not reduce it beyond a threshold. Therefore, our results do not contradict the findings of RankMe, LiDAR, and $\alpha$-ReQ. Instead, our work presents a complementary theoretically grounded explanation of how these design factors affect representation learning in SSL frameworks and under what conditions can good representations be learned using projectors with fewer units. But, in line with the reviewer’s question, RankMe, and LiDAR do indicate why one would not want to reduce the projector dimensionality down to extremely low values (e.g. 16 units). We will clarify this point in the updated manuscript.

We would like to thank the reviewer again for their thoughtful comments and excellent questions. Please let us know if you have any other feedback or concerns.

Best regards,
The Authors

[1] Garrido et al. On the duality between contrastive and non-contrastive self-supervised learning. ICLR, 2023.

[2] Zhai et al. Understanding augmentation-based self-supervised representation learning via rkhs approximation and regression. ICLR, 2024.

[3] Schaeffer et al. Towards an Improved Understanding and Utilization of Maximum Manifold Capacity Representations. arXiv, 2024.

[4] Agrawal et al. $\alpha$-req: Assessing representation quality in self-supervised learning by measuring eigenspectrum decay. NeurIPS, 2022.

We would like to thank the reviewer for their thorough and insightful comments. We are also glad the reviewer found our proposal theoretically motivated and the empirical validation satisfactory. Below, we address the questions raised by the reviewer:

1. Definitions of weak/strong orthogonality constraints: The reviewer is correct in pointing out that we did not provide explicit definitions in the original manuscript as to what constitutes ”weak” vs “strong” orthogonality constraints. In the extremes, the weakest possible orthogonality constraint is a $\beta = 0$, and the higher the value of $\beta$, the stronger it is. Thus, we generally use these terms to describe the relative values of the hyperparameter, $\beta$, i.e., we generally use these words to indicate differences between different models. But, we do not have a specific threshold that universally defines “weak” vs. “strong”. Based on reference [31], we generally consider $\beta = 1$ a “strong” constraint. We will clarify this usage in the updated text and make it clear that there is no explicit threshold.

2. Experiments on other architectures like Vision Transformers (ViTs) or other CNNs: We thank the reviewer for this suggestion. In this work, we train Resnet-50 and Resnet-18 (Imagenet experiments) following previous literature in non-contrastive SSL. To our knowledge, self-supervised ViTs are generally trained using either the MAE or self-distillation objectives (JEPA-style). While the insight regarding projector dimensionality and orthogonalization constraint is not directly transferable to such settings, the loss formulation involving multiple augmentations can improve convergence in self-distillation settings. We can, however, test on other CNN architectures if the reviewer feels this would strengthen the empirical results. We will do this over the discussion period, but we are confident we will see similar results for other architectures.

3. Core contributions of our work: We are glad that the reviewer found our work to be theoretically well motivated with sufficient empirical evaluation. The reviewer's point is well taken that it is not surprising that additional augmentations can help pretraining. However, we would like to highlight that while existing works and heuristics help with improving the generalization performance, they still require long pretraining, usually 800-1000 epochs for Imagenet. Instead, our work focusses on improving the convergence speed while maintaining performance. Moreover, we provide theoretical justification for how our recommendations impact the SSL pretraining, and such theoretical grounding helps to move machine learning towards more principled practice.

4. Low-dimensional projector experiments beyond CIFAR-10 and STL-10: We apologize for the oversight in presenting the results from Imagenet-100 in Table 1. In the rebuttal document, we have presented the results for Resnet-18 pretrained using BarlowTwins on Imagenet-100. We will also modify Table 1 to include all intermediate dimension sizes (see the rebuttal document for the updated Table 1 to be put in the revised manuscript). We hope these fleshed-out results will be more convincing for the reviewer.

We thank the reviewer again for their thoughtful comments and excellent questions. We will carefully consider your points as we revise and extend this work. Please let us know if you have any other feedback or concerns.

Best regards,

The Authors

Thank you for your thorough and insightful review of our work. We greatly appreciate your time and effort in assessing our paper and providing constructive feedback. We are delighted you found the paper well-written, clearly presented, and understandable to a broad audience. It's delightful to hear that you find our approach new, original, and well-supported, theoretically and empirically. Your acknowledgment of the range and importance of our suggestions is very encouraging.

1. Avoiding collinearity of the encoder features: You raise an intriguing point about avoiding feature collinearity after dropping the projector. In practice, enforcing an orthogonality among the projector output features implicitly avoids collinearity in the feature space of the encoder. An intuitive explanation of this phenomenon is as follows. Since the projector output is a piecewise nonlinear function of the encoder output (owing to the fully connected layers with ReLU in the projector network), a collinearity in the encoder output space would imply a collinearity in the projector output space. Therefore, by enforcing an orthogonality constraint in the projector output space, we are able to avoid a collapse or collinearity in the feature space. An in-depth analysis of the learning dynamics is, however, an open research problem with some recent progress being made in simplified settings [1].

2. Metrics of efficiency: The suggestion to include additional efficiency measures such as iterations per second and memory consumption is well-taken. We will incorporate these metrics in a revision to provide a more comprehensive picture of the efficiency improvements. Disentangling the contributions of each component (reformulation, orthogonality constraint, reduced projector) to the overall speedup is also a valuable direction for further analysis.

3. Other questions:

• Regarding downstream performance, our proposed formulation achieves comparable best-case results to the original formulations while being more efficient. We believe there is potential for further performance gains with additional tuning and optimization.
• The depth and width of the projector do impact the observed performance. We found that relatively shallow projectors (1-2 layers) with moderate width (e.g., 128-512 dimensions) worked well, but the optimal configuration likely depends on the specific dataset and backbone architecture. More systematic ablations on projector design would be informative.

Thank you again for your excellent suggestions and questions. We will carefully consider your points as we revise and extend this work. Please let us know if you have any other feedback or concerns.

Best regards,

The Authors

[1] Y. Xue, E. Gan, J. Ni, S. Joshi, and B. Mirzasoleiman, Investigating the Benefits of Projection Head for Representation Learning. ICLR 2024.