Paper: Augmentation-aware Self-Supervised Learning with Conditioned Projector (ICLR-2024)

Reviewer comment (Novelty):
[Garrido et al.] would be one of the most recent work among prior works but missed in this paper.

Author response:
**ad. W1: Comparison to the work of Garrido et. al.**

Thank you for this question. In their recent work, Garrido et. al propose to extend VicReg with a hypernetwork-based predictor to learn representations that are equivariant to transformations. Similarly to one of the variants of CASSLE, their hypernetwork predicts the parameters of the predictor based on the transformation descriptors. In our work, we show that there exist methods of conditioning superior to hypernetworks. We also evaluate our approach on real-life image data, as well as a wider range of SSL methods.

Reviewer comment (Theory):
The proposed method simply provides the additional information about augmentations together with the encoder outputs, and it is not clear how it helps to "preserve more information about augmentations" in representations. Figure 3 shows that injecting information of random augmentations results in reduced cosine similarities. This implies that the projector relies on the given information about augmentations, which is not directly related to the learned representations (the output of the encoder), i.e., learned representations do not have to be changed regardless of whether the projector relies on the additional information about augmentations or not. Any theoretical justification on the effect of the proposed method to the learned representations would be welcome.

Author response:
**ad. W2: request for additional theoretical justification:**

We have expanded the discussion of Figure 3 in Section 3.2, showing that the conditional probability of feature extractor representations on the condition of correct augmentation information is larger than on the condition of any random augmentation information. This implies that the learned representations of images are correlated, to an extent, with the parameters of augmentations applied to them.

---

1. We believe our paper sufficiently proves experimentally that CASSLE leads to preserving information about perturbed features.
 - Figure 3 indicates that falsified augmentation information has a negative effect on CASSLE's projector performance. If the augmented features were not preserved in feature extractor representation, the projector would be indifferent to augmentation information used for conditioning. 
 - Figure 4 shows that CASSLE's representaion is more subject to augmentation-induced noise, compared to other models.
 - Table 6 shows that, in case of SimCLR and Barlow Twins, CASSLE model performs best in the task of rotation prediction - prediction of augmentation which was not used during pretraining.

Reviewer comment (Evaluation):
The performance gain is overall minor and often it underperforms previous methods.

Author response:
**ad. W3: The performance gain is overall minor and often it underperforms previous methods**

While it is true that CASSLE does not always achieve the best downstream performance, it does so in the majority of cases. We have compiled the linear evaluation results of different approaches on different downstream tasks (Tables 1 and 7) and ranked each approach from best to worst-performing in each downstream task (See Figure 5 in the revised manuscript). 

We find that:

* CASSLE performs best in 54 out of 91 downstream tasks, i.e. 59.34% of cases where it was evaluated.
* AugSelf performs best in 31 out of 91 downstream tasks, i.e. 34.07% of cases where it was evaluated.
* AI performs best in 5 out of 16 downstream tasks, i.e. 31.25% of cases where it was evaluated.
* IFM performs best in 1 out of 13 downstream tasks, i.e. 7.69% of cases where it was evaluated.
* Vanilla performs best in 0 out of 91 downstream tasks, i.e. 0.00% of cases where it was evaluated.

(we excluded LooC from the analysis, see the below answer to W4).

---

2. In our previous answer, we focused on answering the concern that "The performance gain is overall minor and often it underperforms previous methods.". We provided a measure that shows that CASSLE outperforms numerous other methods in the majority of cases. As for the discussion of BYOL and SimSiam, we suspect that this may be caused by the asymmetric representation used by those methods, which distill the representations between the projector and additional predictor network. Observe that CASSLE and AugSelf also perform comparably on MoCo-v3, which also utilizes an additional predictor. It is worth noticing that CASSLE improves the performance of the vanilla versions of BYOL and SimSiam as well. We have added this remark to the revised version of the paper.

