DIVA: Domain Invariant Variational Autoencoders

Abstract: We consider the problem of domain generalization, namely, how to learn representations given data from a set of domains that generalize to data from a previously unseen domain. We propose the Domain Invariant Variational Autoencoder (DIVA), a generative model that tackles this problem by learning three independent latent subspaces, one for the domain, one for the class, and one for any residual variations. We highlight that due to the generative nature of our model we can also incorporate unlabeled data from known or previously unseen domains. To the best of our knowledge this has not been done before in a domain generalization setting. This property is highly desirable in fields like medical imaging where labeled data is scarce. We experimentally evaluate our model on the rotated MNIST benchmark and a malaria cell images dataset where we show that (i) the learned subspaces are indeed complementary to each other, (ii) we improve upon recent works on this task and (iii) incorporating unlabelled data can boost the performance even further.

• The paper presents DIVA, which partitions the latent space into domain, class, and residual subspaces to enable robust domain generalization.
• It employs auxiliary classifiers to disentangle representations, effectively integrating unlabeled data and enhancing performance.
• Experimental results on rotated MNIST and malaria cell datasets demonstrate DIVA’s superior accuracy compared to state-of-the-art methods.

Summary of "DIVA: Domain Invariant Variational Autoencoders"

The paper introduces the Domain Invariant Variational Autoencoder (DIVA), a novel generative model designed to address the domain generalization problem. Domain generalization involves learning representations that generalize across multiple domains, particularly when the model encounters previously unseen domains during testing. DIVA extends the traditional variational autoencoder (VAE) framework by partitioning the latent space into three independent subspaces, each representing distinct aspects: domain, class, and residual variations.

Key Contributions and Methodology

DIVA distinguishes itself from existing domain generalization methods by employing a generative model to achieve domain invariance. This is realized through the introduction of three orthogonal latent subspaces:

1. Domain-specific latent space ($z_d$): Captures variations specific to each domain.
2. Class-specific latent space ($z_y$): Encodes information pertinent to the class labels.
3. Residual latent space ($z_x$): Captures any remaining variations not explained by domain or class.

The model benefits from a unique ability to incorporate unlabeled data from known or novel domains, which is particularly advantageous in fields with limited labeled data, such as medical imaging.

The DIVA framework encourages disentanglement within the latent space by utilizing auxiliary classifiers during training to predict domain labels from $z_d$ and class labels from $z_y$. The learned representations are evaluated on the rotated MNIST benchmark and a malaria cell images dataset, demonstrating enhanced generalization performance compared to existing approaches.

Experimental Results

The paper's empirical evaluation reports significant performance improvements over state-of-the-art domain generalization methods, evidenced by strong numerical results:

• On the rotated MNIST dataset, where the task involves classifying digits across multiple image rotation angles, DIVA consistently achieves higher classification accuracy than domain adversarial neural networks (DA), and other recent approaches like LG, HEX, and ADV.
• Importantly, DIVA’s architecture facilitates effective usage of additional unlabeled data, further boosting its performance.
• On the malaria cell images dataset, DIVA showcases its strength in learning disentangled representations, effectively separating variations due to domain-specific attributes (e.g., cell staining) and class-specific attributes (e.g., presence of parasites).

Implications and Future Directions

The paper's findings underline the effectiveness of structured generative models in capturing domain-invariant characteristics, showing promise for applications in medical image classification and beyond. The capability to leverage unlabeled data broadens DIVA's applicability in scenarios where labeled data is sparse.

As the research community continues to explore domain generalization, potential future developments could include:

• Extending the DIVA framework to other challenging datasets or applications that involve complex domain shifts.
• Investigating the scalability of DIVA in scenarios with a significant increase in the number of domains or classes.
• Enhancing the model by integrating more sophisticated generative structures or advanced inference techniques to improve efficiency and robustness.

In summary, DIVA represents a significant step forward in domain generalization through its innovative use of disentangled generative modeling, with demonstrated efficacy in both supervised and semi-supervised learning contexts. The model's ability to integrate unlabeled data opens new avenues for its deployment in data-constrained environments, holding promise for practical applications in diverse AI fields.