Views Can Be Deceiving: Improved SSL Through Feature Space Augmentation

We thank the reviewer for their constructive feedback on our work. We appreciate the time taken to provide thoughtful comments to help strengthen our paper. We respond to each point below.

W1. The paper's structure and content are not easy to follow. Specifically, the use of symbols like alpha, beta, and gamma appears to be overloaded, making it challenging to distinguish between the connectivity and the P matrix.

Thank you for pointing out the notation overload. To address this issue, we keep the symbol $\alpha$ at section 3.2 to mean the expected spurious connectivity across dataset, and have revised the notation when introducing the toy setup to clarify the relation between $\alpha_{\text{toy}}$, and $\alpha$. We agree that overload of these symbols made parts of the theory section difficult to follow, and have clarified this in the revised draft. If you have any additional suggestions for paper structure, please let us know.

W2. It's unclear how the correlation in Table 1 is computed. It seems spurious attributes and core features are not defined for real data.

We originally discussed how the connectivity error rates in Table 1 are computed in Appendix E. We have added a short summary of this procedure to Section 4.2. To summarize, we utilize the spurious attributes given as metadata in each of the datasets, along with the label (i.e. the core attribute), to partition the dataset into different groups. We then compute the connectivity terms by training classifiers to distinguish augmented data from each combination of the two groups and reporting their error rates.

W3. The theoretical foundation leans heavily on Spectral Contrastive Learning, yet it is not used in the experimental section. It begs the question of its relevance or applicability. And whether the empirical results are well justified by theory.

HaoChen et al. (2021) [1] shows that spectral contrastive learning achieves similar empirical performance to other contrastive learning methods, but enables easier theoretical analysis. Theoretically, the spectral contrastive loss is similar to many popular contrastive losses. For instance, the contrastive loss in SimCLR can be re-written as an expression similar to the spectral contrastive loss. We have added a short clarification to Appendix C.

W4. The presentation would benefit significantly from the inclusion of an algorithmic representation of the proposed Late TVG method, enhancing its comprehensibility for readers.

Per the reviewer's suggestion, we have included the Late-TVG algorithm box including details in Appendix B to improve comprehension of the proposed method.

Q1. Which training parameters were utilized in the LateTVG method?

The training parameters used for Late-TVG are identical to the SSL-Base model. The Late-TVG encoder parameters are modified via pruning during training. We have clarified this in Appendix D.2.

Q2. How do the authors identify spurious features? The paper does not provide clear definitions distinguishing between spurious attributes and core features.

The spurious features are provided as metadata in each dataset we use, thus are known apriori (e.g. background in Waterbirds). We do not attempt to discover unknown spurious correlations. We have added text to clarify this assumption in Appendix D.1.

We believe these changes address the major points raised by the reviewer. Please let us know if any part of our response remains unclear, or if you have any other suggestions to further improve the paper. We greatly appreciate you taking the time to provide feedback on our work.

[1] HaoChen, Jeff Z., et al. "Provable guarantees for self-supervised deep learning with spectral contrastive loss." Advances in Neural Information Processing Systems 34 (2021): 5000-5011.

We thank the reviewer for the positive feedback and constructive suggestions.

W1. The intuition / reasoning behind the choice to prune the later layers to forget the spurious feature is not well-explained / discussed sufficiently.

We have clarified the paragraph on "Intuition for lateTVG" at the end of Section 5.1 in the updated version. To summarize,

1. Existing SSL augmentations cause undesirable invariances in the representation space, which we demonstrate in the result that the spurious connectivity is greater than the invariant connectivity. This means that images with the same spurious feature are more likely to be mapped closer to each other in the representation space.

2. One approach would be to modify the augmentations in the image space, but here, we take the approach of augmenting the representation space.

3. Prior works [3, 4] have shown that, for a supervised network, spurious features are encoded in the final layers, and that perturbing the final layers is effective in removing the spurious feature. Furthermore, pruning has been shown to be more effective on minority examples [5].

4. If this intuition also holds for an SSL model, then the pruned network (which encodes invariant features more than spurious ones) provides a view of the samples in the representation space that do not include the spurious feature. In this case, the main encoder is forced to use the core feature to map representations to these views.

Q1. The problem of spurious correlations in SSL is actually equivalent to feature suppression. It will be useful to discuss the equivalence of this problem as well as other relevant literature.

We had cited the suggested work in the related work section when discussing representation learning under dataset imbalances and shortcuts (Chen et al., 2021). We are happy to cite any other suggestions for additional relevant prior work in a future revision.

Q2. In supervised learning, the models trained on Waterbirds etc. are initialized from ImageNet pretrained weights. Is that the case here as well?

We use randomly initialized networks for our experiments, so spurious cues from ImageNet pretraining [e.g. 1] would not impact our evaluations. We explain this decision further in “The role of Initialization” paragraph at the end of Section 4.1.

Q3. Experiments on some larger scale datasets provided in packages such as WILDS or SpuCo can strengthen the empirical success of the method.

We thank the reviewer for the suggestion. We found that the only dataset in WILDS with explicit spurious correlation is the CivilComments, which is a natural language dataset, and is outside the scope of this work. The other WILDS datasets do not appear to contain strong spurious cues to the best of our knowledge. SpuCo is certainly a relevant dataset, and we have cited it in our work. We hope to conduct experiments on this dataset for a potential camera-ready, but did not prioritize this as it includes subsets that are similar to our current datasets, such as waterbirds and CMNIST.

