Invariant Causal Mechanisms through Distribution Matching

1 INTRODUCTION . Learning structured representations which capture the underlying causal mechanisms generating data is of central importance for training robust machine learning models ( Bengio et al. , 2013 ; Schölkopf et al. , 2021 ) . One particular structure the learned representation should capture is invariance to changes in nuisance variables . For example , we may want the representation to be invariant to sensitive attributes such as the race or gender of an individual in order to avoid discrimination or biased decision making in a downstream task ( Creager et al. , 2019 ; Locatello et al. , 2019 ; Träuble et al. , 2021 ) . While learning invariant representations is thus highly important for fairness applications , it also appears in seemingly unrelated tasks such as domain adaptation ( DA ) and domain generalization ( DG ) , where one aims to be invariant across the different domains ( Muandet et al. , 2013 ; Zemel et al. , 2013 ; Ganin et al. , 2016 ; Peters et al. , 2016 ) . For tasks such as DA and DG invariance across domains or environments implies to being invariant to the domain index , which thus is the `` sensitive attribute '' in this case and typically implies a change in the distribution of the data generating process . Being invariant to the domain index is thus a proxy to being invariant to latent unobserved factors that can change in distribution . Established approaches for enforcing invariance in the learned representation usually aim to learn a representation whose statistical distribution is independent to the sensitive attribute e.g. , by including an adversary during training ( Ganin et al. , 2016 ; Xie et al. , 2017 ) . As an adversary is essentially a parametric distributional distance , other approaches minimize different distribution distances , such as maximum mean discrepancy ( MMD ) ( Louizos et al. , 2017 ; Li et al. , 2018b ) , or optimal transport ( OT ) based distances ( Shen et al. , 2018 ; Damodaran et al. , 2018 ) . To enforce independence , these methods add a regularizer to the loss that consists in the pairwise distributional distance between all possible combination of the sensitive attribute , i.e. , dist ( p ( z|d ) , p ( z|d′ ) ) ∀d , d′ ∈ D. As such , the complexity of the loss grows quadratically in the size of the support of the sensitive attribute , which can limit the applicability of these models when the support of D is large ( Koh et al. , 2021 ) . Despite the importance of learning invariant representations and their potential societal impact in the medical domain or fair decision making , most established approaches are still based on heuristics and specialized for different tasks at hand . We take first steps towards a unifying framework by viewing invariance as a property of a causal process ( Pearl , 2009 ; Peters et al. , 2017 ) and our key contributions can be summarized as follows : • We introduce a unifying causal framework for invariant representation learning , which allows us to derive a new algorithm to enforce invariance through distribution matching . One advantage of our algorithm is that only one distributional distance between two batches needs to be computed at each step , irrelevant of the size of the support of D. • We define the notion of style variable and present some necessary and sufficient conditions under which being invariant to the domain index actually leads to invariance to the style variables . We argue that our proposal naturally captures most of the existing invariant representation learning tasks and datasets . • Finally , we conduct a large number of experiments across different tasks and datasets , demonstrating the versatility of our framework . We obtain competitive results on the task of learning fair representations and we are able to significantly boost the performance of existing models using our proposed algorithm for the task of DG . 2 INVARIANT REPRESENTATION LEARNING ACROSS TASKS . In this section , we highlight how the learning of an invariant representation is a goal that is ( implicitly ) pursued in a large spectrum of machine learning tasks . Domain Adaptation The range of techniques used in Domain Adaptation and the different assumptions followed are vast ( see Wilson and Cook ( 2020 ) for a more in depth review ) . Thus , we here concentrate only on a subset of the literature . A direction that is widely followed in DA , and which is the closest to our framework , is the alignment of the latent distribution of the source and target datasets . Under the covariate shift assumption , which assume that the labeling function P ( Y |X ) is fixed , and that only P ( X ) varies across environments , the goal is then to learn a representation h ( X ) that is invariant across source and target and that remains useful to learn a discriminator on the source dataset . Ganin et al . ( 2016 ) uses a domain adversarial network to align the two latent spaces , whereas others uses distributional divergences directly , such as MMD ( Baktashmotlagh et al. , 2016 ) , Wasserstein and optimal transport in general ( Shen et al. , 2018 ; Damodaran et al. , 2018 ; Redko et al. , 2017 ) . DA under different assumptions , such as the case where both P ( Y ) and P ( X|Y ) , have also been studied ( Gong et al. , 2016 ) . Domain Generalization Though very similar to DA , DG differs in one significant way : the test domain is not observed at training time . As such , it is a way harder task as the test domain could exhibit arbitrary shifts in distribution , and the learned model is supposed to handle any reasonable shifts in distribution . Without any assumptions , there is little hope to obtain models that actually generalizes . Nevertheless , many inductive biases and models have been proposed , which have stronger assumptions than classical empirical risk minimization ( ERM ) ( Vapnik , 1998 ) . Given its similarity to DA , similar models have been proposed , and most models work for both tasks . Nevertheless , until recently ( Albuquerque et al. , 2019 ; Deng et al. , 2020 ) , theoretical justification , e.g. , for minimizing the distance between pairs of latent variables coming from different domains , was missing , as results from domain adaptation assumes that the test domain is observed . Without some assumptions , there exists no theoretical reasons to infer that a constant distribution of the latent variables across the training domains leads to better generalization on the test domains . Indeed , many benchmarks ( Gulrajani and Lopez-Paz , 2020 ; Koh et al. , 2021 ) show that it is difficult to create algorithms that consistently beat ERM across different tasks . Invariant representations for DG was first proposed by Muandet et al . ( 2013 ) . This idea was then extended to use other distributional distances , such as MMD ( Li et al. , 2018b ) , Adversarial ( Li et al. , 2018d ; Deng et al. , 2020 ; Albuquerque et al. , 2019 ) , and Optimal Transport ( Zhou et al. , 2020 ) ( see Table 1 ) . On the theoretical side , both Albuquerque et al . ( 2019 ) and Deng et al . ( 2020 ) attempt to give theoretical grounding to the use of an adversarial loss by deriving bounds similar to what exists in DA . Domain Generalization and Causal Inference Many links between causal inference and domain generalization have been made , arguing that domain generalization is inherently a causal discovery task . In particular , causal inference can be seen as a form of distributional robustness ( Meinshausen , 2018 ) . In regression , one way of ensuring interventional robustness is by identifying the causal parents of Y , whose relation to Y is stable . This can be achieved by finding a feature representation such that the optimal classifiers are approximately the same across domains ( Peters et al. , 2016 ; Rojas-Carulla et al. , 2018 ; Arjovsky et al. , 2019 ) . Unfortunately , most of these models do not really apply to classification of structured data such as images , where the classification is predominantly Table 1 : Review of invariance across different tasks . Note that the general loss is defined as 1 n ∑n i=0 L ( xi , yi ) + λ · ( dist ( zn1 , z N n+1 ) ) . Task Adversarial MMD Wasserstein Distance equation 1n ∑n i=0 log 1 Gd ( zi ) + 1n′ ∑N i=n+1 log 1 1−Gd ( zi ) ∥∥∥ 1n∑ni=0 φ ( zi ) − 1n′ ∑Ni=n+1 φ ( zi ) ∥∥∥H 1n∑ni=0Gd ( zi ) − 1n′ ∑Ni=n+1Gd ( zi ) Domain Adaptation Ganin et al . ( 2016 ) Baktashmotlagh et al . ( 2016 ) Shen et al . ( 2018 ) Hoffman et al . ( 2017 ) Damodaran et al . ( 2018 ) Domain Generalization Ganin et al . ( 2016 ) ; Albuquerque et al . ( 2019 ) Li et al . ( 2018b ) Zhou et al . ( 2020 ) Li et al . ( 2018d ; c ) ; Deng et al . ( 2020 ) Fair Representation Learning Edwards and Storkey ( 2015 ) ; Xie et al . ( 2017 ) Louizos et al . ( 2017 ) Jiang et al . ( 2020 ) Roy and Boddeti ( 2019 ) anti-causal and where the wanted invariance is not toward the pixels themselves but towards the unobserved generating factors . In a similar setting to ours , Heinze-Deml and Meinshausen ( 2021 ) tackles the task of image classification and propose a new model . A significant difference to our work is that they rely on the observation of individual instances across different views , i.e. , the images are clustered by an ID . Fair Representation Learning Fair representation learning can also be viewed as an invariant representation learning task . This task consists in learning a representation that maximizes usefulness towards predicting a target variable , while minimizing information leakage of a sensitive attribute ( e.g. , gender , age , race ) . The seminal work in this field is Zemel et al . ( 2013 ) , which aims at learning a multinomial random variable Z , with associated vectors vk , such as the representation Z is fair . More recent work directly learns a continuous variable Z that has minimal information about the sensitive attribute , either through minimizing the MMD distance ( Louizos et al. , 2017 ) , through adversarial training ( Edwards and Storkey , 2015 ; Xie et al. , 2017 ; Roy and Boddeti , 2019 ) , or through a Wasserstein distance ( Jiang et al. , 2020 ) . 3 INVARIANCE AS THE PROPERTY OF A CAUSAL PROCESS . In this section , we will first consider the assumptions for the causal process underlying the data generating mechanism using a structural causal model ( SCM ) type graph from Causality theory ( Pearl , 2009 ) and following the causal view of learning disentangled representations ( Suter et al. , 2019 ) , as illustrated in Figure 1 . G1 toGk represents all the factors of variation that generate the data , i.e. , there exists a ( one-to-one ) function such that given all the factors , X is fixed : X ←− g ( G1 , . . . , Gk ) Y is a target value that we may want to predict in a downstream task and is either known ( supervised setting ) or unobserved ( unsupervised ) . D is another confounder that we want to be invariant to . It can be a domain index , such as in DA and DG , or a sensitive attribute such as in fairness . We will assume for now that D does not have an effect on Y . Lastly , the generative factors Gi are assumed to not have any causal relations between them , and any correlation between some factors may only come from a hidden confounder . This assumption is similar to the assumptions in Suter et al . ( 2019 ) . Furthermore , in this work , we assume that the label Y and D directly have an effect on the latent generating factors . In this setting , Y and D are thus independent . Given our data generating framework , we can now give some definitions , especially the notion of style generating factors . Definition 3.1 . We call style variables the set of variables G that are children of D in the DAG . We denote this set S. Observation 3.1 . X and Z are independent from D given S , as they are d-separated from D by the set S in the graph . To the best of our knowledge there is no consistent and widely accepted definition of an invariant representation , yet . Using the above framework , we propose the following definition : Definition 3.2 . We say that a representation Z is invariant to a variable D if and only if D has no total causal effect1on Z . This definition of invariance is very robust since it guarantees that no intervention on the variable D can break the independence between Z and D. This is particularly relevant in application such as fair and private representation learning , as we may not want that intervening on the distribution of a sensitive variable breaks the property of fairness or privacy of a representation . The goal of invariant representation learning can then be described as creating a new variable Z = f ( X ) such that D has no total causal effect on Z . In a way , we can view it as adding a new variable in the SCM and learning its structural equation . If we assume that our distribution follows our proposed SCM ( Figure 1 ) , then absence of total causal effect is equivalent to independence , as D has no parent in the DAG . We use this assumption of D having no causal parents in Theorem 3.1 . Theorem 3.1 . Under the assumption of the graph in Figure 1 , we have that : Z is independent from D ( equivalently , D has no total causal effect on Z , or p ( z|d ) = p ( z|d′ ) for all d , d′ ) ⇐⇒ p ( z|do ( d = Nd ) ) = p ( z ) for all Nd ( intervention on the distribution of D ) . In summary , Theorem 3.1 states that no total causal effect is equivalent to independence given the right assumptions , and that having a constant marginal distribution of Z under different mixtures of D leads to independence . The proof is presented in Appendix B . In the case where the full support of D is observed during training , we have the guarantee that independence in training I ( Z ; D ) = 0 will hold in any test setting . However , in DG for example , we observe a new value of D at test time . Having p ( z|dtraini ) = p ( z|dtrainj ) ∀i , j does not guarantee that p ( z|dtest ) = p ( z|dtraini ) . Indeed , in this setting , the variable D works more like an index , where each value indicates a domain where the distribution of X has changed . Invariance to the variable D is thus a proxy to being invariant to the unobserved style variables ( see Definition 3.1 ) .

Summary: This work provides a causal perspective and new algorithm for learning invariant representations from multiple domain datasets. This work introduces the notion of style variables and shows theoretically that being invariant to the domain index actually leads to invariance to the style variables. The extensive experiments demonstrate the usefulness of the proposed algorithm in various tasks.