Reviewer comment (Presentation):
Why does MoCo-v2 in Table 1 contain only one performance of LooC? It looks quite not informative.

Author response:
**ad. W4: Why does MoCo-v2 in Table 1 contain only one performance of LooC? It looks quite not informative.**

Out of the datasets we are transferring to in Table 1, the authors of LooC have conducted experiments only with transferring to the CUB dataset, which we re-reported in our work. To the best of our knowledge, neither the codebase nor network checkpoints trained by LooC are publicly available. As such, we were not able to benchmark this method on more data.
It is also worth noting that both we and the authors of AugSelf observed lower performance of baseline MoCo-v2 transferred to CUB, compared to the authors of LooC.
We also note that other works on augmentation-aware SSL [1,2] do not reimplement LooC and also re-reported their results.

Reviewer comment (Experiments):
Why does MoCo-v3 in Table 1 miss the performance of "AI by [Chavhan et al.]," while the original paper presents its performance?

Author response:
**ad. W5: Why does MoCo-v3 in Table 1 miss the performance of "AI by [Chavhan et al.]," while the original paper presents its performance?**

This is because we trained a different backbone (ViT-Small, which is 4x smaller than ViT-base), compared to the one used by Chavhan et. al (ViT-Base).

Reviewer comment (Presentation):
The reference section requires thorough proofreading, as there are many incomplete/inaccurate references. For example, the closest prior work by [Lee et al.] is published in NeurIPS'21, but its arXiv version is cited. Also, many references miss the name of the published venue.

Author response:
**ad. W6: updating the bibliography**

Thank you for pointing this out. We have updated the bibliography in the revised version of the paper in order to refer to journal /conference versions of the cited papers.

Reviewer comment (Experiments):
I think your experiments can only be used to "indirectly" justify if the encoder preserves the information about augmentations. As I suggested, any theoretical analysis or a more direct experiment that checks if the proposed method better retains the augmentation information compared to the baseline would be helpful.

Author response:
1. We believe our paper sufficiently proves experimentally that CASSLE leads to preserving information about perturbed features.
 - Figure 3 indicates that falsified augmentation information has a negative effect on CASSLE's projector performance. If the augmented features were not preserved in feature extractor representation, the projector would be indifferent to augmentation information used for conditioning. 
 - Figure 4 shows that CASSLE's representaion is more subject to augmentation-induced noise, compared to other models.
 - Table 6 shows that, in case of SimCLR and Barlow Twins, CASSLE model performs best in the task of rotation prediction - prediction of augmentation which was not used during pretraining.

Reviewer comment (Experiments):
Only ResNet-50 is used throughout experiments, so my concern on the scalability still remains. Note that ViT is just one option to resolve this concern; you can use different number of layers or other types of CNN architectures. Generally speaking, ViT becomes prevalent in the last ~ 2 years, so experiments with this type of architecture would strengthen your contribution.

Author response:
3. Our work focuses primarily on convolutional architectures in augmentation-aware SSL. There are two main reasons for that: 1) convolutional architectures are the most popular ones in prior works on augmentation-aware SSL, such as AugSelf, AI and LooC; 2) following the main research line of augmentation-aware SSL we can fairly compare our methods and ideas against most of the existing methods using convolutional architectures. 
Please note that prior to our work, the only augmentation-aware method to have been benchmarked on ViTs was AI, where the ViT variant of the method worked because the authors trained it on top of pre-trained ViT. Moreover, in the case of ViT, their method was often outperformed by fine-tuning and not compared to AugSelf, leaving the question of scaling largely unanswered. On the other hand, we compare AugSelf and CASSLE with MoCo-v3, where all three models are trained from scratch. We believe that CASSLE should scale to larger models. 
While the question of whether augmentation-aware models scale to larger architectures is a valid one, it requires a more rigid analysis of all proposed techniques. For instance, larger models may benefit from different hyperparameter values of augmentation-aware methods than those established for ResNets. Such a detailed analysis is out of the scope of our paper.

