DIVA: D OMAIN I NVARIANT V ARIATIONAL A UTOENCODER
A BSTRACT
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 VAE (DIVA), a generative model that tackles this problem by learning three independent latent subspaces, one for the class, one for the domain and one for the object itself.
In addition, we highlight that due to the generative nature of our model we can also incorporate unlabeled data from known or previously unseen domains.
This property is highly desirable in ﬁelds like medical imaging where labeled data is scarce.
We experimentally evaluate our model on the rotated MNIST benchmark 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
I NTRODUCTION
Deep neural networks (DNNs) led to major breakthroughs in a variety of areas like computer vision and natural language processing.
Despite their big successes recent research shows that DNNs learn the bias present in the training data.
As a result they are not invariant to cues that are irrelevant to the actual task (Azulay & Weiss, 2018).
This leads to a dramatic performance decrease when tested on data from a different distribution with a different bias (Torralba & Efros, 2011).
In domain generalization the goal is to learn representations from a set of similar distributions, here called domains, that can be transferred to a previously unseen domain during test time.
A common motivating application, where domain generalization is crucial, is medical imaging (Blanchard et al., 2011; Muandet et al., 2013).
For instance, in digital histopathology a typical task is the classiﬁcation of benign and malignant tissue.
However, the preparation of a histopathology image includes the staining and scanning of tissue which can greatly vary between hospitals.
Moreover, a sample from a patient could be preserved in different conditions.
As a result, each patient data could be treated as a separate domain (Lafarge et al., 2017; Ciompi et al., 2017).
Another problem commonly encountered in medical imaging is class label scarcity.
Annotating medical images is an extremely time consuming task that requires expert knowledge.
However, obtaining domain labels is surprisingly cheap, since hospitals generally store information about the patient (e.g., age and sex) and the medical equipment (e.g., manufacturer and settings).
Therefore, we are interested in extending the domain generalization framework to be able to deal with additional unlabeled data.
We hypothesize that additional unlabeled data can lead to better domain generalization.
In this paper, we propose to tackle domain generalization via a new deep generative model that we refer to as the domain invariant variational autoencoder (DIVA).
We extend the variational autoencoder (VAE) framework (Kingma & Welling, 2013) by introducing independent latent representations for an object (e.g., an image), a class label and a domain label.
Such partitioning of the latent space will encourage and guide the model to disentangle these sources of variation.
Finally, by virtue of having a generative model we can also naturally handle the semi-supervised scenario, similarly to Kingma et al.
We evaluate our model on a version of the MNIST dataset where each domain corresponds to a speciﬁc rotation angle of the digits.
[Figure]
According to the graphical model we obtain
We propose to factorize the variational represent auxiliary classiﬁers.
T OWARDS DOMAIN GENERALIZATION WITH GENERATIVE MODELS 2. D OMAIN GENERALIZATION
We follow the domain generalization deﬁnitions used in (Muandet et al., 2013).
A domain is deﬁned as a joint distribution p ( x, y ) on X × Y , where X denotes an input space and Y an output space.
Let P X×Y be the set of all domains.
We assume that we cannot observe domains directly.
Instead, the training set consists of samples S taken from N domains, S = { S ( d = i ) } Ni =1 .
The i th domain p ( d = i )
( x, y ) is represented by n i samples, S ( d = i ) =
{ ( x ( d = i ) k , y ( d = i ) k ) } n i k =1 .
Here each of the N distributions p ( d =1) ( x, y ) , . . .
, p ( d = i ) ( x, y ) , . . .
, p ( d = N ) ( x, y ) is sampled from P X×Y .
We further assume that p ( d = i ) ( x, y ) (cid:54) = p ( d = j ) ( x, y ) , therefore, the samples in S are non-i.i.d.
During test time we are presented with samples S ( d = N +1) from a previously unseen domain p ( d = N +1)
( x, y ) .
We are interested in learning representations that generalize from p ( d =1) ( x, y ) , . .
, p ( d = N ) ( x, y ) to this new domain.
Training data are given as tuples ( d, x, y ) in the case of supervised data or as ( d, x ) in the case of unsupervised data.
2. DIVA: D OMAIN I NVARIANT VAE
Assuming a perfectly disentangled latent space (Higgins et al., 2018), we hypothesize that there exists a latent subspace that is invariant to changes in d , i.e., that is domain invariant.
We propose a generative model with three independent sources of variation; z x , which is object speciﬁc, z d , which is domain speciﬁc and ﬁnally z y , which is class speciﬁc.
While z x keeps an independent Gaussian prior p ( z x ) , z d and z y have conditional priors p θ
d
( z d | d ) , p θ y ( z y | y ) with learnable parameters θ d , θ y .
This will encourage information about the domain d and label y to be encoded into z d and z y respectively.
Furthermore, as z d and z y are marginally independent by construction, we argue that the model will learn representations z y that are invariant with respect to the domain d .
All three of these latent variables are then used by a single decoder p θ ( x | z d , z x , z d ) for the reconstruction of x .
Since we are interested in using neural networks for p θ ( x | z d , z x , z d ) , exact inference will be intractable.
For this reason, we perform amortized variational inference with an inference network (Kingma & Welling, 2013), i.e., we employ a variational autoencoder (VAE) framework.
We introduce three separate encoders q φ d ( z d | x ) , q φ x ( z x | x ) and q φ y ( z y | x ) that serve as variational posteriors over these latent variables.
Notice that we do not share their parameters as we empirically found that sharing parameters leads to a decreased generalization performance.
For the prior and variational posterior distributions over the latent variables z x , z d , z y we assume fully factorized Gaussians with parameters given as a function of their input.
We coin the term domain invariant VAE (DIVA) for our overall model, which can be seen in Figure 1.
Given a speciﬁc dataset, all of the aforementioned parameters can then be optimized by maximizing the following variational lower bound:
[Equation]
Notice that we have introduced a weigting term, β .
This is motivated by the β -VAE (Higgins et al., 2017) and serves as a constraint that controls the capacity of the latent spaces of DIVA.
Larger values of β limit the capacity of each z and in the ideal case each dimension of z captures one of the conditionally independent factors in data.
The β -VAE framework offers a trade-off between the information preservation, i.e., how well one can reconstruct x from the z ’s, and the capacity, i.e., how well the z ’s compress information about x .
In order to further encourage separation of z d and z y into domain and class speciﬁc information respectively, we add two auxiliary objectives.
During training z d is used to predict the domain d and z y is used to predict the class y for a given input x
[Equation]
where α d , α y are weighting terms for each of these auxiliary objectives.
Since our main goal is a domain invariant classiﬁer, during test time we are left with the encoder q φ y
( z y | x ) and the auxiliary classiﬁer q ω y ( y | z y ) .
For predicting the class y of a new input x we only use the mean of z y .
2. G UIDED DISENTANGLEMENT
Locatello et al.
and Dai & Wipf (2019) claim that learning a disentangled representation, i.e., q φ ( z ) =
(cid:81)
i q φ ( z i ) , in an unsupervised fashion is impossible for arbitrary generative models.
Inductive biases, e.g., some form of supervision or constraints on the latent space, are necessary to ﬁnd a speciﬁc set of solutions that matches the true generative model.
Consequently,
DIVA is using domain labels d and class labels y in addition to input data x during training.
Furthermore, we enforce the factorization of the marginal distribution of z in the following form:
[Equation]
, which prevents the impossibility described at Locatello et al.
We argue that the strong inductive biases in DIVA make it possible to learn disentangled representations that match the ground truth factors of interest, namely, the domain d and class y .
2. S EMI - SUPERVISED DIVA
In (Kingma et al., 2014) an extension to the VAE framework was introduced that allows to use labeled as well as unlabeled data during training.
While Kingma et al.
(2014) introduced a two step procedure, Louizos et al.
(2015) presented a way of optimizing the decoder of the VAE and the auxiliary classiﬁer jointly.
We use the latter approach to learn from supervised data { ( d n , x n , y n ) } as well as from unsupervised data { ( d m , x m ) } .
Analogically to (Louizos et al., 2015), we use q ω y
( y | z y ) to marginalize out y :
[Equation]
where we use Monte Carlo sampling with the reparametrization trick (Kingma & Welling, 2013) for the continuous latents z d , z x , z y and explicitly marginalize over the discrete variable y .
The ﬁnal objective combines the supervised and unsupervised variational lower bound as well as the two auxiliary objectives.
By assuming N labeled and M unlabeled data tuples we arrive at the following objective
[Equation]
R ELATED W ORK
Domain generalization The majority of recently proposed deep learning methods for domain generalization falls into one of two categories: 1) Learning a single domain invariant representation, e.g., using adversarial methods (Carlucci et al., 2018; Ghifary et al., 2015; Li et al., 2018; 2017;
Motiian et al., 2017; Shankar et al., 2018; Wang et al., 2019).
Ensembling models, each trained on an individual domain from the training set (Ding & Fu, 2018; Mancini et al., 2018).
Multi-task learning Zamir et al.
(2018) show that sharing parameters among multiple tasks can lead to better performance and can decrease the amount of labeled examples necessary for each task.
Even though we make use of multiple tasks, our goal is the opposite.
The three encoders of DIVA are forced to learn complementary features.
For a general introduction to multi-task learning we refer to Ruder (2017).
Fairness The goal of fair classiﬁcation is to learn good representation that at the same time cannot be used to associate a data sample to a certain group (Zemel et al., 2013).
The only difference to domain generalization is the intention behind that goal, e.g., to protect groups of individuals vs. being robust to technical variations.
Consequently, DIVA is closely related to the fair VAE (Louizos et al., 2015).
The major difference lies in the fashion we replace the hierarchical latent space of the fair VAE with a partitioned latent space.
Latent subspaces A series of very recent papers demonstrate the beneﬁts of having a partitioned latent space in a VAE.
Klys et al.
(2018) show that two latent supaces can be better disentangled using binary labels and adversarial training.
In (Bouchacourt et al., 2018) two latent subspaces exist as well.
Here one is used for individual samples and one is shared among samples from the same group.
Last, Hsu & Glass (2018) show that latent subspaces can be used to capture information from different modalities (e.g., image and audio).
E XPERIMENTS 4. R OTATED MNIST
The construction of the rotated MNIST dataset follows (Ghifary et al., 2015).
We sample 100 images from each of the 10 classes from the original MNIST training dataset (in contrast to (Ghifary et al., 2015)
we do not resize the images).
This set of images is denoted M 0 ◦ .
To create ﬁve additional domains the images in M 0 ◦ are rotated by 15, 30, 45, 60 and 75 degrees.
In order to evaluate their domain generalization abilities models are trained on ﬁve domains and tested on the remaining 6th domain, e.g., train on M 0 ◦ , M 15 ◦ , M 30 ◦ , M 45 ◦ and M 60 ◦ , test on M 75 ◦ .
The evaluation metric is the classiﬁcation accuracy on the test domain.
All experiments are repeated 10 times with 10 different seeds, resulting in 10 different datasets.
Detailed information about hyperparameters, architecture and training schedule can be found in the Appendix.
4. Q UALITATIVE DISENTANGLEMENT
First of all, we visualize the three latent spaces z d , z x and z y , to see if DIVA is able to successfully disentangle them.
In addition, we want to see if DIVA utilizes z x in a meaningful way, since it is not directly connected to any downstream task.
For the following visualizations we restrict the size of each latent space z d , z x and z y to 2 dimensions.
Therefore, we can plot the latent supspaces directly without applying dimensionality reduction.
DIVA is trained on 5000 images from ﬁve domains ( M 0 ◦ , M 15 ◦ , M 30 ◦ , M 45 ◦ and M 60 ◦ ).
Figure 2 shows 5000 embeddings z y encoded by q φ y ( z y | x ) .
In Figure 2 (left) the colors indicate the 10 different classes of the MNIST dataset.
We observe 10 well separated clusters, each corresponding to one of the 10 classes.
In stark contrast, Figure 2 (right), where the colors indicate the ﬁve different training domains, shows no such clustering.
It appears that z y is indeed capturing all necessary information for predicting the class y while containing very little information about the domain d .
In Figure 3 we visualize the two-dimensional latent space for z d .
Each of the 5000 training images is encoded by q φ d ( z d | x ) .
Figure 3 (left) shows very little clustering according to the class label y .
However, Figure 3 (right), where each color represents a different domain, shows strong clustering.
Each cluster corresponds to one of the ﬁve domains.
In contrast to z y and z d , z x is only used for reconstructing x and has an independent Gaussian prior.
In Figure 4 (left) we can see that there is a certain amount of clustering.
We ﬁnd that narrow numbers with thin lines, e.g., ’1’ and ’7’, cluster in the bottom left half of the plot.
Whereas round numbers with thicker lines, e.g., ’0’ and ’6’ appear to cluster in the top right half.
We conclude that z x models the remaining variations that are not captured by z d and z y .
In Figure 4 (right) we do not notice any apparent clustering.
From these initial qualitative results we conclude that DIVA is disentangling the information contained in x as intended, as z y is only containing information about y and z d only information about d .
In the case of the rotated MNIST dataset z x captures information about line thickness and digit width, two factors of variation that are not correlated with either the class y or the domain d .
We
[Table]
also can perform conditional reconstructions with DIVA.
The results along with more details can be found in the Appendix.
4. C OMPARISON TO OTHER METHODS
We compare DIVA against the well known domain adversarial neural networks (DA) (Ganin et al., 2015) as well as three recently proposed methods: LG (Shankar et al., 2018), HEX (Wang et al., 2019) and ADV (Wang et al., 2019).
For the ﬁrst half of Table 1 (until the vertical line) we only use labeled data.
The ﬁrst column indicates the rotation angle of the test domain.
We report test accuracy on y for all methods.
For
DIVA we report the mean and standard error for 10 repetitions.
DIVA achieves the highest accuracy across all test domains.
In addition we achieve the highest average test accuracy among all proposed methods.
The second half of Table 1 showcases the ability of DIVA to use unlabeled data.
For this experiment we add: The same amount (+1) of unlabeled data as well as three (+3), ﬁve (+5) and nine (+9) times the amount of unlabeled data to our training set.
Here, we ﬁrst add the unlabeled data to M 0 ◦ and create the data for the other domains as described in Section 4.1.
In Table 1 we can clearly see a performance increase when unlabeled data is added to the training set.
The effect seems to become smaller when the amount of unlabeled data is much larger than the amount of labeled data as seen in the last two columns of Table 1.
4. A DDITIONAL UNLABELED DOMAINS
In Section 4.3 we show that the performance of DIVA increases when it is presented with additional unlabeled data for each domain.
As a result each training domain consists of labeled and unlabeled examples.
In this section we investigate a more challenging scenario: We add an additional domain to our training set that consists of only unlabeled examples.
Coming back to our introductory example of medical imaging, here we would we add unlabeled data from a new patient or hospital to the training set.
In contrast to the experiment in Section 4.3 where we would add unlabeled data from each known patient or hospital to the training set.
In the following, we are looking at two different experimental setups: 1)
The additional domain is dissimilar to the test domain, e.g. M 30 ◦ and M 75 ◦ .
The additional domain is similar to the test domain, e.g. M 60 ◦ and M 75 ◦ .
In both cases we show that DIVA improves when trained with unlabeled data from an unseen domain.
4.4. D ISSIMILAR DOMAIN
For the ﬁrst experiment we choose the domains M 0 ◦ , M 15 ◦ , M 45 ◦ and M 60 ◦ to be part of the labeled training set.
In addition, unlabeled data from M 30 ◦ is used.
The test domain is M 75 ◦ .
In
Table 2 we can see that even in the case where the additional domain is dissimilar to the test domain
DIVA is able to slightly improve.
[Table]
S IMILAR DOMAIN
For the second experiment we choose the domains M 0 ◦ , M 15 ◦ , M 30 ◦ and M 45 ◦ to be part of the labeled training set.
In addition, unlabeled data from M 60 ◦ is used.
The test domain is M 75 ◦ .
When comparing the results in Table 3 to the results in Table 1 and 2 we notice a drop in accuracy of about 20 % for DIVA trained with only labeled data.
However, when trained with unlabeled data from M 60 ◦ the drop in accuracy is only about 13 % .
[Table]
C ONCLUSION
We have proposed DIVA as a generative model with three independent sources of variation.
We demonstrate through quantitative and qualitative experiments on rotated MNIST that our model successfully learns representations z y that are invariant with respect to the domain d .
In future work, we want to evaluate DIVA on a more complex dataset.
A CKNOWLEDGMENTS
The authors thank Patrick Forr´e, Rianne van den Berg, Marco Federici, Daniel Worrall and Bas
Veeling for helpful discussions and comments.
Maximilian Ilse was funded by the Nederlandse Organisatie voor Wetenschappelijk Onderzoek (Grant DLMedIa: Deep Learning for Medical Image Analysis)
A PPENDIX T RAINING PROCEDURE
All DIVA models are trained for 500 epochs.
The training is terminated if the classiﬁcation loss for y has not improved for 100 epochs.
As proposed in Burgess et al.
(2018), we linearly increase β from 0.0 to 1.0 during the ﬁrst 100 epochs of training.
We set α d = 2000.
As seen in (Maale et al., 2019), we adjust α y according to the ratio of labeled (N) and unlabeled data (M),
[Equation]
(5) where we set γ = 3500 .
Last, z d , z x and z y each have 64 latent dimensions.
All hyperparameters were determined by training DIVA on M 0 ◦ , M 15 ◦ , M 30 ◦ , M 45 ◦ and testing on M 60 ◦ .
All models were trained using ADAM (Kingma & Ba, 2014) (with default settings), a pixel-wise cross entropy loss and a batch size of 100.
A RCHITECTURE
To enable a fair experiment, the encoder q φ y ( z y | x ) and auxiliary classiﬁer q ω y ( y | z y ) form a DNN with the same number of layers and weights as described in Wang et al.
[Table]
Linear
[Table]
), BatchNorm1d,
[Table] [Table]
[Table] [Table]
, BatchNorm2d,
[Table]
, BatchNorm2d,
[Table] [Table]
[Table] [Table]
C ONDITIONAL GENERATION
Yet another way to gain insight into the disentanglement abilities of DIVA is conditional generation.
We ﬁrst train DIVA with β = 10 using M 0 ◦ , M 15 ◦ , M 30 ◦ , M 45 ◦ and M 60 ◦ as training domains.
After training we perform two experiments.
In the ﬁrst one we are ﬁxing the class and varying the domain.
In the second experiment we are ﬁxing the domain and varying the class.
Change of class
The ﬁrst row of Figure 5 (left) shows the input images x for DIVA.
First, we generate embeddings z d , z x and z y for each x using q φ d ( z d | x ) , q φ x ( z x | x ) and q φ y ( z y | x ) .
Second, we replace z y with a sample z (cid:48)
y from the conditional prior p θ y ( z y | y ) .
Last, we generate new images from z d , z x and z (cid:48)
y using the trained encoder p θ ( x | z d , z x , z y ) .
In Figure 5 (left) rows 2 to 11 correspond to the classes ’0’ to ’9’.
We observe that the rotation angle (encoded in z d ) and the line thickness (encoded in z x ) are well preserved, while the class of the image is changing as intended.
Change of domain We repeat the experiment from above but this time we keep z x and z y ﬁxed while changing the domain.
After generating embeddings z d , z x and z y for each x in the ﬁrst row of Figure 5 (right), we replace z d with a sample z (cid:48)
d from the conditional prior p θ d ( z d | d ) .
Finally, we generate new images from z (cid:48)
d , z x and z y using the trained encoder p θ ( x | z d , z x , z y ) .
In Figure(right) rows 2 to 6 correspond to the domains M 0 ◦ to M 60 ◦ .
Again, DIVA shows the desired behaviour: While the rotation angle is changing the class and style of the original image is maintained.
Q UALITATIVE D ISENTANGLEMENT : T EST DOMAIN
In this section, we visualize the z d and z y for data points x from the test domain M 75 ◦ for the model trained in Section 4.2.
Figure 6 shows 1000 embeddings z y encoded by q φ y ( z y | x ) . Figure 7 shows 1000 embeddings z d encoded by q φ d ( z d | x ) .