Paper: SASSL: Enhancing Self-Supervised Learning via Neural Style Transfer (ICLR-2024)

Reviewer comment (Experiments):
While the 2% improvement for MoCo V2 is noteworthy, the generalizability of the proposed augmentation technique necessitates evaluation across various methods, such as BYOL and SimSiam.

Author response:
>### 1. SASSL performance on other SSL techniques.

As suggested, we conducted additional experiments to assess SASSL's performance across alternative SSL methods, namely SimCLR and BYOL. Our results consistently show performance improvements across all methods. For details on these experiments, please refer to our General Response (item 2).

Reviewer comment (Novelty):
The concept of controlling augmentation to preserve semantic information is not new. It's crucial for the authors to position their work within the context of existing literature.

Author response:
>### 3. Comparison to work on data augmentation

In what follows, we compare our proposed method and the work suggested by the reviewer. Please refer to Appendix Sec. A.9 for a more in-depth comparison.

Purushwalkam et al [1] suggest an augmentation strategy based on natural (temporal) transformations that occur in videos as an alternative to learning from occluded (cropped) images from object-centric datasets. In contrast, our approach utilizes neural style transfer to decouple images into content and style, allowing us to apply transformations exclusively to the style of pre-training samples while preserving the content information. Their approach relies on datasets generated on video frames to incorporate temporal invariance, while SASSL can be used on standard image datasets, enabling its integration to default augmentation pipelines.

Lee et al. [2] propose to include an auxiliary loss for SSL methods to capture the difference between augmented views, leading to better downstream performance on datasets where the task relies on information that may have been lost due to invariances introduced by aggressive augmentation. While their method modifies the pre-training loss and requires the use of data augmentation parameters to learn representations, SASSL does not require any additional loss components or auxiliary information from each applied augmentation to function. Our method complements the default data augmentation pipeline with a content-preserving transformation to obtain more general representations.

Bai et al. [3] propose an alternative SSL augmentation pipeline to prevent the loss of semantic information in the learned representations by encouraging learning from weakly augmented views and gradually transitioning to strongly augmented views. The method avoids semantic shifts caused by aggressive image transformations by leveraging the tendency of SSL models to memorize clean views during early pre-training stages and by proposing a multi-stage augmentation pipeline to generate weak and strongly augmented views. Differently, our approach does not require additional augmented views or dynamic control over the loss weights during training. SASSL utilizes a content-preserving transformation that only modifies stylistic characteristics (color and texture), leaving the natural image structure and objects unchanged. Our approach incorporates a transformation into the default augmentation pipeline without altering the loss function or instance discrimination strategy.

Reviewer comment (Evaluation):
The enhancements observed during fine-tuning are not as pronounced as those in linear probing. Could the authors elucidate the factors contributing to this disparity?

Author response:
>### 4. Linear Probing vs. Fine-Tuning improvement gap

The difference in performance between linear probing and fine-tuning arises from how the representation model is treated. Linear probing involves training a classification head atop representations extracted through SASSL. In contrast, fine-tuning entails fully training both the representation and classification models on the target dataset. Consequently, fine-tuning diminishes the impact of style transfer employed during pre-training by retraining the backbone, lessening its effect on downstream performance.

Despite this performance gap, SASSL demonstrates improved performance in both linear probing and fine-tuning. This can be attributed to (i) more robust image representations (evidenced by the linear probing results) and (ii) better model initialization (supported by the fine-tuning results). It's noteworthy that alternative SSL methods also show less improvement over the baseline during fine-tuning. For instance, BYOL exhibits this behavior compared to SimCLR, as shown in Table 3 of BYOL's paper [r3].

[r3] Grill, Jean-Bastien, et al. "Bootstrap your own latent-a new approach to self-supervised learning." Advances in neural information processing systems 33 (2020): 21271-21284.