Reviewer comment (Writing):
The reference section is still not proofread after being pointed out twice; at this point, I am not sure if authors are willing to show proper respect for previous works. MoCo v2 is an arXiv preprint and MoCo v3 is published in ICCV'21, but both miss their venues.

Author response:
4. In our manuscript we consistently use only the `\citep` command. We recognize that `\citet` would be preferable in a few places of the manuscript. We do not quite understand what you mean by "too verbose" references.  Thank you for pointing out those additional papers with missing metadata. We will take that into account in the revision of our paper.

Reviewer comment (Evaluation):
Based on Table 7, the proposed method seems to be less effective for SimSiam and BYOL compared to InfoNCE-based methods. The manuscript currently claims that CASSLE is applicable to all joint-embedding architectures, but the current experimental results do not demonstrate this.

Author response:
**ad. W1: applicability to SimSiam and BYOL.**

When talking about applicability, the meaning is that the design of CASSLE is compatible with a wide range of joint-embedding approaches and thus, can be applied to them. We do not claim that it is guaranteed to outperform approaches such as AugSelf. Nevertheless, while AugSelf outperforms CASSLE on SimSiam and BYOL, CASSLE still offers a performance boost compared to the vanilla versions of those methods, confirming its applicability in the sense of performance.

Reviewer comment (Experiments):
The work does not address the large body of work surrounding “feature suppression”, an important issue of contrastive models becoming invariant to features important for downstream tasks. The work can be strengthened by including comparisons to methods proposed to address feature suppression [2], as well as evaluation on some feature suppression benchmarks [1].

Author response:
**ad. Q1/ W3: How does CASSLE relate to feature suppression [1] and shortcut solutions [2] in contrastive learning?**

Thank you for pointing this out. We believe that invariance to augmentation can be regarded as a consequence of feature suppression, and have added an according reference to the above works in the Related Work section of the revised version of our manuscript. It is also interesting to verify how transferable are methods designed for mitigating feature suppression. We also added a comparison to MoCo-v2 + implicit feature modification (IFM) [2] to find how this established method of mitigating feature suppression transfers to downstream tasks.

Reviewer comment (Experiments):
Current experiments do not demonstrate the effectiveness of CASSLE with augmentation-free approaches to self-supervised learning. This limits the modalities in which it can be applied to those where augmentations can be selected a priori.

Author response:
**ad. W4: comparison with augmentation-free approaches**

Our research problem focuses on how augmentation-based SSL methods become invariant to augmentations. We are not aware of any evidence in the literature that MIM-based methods also suffer from this issue. As such, a comparison to them is not relevant in the context of CASSLE.

Reviewer comment (Evaluation):
For Table 1, and other similar tables, the authors could add a column denoting mean improvement, taken over datasets, over the vanilla baseline to more easily compare each of the methods to CASSLE. It does not have to specifically be an additional column, but it would be nice to have an aggregate metric of performance in comparison to the baseline.

Author response:
**ad. W5: aggregate metric of improvement over baseline**

Thank you for this valuable suggestion. We have compiled the linear evaluation results of different approaches on different downstream tasks (Tables 1 and 7) and summarized them in Figure 5 (appendix C) in the revised manuscript. CASSLE is usually ranked the best in terms of performance on downstream tasks compared to AugSe and Vanilla approaches. AugSelf and CASSLE improve over Vanilla approaches by a comparable margin. Finally, CASSLE achieves the best performance in the largest number of downstream tasks.

Reviewer comment (Novelty):
How does CASSLE relate to feature suppression [1] and shortcut solutions [2] in contrastive learning?

Author response:
**ad. Q1/ W3: How does CASSLE relate to feature suppression [1] and shortcut solutions [2] in contrastive learning?**

