Memory-Efficient Self-Supervised Contrastive Learning with a Supervised Loss

Authors: Eric Gan, Baharan Mirzasoleiman
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: contrastive learning, self-supervised learning, representation learning, machine learning theory

Abstract:
Contrastive Learning (CL) is among the most popular methods for self-supervised representation learning. However, CL requires a large memory and sample size and careful hyperparameter tuning.
These factors make it difficult to
learn high-quality representations with limited amount of memory. In this work, we theoretically analyze a recently proposed \textit{supervised} approach, DIET, for self-supervised representation learning. DIET labels every example by its datum index and trains on the labeled data with a supervised loss. DIET does not require a large sample size 
or hyperparameter tuning. However, it falls short when using smaller encoders and is memory intensive due to its massive classifier head.
Given its remarkable simplicity, it is not obvious whether DIET can match the performance of CL methods, which explicitly model pairwise interactions between augmented examples. We prove that, perhaps surprisingly, for a linear encoder DIET with MSE loss is equivalent to spectral contrastive loss. Then, we prove that DIET is prone to learning less-noisy features and may not learn all features from the training data. We show feature normalization can provably address this shortcoming and use of a projection head can further boost the performance. Finally, we address the scalability issue of DIET by reducing its memory footprint.
The modified approach, namely S-DIET, substantially improves on the linear probe accuracy of DIET across a variety of datasets and models and 
outperforms other SSL methods,
all with limited memory and without extensive hyperparameter tuning. This makes S-DIET a promising alternative for simple, effective, and memory-efficient representation learning.

Decision: desk reject
Reviewer scores: 3, 5, 3, 5

Reviews and discussion:
<thread>
creation date:
1730194834025

thread_type:
official_review

<thread_object>
creator: reviewer_cxwx

<content>
summary:
This study provides a theoretical analysis of DIET, a recently proposed supervised approach for self-supervised learning. DIET, as a menthod of CL, labels each example by its datum index and employs a supervised loss for training. This work obtains several conclusions, including (i) for linear encoders, DIET with MSE loss is equivalent to spectral contrastive loss; (ii) DIET tends to learn features with less noise but may not capture all relevant aspects of the training data; (iii) feature normalization can help mitigate this issue, while incorporating a projection head can further enhance performance. This work further introduces SCALED-DIET (S-DIET) to improve the model's linear probe accuracy.

strengths:
- This paper explore the limitation of DIET, an important method of CL, and obtain several conclusions: (i) for linear encoders, DIET with MSE loss is equivalent to spectral contrastive loss; (ii) DIET tends to learn features with less noise but may not capture all relevant aspects of the training data; (iii) feature normalization can help mitigate this issue, while incorporating a projection head can further enhance performance. 

- This work further introduces SCALED-DIET (S-DIET) to improve the model's linear probe accuracy, i.e., use batch cross entropy and the multistep update formula for AdamW.

- Some experiments demonstrate the effectiveness of the proposed S-DIET.

weaknesses:
- The motivation behind this paper is unclear. According to P2 in the Introduction, DIET's advantage lies in its ability to mitigate CL's reliance on large datasets, while requiring a smaller parameter dimension to balance with sample size. The claim that smaller encoder dimensions are a key drawback for handling large data is not well justified. Furthermore, the authors state, "not clear whether DIET can capture the pairwise similarities between views...SSL," yet DIET, as a CL algorithm utilizing supervised loss, does not explicitly depend on pairwise similarities; this is merely an implementation choice rather than a fundamental mechanism. I fail to see how this motivation strongly connects DIET with contrastive loss, is pairwise similarity closely related to memory? While exploring CL from an efficiency perspective do be valuable, a thorough reading of the paper reveals a lack of such information. Perhaps I overlooked some details, and I hope the authors can clarify their insights.

- The key conclusions in this work need further explanation in relation to the core ides, i.e., MEMORY-EFFICIENT. For instance, the relationships between memory and features, encoder parameter dimensions, and projection heads are not adequately described, leading to fragmented conclusions.

- The choice to study DIET is justified by its independence from large training data, but there are other CL algorithms that also do not rely on large datasets, such as few-shot SSL. Additionally, DIET requires labeled information. Can the analyses in this work be applied to these other methods? If so, what differentiates this work? A broader exploration of algorithms and their mechanisms might enhance the reliability of this study.

