Disentangling Masked Autoencoders for Unsupervised Domain Generalization

Abstract

Domain Generalization (DG), designed to enhance out-of-distribution (OOD) generalization, is all about learning invariance against domain shifts utilizing sufficient supervision signals. Yet, the scarcity of such labeled data has led to the rise of unsupervised domain generalization (UDG)—a more important yet challenging task in that models are trained across diverse domains in an unsupervised manner and eventually tested on unseen domains. UDG is fast gaining attention but is still far from well-studied.

To close the research gap, we propose a novel learning framework designed for UDG, termed the Disentangled Masked AutoEncoder (DisMAE), aiming to discover the disentangled representations that faithfully reveal the intrinsic features and superficial variations without access to the class label. At its core is the distillation of domain-invariant semantic features, which can not be distinguished by domain classifier, while filtering out the domain-specific variations (for example, color schemes and texture patterns) that are unstable and redundant. Notably, DisMAE co-trains the asymmetric dual-branch architecture with semantic and lightweight variation encoders, offering dynamic data manipulation and representation level augmentation capabilities. Extensive experiments on four benchmark datasets (i.e. DomainNet, PACS, VLCS, Colored MNIST) with both DG and UDG tasks demonstrate that DisMAE can achieve competitive OOD performance compared with the state-of-the-art DG and UDG baselines, which shed light on potential research line in improving the generalization ability with large-scale unlabeled data. Our codes are available at https://github.com/rookiehb/DisMAE.

A. Zhang and H. Wang—These authors contribute equally to this work.

Appendices

A Algorithm

Algorithm 1 depicts the detailed procedure of DisMAE.

B Discussion About Differences

We argue that DisMAE is novel and significantly different from prior studies w.r.t. three aspects. 1) Scope. Transitioning to UDG is non-trivial. Previous disentangled methods like DADA [38], DIDN [32], and DIR [19], while effective in DG, struggle with unsupervised data due to their high dependence on class labels to encapsulate semantic attributes. 2) Disentangled Targets. Without class label guidance, achieving a domain-invariant semantic encoder is challenging. Many UDG methods, such as DiMAE [54] and CycleMAE [53], can only separate domain styles using multiple decoders but fall short in disentangling domain-invariant semantics from variations. 3) Disentangle Strategy. DisMAE is grounded in disentanglement and invariance principles, uniquely combining adaptive contrastive loss with reconstruction loss collaboratively. The adaptive contrastive loss, in particular, is designed by seamlessly leveraging the domain classifier and intra-domain negative sampling. The differences are summarized in Table 4.

Table 4. Comparison with previous works

C Experiments

1.1 C.1 Experimental Settings

Baseline Hyperparameter Tuning. For a fair comparison, we uniformly substitute the backbones of all baselines with the same ViT-B/16 and rerun the experiment using UDG and DG open-source codebases. And we provide the default hyperparameters for UDG baselines in Table 5. And the search distribution for each hyperparameter in each DG baseline is detailed in Table 6.

Table 5. Hyperparameters for baselines in UDG. BS represents the batch size, and WD denotes weight decay.

Table 6. Default hyperparameters and random search distribution for baselines in DG

1.2 C.2 Overall Performance

Unsupervised Domain Generalization. Due to limited space in the paper, we show the rest UDG results in Table 7. We employ Clipart, Infograph, and Quickdraw as training domains and Painting, Real, and Sketch as test domains. Following the same all correlated settings and protocols in DARLING, we find that our DisMAE could achieve 1.14%, 1.19%, 4.40%, and 5.45% gains for average accuracy over the second-best baselines across 1%, 5%, 10%, and 100% fraction setting respectively.

Table 7. Unsupervised domain generalization results on DomainNet. We employ Clipart, Infograph, and Quickdraw as training domains and Painting, Real, and Sketch as test domains. All the models are unsupervised pre-trained before fine-tuning on the labeled data. Overall and Avg. are the overall test data accuracy and the arithmetic mean of individual domain accuracy respectively. Note that they are different because the size of each test domain isn’t equal. Bold = best, underline = second best.

Domain Generalization. Aligning with the training-domain validation setup in DomainBed, we achieve 0.50% gains for the average accuracy in PACS datasets, as shown in Table 8.

Table 8. Domain generalization results on PACS. Bold = best, underline = second best.

1.3 C.3 Discussion of the Invariance Principle