Thank you for pointing this out. We believe that invariance to augmentation can be regarded as a consequence of feature suppression, and have added an according reference to the above works in the Related Work section of the revised version of our manuscript. It is also interesting to verify how transferable are methods designed for mitigating feature suppression. We also added a comparison to MoCo-v2 + implicit feature modification (IFM) [2] to find how this established method of mitigating feature suppression transfers to downstream tasks.

Reviewer comment (Experiments):
Table 4 indicates that many of the methods were trained with a batch size of 256. Can the authors clarify why this was set so low? In the SimCLR paper it is shown that contrastive methods perform much worse when trained with a smaller batch size. Does CASSLE scale to larger batch sizes? Does CASSLE still perform well with a large batch size?

Author response:
**ad. Q2: Question about batch size of 256** 

We have set this batch size for all methods in compliance with the hyperparameters of AugSelf [4]. We note that the batch size is mainly important for SimCLR, as approaches such as MoCo decouple the batch size from the number of negatives, and Barlow Twins, SimSiam and BYOL work well with batch sizes of 256. CASSLE is unrelated to the batch sizes and we expect it to perform equally well with larger batch sizes.

--------
2. In order to fully answer your question about the batch size in SimCLR, we ran experiments with this batch size for vanilla SimCLR, AugSelf, and CASSLE. We present the linear evaluation in the table below.

| SimCLR   | C10       | C100      | Food      | MIT       | Pets      | Flowers   | Caltech   | Cars      | FGVCA     | DTD       | SUN       | CUB       | 300W      |
|----------|-----------|-----------|-----------|-----------|-----------|-----------|-----------|-----------|-----------|-----------|-----------|-----------|-----------|
| Vanilla  |     84.41 |     61.77 |     57.48 |     63.10 |     71.60 |     83.37 | **79.67** |     35.14 |     40.03 |     65.59 |     46.92 |     30.98 |     88.59 |
| AI       |     83.90 |     63.10 | --        | --        |     69.50 |     68.30 |     74.20 | --        | --        |     53.70 | --        | **38.60** |     88.00 |
| AugSelf  |     84.45 |     62.67 |     59.96 |     63.21 |     70.61 | **85.77** |     77.78 |     37.38 |     42.86 |     65.53 | **49.18** |     34.24 |     88.27 |
| CASSLE   | **86.31** | **64.36** | **60.67** | **63.96** | **72.33** |     85.22 |     79.62 | **39.86** | **43.10** | **65.96** |     48.91 |     33.21 | **88.88** |

We observe that while most of the results are indeed better than for SimCLR variants trained with the batch size of 256, the relative downstream performance of the compared extensions remains relatively unchanged. We maintain that CASSLE is the best-performing extension of SimCLR among the compared methods. 

However, a full benchmark of the new SimCLR variants would require us to re-run all the evaluations reported in our manuscript, in particular the lengthy object detection evaluation, which we do not have the capacity for within the remaining discussion timeframe. In order to not cause confusion about which SimCLR model was used in which evaluation, we choose to keep the original results in the manuscript for now. We will conduct a full evaluation of the models trained with larger batch sizes for the camera-ready version.

Reviewer comment (Novelty):
Can CASSLE be applied to masked self-supervision? There seems to be a connection between CASSLE and the MAE, where the latter conditions on mask tokens to reconstruct masked patches.

Author response:
**ad. Q3: Can CASSLE be applied to masked self-supervision?**

Indeed, there seems to be such a connection. While CASSLE has been designed with joint-embedding approaches in mind, applying it to MIM-based approaches could be an interesting idea for future work. However, we are not aware of any works that indicate that MIM-based methods also suffer from the issue of augmentation invariance.

Reviewer comment (Experiments):
Have the authors tried performing feature inversion like in [3]? It would be interesting to see if CASSLE results in inverted features that are more reconstructive of attributes like color compared to vanilla contrastive features.

Author response:
**ad. Q4: Have the authors tried performing feature inversion like in [3]?**

