DIVA: Domain Invariant Variational Autoencoder

1 INTRODUCTION . Deep neural networks ( DNNs ) led to major breakthroughs in a variety of areas like computer vision and natural language processing . Despite their big success , 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 . 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 classification 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 ( Ciompi et al. , 2017 ) . As a result , each patient ’ s data could be treated as a separate domain ( Lafarge 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 , as we hypothesize that it can help to improve performance . 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 ; Rezende et al. , 2014 ) by introducing independent latent representations for a domain label , a class label and any residual variations in the input x . 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 naturally handle the semi-supervised scenario , similarly to Kingma et al . ( 2014 ) . We evaluate our model on a version of the MNIST dataset where each domain corresponds to a specific rotation angle of the digits , as well as on a Malaria Cell Images dataset where each domain corresponds to a different patient . An implementation of DIVA can be found under ( URL was removed to preserve anonymity ) . 2 TOWARDS DOMAIN GENERALIZATION WITH GENERATIVE MODELS . We follow the domain generalization definitions used in Muandet et al . ( 2013 ) . A domain is defined as a joint distribution p ( x , y ) on X × Y , where X denotes the input space and Y denotes the output space . Let PX×Y be the set of all domains . The training set consists of samples S taken from N domains , S = { S ( d=i ) } Ni=1 . Here , the ith domain p ( d=i ) ( x , y ) is represented by ni samples , S ( d=i ) = { ( x ( d=i ) k , y ( d=i ) k ) } ni k=1 . Each of the N distributions p ( d=1 ) ( x , y ) , . . . , p ( d=i ) ( x , y ) , . . . , p ( d=N ) ( x , y ) are sampled from PX×Y . We further assume that p ( d=i ) ( x , y ) 6= 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.1 DIVA : DOMAIN INVARIANT 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. , it is domain invariant . We propose a generative model with three independent sources of variation ; zd , which is domain specific , zy , which is class specific and finally zx , which captures any residual variations left in x . While zx keeps an independent Gaussian prior p ( zx ) , zd and zy have conditional priors pθd ( zd|d ) , pθy ( zy|y ) with learnable parameters θd , θy . This will encourage information about the domain d and label y to be encoded into zd and zy , respectively . Furthermore , as zd and zy are marginally independent by construction , we argue that the model will learn representations zy that are invariant with respect to the domain d. All three of these latent variables are then used by a single decoder pθ ( x|zd , zx , zy ) for the reconstruction of x . Since we are interested in using neural networks to parameterize pθ ( x|zd , zx , zy ) , exact inference will be intractable . For this reason , we perform amortized variational inference with an inference network ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) , i.e. , we employ a VAE-type framework . We introduce three separate encoders qφd ( zd|x ) , qφx ( zx|x ) and qφy ( zy|x ) that serve as variational posteriors over the 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 zx , zd , zy 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 is depicted in Figure 1 . Given a specific dataset , all of the aforementioned parameters can be optimized by maximizing the following variational lower bound per input x : Ls ( d , x , y ) = Eqφd ( zd|x ) qφx ( zx|x ) , qφy ( zy|x ) [ log pθ ( x|zd , zx , zy ) ] − βKL ( qφd ( zd|x ) ||pθd ( zd|d ) ) − βKL ( qφx ( zx|x ) ||p ( zx ) ) − βKL ( qφy ( zy|x ) ||pθy ( zy|y ) ) . ( 1 ) 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 x . To further encourage separation of zd and zy into domain and class specific information respectively , we add two auxiliary objectives . During training zd is used to predict the domain d and zy is used to predict the class y for a given input x : FDIVA ( d , x , y ) : = Ls ( d , x , y ) + αdEqφd ( zd|x ) [ log qωd ( d|zd ) ] + αyEqφy ( zy|x ) [ log qωy ( y|zy ) ] , ( 2 ) where αd , αy are weighting terms for each of these auxiliary objectives . Since our main goal is a domain invariant classifier , during inference we only use the encoder qφy ( zy|x ) and the auxiliary classifier qωy ( y|zy ) . For the prediction of the class y for a new input x we use the mean of zy . Consequently , we regard the the variational lower bound Ls ( d , x , y ) as a regularizer . Therefore , evaluating the marginal likelihood p ( x ) of DIVA is outside the scope of this paper . Locatello et al . ( 2018 ) and Dai & Wipf ( 2019 ) claim that learning a disentangled representation , i.e. , qφ ( z ) = ∏ i qφ ( zi ) , 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 find a specific 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 . Recent work by Khemakhem et al . ( 2019 ) shows that conditional priors , like pθd ( zd|d ) and pθy ( zy|y ) in DIVA , lead to identifiability guarantees in VAEs . Furthermore , we enforce the factorization of the marginal distribution of z in the following form : qφ ( z ) = qφd ( zd ) qφx ( zx ) qφy ( zy ) , which prevents the impossibility described in Locatello et al . ( 2018 ) . 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 factors zd and class factors zy . To highlight the importance of a partitioned latent space we compare DIVA to a VAE with a single latent space , the results of this comparison can be found in the Appendix . 2.2 SEMI-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 classifier jointly . We use the latter approach to learn from supervised data { ( dn , xn , yn ) } as well as from unsupervised data { ( dm , xm ) } . Analogically to Louizos et al . ( 2015 ) , we use qωy ( y|zy ) to impute y : Lu ( d , x ) = Eqφd ( zd|x ) qφx ( zx|x ) qφy ( zy|x ) [ log pθ ( x|zd , zx , zy ) ] − βKL ( qφd ( zd|x ) ||pθd ( zd|d ) ) − βKL ( qφx ( zx|x ) ||p ( zx ) ) + βEqφy ( zy|x ) qωy ( y|zy ) [ log pθy ( zy|y ) − log qφy ( zy|x ) ] + Eqφy ( zy|x ) qωy ( y|zy ) [ log p ( y ) − log qωy ( y|zy ) ] , ( 3 ) where we use Monte Carlo sampling with the reparametrization trick ( Kingma & Welling , 2013 ) for the continuous latents zd , zx , zy and explicitly marginalize over the discrete variable y . The final objective combines the supervised and unsupervised variational lower bound as well as the two auxiliary losses . Assuming N labeled and M unlabeled examples , we obtain the following objective : FSS-DIVA = N∑ n=1 FDIVA ( xn , yn , dn ) + M∑ m=1 Lu ( xm , dm ) + αdEqφd ( zd|xm ) [ log qωd ( dm|zd ) ] . ( 4 )

Summary: In this work, the authors propose a domain invariant variational autoencoder for domain generalization problem. Specifically, the data is assumed to be constructed from three independent variables, one for the domain, one for the class and one for the residual variations. The method can be used for both unsupervised and semi-supervised cases. Experimental studies on rotated MNIST dataset and a malaria cell images dataset verify the effectiveness of the proposed method.