Reviewer comment (Experiments):
Have the authors employed strong augmentations during the linear probing phase? Typically, in settings of SimCLR and MoCo V2, only week augmentations are utilized. It would be help present the linear probing results on ImageNet-1K without strong augmentations.

Author response:
>### 5. Data augmentation during linear probing

SASSL employs the same data augmentation and optimizer settings for both linear probing and fine-tuning as MoCo v2, which correspond to **weak data augmentations**. For a more detailed description of these settings, please refer to Appendix Sec. A.4.

Please note that SASSL utilizes style transfer only during pre-training, not during downstream training.

Reviewer comment (Experiments):
One major concern is the lack of a detailed discussion about relevant data augmentation strategies, especially those manipulating image textures like [1]. Regarding the experiments, no existing data augmentation algorithm is compared against SASSL.

Author response:
>### 1. Discussion on previous texture-based techniques

We thank the reviewer for pointing out related work on image texture and its effect on perceptual tasks. While Geirhos et al. [1] studies how CNNs are biased towards texture in the context of supervised learning, we focus on preserving semantic information in self-supervised learning by exclusively distorting the style component of pre-training samples. While their main hypothesis is that CNN classifiers trained via supervised learning are biased towards texture, ours is that decoupling content and style in order to exclusively distort the style leads to stronger self-supervised image representations.

While both works rely on NST, their method uses it to create a stylized dataset (Stylized-ImageNet), whereas our approach is to incorporate style transfer in the augmentation pipeline to create stylized images on-the-fly while controlling their effect via blending and interpolation hyperparameters. Experimentally, we also cover a wider set of style datasets (beyond Painter by Numbers) and downstream datasets. We emphasize that Geirhos's approach is designed for supervised learning, while SASSL focuses on representation learning using SSL. Consequently, these methods are not directly comparable (See Appendix Sec. A.9 for a more detailed comparison between techniques).

Reviewer comment (Experiments):
Another potential weakness of the paper is the lack of a comprehensive ablation study to analyze the effect of the number of layers in the stylization network on the performance. Conducting such an ablation study would provide deeper insights into the strengths and limitations of SASSL and help guide future improvements.

Author response:
>### 2. Ablation study on stylized layers

At the reviewer's suggestion, we conducted an ablation study to investigate the impact of the number of layers stylized during pre-training. We examined three scenarios: (i) stylizing the first two residual blocks (four layers corresponding to Residual blocks 1 and 2), (ii) stylizing the first four residual blocks (eight layers corresponding to Residual blocks 1 to 4), and (iii) stylizing all residual layers (ten layers corresponding to Residual layers 1 to 5).

For each scenario, we pre-trained and linearly probed a ResNet-50 on ImageNet using SASSL + MoCo v2 with its optimal configuration (blending and interpolation between 0.1 and 0.3, using Painter by Numbers as the external style dataset). To provide a clear comparison, we compared the classification accuracy of these three models to the model trained without feature alignment (using the stylization network as an autoencoder), and our complete SASSL + MoCo v2 model (stylizing all layers). Note that the latter stylizes all residual and upsampling blocks (a total of thirteen layers).

Results indicate a gradual improvement in accuracy as the number of stylized layers increases. Stylizing only the first four layers yielded no significant improvements, resulting in approximately the same accuracy as the model pre-trained without feature alignment. Stylizing the first eight layers enhanced top-1 classification accuracy by 0.32% over the model with no feature alignment. Finally, stylizing the first ten layers improved accuracy by 0.5%. These findings suggest that deeper layers have a more pronounced impact on accuracy improvement. Additionally, the model pre-trained with full stylization achieved a 1.61% improvement over the model with no feature alignment, indicating that the stylization occurring in the upsampling layers plays a dominant role in the gains provided by SASSL.

Our study suggests that significant improvements can be achieved by solely stylizing deeper layers, potentially reducing the computational demands of SASSL while still enabling the acquisition of robust image representations (See Appendix Sec. A.7 for more details).