- A code that can reflect the idea of ​​algorithm implementation is encouraged, since it is currently only described in 6 lines. At the same time, the introduction of these modules will increase the computational overhead, and related experiments are also necessary, after all, the focus is on memory.

- (Minor) The paper's template appears to differ from the one provided on the official website, such as in the line numbering. Please consider making further corrections.

**I would be happy to reconsider my score if these concerns can be addressed.**

questions:
Please see **Weaknesses**.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730388718446

thread_type:
official_review

<thread_object>
creator: reviewer_umd8

<content>
summary:
The manuscript studied the property of a supervised representation learning, called DIET, and then propose an improved version of this method. Specifically, the authors show the equivalence between DIET and spectral contrastive loss proposed by HaoChen \& Ma under a linear case. In addition, the improvement is motivated by the insight derived from the model introduced in the setting of Section 5.1. Although it looks interesting when all the strong assumptions are true, it is unclear if the results presented in the manuscript can provide guidance in the practical settings.

strengths:
The connection between DIET and spectral contrastive loss is interesting.

weaknesses:
The analysis is built on some unrealistic assumptions. It is unclear whether the results and development in the manuscript can hold in practice.  It could be more meaningful to make assumptions more carefully.

questions:
1) The review in the manuscript is clearly incomplete. Some important literature are missing. The author may want to have a more comprehensive literature review. 

2) The analysis in the Section 4 is not surprising due to the linear setting. Can the analysis be extended to the nonlinear case? The author may read the recent development under the nonlinear setting in Wang 2023. 

3)Assuming $W_H$ is an isometry seems a very strong assumption. Do the authors put a constraint for $W_H$ in the loss function to enforce such isometry? Otherwise, the authors may consider removing this assumption.

4)The form of training example introduced in Section 5 is unrealistic. When does this assumption hold (at least approximately) in piratical example? Why are there two features, one is low noise and the other is high noise?

5)Can we generalize the results in Theorem 5.1 to a more realistic setting?

6) The idea in memory-efficient DIET is straightforward.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730621728692

thread_type:
official_review

<thread_object>
creator: reviewer_nqvr

<content>
summary:
This paper presents S-DIET, a memory-efficient modification of DIET for self-supervised contrastive learning. It proves that DIET with a linear encoder and MSE loss is theoretically equivalent to spectral contrastive loss, and proposes feature normalization and projection head use to enhance performance. S-DIET significantly reduces DIET's memory requirements and achieves state-of-the-art performance without extensive hyperparameter tuning.

strengths:
1. The paper provides comprehensive and rigorous theoretical proofs.
2. Addressing the high memory demand of DIET is a well-motivated objective with strong practical significance.
3. The experimental results presented in Table 5 demonstrate promising improvements.
4. The paper conducts a detailed and insightful ablation study.

weaknesses:
1. The paper aims to address three issues: why DIET can perform comparably to CL, DIET's failure to learn all features, and its high memory demand. However, these issues are addressed in a fragmented manner without clear logical connections between them, making it difficult for readers to grasp the paper's central thesis.
2. The paper does not provide code or pseudocode, which hinders understanding of the proposed method and limits the ability to verify its effectiveness.
3. It is well-known that MSE is not typically used as a loss function for classification tasks. In the original DIET paper, each sample is treated as a separate class in a classification problem using cross-entropy loss. Why, then, is MSE employed as the loss function in Section 4? Does Theorem 4.5 hold if cross-entropy loss is used instead?
4. Due to the use of W1, it is unclear how Theorem 4.3 is related to the proposed method.
5. The theoretical analysis in the paper relies entirely on linear assumptions, while the proposed method (Equation 5) is based on empirical assumptions. These assumptions raise concerns about the rigor of this work.

questions:
Please refer to the weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730715363189

thread_type:
official_review

<thread_object>
creator: reviewer_t8fx

<content>
summary:
This paper studies DIET, which is a method that has rather few adoptions. This paper claims that DIET uses fewer memory than common contrastive learning approaches, and proved some theoretical results showing that DIET and spectral contrastive learning share the same solutions. Moreover, this paper proposes a new alternative S-DIET to further improve its performance.