While this could be interesting, we did not perform such an evaluation. We are interested in whether features learned by CASSLE are more transferable to downstream tasks and harder to match together, whereas the quality of reconstruction of particular attributes such as color could be subject to large variance.

Reviewer comment (Evaluation):
CASSLE performs better (compared to AugSelf) for contrastive methods and BarlowTwins than others, i.e. BYOL and SimSiam. A discussion on why this happens can be interesting.

Author response:
Thank you for this question. We suspect that this may be caused by the asymmetric representation used by BYOL and SimSiam, which distill the representations between the projector and additional predictor network. Observe that CASSLE and AugSelf also perform comparably on MoCo-v3, which also utilizes an additional predictor. It is worth noticing that CASSLE improves the performance of the vanilla versions of those methods as well.  
We have added this remark to the revised version of the paper.

Reviewer comment (Evaluation):
Experiments on object detection task demonstrate a marginal improvement.

Author response:
While this improvement is small, CASSLE consistently improves over the performance of Vanilla and AugSelf models for both MoCo-v2 and SimCLR.

Reviewer comment (Evaluation):
The paper does not report confidence intervals of their results.

Author response:
We did not report the confidence intervals of individual downstream task performance, as this would limit the readability of (already massive) Tables 1 and 7. We have compiled the linear evaluation results of different approaches on different downstream tasks in Figure 5 (Appendix C) of the revised manuscript and report the mean rank of each model (when ranked from worst to best performance on downstream tasks) and mean improvements of CASSLE and AugSelf over Vanilla approaches, both with 95% confidence intervals.

Reviewer comment (Novelty):
Missing relevant approach to CASSLE is [1]. They provide with a method which can be perceived as a kind of conditioning to augmentation information.

Author response:
The authors of [1] use a mapping function that, similarly to our conditioned projector, acts on the combined information of image embeddings and augmentation parameters. This mapping allows for supplementing the training objective with equivariance regularization and leads to learning representations that are equivariant to augmentations applied to the images. There are two key differences between our work and [1]:  
CASSLE does not modify the objective function of the trained model. While we do not optimize for equivariance, we observe that increased sensitivity to augmentations emerges in CASSLE due to injecting augmentation information into the projector.  
[1] is evaluated in the supervised learning setting, whereas CASSLE is designed for self-supervised learning. We have included a reference to [1] in our related work section.

Reviewer comment (Writing):
In *Related Work*, contrastive learning objectives usually refer to methods which prevent representational collapse by contrasting against negative pairs. Please clarify this distinction.

Author response:
While there obviously exists a distinction between contrastive, distillation-based, and CCA-based approaches, their objectives are often collectively referred to as contrastive objectives [2,3]. We have adopted this term in our work since CASSLE is applicable to different joint-embedding models regardless of their objective. We are open to your suggestions for a more accurate term describing those three kinds of objectives.

Reviewer comment (Reproducibility):
In *Section 4.2*, an analysis of activation invariance is presented based on the InfoNCE loss. Which similarity function was used to compare earlier representations?

Author response:
For all stages of the network, we measure the cosine similarity of representations and calculate the InfoNCE loss subsequently.

Reviewer comment (Evaluation):
In *Table 3*, how is the rank computed exactly?

Author response:
We compare different variants of MoCo-v2+CASSLE on the same classification and regression tasks as in Section 4.1. We rank the models from best to worst performance on each task and report the average ranks in Table 3. The full results are in Table 12.

Reviewer comment (Experiments):
Lack of comparison with recent augmentation-free SSL methods. Recently, there have been proposed many augmentation-free self-supervised learning methods, including data2vec [1-2], I-JEPA [3], and Masked Image Modeling (MIM) [4-5]. The augmentation-free SSL methods do not use augmentation, in other words, they aim to learn full information about original images, rather than learning augmentation-invariant representations. Also, since they are often better than MoCo-v2 and SimCLR in various benchmarks (e.g., linear evaluation, fine-tuning, scalability), the authors should compare the proposed method with the methods.