In Figs. 2 and 4, we visualize the representation acquired through MAE, our semantics encoder, and the variational encoder via t-SNE. We find that: (1) The representations generated by MAE for each domain showcase a degree of overlap at the center of the picture, accompanied by slight variations within each distribution. This suggests that MAE captures both semantics and domain-variant information but fails to disentangle them effectively. (2) Our semantic representations in each domain distribute uniformly. This justifies that DisMAE could learn domain-invariant representations from each domain. (3) The variation representations in each domain has their specific distribution. Clusters of similar variation data further emphasize domain-specific characteristics.

1.4 C.4 More Ablation Study

Effects of Decoder Depth. The efficacy of the adaptive contrastive loss hinges on the output of decoders. This prompts the inquiry: how many decoder layers are optimal for achieving peak performance? As shown in Table 9, a deeper decoder may lead to overfitting in reconstruction and subsequently diminish the effect of our contrastive loss. Thus, adopting a lightweight decoder could both accelerate the training and guarantee robustness.

Table 9. Hyperparameter Analysis of the decoder layer

Effects of Mask Ratios. In Table 10, we set different mask ratios to test the robustness of our model. And we found that the 80 percentile of the mask ratio reaches the optimal result. We set it as our default protocol.

Table 10. Hyperparameter Analysis of the mask ratio

1.5 C.5 Failure Cases

Some failure cases of our proposed DisMAE are in Fig. 6. Our approach struggles with reconstruction containing intricate details and lines. It frequently fails to generate images that possess sufficient detail while simultaneously providing clear augmented variations. We attribute these failures to two primary reasons: 1) The MAE backbone operates on patches, making pixel-level reconstruction difficult, and our method heavily relies on the MAE model’s reconstruction outcomes. 2) Our disentanglement lacks granularity, often capturing broad color regions and background information rather than nuanced details. In the context of UDG, reconstructing images with fine granularity, high resolution, and authenticity remains a challenging and crucial research direction. We are also keenly interested in exploring the potential integration of the diffusion model within the UDG framework.

Fig. 6.

Some failure cases of reconstructed images generated by DisMAE.

1.6 C.6 Qualitative Reconstructions

Additional visualization results of image reconstruction, spanning both colored MNIST and DomainNet, can be observed in Fig. 7.

DisMAE differentiates between the foreground and background of an image. Remarkably, DisMAE can discern domain styles and fuse domain-specific elements across them—a notable instance is superimposing the sun from a sketch onto a painting. Such disentanglement ability endows DisMAE with the flexibility to generate controllable images by manipulating semantic and variation factors through swapping.

Fig. 7.

Illustrative reconstructed images generated by DisMAE.

1.7 C.7 Detailed Implementation of DisMAE

We conduct all the experiments in Pytorch on a cluster of 8 NVIDIA Tesla A100 GPUs with 40GB each. Our default backbone consists of 12 blocks of semantic encoder, 6 blocks of variation encoder, and a transform-based decoder. We utilize ViT-B/16 as our default backbone for both visualization and main experiments. And we use ViT-Tiny/16 in our ablation study. We let margin $\gamma =0.008$ and $\tau =0.4$. In UDG, we choose the AdamW optimizer for the main branch and set the learning rate as 1e-4 and betas as (0.9, 0.95) for pre-training. As for finetuning, we adopt the learning rate as 0.025, 0.05, 5e-5, 5e-5 and batch size as 96, 192, 36, and 36 in the label fraction 1%, 5%, 10%, and 100% experiments respectively. And we finetune all the checkpoints for 50 epochs. In DG, $\lambda _1$ is selected within {5e-4, 1e-3, 5e-3, 1e-2} and $\lambda _2$ is selected within {0.1, 0.5, 1.0, 2.0}. The detailed hyperparameters for UDG and DG are listed in Table 11.

Table 11. Hyperparameters Selection of DisMAE.

Details About Training the Domain Classifier. As for the domain classifier, we use the SGD optimizer with a learning rate of 0.0005, momentum of 0.99, and weight decay of 0.05. We choose adaptive training intervals $T_{ad}$ as 15 and maximum adaptive training epochs $E_{ad}$ as 100 in the UDG setting. We only update the domain classifier by minimizing the cross-entropy loss while freezing backbones when $e \bmod T_{ad} == 0$ and $e \le E_{ad}$, where e is the current training epoch. The detailed algorithm can be found in Appendix A.