strengths:
1. This paper has theoretical claims that connects DIET to CL.
2. This paper has some empirical evidence that the proposed S-DIET can match the performance on benchmarks like CIFAR and ImageNet-100.

weaknesses:
The method DIET studied in this paper has a fatal limitation, which is that the labels are essentially the sample index. As the dataset size increases, the classification head will need to increase linearly as well, which makes it impractical to use in large scale dataset training. Even though the proposed S-DIET does not require the classification head to be always loaded into the memory, it is still unnecessary to store such a large head, especially when one is training on millions of data.

questions:
None
</content>
</thread_object>

</thread>
<thread>
creation date:
1732146709873

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We would like to highlight the main motivation and contribution of our work, which is not fully captured in all the reviews. Contrastive Learning (CL) is one of the most popular and successful methods for representation learning. However, CL requires a large encoder and crucially relies on a large batch size to learn high-quality representation. Therefore, CL methods require a considerably large GPU memory. For example, the well-known SimCLR method trained their best model with CloudTPUs, using 128 cores and a batch size of 8192 to train encoders 4x larger than ResNet50 (Section 2.2 in [1])! In our work, we investigated an alternative “supervised” approach for representation learning (S-DIET), with considerably lower memory requirements. We theoretically proved the promise of S-DIET to learn high-quality representations, and empirically confirmed its effectiveness:
1. We proved a precise equivalence between solutions of S-DIET and spectralCL, which confirms that supervised approaches can learn high-quality representations. To our knowledge, this is the first fully rigorous connection between supervised and contrastive learning.
2. We showed that S-DIET can be implemented with much lower memory requirements, compared to CL.
3. To further boost the S-DIET performance, we proved in a simple theoretical example that normalizing embeddings enables S-DIET to learn noisier features in addition to the less-noisy ones. 
4. We empirically showed that S-DIET matches the performance of state-of-the-art CL methods while using substantially less memory than all the other methods.
We believe our contributions show the promise of simpler and more efficient approaches for representation learning.

[1] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations, 2020.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper studies DIET, a method for self-supervised representation learning by labeling each sample as a distinct class and training a classifier as in standard classification networks. The main contributions include

- Theoretical Insight: The authors prove that DIET with a linear encoder trained with MSE loss is equivalent to the spectral contrastive loss, connecting supervised losses to contrastive learning frameworks.
- Feature Learning Limitation: DIET tends to prioritize learning less noisy features but may miss others. The authors demonstrate that feature normalization mitigates this issue, and a projection head further enhances performance.
- Memory Efficiency: S-DIET addresses DIET's memory inefficiency by storing the classification head on disk instead of memory. This significantly reduces GPU memory consumption.
- Empirical Results: S-DIET achieves state-of-the-art performance on benchmarks like CIFAR-10 and ImageNet-100 with lower memory usage and no extensive hyperparameter tuning.

Reviewers generally recognize the importance of the problem, appreciate the theoretical insights and the practicality demonstrated through experiments. The main concerns are the following:

- Unrealistic Assumptions: The theoretical analysis relies on linear model setting and specific data models, questioning the applicability to more general, practical cases. 
- Fragmented Presentation: Reviewer nqvr criticized the paper for addressing multiple issues (theoretical equivalence, feature learning, and memory efficiency) in a "fragmented manner," which weakens the paper’s logical flow. Similar concern was raised by Reviewer cXwx who complained about fragmented conclusions and lack of motivation.

These concerns indicate that the paper requires improved writing to make the motivation and conclusion clearer to the audience of the conference and possibly more effort on the theoretical study. 

In particular, while simplifying assumptions are necessary, the particular assumption of a linear network represents a significant departure from the practical use cases of contrastive learning. 

With regard to the MSE loss, while the authors use Ref [1] (which provides a neural collapse analysis of MSE) to justify their choice, please be noted that neural collapse analysis has also been conducted for CE in earlier work. 

Given these limitations, I’d place this interesting and potentially impactful study in a marginally below borderline area. I strongly encourage the authors to improve their manuscript addressing the clarity and theoretical assumption issues above for future submissions.