PaperlayerView exact revision on arXiv

Source-linked AI summary

DIVA: Domain Invariant Variational Autoencoders

Maximilian Ilse, Jakub M. Tomczak, Christos Louizos, Max Welling

arXiv:1905.10427v2stat.MLcs.LG

TL;DR

Domain generalization requires representations that transfer from observed domains to previously unseen domains despite domain-specific biases. DIVA addresses this with a generative model that separates domain, class, and residual variation, and reports disentangled representations, improved performance over other methods, and further gains from unlabeled data.

  • Problem

    Domain generalization seeks representations that transfer to previously unseen domains, while training-data bias and scarce medical-image labels limit ordinary supervised approaches.

  • Method

    DIVA is a generative VAE with independent latent subspaces for domain-specific, class-specific, and residual variation, supporting supervised and unsupervised data.

  • Results

    DIVA learns disentangled representations and improves upon other methods on rotated MNIST and malaria cell images; unlabeled samples can boost performance further.

  • Takeaways & Limitations

    Unlabeled samples, including samples from entirely new domains, can improve DIVA’s domain-generalization performance, which is desirable when labeling is costly.

  • Takeaways & Limitations

    DIVA does not evaluate the marginal likelihood p(x), because the variational lower bound is used as a regularizer during inference.

Abstract

from arXiv · show

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.

1 INTRODUCTION

DIVA addresses domain generalization by learning representations that transfer from observed domains to unseen domains. It uses a generative model with separate latent representations for domain, class, and residual variation, and evaluates this approach on rotated MNIST and malaria cell images.

  • Training-data bias can make deep networks sensitive to irrelevant cues and reduce performance on data from a different distribution.
  • Domain generalization seeks representations learned from several domains that transfer to a previously unseen domain.
  • DIVA extends VAEs with independent latent representations for domain labels, class labels, and residual input variation.
  • The generative formulation enables DIVA to incorporate unlabeled data from known or previously unseen domains.
  • DIVA is evaluated on rotated MNIST, where domains are digit rotation angles, and on malaria cell images, where domains are patients.

2 TOWARDS DOMAIN GENERALIZATION WITH GENERATIVE MODELS

DIVA factorizes variation into domain, class, and residual latent subspaces, using separate encoders and auxiliary prediction objectives to encourage disentanglement. Its generative VAE framework also supports supervised and unsupervised training, while inference uses only the class representation and classifier.

  • DIVA: DOMAIN INVARIANT VAE: DIVA models three independent sources of variation: domain-specific zd, class-specific zy, and residual zx.
  • DIVA: DOMAIN INVARIANT VAE: Separate encoders qφd(zd|x), qφx(zx|x), and qφy(zy|x) approximate the posterior without shared parameters.
  • DIVA: DOMAIN INVARIANT VAE: The β weighting term controls latent-space capacity, with larger values limiting the capacity of each latent variable.
  • DIVA: DOMAIN INVARIANT VAE: Auxiliary objectives train zd to predict domain d and zy to predict class y, strengthening their respective information assignments.
  • DIVA: DOMAIN INVARIANT VAE: At inference, DIVA uses only qφy(zy|x) and qωy(y|zy), treating the variational lower bound as a regularizer rather than evaluating p(x).
  • SEMI-SUPERVISED DIVA: Semi-supervised training combines supervised and unsupervised variational bounds with auxiliary losses, using labeled tuples and unlabeled pairs.

3 RELATED WORK

The paper positions DIVA against invariant-representation and ensemble approaches, emphasizing its complementary latent subspaces and inference cost independent of the number of training domains. It also distinguishes DIVA from related fairness and information-bottleneck models.

  • Domain-generalization methods commonly learn a single invariant representation or ensemble models trained separately for each domain.
  • Unlike adversarial approaches, DIVA learns complementary zd, zx, and zy representations through a generative architecture rather than explicitly regularizing zy with d.
  • DIVA’s model size is independent of the number of training domains, unlike domain-specific ensembles whose size scales linearly with that number.
  • Unlike the fair VAE and CVIB, DIVA does not require the domain label during inference.

4 EXPERIMENTS

The experiments evaluate whether DIVA learns disentangled factors and improves classification in supervised and semi-supervised settings on rotated MNIST and malaria cell images. Rotated MNIST matches DIVA’s graphical model, whereas malaria images present a more challenging setting.

  • DIVA is evaluated on rotated MNIST and malaria cell images for disentanglement and supervised or semi-supervised classification.
  • Rotated MNIST matches DIVA’s graphical model, while the malaria cell images dataset provides a more challenging evaluation setting.

4.1 ROTATED MNIST

On rotated MNIST, DIVA separates domain, class, and residual information across latent subspaces, achieves the highest accuracy among compared methods, and benefits from unlabeled data, including from an additional domain.

  • Dataset and evaluation: Rotated MNIST creates six domains by rotating digits from 0° to 75°; models train on five domains and test on the remaining domain using classification accuracy.Experiments are repeated 10 times.
  • Qualitative disentanglement: The domain latent zd forms domain-specific clusters, while the class latent zy forms ten class-specific clusters with no visible cross-factor clustering.The residual latent zx captures rotation angle, line thickness, and digit width correlations.
  • Comparison to other methods: DIVA achieves the highest accuracy across all test domains and the highest average test accuracy among the compared domain-generalization methods.The comparison includes DA, LG, HEX, and ADV using labeled data.
  • Semi-supervised learning: Adding unlabeled data to the training set generally increases DIVA’s performance, including when the additional data are from a previously unseen domain.With an unlabeled M60° domain, DIVA improves by about 7% despite a roughly 20% supervised-only accuracy drop relative to Table 1.
  • Semi-supervised learning: DIVA can learn from samples of an additional domain without labels, although balancing loss terms becomes increasingly challenging when unlabeled data greatly outnumber labeled data.This imbalance can cause declining performance in the last two columns of Table 1.