Following the suggestion from Reviewer H1ZX, we have conducted experiments on Hard ImageNet [2], which have been added to Appendix F.1., and which we reproduce below:

We train an SSL model on the full Hard ImageNet train set and a linear classifier on the balanced subset (obtained by spurious rankings), then evaluate the classifier on the original test split and three additional splits with stronger color, shape, and texture spurious cues. This tests the model's ability to learn useful representations without relying on spurious correlations, which should result in lower accuracy on the manipulated splits, which we observe in the above table.

[1] Meehan, Casey, et al. "Do SSL Models Have D'eja Vu? A Case of Unintended Memorization in Self-supervised Learning." arXiv preprint arXiv:2304.13850 (2023) [2] Moayeri, Mazda, Sahil Singla, and Soheil Feizi. "Hard imagenet: Segmentations for objects with strong spurious cues." Advances in Neural Information Processing Systems 35 (2022): 10068-10077.

[3] Lee, Yoonho, et al. "Surgical fine-tuning improves adaptation to distribution shifts." arXiv preprint arXiv:2210.11466 (2022).

[4] Kirichenko, Polina, Pavel Izmailov, and Andrew Gordon Wilson. "Last layer re-training is sufficient for robustness to spurious correlations." arXiv preprint arXiv:2204.02937 (2022).

[5] Hooker, Sara, et al. "What do compressed deep neural networks forget?." arXiv preprint arXiv:1911.05248 (2019).

Thank you for the review and for recognizing the theoretical insights, empirical analysis, and proposed method as strengths of our work.

W1. The paper is missing references to some important works on spurious feature learning

Thanks for pointing out the missing references, we have cited these papers in the revised draft.

W2, Q1. Most of the results given in the paper are using smaller datasets. I believe the analysis of Section 4 could have been carried out a large number of publicly available SSL trained models. I would have liked to see some results using models trained on large datasets.

We appreciate the suggestions of the three datasets. We note that WILDS does not contain strong spurious correlations except for the natural language dataset CivilComments. The image datasets in WILDS exhibit natural distribution shifts, but do not have the type of spurious correlations studied in our work to the best of our knowledge. FOCUS is definitely an interesting dataset to build upon, but for initial analysis on a larger scale dataset, it is unclear which attribute we should have chosen as the spurious attribute. We can add results on this dataset to a future version of this paper.

Per your suggestion, we have added experiments on Hard ImageNet in Appendix F.1 to demonstrate our approach on a larger dataset. For the analysis in Section 4, we focused on datasets that contain the spurious attribute as a metadata, since our goal was to provide an in-depth empirical evaluation of spurious feature learning in SSL.

We train an SSL model on the full Hard ImageNet train set and a linear classifier on the balanced subset (obtained by spurious rankings), then evaluate the classifier on the original test split and three additional splits with stronger spurious cues. This tests the model's ability to learn useful representations without relying on spurious correlations, which should result in lower accuracy on the manipulated splits.

The summarized result table for Hard ImageNet is as follows:

As shown in this table, Late-TVG reduces accuracy on three Gray, Gray BBox, and Tile splits, meaning that the downstream classifier trained on Late-TVG representations relies less on the spurious (i.e. non-object) features. We have added more details and results on this experiment to Appendix F.1, and will include it in the revised main paper.

We believe the added experiments on Hard ImageNet help address your concern about evaluation on larger datasets, and we hope to add the other large-scale dataset in a following version of this paper.

Thank you for your positive feedback on our work and the suggestions for improvement. We appreciate that you recognize the originality, quality, and significance of our work.

W1, Q1. Scalability and Efficiency

The SSL-LateTVG model updates the same number of parameters as SSL-Base during training, with the forward pass keeping both the original and pruned encoder (See Algorithm 1 in Appendix B). The pruning operation is cost O(n) where n is the number of parameters. So any FLOPs used for the extra pruning mechanism will be very small compared to a single forward pass.

In summary, the training parameters are identical between SSL-Base and LateTVG, with the LateTVG encoder modified via pruning. For the backward pass, only the original encoder and other SSL parameters like the projection head are updated. We have also clarified this computational process in Appendix B, and parameters in Appendix D.2.

We provide runtimes on the MetaShift dataset below as an example:

SimSiam: SSL-Base: 4h 26m, SSL-LateTVG: 5h 14m

SimCLR: SSL-Base: 4h 39m SSL-LateTVG: 4h 50m

W2. Domain Generalization

We appreciate the reviewer raising this point. We focus on the image modality in this work because SSL methods we use as base implementations are designed for images. We agree that evaluating our approach on non-visual modalities is an important direction for future work.

W3. Comparison with State-of-the-Art

To the best of our knowledge, this is the first work evaluating visual self-supervised representation learning on datasets containing known spurious correlations. If the reviewer could suggest recent SSL papers addressing similar challenges, we would be happy to include comparisons.

Q2. Long-Term Effects of Pruning

While we agree that this is an interesting direction, our focus is on mitigating spurious correlations for self-supervised representation learning itself. We leave investigations into pruning's effects on adapting to new data or tasks to future work.

Q3. Robustness to Adversarial Attacks

We believe that evaluating adversarial robustness is orthogonal to the problem studied in this paper. Our method targets removing spurious correlations, not increasing robustness to adversarial examples. While an interesting direction, analyzing LateTVG's robustness to adversarial attacks is out of scope for this work. We leave this to future work and maintain our focus on addressing spurious correlations in self-supervised learning.