Adapting Self-Supervised Representations to Multi-Domain Setups

TL;DR

This paper introduces a lightweight Domain Disentanglement Module (DDM) that enhances self-supervised models' ability to generalize across multiple diverse domains by disentangling domain-specific and domain-invariant features, improving performance on unseen domains.

Contribution

The paper proposes a novel DDM that can be integrated into any self-supervised encoder to improve multi-domain generalization, even without domain labels, by disentangling representations.

Findings

Up to 3.5% improvement in linear probing accuracy on multi-domain benchmarks.

7.4% better generalization to unseen domains.

Effective with multiple self-supervised models like SimCLR, MoCo, BYOL, DINO, SimSiam, Barlow Twins.

Abstract

Current state-of-the-art self-supervised approaches, are effective when trained on individual domains but show limited generalization on unseen domains. We observe that these models poorly generalize even when trained on a mixture of domains, making them unsuitable to be deployed under diverse real-world setups. We therefore propose a general-purpose, lightweight Domain Disentanglement Module (DDM) that can be plugged into any self-supervised encoder to effectively perform representation learning on multiple, diverse domains with or without shared classes. During pre-training according to a self-supervised loss, DDM enforces a disentanglement in the representation space by splitting it into a domain-variant and a domain-invariant portion. When domain labels are not available, DDM uses a robust clustering approach to discover pseudo-domains. We show that pre-training with DDM can show up to 3.5% improvement in linear probing accuracy on state-of-the-art self-supervised models including SimCLR, MoCo, BYOL, DINO, SimSiam and Barlow Twins on multi-domain benchmarks including PACS, DomainNet and WILDS. Models trained with DDM show significantly improved generalization (7.4%) to unseen domains compared to baselines. Therefore, DDM can efficiently adapt self-supervised encoders to provide high-quality, generalizable representations for diverse multi-domain data.