4.2 MALARIA CELL IMAGES

The malaria experiments treat patient identity as a domain because cell appearance varies across patients, then evaluate DIVA’s disentanglement and domain-generalization performance.

  • Dataset: Patient-specific staining creates appearance variation that can hinder machine-learning generalization, so patient ID is used as the domain label.Cell color is relatively constant within a patient but can vary greatly between patients; patient ID is treated as a proxy for appearance variability.
  • Dataset: The dataset contains 5922 cell images from 10 patients, with 400–700 cells per patient and similar marginal label distributions.The subset selection was motivated by comparable cell counts and label distributions, described as necessary for successful domain generalization.
  • Qualitative disentanglement: DIVA reconstructs the original cell images with all three latent subspaces while assigning cell color to zd, cell shape to zx, and parasite location to zy.Reconstructions using individual subspaces support the reported correspondence between latent factors and appearance or parasite-related properties.
  • Qualitative disentanglement: Classifiers trained on zd or zx perform worse than a majority-class predictor, supporting separation of class information into zy.This classifier result is reported as an additional qualitative disentanglement check.
  • Supervised case: Supervised experiments train on nine patient domains, hold out the tenth, evaluate ROC AUC, and repeat every experiment five times.DIVA significantly improves over the baseline and DA in five test domains and improves on average, although the average gains are within standard error.
  • Semi-supervised case: The semi-supervised setup tests on C116P77 and adds unlabeled C59P20 data, using ROC AUC to assess whether unlabeled neighboring-domain information helps.The supplied passages specify the comparison table and explain that DIVA uses unlabeled data to update latent-variable components and the decoder.

5 CONCLUSION

The paper concludes that DIVA learns disentangled, domain-invariant representations on rotated MNIST and malaria images, improves over other methods, and benefits from unlabeled data, including from new domains.

  • Conclusion: DIVA learns three latent subspaces whose factors match domain, class, and residual variation, with domain-invariant zy improving performance on both datasets.The conclusion identifies class y and domain d as the ground-truth factors of interest.
  • Conclusion: Unlabeled samples can further boost DIVA’s performance, including samples from entirely new domains without labeled examples.The paper highlights this capability as desirable for medical imaging because labeling is time-consuming and costly.
  • Interpolation and extrapolation: Performance appears better for interpolation domains near the center of an interval than for extrapolation domains at its ends.The authors argue that unlabeled data from a domain close to the test domain can improve extrapolation performance.

5.1 ROTATED MNIST

On rotated MNIST, DIVA is evaluated with controlled training procedures and conditional-generation visualizations that test whether its latent subspaces preserve distinct factors. The experiments report that class and domain can be varied independently while other image characteristics remain stable.

  • Training procedure: DIVA models are trained for 500 epochs with early stopping, KL warm-up, αd = 2000, and 64 dimensions for each latent space.The models use ADAM, pixel-wise cross entropy, and batches of 100.
  • Training procedure: The rotated-MNIST setup uses matched encoder and classifier capacity for qφy(zy|x) and qωy(y|zy) relative to Wang et al. (2019).
  • Generative samples: DIVA generates rotated-MNIST samples that are described as almost indistinguishable from real datapoints when all three latent variables are sampled from their priors.The generated examples are shown in Figure 4.
  • Conditional generation: When zy is replaced to change class, the rotation angle and line thickness remain preserved while the image class changes as intended.Rows 2 to 11 in Figure 5 correspond to classes 0 through 9.
  • Conditional generation: When zd is replaced to change domain, the rotation angle changes while the original image’s class and style are maintained.Rows 2 to 6 in Figure 5 correspond to domains M0° through M60°.
  • Ablation study: An ablation compares DIVA with a VAE using a single latent space, a standard Gaussian prior, and two auxiliary tasks; Table 10 reports classification accuracy.The passage states that the partitioned latent space benefits DIVA.

5.2 MALARIA CELL IMAGES

The malaria experiments use supervised and semi-supervised training procedures with a shared feature-extractor architecture across models. They assess whether the learned latent embeddings support class prediction on a held-out test domain.

  • Training and architecture: Malaria experiments train models for 500 epochs with early stopping based on validation accuracy for y, using 64 dimensions for zd, zx, and zy.The hyperparameter search varies classifier weights, latent dimensionality, and βmax.
  • Training and architecture: The supervised and semi-supervised experiments use matched training procedures, while domain adaptation follows the procedure of Ganin et al. (2015).The semi-supervised epoch count is adjusted to match DIVA’s number of parameter updates.
  • Training and architecture: The baseline uses qφy(zy|x) followed by qωy(y|zy), domain adaptation uses qφy(zy|x) as a feature extractor, and all models use seven ResNet blocks.Class and domain classifiers consist of two linear layers.
  • Latent-space evaluation: A two-layer MLP is trained separately on zd, zx, and zy embeddings and evaluated on the held-out test domain to assess their class predictiveness.The passage introduces Table 15 as reporting the resulting classification accuracies.
Loading 1905.10427v2…