| Augmentation Settings |              Stylization Mode              | Top-1. Acc (%) | Top-5 Acc. (%) |   |
|:---------------------:|:------------------------------------------:|:--------------:|:--------------:|---|
|        Default        |                      -                     |      72.97     |      90.86     |   |
|         SASSL (Ours)        | Image encoding only (no feature alignment) |      73.77     |      91.64     |   |
|         SASSL (Ours)        |          Stylize first four layers         |      73.75     |      91.58     |   |
|         SASSL (Ours)        |         Stylize first eight layers         |      74.09     |      91.76     |   |
|         SASSL (Ours)        |          Stylize first ten layers          |      74.27     |      91.74     |   |
|         **SASSL (Ours)**        |              Full stylization              |      **75.38**     |      **92.21**     |   |

Reviewer comment (Experiments):
The paper could provide more insights into the computational efficiency of the SASSL technique and comparison with other existing methods that also leverage style transfer for data augmentation. This would provide a clearer understanding of the computational requirement and unique contributions and advantages of SASSL compared to previous approaches.

Author response:
>### 3. Computational Requirements

We ran additional experiments to compare the runtime of our proposed method to the default augmentation pipeline, as suggested by the reviewer. We calculated the throughput (augmented images per second) of SASSL relative to MoCo v2's data augmentation. The throughput was determined by averaging 100 independent trials on 128 x 128-pixel images with a batch size of 2048. We also report the relative change, which is the percentage decrease in throughput compared to the default data augmentation. All experiments were carried out on a single TPU.

The following table summarizes the throughput comparison. SASSL reduces throughput by about 20% due to the computational expense of stylizing large batches, which involves running a forward pass of the stylization model. However, SASSL achieves up to a 2% top-1 classification accuracy improvement, as shown empirically. We believe this to be a favorable tradeoff between performance and runtime (See Appendix Sec. A.8 for more details).

| Augmentation Method | Throughput (Images/second) | Relative change (%) |   |   |
|:-------------------:|:--------------------------:|:-------------------:|:-:|---|
|    SimCLR (default)   |            37.45           |          -          |   |   |
|    **SASSL + SimCLR (Ours)**   |            29.48           |        21.28        |   |   |

Reviewer comment (Experiments):
The downstream task is only about classification. Maybe it can happen to be successful.

Author response:
We thank the reviewer for appreciating our paper’s presentation and for the insightful feedback.

>### 1. Additional downstream tasks

We appreciate the reviewer's recommendation to evaluate our method on additional downstream tasks. In response, we conducted further experiments on few-shot learning. The results from one-shot and ten-shot experiments on ImageNet demonstrate that the representations learned using SASSL enhance performance in both scenarios. Please refer to our General Response (item 3) for the details.

Reviewer comment (Experiments):
Can you or did you do some other downstream task to demonstrate the success of data augmentation (e.g. object detection?)

Author response:
We thank the reviewer for appreciating our paper’s presentation and for the insightful feedback.

>### 1. Additional downstream tasks

We appreciate the reviewer's recommendation to evaluate our method on additional downstream tasks. In response, we conducted further experiments on few-shot learning. The results from one-shot and ten-shot experiments on ImageNet demonstrate that the representations learned using SASSL enhance performance in both scenarios. Please refer to our General Response (item 3) for the details.

Reviewer comment (Experiments):
The proposed idea (NST for SSL data augmentation) is only evaluated on CSN of NST and SimCLR of SSL. More different trials, such as adaptive instance normalization (AdaIN) of NST and more SSL frameworks should be evaluated to demonstrate the effectiveness of this idea. Furthermore, why would you call the normalization method as CSN? This method was noted as Conditional Instance Normalization (CIN) in the original paper, and this notion CIN has been widely used in current NST literature.

