SASSL: Enhancing Self-Supervised Learning via Neural Style Transfer

We thank the reviewer for the valuable comments and feedback.

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).

2. Additional downstream tasks

We value the reviewer's suggestion to assess our method on additional downstream tasks. We conducted further experiments on few-shot learning. The results from one-shot and ten-shot experiments on ImageNet reveal that the representations acquired using SASSL improve performance in both scenarios. For the details, please refer to our General Response (item 3).

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.

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.

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.

We appreciate the reviewer's acknowledgment of our paper's presentation and for sharing insightful feedback.

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).

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).

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).

4. Effect of feature matching

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).

5. Hyperparameter selection

The intervals for alpha and beta (0.1 to 0.3 in both cases) were selected by sweeping their minimum and maximum values and measuring their downstream performance on ImageNet. Similar to the other augmentations in MoCo v2, e.g., random cropping, color jittering and solarization, we fixed these intervals and used them in all our transfer-learning experiments.

We respectfully disagree with the reviewer's assessment of the practical limitations of our proposed method. While further performance gains might be possible with different alpha and beta intervals, our experiments show that our method consistently enhances performance across datasets without the need for hyperparameter tuning. This aligns with most augmentations in MoCo v2 and other SSL methods, which rely on multiple hyperparameters yet still improve representation learning with fixed hyperparameter values. Our NST augmentation functions similarly and can be seamlessly integrated into the default augmentation pipeline to further enhance downstream performance.

6. Details on ablation study (Table 3)

Table 3 demonstrates the significance of combining the intermediate stylized and content images using an interpolation factor beta. While preserving the entire intermediate stylized image negatively impacts performance, it is important to note that this is not our final proposed technique but rather the scenario where the interpolation step is disabled. As a result, fine-grained content details are lost due to the bottleneck compression introduced by the stylization network.

Several NST methods require blending and interpolation because feature alignment techniques tend to overemphasize style attributes (as seen in the Universal Style Transfer implementation [r3]). So, it is expected that fully utilizing the stylized image hinders performance. The performance drop caused by not having control over the pixel interpolation is comparable to that observed in other SSL methods like BYOL, where removing a single augmentation can lead to a decline in top-1 accuracy of over 8% (refer to BYOL's paper, Appendix F.3).

[r3] Li, Yijun, et al. "Universal style transfer via feature transforms." Advances in neural information processing systems 30 (2017).

7. Math notation

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

We thank the reviewer for the helpful comments and feedback.

1. Comparison to NST work on supervised learning

We appreciate the reviewer bringing these recent work to our attention. In what follows, we will address their key differences from our proposed method.

In their work [r2], Jackson et al. propose employing NST as a data augmentation technique for supervised learning, where the style reference is drawn from a normal distribution. However, our experiments demonstrate that this approach hinders performance in the context of SSL (as shown in item 4). Instead, SASSL utilizes precomputed style representations from external datasets or samples from the content datasets themselves (in-batch stylization), resulting in improved or on-par downstream performance across style datasets.

Furthermore, their approach solely considers the blending of style representations, overlooking the distortions introduced by the stylization network itself. Our work reveals that the compression generated by the stylization network, even without using a style representation, leads to a loss of information about the structure (edges and small image details) of pre-training images, deteriorating downstream performance. This is the rationale behind SASSL's reliance on both pixel interpolation and feature blending. While both techniques employ NST as a content-preserving transformation, SASSL enables its use by incorporating it into the default data augmentation pipeline of various SSL methods (as shown in our General Response, item 2), which differs from utilizing NST for domain transfer tasks in supervised learning.

As mentioned in Section 2.2, Zheng et al. [r1] utilize NST to create additional samples for supervised learning, using only eight style images. It's important to note that their approach does not incorporate NST as data augmentation; instead, the stylized images are fixed and added to the original training dataset. Additionally, their experiments are limited to a single classification dataset (Caltech). This contrasts sharply with our method, which proposes different stylization strategies (external and in-batch) and ways to control the stylization effect (via ablation and interpolation) to augment data on-the-fly in order to enhance SSL performance. While both methods share the use of style NST, their tasks and approaches are distinct.

2. Details on normalization

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 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.

3. Style dataset selection

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).

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.

3. Additional Tasks

To further demonstrate the enhanced representation learning capabilities of SASSL, we conducted additional experiments on another downstream task: few-shot classification. Due to computational constraints, semantic segmentation results using SASSL are not yet available.

We evaluated the MoCo v2 representations obtained through SASSL against the default MoCo v2 model in the context of few-shot learning, specifically on ImageNet's one-shot and ten-shot settings.

The table below presents the classification accuracy in both scenarios. The results reveal that SASSL boosts few-shot classification accuracy by over 1% in terms of top-1 accuracy in both cases. This aligns with our classification experiments, suggesting that SASSL promotes more general image representations (See Appendix Sec. 6 for more details).