Domain Invariant Representation Learning with Domain Density Transformations

1 Introduction . Domain generalization refers to the machine learning scenario where the model is trained on multiple source domains so that it is expected to generalize well to unseen target domains . The key difference between domain generalization [ 25 , 37 , 18 ] and domain adaptation [ 49 , 48 , 14 , 45 ] is that , in domain generalization , the learner does not have access to data of the target domain , making the problem much more challenging . One of the most common domain generalization approaches is to learn an invariant representation across domains , aiming at a good generalization performance on target domains . For instance , in the representation learning framework , the prediction function y = f ( x ) , where x is data and y is a label , is obtained as a composition y = h ◦ g ( x ) of a deep representation network z = g ( x ) , where z is a learned representation of data x , and a smaller classifier y = h ( z ) , predicting label y given representation z , both of which are shared across domains . With this framework , we can aim to learn an “ invariant ” representation z across the source domains with the “ hope ” of a better generalization to the target domain . Most existing “ domain-invariance ” -based methods in domain generalization focus on the marginal distribution alignment [ 37 , 1 , 44 , 43 , 32 ] , which are still prone to distributional shifts when the conditional data distribution is not stable . In particular , the marginal alignment refers to making the 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . representation distribution p ( z ) to be the same across domains . This is essential since if p ( z ) for the target domain is different from that of source domains , the classification network h ( z ) would face out-of-distribution data at test time . Conditional alignment refers to aligning the conditional distribution of the label given the representation p ( y|z ) to expect that the classification network ( trained on the source domains ) would give accurate predictions at test time . The formal definitions of these two types of alignment are discussed in Section 3 . In Figure 1 we illustrate an example where the representation z satisfies the marginal alignment but not the conditional alignment . Specifically , x is distributed uniformly on the circle with radius 2 ( and centered at the origin ) for domain 1 and distributed uniformly on the circle with radius 3 ( centered at the origin ) for domain 2 . The representation z defined by the mapping z = g ( x ) = x/||x||2 will align the marginal distribution p ( z ) , i.e. , z is now distributed uniformly on the unit circle for both domains . However , the conditional distribution p ( y|z ) is not aligned between the two domains ( y is represented by color ) , which means using this representation for classification is suboptimal , and in this extreme case would lead to 0 % accuracy in the target domain 2 . This is an extreme case of misalignment but it does illustrate the importance of the conditional alignment . Therefore , we need to align both the marginal and the conditional distributions for a domain-invariant representation . Recently , there have been several attempts [ 33 , 34 , 50 ] to align the joint distribution of the representation and the label p ( y , z ) in a domain generalization problem by aligning the distribution of z across domains for each class , i.e. , p ( z|y ) ( given that the label distribution p ( y ) is unchanged across domains ) . However , the key drawbacks of these methods are that they either do not scale well with the number of classes or have limited performance in real-world computer vision datasets ( see Section 5 ) . In this paper , we focus on learning a domain-invariant representation that aligns both the marginal and the conditional distributions in domain generalization problems . We present theoretical results regarding the necessary and sufficient conditions for the existence of a domain-invariant representation ; and subsequently propose a method to learn such representations by enforcing the invariance of the representation network under domain density transformation functions . A simple intuition for our approach is that if we enforce the representation to be invariant under the transformations among the source domains , the representation will become more robust under other domain transformations . Furthermore , we introduce an implementation of our method in practice , in which the domain transformation functions are learned through the training process of generative adversarial networks ( GANs ) [ 20 , 12 ] . We conduct extensive experiments on several widely used datasets and observe a significant improvement over relevant baselines . We also compare our methods against other state-of-the-art models and show that our method achieves competitive results . Our contribution in this work is threefold : • We revisit the domain invariant representation learning problem and shed some light by providing several observations : a necessary and sufficient condition for the existence of a domain-invariant representation and a connection between domain-independent representation and a marginally-aligned representation . • We propose a theoretically grounded method for learning a domain-invariant representation based on domain density transformation functions . We also demonstrate that we can learn the domain transformation functions by GANs in order to implement our approach in practice . • We empirically show the effectiveness of our method by performing experiments on widely used domain generalization datasets ( e.g. , Rotated MNIST , VLCS and PACS ) and compare our method with relevant baselines ( especially CIDG [ 33 ] , CIDDG [ 34 ] and DGER [ 50 ] ) . 2 Related Work . Domain generalization : Domain generalization is an seminal task in real-world machine learning problems where the data distribution of a target domain might vary from that of the training source domains . Therefore , extensive research has been developed to handle that domain-shift problem , aiming at a better generalization performance in the unseen target domain . A predominant approach for domain generalization is domain invariance [ 37 , 33 , 34 , 3 , 47 , 2 , 24 , 50 , 1 , 32 , 44 , 43 ] that learns a domain-invariant representation ( which we define as to align the marginal distribution of the representation or the conditional distribution of the output given the representation or both ) . We are particularly interested in CIDG [ 33 ] , CIDDG [ 34 ] and DGER [ 50 ] , which also learn a representation that aligns the joint distribution of the representation and the label given that the class distribution is unchanged across domains . It should be noted that Zhao et al . [ 50 ] assume the label is distributed uniformly on all domains , while our proposed method only requires an assumption that the distribution of label is unchanged across domains ( and not necessarily uniform ) . We also show later in our paper that the invariance of the distribution of class label across domains turns out to be the necessary and sufficient condition for the existence of a domain-invariant representation . Moreover , we provide a unified theoretical discussion about the two types of alignment , and then propose a method to learn a representation that aligns both the marginal and conditional distributions via domain density transformation functions for the domain generalization problem . Note that there exist several related works , such as Ajakan et al . [ 1 ] , Ganin et al . [ 17 ] , that use adversarial loss with a domain discriminator to align the marginal distribution of representation among domains , but they are different from our approach . In particular , our method only uses GANs or normalizing flows to learn the transformation among domains , and learn a representation that is invariant under these functions , without using an adversarial loss on the representation ( which can lead to very unstable training [ 19 , 27 ] ) . There also exist works [ 35 , 23 , 7 , 40 ] in the domain adaptation literature that use generative modeling to learn a domain transformation function from source to target images , and use the transformed images to train a classifier . Our method differs from these by enforcing the representation to be invariant under the domain transformation , and we show theoretically that the representation learned that way would be domain-invariant marginally and conditionally . Meanwhile , the above works use the domain transformation to transform the images and train the classifier directly on the transformed data , and are not effective or applicable for domain generalization . Another line of methods that received a recent surge in interest is applying the idea of meta-learning for domain generalization problems [ 16 , 4 , 31 , 5 ] . The core idea behind these works is that if we train a model that can adapt among the source domains well , it would be more likely to adapt to unseen target domains . Recently , there are approaches [ 15 , 9 , 41 ] that make use of the domain specificity , together with domain invariance , for the prediction problem . The argument here is that domain invariance , while being generalized well between domains , might be insufficient for the prediction of each specific domain and thus domain specificity is necessary . We would like to emphasize that our method is not a direct competitor of meta-learning based and domain-specificity based methods . In fact , we expect that our method can be used in conjunction with these methods to get the best of both worlds for better performance . Density transformation between domains : Since our method is based on domain density transformations , we will review briefly some related works here . To transform the data density between domains , one can use several types of generative models . Two common methods are based on GANs [ 51 , 12 , 13 ] and normalizing flows [ 21 ] . Although our method is not limited to the choice of the generative model used for learning the domain transformation functions , we opt to use GAN , specifically StarGAN [ 12 ] , for scalability . This is just an implementation choice to demonstrate the use and effectiveness of our method in practice , and it is unrelated to our theoretical results . Connection to contrastive learning : Our method can be interpreted intuitively as a way to learn a representation network that is invariant ( robust ) under domain transformation functions . On the other hand , contrastive learning [ 10 , 11 , 36 ] is also a representation learning paradigm where the model learns images ’ similarity . In particular , contrastive learning encourages the representation of an input to be similar under different transformations ( usually image augmentations ) . However , the transformations in contrastive learning are not learned and do not serve the purpose of making the representation robust under domain transformations . Our method first learns the transformations between domains and then uses them to learn a representation that is invariant under domain shifts . 3 Theoretical Approach . 3.1 Problem Statement . Let us denote the data distribution for a domain d ∈ D by p ( x , y|d ) , where the variable x ∈ X represents the data and y ∈ Y is its corresponding label . The graphical model for our domain generalization framework is depicted in Figure 2 , in which the joint distribution is presented as follows : p ( d , x , y , z ) = p ( d ) p ( y ) p ( x|y , d ) p ( z|x ) . ( 1 ) In the domain generalization problem , since the data distribution p ( x , y|d ) varies between domains , we expect the changes in the marginal data distribution p ( x|d ) or the conditional data distribution p ( y|x , d ) or both . In this paper , we assume that p ( y|d ) is invariant across domains , i.e. , the marginal distribution of the label y is not dependent on the domain d—this assumption is shown to be the key condition for the existence of a domain-invariant representation ( see Remark 1 ) . This is practically reasonable since in many classification datasets , the class distribution can be assumed to be unchanged across domains ( usually uniform distribution among the classes , e.g. , balanced datasets ) . Our aim is to find a domain-invariant representation z represented by the mapping p ( z|x ) that can be used for the classification of label y and be generalized among domains . In practice , this mapping can be deterministic ( in that case , p ( z|x ) = δgθ ( x ) ( z ) with some function gθ , where δ is the Dirac delta distribution ) or probabilistic ( e.g. , a normal distribution with the mean and standard deviation outputted by a network parameterized by θ ) . For all of our experiments , we use a deterministic mapping for an efficient inference at test time , while in this section , we present our theoretical results with the general case of a distribution p ( z|x ) . In most existing domain generalization approaches , the domain-invariant representation z is defined using one of the two following definitions : Definition 1 . ( Marginal Distribution Alignment ) The representation z is said to satisfy the marginal distribution alignment condition if p ( z|d ) is invariant w.r.t . d. Definition 2 . ( Conditional Distribution Alignment ) The representation z is said to satisfy the conditional distribution alignment condition if p ( y|z , d ) is invariant w.r.t . d. However , when the joint data distribution varies between domains , it is crucial to align both the marginal and the conditional distribution of the representation z . To this end , this paper aims to Domain 1 Domain 2 learn a representation z that satisfies both the marginal and conditional alignment conditions . We justify our assumption of independence between y and d ( thus p ( y|d ) = p ( y ) ) by the following remark , which shows that this assumption turns out to be the necessary and sufficient condition for learning a domain-invariant representation . Note that this condition is also used in several existing works [ 50 , 33 , 34 ] . Remark 1 . The invariance of p ( y|d ) across domains d is the necessary and sufficient condition for the existence of a domain-invariant representation ( that aligns both the marginal and conditional distributions ) . Proof . provided in the appendix . It is also worth noting that methods which learn a domain independent representation , for example , [ 24 ] , only align the marginal distribution . This comes directly from the following remark : Remark 2 . A representation z satisfies the marginal distribution alignment condition if and only if I ( z , d ) = 0 , where I ( z , d ) is the mutual information between z and d. Proof . provided in the appendix . The question still remains that how we can learn a non-trivial domain invariant representation that satisfies both of the distribution alignment conditions . This will be discussed in the following subsection .

Summary: This work focuses on solving domain generalization. Given labeled data of multiple source domains, the goal is to predict the labels of target domain data. In domain generalization, the target domain data are not provided at training time. The main challenge of domain generalization is to learn a domain-invariant representation. To solve the problem, most existing approaches focus on the marginal distribution alignment p(z|d) without considering the conditional distribution p(y|z,d). However, it is needed to align both the marginal and the conditional distributions to classify the target data. To address the challenge, they propose a theoretically grounded method for learning a domain-invariant representation based on domain density transformation functions. They enforce the representation to be invariant under the transformations among the source domains to the representation become more robust under other domain transformations. It aligns both the marginal and the conditional distributions. The transformation functions are based on GANs in practice.