Author response:
Due to computational constraints, we are unable to perform additional experiments on alternative NST methods with different feature matching techniques. Instead of focusing on how features are matched using alternative alignment techniques (e.g. Whitening and coloring transforms, histogram matching, optimal transport, adaptive instance normalization, etc.), we provide a more in-depth examination of how the feature matching itself affects downstream performance. We conducted additional ablation studies to gain a better understanding of the impact of feature alignment, as described next.

We ran additional experiments to analyze the downstream task performance of models pre-trained via MoCo v2 under four different scenarios: (i) default augmentations, (ii) SASSL with no conditional instance normalization (no feature alignment, only image encoding), (iii) SASSL using Gaussian noise instead of style representations for conditional instance normalization, and (iv) SASSL using external style representations from Painter by Numbers.

The following table shows that the performance boost obtained via NST relies on both the image encoding and feature alignment processes, where feature alignment has a greater impact on performance. Without aligning the content representation to those of a style reference, performance improves by 0.8%. Aligning content representations using noise as a style reference decreases performance compared to using no alignment. Finally, using an external dataset as style representation improves performance by an additional 1.61%. This suggests that feature alignment is an important part of the stylization process, but it is not the only factor. The encoding component of NST, which is common across stylization methods, also contributes to improved representation learning (See Appendix Sec. A.7 for more details).

| Augmentation Settings |              Stylization Mode              | Top-1. Acc (%) | Top-5 Acc. (%) |   |
|:---------------------:|:------------------------------------------:|:--------------:|:--------------:|---|
|        Default        |                      -                     |      72.97     |      90.86     |   |
|         SASSL (Ours)        | Image encoding only (no feature alignment) |      73.77     |      91.64     |   |
|        SASSL (Ours)         |   Gaussian noise as style representation  |      73.21     |      91.17     |   |
|         **SASSL (Ours)**         |    **External dataset (Painter by Numbers)**   |      **75.38**     |      **92.21**     |   |


In Section 3.2, we use the name Conditional Style Normalization to follow Fast Style Transfer’s official implementation. Please note that the [official codebase]((https://github.com/magenta/magenta/blob/main/magenta/models/image_stylization/ops.py)) makes a distinction between Conditional Instance Normalization, Weighted Instance Normalization, and Conditional Style Normalization. Following the reviewer’s suggestion, we have renamed the normalization method to Conditional Instance Normalization in our paper to prevent any confusion with the NST literature.

Reviewer comment (Experiments):
Table 2 shows that not all style datasets can outperform the original MoCo V2 for Retinopathy and DTD. The optimal style dataset also varies across different target datasets. The choice of style dataset would be a problem.

Author response:
When working with a target dataset where style/texture is not meaningful for classification (e.g., ImageNet and iNat21), all the evaluated style datasets enhance downstream performance. This means that **selecting any of these style datasets leads to better performance than using the baseline**. However, for target datasets where semantic information is intertwined with stylistic attributes, employing NST as data augmentation does not guarantee improved performance.

Retinopathy and DTD datasets embed semantic information within their textures and local patterns. Retinopathy images employ local abnormalities and textures to determine the stage of diabetes. Similarly, DTD relies on patterns and appearance to distinguish between texture categories. In both cases, the semantic information is intrinsic to the texture itself. As a result, SASSL produces image representations that achieve performance comparable to the baseline (within one standard deviation).

To provide a clearer understanding of these scenarios, we have extended our discussion and included a visual analysis that illustrates the differences between styles in different datasets (See Appendix Sec. A.2 for more details).

Reviewer comment (Writing):
The notations in this paper are somewhat confusing. 1) $\beta^{(k)}$ denotes the trainable function in conditional style normalization, while $\beta$ denotes the interpolation factor of content and style images in Eq. 11. The authors should reconsider the notations for either of them; 2) Should $f$ be $h$ in Eq. 2? 3) In Eq. 7 $\mathcal{T}$ is a stylization network, while in Eq. 1 $\mathcal{T}$ denotes all possible augmentations.

Author response:
We thank the reviewer for the suggested changes. We have updated the mathematical notation on the paper accordingly.