Author response:
**ad. W1: Lack of comparison with recent augmentation-free SSL methods.**

Our research problem focuses on how augmentation-based SSL methods become invariant to augmentations. We are not aware of any evidence in the literature that MIM-based methods also suffer from this issue. As such, a comparison to them is not relevant in the context of CASSLE.

Reviewer comment (Evaluation):
Experimental results are not convincing. The performance improvement of CASSLE over AugSelf is marginal.

Author response:
**ad. W2: The performance improvement of CASSLE over AugSelf is marginal.**

While it is true that CASSLE does not always achieve the best downstream performance, it does so in the majority of cases. We have compiled the linear evaluation results of different approaches on different downstream tasks (Tables 1 and 7) and ranked each approach from best to worst-performing in each downstream task (See Figure 5 in the revised manuscript). We find that:

* CASSLE performs best in 54 out of 91 downstream tasks, i.e. 59.34% of cases where it was evaluated.
* AugSelf performs best in 31 out of 91 downstream tasks, i.e. 34.07% of cases where it was evaluated.
* AI performs best in 5 out of 16 downstream tasks, i.e. 31.25% of cases where it was evaluated.
* IFM performs best in 1 out of 13 downstream tasks, i.e. 7.69% of cases where it was evaluated.
* Vanilla performs best in 0 out of 91 downstream tasks, i.e. 0.00% of cases where it was evaluated.

Reviewer comment (Novelty):
Lack of novelty. The proposed method is neither novel nor interesting. First, the goal of this paper has been widely studied via augmentation-aware objectives (e.g., AugSelf) and augmentation-free SSL methods (e.g., I-JEPA). Also, it is hard to find a strong advantage of the proposed idea compared to AugSelf. In the reviewer’s opinion, the choice between injection and prediction cannot make meaningful novelty.

Author response:
**ad. W3a: Lack of novelty. First, the goal of this paper has been widely studied via augmentation-aware objectives (e.g., AugSelf) and augmentation-free SSL methods (e.g., I-JEPA).**

We disagree that the existence of augmentation-free methods such as MIM and I-JEPA exhausts the need for augmentation-based methods and therefore, for research on augmentation awareness. We see several reasons for research into augmentation-aware approaches:
* recent augmentation-based SSL approaches such as DINO-v2 [9] achieve state-of-the-art against augmentation-free methods
* augmentation-based and augmentation-free methods learn different kinds of features - while augmentation-based methods lack spatial sensitivity which requires modeling the local structure within each image, MIM does not have good semantic alignment [8].
* it should also be highlighted that augmentation-free models focus primarily on pretraining Vision Transformers (with some rare exceptions [7]). On the other hand, contrastive SSL continues to be the mainstream family of approaches for pretraining convolutional architectures such as ResNets, which continue to be widely used.

Thus, better understanding and improving these methods remains an important research problem.

**ad. W3b: Also, it is hard to find a strong advantage of the proposed idea compared to AugSelf.**

We believe our approach has several advantages over AugSelf:
* a stronger theoretical justification, which we added in the revised Section 3.2 of the paper
* AugSelf is a multi-task learning method. It requires the user to choose the proportions of the SSL, and augmentation prediction losses. In fact, the authors of AugSelf show that different proportions are needed for different types of backbones and that the choice of augmentations whose parameters to predict during pretraining is not trivial. 
* On the other hand, CASSLE does not introduce new losses to the training. We also demonstrate that utilizing all possible augmentation information during pretraining yields the best downstream performance. 
* Finally, during linear evaluation, there are 12 cases of AugSelf performing worse than the vanilla approach (in MoCo-v1, MoCo-v3, SimCLR, SimSiam, Barlow Twins), whereas CASSLE does so in only 4 cases which are isolated to the BYOL model.

Thus, CASSLE requires simpler modifications to the extended SSL models, while offering a larger performance boost than AugSelf in the majority of cases.