Towards domain-invariant Self-Supervised Learning with Batch Styles Standardization

TL;DR

This work introduces Batch Styles Standardization (BSS), a Fourier-based batch-level style harmonization technique designed to improve domain generalization for self-supervised learning (SSL) without domain labels or domain-specific components. By transferring low-frequency amplitude components from randomly chosen batch images to all batch images, BSS reduces spurious correlations and encourages semantic rather than style-based distinctions, enabling stronger domain-invariant SSL representations. When integrated with SimCLR, SWaV, and MSN, BSS yields consistent improvements on unseen domains across PACS, DomainNet, and Camelyon17 WILDS, often surpassing or rivaling existing unsupervised domain generalization methods. Ablation studies reveal that batch-wise color jitter and FA, when combined with BSS, maximize gains, while mechanistic analyses show reductions in domain purity and improved hardness of negatives, with reduced reliance on large batch sizes. The approach thus offers a practical, adaptable pathway to robust cross-domain SSL without requiring domain annotations, and invites exploration of other style-transfer strategies for further gains.

Abstract

In Self-Supervised Learning (SSL), models are typically pretrained, fine-tuned, and evaluated on the same domains. However, they tend to perform poorly when evaluated on unseen domains, a challenge that Unsupervised Domain Generalization (UDG) seeks to address. Current UDG methods rely on domain labels, which are often challenging to collect, and domain-specific architectures that lack scalability when confronted with numerous domains, making the current methodology impractical and rigid. Inspired by contrastive-based UDG methods that mitigate spurious correlations by restricting comparisons to examples from the same domain, we hypothesize that eliminating style variability within a batch could provide a more convenient and flexible way to reduce spurious correlations without requiring domain labels. To verify this hypothesis, we introduce Batch Styles Standardization (BSS), a relatively simple yet powerful Fourier-based method to standardize the style of images in a batch specifically designed for integration with SSL methods to tackle UDG. Combining BSS with existing SSL methods offers serious advantages over prior UDG methods: (1) It eliminates the need for domain labels or domain-specific network components to enhance domain-invariance in SSL representations, and (2) offers flexibility as BSS can be seamlessly integrated with diverse contrastive-based but also non-contrastive-based SSL methods. Experiments on several UDG datasets demonstrate that it significantly improves downstream task performances on unseen domains, often outperforming or rivaling with UDG methods. Finally, this work clarifies the underlying mechanisms contributing to BSS's effectiveness in improving domain-invariance in SSL representations and performances on unseen domain.

Figures (5)

• Figure 1: (a) BSS: Fourier Transform $\mathcal{F}$ is applied on all batch images then low-frequency components of the amplitudes $\mathcal{A}$ (determined by the areas ratio $r$) are replaced by those of a randomly chosen image (the first one in this case). Finally, inverse Fourier transform $\mathcal{F}^{-1}$ is applied to the altered Fourier transforms to build images with standardized styles. (b) Augmented images with standard independent FA. (c) Augmented images with BSS.
• Figure 2: SimCLR experiments with standard augmentation, FA or BSS on Camelyon17 WILDS.
• Figure 3: Averaged homogeneity scores for representations assigned to the same prototype.
• Figure 4: t-SNE plots of the backbone representations for different UDG methods on Camelyon17 WILDS. Colors and markers correspond respectively to different domains and classes. On the target domains (hospital_1, hospital_2), our method (SimCLR w/ BSS) shows better domain confusion while keeping better class separability. Zoom on pdf for better visualization.
• Figure 5: Impact of hyperparameter $r$ on augmented images with BSS