[Paper Review] DIVA: Domain Invariant Variational Autoencoders

Maximilian Ilse, Jakub M. Tomczak+2 more|arXiv (Cornell University)|May 24, 2019

Domain Adaptation and Few-Shot Learning43 references66 citations

TL;DR

DIVA learns three independent latent subspaces (domain, class, residual) within a VAE to achieve domain generalization and can leverage unlabeled data to boost performance.

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.

Motivation & Objective

• Motivate domain generalization where training domains differ from unseen test domains, with applications in medical imaging.
• Propose a generative model (DIVA) that disentangles domain, class, and residual factors into separate latent subspaces.
• Enable semi-supervised learning by leveraging unlabeled data from known or unseen domains.
• Demonstrate disentanglement and improved generalization on rotated MNIST and malaria cell image datasets.
• Show that unlabeled data can boost performance and discuss interpolation vs extrapolation in domain spaces.

Proposed method

• Introduce three independent latent variables z_d (domain), z_y (class), and z_x (residual) with respective priors p(z_d|d), p(z_y|y), and p(z_x).
• Use three separate encoders q_phi_d, q_phi_y, q_phi_x to infer z_d, z_y, z_x from x, with a shared decoder p_theta(x|z_d,z_x,z_y).
• Optimize a beta-VAE-like lower bound with a reconstruction term and KL penalties for each latent: L_s = E[...] log p_theta(x|z_d,z_x,z_y) - beta[KL(q_phi_d(z_d|x)||p_theta_d(z_d|d)) + KL(q_phi_x(z_x|x)||p(z_x)) + KL(q_phi_y(z_y|x)||p_theta_y(z_y|y))].
• Include auxiliary objectives to predict domain from z_d and class from z_y to encourage disentanglement: F_DIVA = L_s + alpha_d E[log q_omega_d(d|z_d)] + alpha_y E[log q_omega_y(y|z_y)].
• Extend to semi-supervised DIVA by jointly training with labeled (d,x,y) and unlabeled (d,x) data, marginalizing over y and incorporating an auxiliary classifier for y on z_y and an objective that blends supervised and unsupervised terms as in equation (4).

Experimental results

Research questions

• RQ1 Can a three-branch latent space (domain, class, residual) in a VAE disentangle domain-specific and class-specific information to improve generalization to unseen domains?
• RQ2 Does incorporating unlabeled data from known or unseen domains boost domain generalization performance in DIVA?
• RQ3 How does DIVA compare to domain adversarial and other domain generalization methods on benchmarks like rotated MNIST and malaria cell images?

Key findings

• DIVA achieves higher test accuracy than DA, LG, HEX, and ADV across rotated MNIST test domains.
• Incorporating unlabeled data (+1, +3, +5, +9) generally improves performance, with diminishing gains when unlabeled data dominates labeled data.
• Latent subspaces are disentangled: z_d captures domain, z_y captures class, z_x captures residual variation, enabling conditional reconstructions and sample generation.
• On malaria cell images, DIVA improves ROC AUC over baselines in several domains and benefits from unlabeled data in semi-supervised setups.
• DIVA can learn from unlabeled data in new domains, updating y-predictor and domain encoders to improve generalization.