Year: 2021

Abstract:
Learning representations that capture the underlying data generating process is a key problem for data efficient and robust use of neural networks. One key property for robustness which the learned representation should capture and which recently received a lot of attention is described by the notion of invariance. In this work we provide a causal perspective and new algorithm for learning invariant representations. Empirically we show that this algorithm works well on a diverse set of tasks and in particular we observe state-of-the-art performance on domain generalization, where we are able to significantly boost the score of existing models.

heading: 1 INTRODUCTION
text: 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.

heading: 2 INVARIANT REPRESENTATION LEARNING ACROSS TASKS
text: 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).

heading: 3 INVARIANCE AS THE PROPERTY OF A CAUSAL PROCESS
text: 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).

heading: 4 AN ALGORITHM FOR INVARIANT LATENT VARIABLE DISTRIBUTIONS
text: Necessary Condition for Invariant Representation Learning We first start by studying whether invariance to D may lead to invariance (or at least independence) to the style variables. For simplicity, suppose we are given two datasets drawn from two distributions p1 and p2 on X , i.e., D ∈ {1, 2}. We assume that only the style variables had their distribution changed across the two datasets.
The goal is to learn a representation of X that is invariant to the style variables. We thus need to learn an encoder f ∈ F , z = f(x). Theorem 4.1. Independence to D is a necessary condition for the representation to be invariant to the style variables.
The proof is trivial and can be found in Appendix B. As this theorem shows, having a representation z that is invariant to D is at least necessary for invariance to the style variables. Unfortunately, it is in general not sufficient. It may be sufficient under some assumptions on the distribution of the style variables, but this is a direction we will not follow as it often leads to a very restrictive setting.
1A variable i has a total causal effect on a variable k if and only if: Xi ⊥/ Xk in P do(Xk=Ñk) for some random variable Ñk (Definition 6.12 in Peters et al. (2017)).
Sufficient conditions for Invariant Representation Learning We here present an idealistic setting where invariance to D is sufficient. Obviously, this setting is almost never observed in practice. It nevertheless gives a sense of how difficult the task is and of what kind of direction we should follow.
Theorem 4.2. If we are given a (possibly infinite) number of domains, where each domain exhibits a different possible intervention on the style variables S, then independence to D implies no total causal effect of S on Z.
The proof can be found in Appendix B. We here take the view that the variable D indexes an intervention, which is a common view in causal inference via DG (see Section 2). It also shows that there may always be a test domain where we observe an intervention to which our representation is not invariant, especially if it changes the support of the style variables. For examples, if we want to be invariant to the brightness in an image, we can only do so to a certain extent, as being invariant to full brightness (i.e., a white image) would force a constant and non-informative representation. This type of question is similar to what is studied in adversarial robustness. Consequently, we need assumptions on what types of interventions we may encounter in the wild. One reasonable assumption is that the full possible support of the style variables is observed during training. We then need to be invariant to any combination of possible values of the style variables. If the style variables are not observed, we cannot directly intervene on them. One possibility is thus to use the variable D as a proxy to simulate interventions on the style variables.
Conjecture 4.3. Given a finite number of domains, we can create new domains via mixtures of the given domains, which simulates new types of interventions on the style variables S.
The idea is that having these new created domains may facilitate identifiability of a representation Z invariant to the style variables S, as in a way it gets us closer to the conditions of Theorem 4.2.
A New Algorithm for Invariant Representation Learning Based on the underlying assumptions of Figure 1, the equivalence proved in Theorem 3.1 and Conjecture 4.3, we present a new algorithm to learn a representation invariant to interventions on D. This algorithm could be useful for example when we have a large number of different values of D, where enforcing an invariant p(z|d) is hard to optimize (pairwise distances between distributions). Instead, we change the distribution of D across batches (simulated intervention) and take the distribution distance between pairs of batches. It also follows Conjecture 4.3: as each time we create a batch with a different distribution of D, it is equivalent to drawing a batch from a created domain, which is a mixture of the initially given domains. We formulate the optimization goal as follows:
min Z=f(X)
L(Y, c(Z)), s.t. p(Z) = const ∀Nd.
We relax this constraint by taking the dual formulation and by approximating the maximum distance between two possible interventions on D by the average distance:
min Z=f(X)
L(Y, c(Z)) + λ · ENd,N ′d [dist(p(Z|do(d = Nd), p(Z|do(d = N ′ d))] (1)
where dist is a distance between distributions (see Appendix A for possible distances), and Nd, N ′d are interventions on the distribution of d.
First, this algorithm gives us a new method to learn invariance when our dataset follows the assumption of Figure 1. If the dataset indeed fulfills our assumptions, we theoretically have that the algorithm will work asymptotically. Unfortunately, these assumptions are not testable. There is thus two possible outcomes: either the optimization converges and we cannot reject that the dataset follows our assumptions, or it does not converge and it probably means that the causal relationships of the variables of our dataset are different than assumed.
Second, as stated in Conjecture 4.3, being invariant to different distributions of D may lead to greater invariance to the style variables. This is especially useful in DG. In Section 5, we experimentally show that our algorithm is indeed a viable method to learn invariance, and that it can also be more
favorable in settings such as DG. We present a possible practical implementation of our algorithm 1 below.
Algorithm 1: Our algorithm for invariant representation learning. 1 Let d be the number of domains; 2 Let n > 0 be the number of samples drawn from each class at each step; 3 begin 4 Draw a batch bi of n samples for each domain; 5 B1, B2← ∅, ∅;
// We create two batches B1 and B2 that approximate the interventions Nd and N ′d of eq. (1)
6 for i← 1 to d do 7 s ∼ U(0, n); 8 B1, B2← (B1, bi[: s]), (B2, bi[s :]); // Concatenate B1 and B2 with a slice of bi 9 end
10 Z1, Z2← f(B1), f(B2); 11 loss← dist(Z1, Z2); 12 return loss; 13 end

heading: 5 EMPIRICAL EVALUATION

heading: 5.1 FAIR REPRESENTATION LEARNING
text: In this section, we present some experiments on Fair representation learning. Here, the goal is not necessarily to obtain better results compared to other baselines. More importantly, we want to show that: (i) Fair representation learning is also an invariant representation learning task, and it is covered by our unifying framework; (ii) Our algorithm is applicable to a wide range of tasks, as it also gives competitive results on this task; (iii) We can control the strength of invariance via the hyperparameter λ; (iv) Fair representation learning datasets probably also follow our proposed data generation graph.
In the context of fair representation learning, the variable D we want to be invariant to here corresponds to what is usually referred to as the sensitive variable. A sensitive variable is a variable that should not have an effect on the predictions of a classifier or regressor. Some examples are the sex, the race or the age of an individual. If we can construct a representation that does not contain information about the sensitive variable, there is no way for a model built on top of this representation to base its prediction on the sensitive variable. Unfortunately, in most datasets, the sensitive variable is actually predictive for the target variable, i.e., the value we are trying to predict. This introduces a trade-off between fairness and accuracy of a model.
Datasets We run experiments on two datasets from the UCI ML-repository (Asuncion and Newman, 2007), the Adult and German dataset. The German dataset consist in predicting whether an individual has good or bad credit, while the sensitive attribute is the gender. The Adult dataset consists in predicting whether the annual income of an individual is more or less than 50, 000$, and the sensitive attribute is the gender. See Table 7 in the Appendix for a summary of the datasets. We also report the size of the majority class for the sensitive and target attribute for each dataset. A fair model should have a sensitive accuracy that is close or below the size of the majority sensitive class, while having a target accuracy as high as possible.
Experiment Design To run our experiments, we reuse the code from Roy and Boddeti (2019) and add our model. We also empirically modify the default latent representation size such that it is optimizable using the MMD distance. As for the synthetic experiment, after training of the encoder, we freeze it and learn two discriminators: one for the target and one for the sensitive attribute. The target discriminators is trained for 100 epochs and the adversary discriminator for 150 epochs. We report the best achieved test accuracy. The goal of this setup is to assess how much information can be extracted from the representation regarding the target and sensitive variables. For each value of λ regularization, we run the experiment three times, and report mean and standard deviation. We
also report the best obtained model and compare it to other baselines. The goal is to show what performance we can potentially achieve with our algorithm and see how it compare with existing models.

heading: 5.1.1 ADULT DATASET
text: The encoder is a neural network with one hidden layer of size 7, and a latent representation size of 2. It is trained for 150 epochs using the Adam optimizer (Kingma and Ba, 2015), with learning rate of 1× 10−4 and weight decay of 5× 10−2. The discriminators are two-hidden-layer neural networks, with hidden layers of size 64 and 32. Both are optimized using Adam with learning rate of 0.001 and weights decay of 0.001. The learning rate of the discriminators is adjusted with Cosine Annealing. Train batch size is set to 128 and test batch size to 1000.
Results are summarized in Figure 2, as well as a comparison with other baselines in Table 2. The results are as we expected: stronger regularization leads to stronger invariance towards the sensitive attribute. We also get that there is a trade-off between target accuracy and adversary accuracy, as the sensitive attribute is informative towards predicting the target. Compared to other baselines, our best model performs well, as it has the best target accuracy for a slightly higher adversary accuracy.

heading: 5.1.2 GERMAN DATASET
text: The encoder is a neural network with two hidden layers of size 15 and 8, and a latent representation size of 32. It is trained for 150 epochs using the Adam optimizer, with learning rate of 1× 10−4 and weight decay of 5× 10−2. The discriminators are two-hidden-layer neural networks, with hidden layers of size 10. Both are optimized using Adam with learning rate of 0.001 and weights decay of 0.001. The learning rate of the discriminators is adjusted with Cosine Annealing. Train batch size is set to 64 and test batch size to 100.
Results are summarized in Figure 2, as well as a comparison with other baselines in Table 3. Here, we observe that 1.0 is a clear optimal value for λ, as it gives the highest target accuracy and the lowest adversary accuracy. A bit more surprisingly, we observe that higher regularization can give lesser invariance, which we can interpret as a form of over-regularization. Compared to other methods, we obtain competitive results as we get the smallest adversary accuracy, even below the majority prediction, while still obtaining the second best target accuracy.

heading: 5.2 DOMAIN GENERALIZATION
text: Datasets For this experiments, we test on seven datasets: ColoredMNIST (Arjovsky et al., 2019), RotatedMNIST (Ghifary et al., 2015), VLCS (Fang et al., 2013), PACS (Li et al., 2017), OfficeHome (Venkateswara et al., 2017), TerraIncognita (Beery et al., 2018) and DomainNet (Peng et al., 2019). In the Appendix, Table 8 shows sample images for each dataset under different domains and Table 9 presents each dataset’s characteristics.
Experiment Design We run our experiments with the DomainBed (Gulrajani and Lopez-Paz, 2020) testbed, which is a recent widely used testbed for DG. We choose this setup as it allows for a highly fair and unbiased comparison with other existing models. DomainBed was designed to be reproducible, to give each algorithm the same amount of hyperparameter search, and to accurately estimate the variance in performance. Three model selection methods are considered: training-domain validation (all training models are pooled and a fraction of each of them is used as validaiton set), leave-one domain-out cross-validation (cross validation is performed using a different domain as validation, and the best models is retrained on all training domains) and test-domain validation set (a fraction of the test domain is used as validation set). The first two methods are closer to a realistic setting, whereas oracle validation allows us to evaluate whether there exists headroom for improvement. Training-domain validation assumes that all training domains and the test domain follows a similar distribution, as we pool all the training domains during training. On the other hand, leave-one domain-out cross-validation is closer to our assumption, as it optimizes for generalization to an unseen domain that is assumed to follow a different distribution.
Proposed Models We take two existing models, MMD and CORAL, based on matching distribution across domains, and propose two new models, CausIRL with MMD and with CORAL. These two new algorithm simply consist in changing how the regularization loss is computed according to our proposed algorithm, i.e., instead of taking pairwise distances across domains, we compute distances between batches that follow different domain distributions. We thus want to see if this simple change in the algorithm leads to better performance, which may be due as we conjecture to greater invariance to the style variables, as well as the fact that it may be easier to optimize. The hyperparameter λ is drawn randomly in 10Uniform(−1,1).

heading: 5.2.1 MODEL SELECTION: LEAVE-ONE-DOMAIN-OUT CROSS-VALIDATION
text: We now look at the DG experiment results for the leave-one-domain-out cross-validation model selection method. The results are summarized in Table 4. The result for this model selection method are the most relevant as it closely follows our assumptions on the training and test distributions. The results for the other model selection methods can be found in Appendix E.
Here, the overall performance of CausIRL with CORAL is almost identical to CORAL. Nevertheless, there are some differences when looking at the performance on individual datasets. CausIRL with CORAL overperform CORAL on PACS, TerraIncognita and DomainNet, where CORAL performs better on VLCS and OfficeHome. However we should note that only the overperformance of CausIRL with CORAL over CORAL on DomainNet is statistically significant when looking at the confidence intervals of the average accuracies.
For CausIRL with MMD, we observe a significant boost in the overall performance compared to MMD. CausIRL with MMD performs better on almost all datasets, and we also observe a significant

heading: Algorithm ColoredMNIST RotatedMNIST VLCS PACS OfficeHome TerraIncognita DomainNet Avg
text: leap in performance on DomainNet going from 23.4% to 38.9%. This shows that the inductive bias of our algorithm matches the real inductive bias of DG far more accurately.

heading: 6 REAL-WORLD DOMAIN GENERALIZATION
text: In this section, we run experiments on more realistic distributional shifts. We use the Wilds (Koh et al., 2021) benchmark and run experiments on two datasets: Camelyon17 (Bandi et al., 2018) and RxRx1 (Taylor et al., 2019). Camelyon17 consists in predicting whether a region of tissue contains tumor tissue, while being invariant to the hospitals where the sample was taken. The goal is to obtain a model that generalizes across hospitals, as hospital specific artifacts of the data collection process can vary. RxRx1 consists of cell images, where the cells received some genetic treatment (as well as no treatment). The goal is to predict the genetic treatment among 1, 139 possible treatments. Here, we want to be invariant to the batch the cells come from, as it is a common observation that batch effect can greatly alter the results.
We test our two proposed models, CausIRL with CORAL and with MMD on both datasets. For the RxRx1 dataset, we use the same hyperparameters than for the CORAL model in the Wilds implementation. For Camelyon17, we change the number of group per batch to three and the batch size to 60. The results are summarized in Table 5. As for the DG experiments on DomainBed before, we observe that CausIRL with CORAL performs better than CORAL. Moreover, CausIRL with MMD performs slightly better than CausIRL with CORAL on Camelyon17 and similarly on RxRx1. Unfortunately, all models perform worse than simple ERM. Nevertheless, we again observe that our proposed models work competitively even on a realistic dataset, and that our proposed algorithm to compute the distributional distance regularization is better than how it is usually done.

heading: 7 CONCLUSION AND FUTURE WORK
text: In this work, we provided a causal perspective on invariant representation learning. We defined style variables in the context of our framework and developed theory on what conditions are necessary or sufficient to be invariant towards the style variables. Based on these theoretic insights and the assumptions on the data generating process, we then proposed a new algorithm for enforcing invariance to style variables in the learned representations. We empirically demonstrated that our algorithm is versatile as it works on a diverse set of tasks and datasets. In particular, it performs strongly in DG, where we obtain state-of-the-art performance.

heading: 8 REPRODUCIBILITY STATEMENT
text: Regarding the theoretical results, all results are based on the assumptions of Figure 1, and the proofs can be found in Appendix B. For the fair representation learning experiments, the architectures and training procedures are precisely described in Section 5.1, and we reused the code of Roy and Boddeti (2019) to ensure a fair comparison to existing models. For the DG experiments, we used two benchmarks, DomainBed (Gulrajani and Lopez-Paz, 2020) and Wilds (Koh et al., 2021), which were designed to be reproducible and unbiased. The hyperparameter used can be found in the main text. Regarding the implementation of our algorithm, the pseudo-code in algorithm 1 and the python code snippets in the Appendix should be sufficient to accurately implement our proposed models. For all experiments, our implementation of the MMD and CORAL distance are taken from the DomainBed (Gulrajani and Lopez-Paz, 2020) code. Our code implementations are provided with all the necessary details in the corresponding sections in the appendix. For the implementation details for fair representation learning we refer to the detailed documentation of the code in section D and for domain generalization we refer to section E.

heading: A BACKGROUND
text: In this chapter, we review some necessary theoretical background that is used for our theoretical results, new proposed method as well as previous works. We review the main elements of Causality theory, which is the central theoretical basis used to model the problems we study and to design new methods and algorithms. We also review some distributional distances used in this work, as our main goal is to study invariant representation learning via invariant latent variable distribution.

heading: A.1 CAUSALITY
text: Causality essentially is the study of cause and effects, which goes beyond the study of statistical associations from observational data. This allows to reason about the notion of interventions, such as a treatment in medicine. The expected effect of an intervention is in general not equivalent to statistical conditioning, which calls for a more profound understating of the data generating process that goes beyond correlations between variables. We here focus on Pearl’s view of causality Pearl (2009), which mainly relies on DAGs.
A DAG allows to represent the relations between variables, where each variable is represented by a node in the graph. Consequently, we can interpret directed edges between nodes as the existence of a causal effect from the parent node (the cause) on the child node (the effect).
Let G = (V,E) be a DAG, and P be a distribution. We say that (G,P ) is a causal DAG model if for any W ⊂ V , we have:
p(xV |do(XW = x′W )) = ∏
i∈V \W
p(xi|xpai)I(xW = x′W )
where xpai are the parents of node i in graph G, I is the indicator function and do(XW = x ′ W ) denotes the intervention on the variables XW . As we can see above, one of the properties of a causal DAG model is that the distribution factorizes according to the parents in the associated graph G.

heading: A.1.1 STRUCTURAL CAUSAL MODELS
text: A SCM can be seen as a more expressive version of a causal DAG model. Formally, an SCM consists of a collection S of d structural assignments, one per variable:
Xj ←− fj(Xpaj , Nj)
where Xpaj ∈ X\Xj , and N1 to Nd are called the noise variables (Definition 6.2 of Peters et al. (2017)). The noise variables are assumed to be jointly independent.
For causal DAG models, we defined an intervention by p(XV |do(XW = x′W )) (sometimes also written as pdo(XW=x ′ W )(XV )), where the value of some variables are set to a constant value. With SCMs, we can give a more precise and general definition of interventions. An intervention now consists in replacing a subset of the collection S of structural assignments by new functions. An intervention can thus consist in replacing a variable by a constant, a new random variable or even by changing the function and its arguments (i.e., its parents). The new distribution over the variables entailed by the new intervened SCM is denoted by P do(Xk=f̃(Xp̃ak ,Ñk)) (see Definition 6.8 of Peters et al. (2017) for more details).
With this definition, we are now equipped to reason about interventions, causes and effects. One important notion is the notion of Total Causal Effect. Definition A.1. (Definition 6.12 in Peters et al. (2017)) We say that a variable i has a total causal effect on a variable k if and only if:
Xi ⊥/ Xk in P do(Xk=Ñk)
for some random variable Ñk.
A total causal effect between a variable Xk and Xi may only exist if there is a directed path from i to k in the DAG associated to our SCM. On the other hand, there may be no total causal effect between two variable even though there exists a directed path between them in the graph.

heading: A.2 DISTRIBUTIONAL DISTANCES
text: The main goal of this work is to study how invariance can be enforced by regularizing different latent spaces to have the same distribution. To this end, we thus need a differentiable distance or divergence between distributions that can be minimized during training. We here present the most commonly used distances in the literature.

heading: A.2.1 ADVERSARIAL
text: Adversarial training was first introduced in Goodfellow et al. (2014) as a new method for Generative modeling. Based on game theory, it can intuitively be described as a two player game, where each player is parameterized by a neural network. The Generator is a function that maps its input distribution to an output distribution. We call it the generated distribution and denote it by pg . On the other hand, a Discriminator tries to distinguish between samples coming from the target dataset and samples produced by the Generator. At convergence, the Generator produces data that is distributed similarly to the target distribution, and thus it becomes impossible for the Discriminator to distinguish samples.
Formally, the objective of the two-player minimax game reads:
min G max D V (D,G) = Ex∼pdata(x) [logD(x)] + Ez∼pz(z) [log (1−D(G(z)))] (2)
where z is the input, x comes from the target distributions, and the Discriminator D should output 1 when its input is a samples from the target, and 0 otherwise. If the Discriminator is optimal for a given G, Equation 2 can be rewritten to show that the Generator actually minimizes the Jensen–Shannon divergence (JSD) between the generated and target distribution.
JSD (P‖Q) = 1 2 DKL
( P ∥∥∥∥12 (P +Q) ) + 1 2 DKL ( Q ∥∥∥∥12 (P +Q) ) ,
where DKL is the Kullback-Leibler (KLd) divergence. It also can be shown that if both networks have sufficient capacity, and if the Discriminator is trained to optimality after each optimization step of the Generator, then the distribution of the Generator converges to the target distribution.
Adversarial training can thus be seen as a proxy distributional distance, which corresponds to the JSD at convergence. This concept of adversarial training has been extended to be used as a regularizer for latent spaces. It can for example be used to enforce a prior distribution on the latent space Makhzani et al. (2015). It can also be used to enforce two latent spaces to have the same distribution. Its use is often justified as wanting two latent spaces to seem indistinguishable for an adversary, which is supposed to force the encoder to discard what is not constant across the two input distribution. We argue that adversarial training is theoretically equivalent to minimizing any distributional divergence, and that only their optimization properties differentiate them. We will also later clarify the intuition of trying to discard the idiosyncratic in favor of the universal, and what it actually corresponds to when we look at the data generation process of a given dataset.

heading: A.2.2 MAXIMUM MEAN DISCREPANCY
text: MMD Gretton et al. (2006) is a distance based on empirical samples from two distributions, based on the distance between the means of the two sets of samples mapped into a reproducing kernel Hilbert space (RKHS). Let {X} ∼ P and {X ′} ∼ Q. Then, we have:
MMD(X,X ′)2 = ∥∥∥∥∥∥ 1n n∑ i=1 φ(xi)− 1 n′ n′∑ i=1 φ(x′i) ∥∥∥∥∥∥ = 1
n2 n∑ i,j=1 k(xi, xj) + 1 n′2 n′∑ i,j=1 k(x′i, x ′ j)− 2 n · n′ n∑ i=1 n′∑ j=1 k(xi, x ′ j),
where k(·, ·) is the associated kernel. One commonly used kernel is the Gaussian kernel k(x, x′) = e−λ‖x−x
′‖2 . Asymptotically, for a universal kernel such as the Gaussian kernel, MMD(X,X ′) = 0 if and only if P = Q. Minimizing the MMD distance during training can thus be used to align two distributions.

heading: A.2.3 OPTIMAL TRANSPORT
text: OT has recently gained traction in machine learning research. It is particularly interesting as it defines multiple distances between distributions.
We present the main formulations of the OT problem and some distances following the notations of Peyré and Cuturi (2019). We concentrate on distances defined on discrete distributions.
Originally, OT can be described as the problem of transporting the mass from a set of points {x1, . . . , xn} to a set of destination points {y1, . . . , ym} that have finite capacities. Furthermore, the cost of transporting from a point to another is fixed by a cost matrix. There is obviously a natural analogy to logistics and planning. The OT problem thus is to fulfill this mass transport while minimizing the cost. To recast this as a distance between distribution, the mass of the source points is now described by an histogram a ∈ {n+: ∑ i ai = 1} and the destination capacities as an histogram
b ∈ {m+ : ∑ i bi = 1}. The cost matrix is C ∈ n×m + where Cij describe the cost of transport from xi to yj . The Kantorovich formulation of this problem is:
LC(a,b) := min T∈U(a,b) 〈C, T 〉 := ∑ i,j Ci,jTi,j (3)
where
U(a,b) := { T ∈n×m+ : T1m = a, TT1n = b } (4)
Ti,j describes how much mass flows from xi to yj . The constrain imposes that all the mass leaves a and that b is filled. Given the cost matrix, we hence can compute the divergence between two distributions, and this divergence elegantly comes with an exact description (the coupling matrix T ) of how to go from the first configuration to the other. Furthermore, if xi, yi ∈ X lie in the same metric space with distance d and the cost matrix satisfies Ci,j = dp(xi, yj) for some p ≥ 1, then:
Wp(a,b) := LC(a,b) 1/p (5)
is a distance, called the p-Wasserstein distance. This distance has had a large range of applications, notably for problems related to invariant representation learning (see Section 2).

heading: B PROOFS OF THEOREMS
text: 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).
Proof. =⇒ As z is a descendant of d, the mechanism p(z|d) is invariant to the distribution of d.
pdo(d=Nd)(z) = ∫ pdo(d=Nd)(z, d′)dd′ = ∫ pdo(d=Nd)(z|d′)Nd(d′)dd′ = p(z) ∫ Nd(d ′)dd′
= p(z)
⇐= As p(z) is constant for all distribution of d, then it is also constant for deterministic distributions, i.e δd.
p(z) = ∫ p(z|d′)δd(d′)dd′ = p(z|d)
This holds for all values of d, which implies that I(Z;D) = 0.
Theorem 4.1. Independence to D is a necessary condition for the representation to be invariant to the style variables.

heading: Proof.
text: p(z|d = 1)− p(z|d = 2) = ∫ p(z|S, d = 1)[p(S|d = 1)dS − ∫ p(z|S, d = 2)p(S|d = 2)]dS
= ∫ p(z|S)[p(S|d = 1)− p(S|d = 2)]dS
= p(z) ∫ [p(S|d = 1)− p(S|d = 2)]dS
= 0
Theorem 4.2. If we are given a (possibly infinite) number of domains, where each domain exhibits a different possible intervention on the style variables S, then independence to D implies no total causal effect of S on Z.
Proof. From the definition of total causal effect, let’s suppose by contradiction that there exists an intervention on S such that pdo(S=Ñs)(z) 6= p(z). Let d̃ denote the domain that correspond to this intervention. We then have a value of D such that p(z|d̃) 6= p(z), which is a contradiction to Z being independent to D.

heading: C SYNTHETIC EXPERIMENT
text: We here implement a simple synthetic experiment to verify that our algorithm effectively enforces invariance to D in a setting that exactly follows our assumptions. We also simplify the setting by considering that we directly observe the generative factors.
The distribution is generated by the following set of structural equations:
Y ←− Ny; D ←− Nd; G1 ←− Y +NG1 ; G2 ←− 2 · Y + 2 ·D +NG2 ; G3 ←− D +NG3 ;
where Ny and Nd ∼ Ber(0.5), NGi ∼ N (0, 1). We draw the associated causal DAG in Figure 3. To create a dataset, we draw 1000 samples from our synthetic distribution and use 200 of them as test samples.
We then learn a representation that is invariant to D and that is predictive towards Y using our proposed loss. As a distributional distance, we use the MMD loss with Gaussian kernel. The architecture of the encoder is a neural network with one hidden layer of size 10 and a representation size of 5. The hidden layer is followed by a batch normalization and a ReLU activation. We use a batch size of 64 and train with the Adam optimizer (Kingma and Ba, 2015) for 200 epochs, with a learning rate of 0.001 and weight decay of 5× 10−5. After training of the encoder, we freeze it and train two one layer linear discriminators: one to predict Y and one to predictD. For each discriminators, we report the best achieved test accuracy. We run this experiment three times for each value of lambda regularization , λ ∈ {0.0, 0.1, 0.5, 1.0, 5.0, 10.0}. The results are summarized in Table 6 and Figure 4.
As expected, we can observe a strong correlation between the strength of regularization and the strength of invariance. We achieved perfect invariance with λ = 10.0, where the adversary accuracy is 50%, but target accuracy is only 54.2%. This is expected: as Y and D are strongly correlated, removing information on D in the representation also reduces the predictive power of the representation. There thus is a trade-off between performance and invariance, that can be controlled via the value of λ. Finally, this experiment confirms that our proposed algorithm is a viable new method to learn invariance.

heading: D FAIR REPRESENTATION LEARNING EXPERIMENT SUPPLEMENTS
text: Compute Resources We run the experiements on NVIDIA GEFORCE RTX 2080 TI GPUs.
Implementation Here is the implementation of our model, which is a class that we added to the code2 of Roy and Boddeti (2019) to the train.py file:
class CausIRL_MMD: def __init__(self, data,
train_loader=None, test_loader=None,
2https://github.com/human-analysis/MaxEnt-ARL
total_epoch=200, alpha=0.1, epsilon=0.1, use_cuda=False, resume=False, ckpt_filename=None, resume_filename=None, privacy_flag=True, privacy_option=’maxent-arl’, print_interval_train=10, print_interval_test=10 ):
# data info self.data = data self.train_loader = train_loader self.test_loader = test_loader self.n_sensitive_class = self.data.n_sensitive_class self.n_target_class = self.data.n_target_class
# models self.adv_net = data.adversary_net self.target_net = data.target_net self.discriminator_net = data.discriminator_net
# optimizer self.optimizer = data.optimizer self.discriminator_optimizer = data.discriminator_optimizer self.adv_optimizer = data.adv_optimizer self.target_optimizer = data.target_optimizer
def my_cdist(x1, x2): x1_norm = x1.pow(2).sum(dim=-1, keepdim=True) x2_norm = x2.pow(2).sum(dim=-1, keepdim=True) res = torch.addmm(x2_norm.transpose(-2, -1),
x1, x2.transpose(-2, -1), alpha=-2).add_(x1_norm)
return res.clamp_min_(1e-30)
def gaussian_kernel(x, y, gamma=[0.001, 0.01, 0.1, 1, 10, 100, 1000]):
D = my_cdist(x, y) K = torch.zeros_like(D)
for g in gamma: K.add_(torch.exp(D.mul(-g)))
return K
def mmd(x, y): Kxx = gaussian_kernel(x, x).mean() Kyy = gaussian_kernel(y, y).mean() Kxy = gaussian_kernel(x, y).mean() return Kxx + Kyy - 2 * Kxy
# loss self.kl_loss = nn.KLDivLoss() self.cross_entropy_loss = nn.CrossEntropyLoss() self.entropy_loss = EntropyLoss() self.nll_loss = nn.NLLLoss()
self.mse_loss = nn.MSELoss() self.mmd_loss = mmd
# filename self.log_file_name = ckpt_filename+"_log.txt" self.adv_log_file_name = ckpt_filename+"_adv_log.txt" self.target_log_file_name = ckpt_filename + "_target_log.txt" self.checkpoint_filename = ckpt_filename self.adv_checkpoint_filename = ckpt_filename+"_adv.ckpt" self.target_checkpoint_filename = ckpt_filename + "_target.ckpt"
# algorithm and visualization parameters self.alpha = torch.tensor([alpha*1.0], requires_grad=True) self.resume = resume self.epoch = 0 self.gamma_param = 0.01 self.plot_interval = 10 self.print_interval_train = print_interval_train self.print_interval_test = print_interval_test self.use_cuda = use_cuda self.privacy_flag = privacy_flag self.privacy_option = privacy_option
# local variables self.uniform = torch.tensor(1 / (self.data.n_sensitive_class)) .repeat(self.data.n_sensitive_class) self.target_label = torch.zeros(0, dtype=torch.long) self.sensitive_label = torch.zeros(0, dtype=torch.long) self.sensitive_label_onehot = torch.FloatTensor(0, self.data.n_sensitive_class) self.target_label_onehot = torch.FloatTensor(0, self.data.n_target_class) self.inputs = torch.zeros(0, 0, 0) self.inputs.requires_grad = False self.batch_uniform = torch.FloatTensor(0, self.data.n_sensitive_class) self.epsilon = torch.tensor([epsilon]).float()
if resume: assert os.path.isdir(’checkpoint’), ’Error: no checkpoint directory found!’ if self.use_cuda:
checkpoint = torch.load(os.path.join(’checkpoint/’, resume_filename))
else: checkpoint = torch.load(os.path.join(’checkpoint/’,resume_filename), map_location=lambda storage, loc: storage)
self.net = checkpoint[’net’] self.best_acc = 0 # checkpoint[’acc’] self.start_epoch = 0 # checkpoint[’epoch’] self.total_epoch = total_epoch # + self.start_epoch
for param in self.net.parameters(): param.requires_grad = True
else: self.net = data.net self.best_acc = 0 self.start_epoch = 0 self.total_epoch = total_epoch
if self.use_cuda:
self.net = self.net.cuda() self.discriminator_net = self.discriminator_net.cuda() self.adv_net = self.adv_net.cuda() self.target_net = self.target_net.cuda() self.net = nn.DataParallel(self.net, device_ids= range(torch.cuda.device_count())) self.target_net = nn.DataParallel(self.target_net, device_ids=range(torch.cuda.device_count())) self.discriminator_net = nn.DataParallel(self.discriminator_net, device_ids=range(torch.cuda.device_count())) self.adv_net = nn.DataParallel(self.adv_net, device_ids=range(torch.cuda.device_count())) cudnn.benchmark = True self.inputs = self.inputs.cuda() self.target_label = self.target_label.cuda() self.sensitive_label = self.sensitive_label.cuda() self.sensitive_label_onehot = self.sensitive_label_onehot.cuda() self.target_label_onehot = self.target_label_onehot.cuda() self.uniform = self.uniform.cuda() self.batch_uniform = self.batch_uniform.cuda() self.alpha = self.alpha.cuda()
self.best_loss = 1e16 self.adv_best_acc = 0 self.target_best_acc = 0 self.t_losses, self.t_top1, self.d_losses, self.d_top1 = AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter() self.e_losses, self.losses = AverageMeter(), AverageMeter() self.t_top5, self.d_top5 = AverageMeter(), AverageMeter() self.adv_losses, self.adv_top1, self.adv_top5, self.entropy_losses = AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter() self.target_losses, self.target_top1, self.target_top5, self.target_entropy_losses = AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter()
def perform_epoch(self, epoch, test_flag=False): if test_flag:
self.net.eval() self.discriminator_net.eval() self.target_net.eval() loader = self.test_loader string = "Test" print_interval = self.print_interval_test data_size = len(self.test_loader)
else: self.net.train() self.discriminator_net.train() self.target_net.train() loader = self.train_loader string = "Train" print_interval = self.print_interval_train data_size = len(self.train_loader)
iteration = 0
self.t_losses.reset() self.e_losses.reset() self.losses.reset() self.d_losses.reset()
self.t_top1.reset() self.d_top1.reset() self.t_top5.reset() self.d_top5.reset() self.entropy_losses.reset()
for batch_idx, (inputs, target_label, sensitive_label) in enumerate(loader):
batch_size = inputs.size(0) iteration += 1
self.inputs.resize_(inputs.size()).copy_(inputs) self.target_label.resize_(target_label.size()). copy_(target_label) self.sensitive_label.resize_( sensitive_label.size()).copy_(sensitive_label) self.sensitive_label_onehot.resize_([batch_size, self.data.n_sensitive_class]) self.sensitive_label_onehot.zero_() self.sensitive_label_onehot.scatter_(1, torch.unsqueeze(self.sensitive_label, 1), 1) self.target_label_onehot.resize_([batch_size, self.data.n_target_class]) self.target_label_onehot.zero_() self.target_label_onehot.scatter_(1, torch.unsqueeze(self.target_label, 1), 1) self.batch_uniform.resize_([batch_size, self.data.n_sensitive_class]) self.batch_uniform[:, :] = 1.0/(self.data.n_sensitive_class) self.batch_uniform.scatter_(1, torch.unsqueeze(self.sensitive_label, 1), 0) self.optimizer.zero_grad()
_, z, e_prob = self.net(self.inputs) target_outputs, _, t_prob = self.target_net(z) t_loss = torch.nan_to_num(self.cross_entropy_loss( target_outputs+1e-16, self.target_label)) entropy_loss = torch.tensor(0) s_loss = torch.tensor(0)
if self.privacy_flag: #d_outputs, _, d_prob = self.discriminator_net(z) #entropy_loss = -self.entropy_loss(d_prob) first = None second = None for i in range(self.n_sensitive_class):
ind = self.sensitive_label == i z_ = z[ind] slice = random.randint(0, len(z_)) if first is None:
first = z_[:slice] second = z_[slice:]
else: first = torch.cat((first, z_[:slice]), 0) second = torch.cat((second, z_[slice:]), 0)
if len(first) > 1 and len(second) > 1: s_loss = torch.nan_to_num(self.mmd_loss(first, second))
loss = t_loss + self.alpha*s_loss else:
loss = t_loss
if not test_flag: # update weights self.optimizer.zero_grad() self.target_optimizer.zero_grad() loss.backward() self.optimizer.step() self.target_optimizer.step()
# measure accuracy and record loss for learner t_prec1 = accuracy(t_prob.data, self.target_label.data) t_prec5 = accuracy(t_prob.data, self.target_label.data, topk=(int(np.min([5, self.n_target_class])),)) self.t_losses.update(t_loss.data.item(), batch_size) self.e_losses.update(s_loss.data.item(), batch_size) self.losses.update(loss.data.item(), batch_size) self.t_top1.update(t_prec1[0], batch_size) self.t_top5.update(t_prec5[0], batch_size) self.entropy_losses.update(s_loss.data.item(), batch_size)
if self.privacy_flag: if not test_flag:
self.discriminator_net.train() d_outputs, _, a_prob = self.discriminator_net(z.detach()) d_loss = self.nll_loss(torch.log(a_prob+1e-16), self.sensitive_label)
if not test_flag: self.discriminator_optimizer.zero_grad() d_loss.backward() self.discriminator_optimizer.step()
d_prec1 = accuracy(a_prob.data, self.sensitive_label.data) d_prec5 = accuracy(a_prob.data, self.sensitive_label.data, topk=(int(np.min([5, self.n_sensitive_class])),)) self.d_losses.update(d_loss.data.item(), batch_size) self.d_top1.update(d_prec1[0], batch_size) self.d_top5.update(d_prec5[0], batch_size)
if iteration % print_interval == 0: print(string + ’_Epoch:[{0}][{1}/{2}] |’
’ T_Loss: {3:.2f} |’ ’ E_Loss: {4:.2f} |’ ’ Loss: {5:.2f} |’ ’ T_Prec: {6:.2f} |’ ’ T_Prec5: {7:.2f} |’ ’ D_Loss: {8:.2f} |’ ’ D_Prec: {9:.2f} |’ ’ D_Prec5: {10:.2f} |’ ’ D_Entropy: {11:.2f} |’
.format( epoch, batch_idx, data_size, float(self.t_losses.avg), float(self.e_losses.avg), float(self.losses.avg),float(self.t_top1.avg.item()), float(self.t_top5.avg.item()), float(self.d_losses.avg), float(self.d_top1.avg.item()), float(self.d_top5.avg.item()), float(self.entropy_losses.avg)))
else:
if iteration % print_interval == 0: print(string + ’_Epoch:[{0}][{1}/{2}] |’
’ T_Loss: {3:.2f} |’ ’ T_Prec: {4:.2f} |’ ’ T_Prec5: {5:.2f} |’
.format( epoch, batch_idx, data_size, float(self.t_losses.avg), float(self.t_top1.avg.item()), float(self.t_top5.avg.item())))
return self.losses.avg, self.t_top1.avg, self.t_top5.avg, self.d_losses.avg, self.d_top1.avg, self.d_top5.avg, self.entropy_losses.avg
def train(self): self.logger = Logger(os.path.join(’checkpoint/’, self.log_file_name), title=’Problem’) self.logger.set_names([’LR’, ’Train-Loss’, ’Test-Loss’, ’Train-Acc.’, ’Train-Acc5.’, ’Test Acc.’, ’Test Acc5.’,
’D-Train Loss’, ’D-Test Loss’, ’D-Train Acc.’, ’D-Train Acc5.’, ’D-Test Acc.’, ’D-Test Acc5.’, ’D-Train-Entropy’, ’D-Test-Entropy’])
for epoch in range(self.start_epoch, self.total_epoch): print(’\nEpoch: %d’ % epoch)
train_loss, train_acc, train_acc5, d_train_loss, d_train_acc, d_train_acc5, d_train_entropy = self.perform_epoch(epoch=epoch, test_flag=False)
with torch.no_grad(): test_loss, test_acc, test_acc5, d_test_loss, d_test_acc, d_test_acc5, d_test_entropy = self.perform_epoch(epoch=epoch,
test_flag=True)
self.logger.append([self.optimizer.param_groups[0][’lr’], float(train_loss), float(test_loss), float(train_acc), float(train_acc5), float(test_acc), float(test_acc5), float(d_train_loss), float(d_test_loss), float(d_train_acc), float(d_train_acc5), float(d_test_acc), float(d_test_acc5), float(d_train_entropy), float(d_test_entropy)])
# it is optimum only when we reach the end of the game by optimization, # any other value e.g., current discriminator feedback is non-optimal if (epoch + 1) % 10:
print(’Saving..’) # Save checkpoint. state = {
’net’: self.net.module if self.use_cuda else self.net, ’state_dict’: self.net.state_dict(), ’acc’: test_acc, ’epoch’: epoch,
’optimizer’: self.optimizer.state_dict() } if not os.path.isdir(’checkpoint’):
os.mkdir(’checkpoint’) torch.save(state, ’checkpoint/’ + self.checkpoint_filename + ’.ckpt’) self.best_acc = test_acc self.best_loss = test_loss
self.logger.close() print("Done")
def perform_epoch_adversary(self, epoch, test_flag=False): if test_flag:
self.adv_net.eval() loader = self.test_loader str = "Test" print_interval = self.print_interval_test
else: self.adv_net.train() loader = self.train_loader str = "Train" print_interval = self.print_interval_train
self.net.eval() iteration = 0 self.adv_losses.reset() self.adv_top1.reset() self.adv_top5.reset() self.entropy_losses.reset()
for batch_idx, (inputs, target_label, sensitive_label) in enumerate(loader): batch_size = inputs.size(0) iteration += 1 if self.data.name == ’mnist’:
inputs = torch.unsqueeze(inputs, 1).float()
self.inputs.resize_(inputs.size()).copy_(inputs) self.target_label.resize_(target_label.size()).copy_(target_label) self.sensitive_label.resize_( sensitive_label.size()).copy_(sensitive_label)
with torch.no_grad(): outputs, z, _ = self.net(self.inputs)
d_outputs, _, prob = self.adv_net(z.detach()) d_loss = self.cross_entropy_loss(d_outputs, self.sensitive_label)
with torch.no_grad(): entropy_loss = -self.entropy_loss(prob)
if not test_flag: self.adv_optimizer.zero_grad() d_loss.backward() self.adv_optimizer.step()
d_prec1 = accuracy(prob.data, self.sensitive_label.data) d_prec5 = accuracy(prob.data, self.sensitive_label.data, topk=(int(np.min([5, self.n_sensitive_class])),))
self.adv_losses.update(d_loss.data.item(), batch_size) self.adv_top1.update(d_prec1[0], batch_size) self.adv_top5.update(d_prec5[0], batch_size) self.entropy_losses.update(entropy_loss.data.item(), batch_size)
if iteration % print_interval == 0: print(str + ’ Epoch:[{0}][{1}/{2}] |’
’ T_Loss: {3:.5f} |’ ’ T_Prec: {4:.2f} |’ ’ T5_Prec: {5:.2f} |’ ’ Entropy: {6:.3f} |’
.format( epoch, batch_idx, len(self.train_loader), float(self.adv_losses.avg), float(self.adv_top1.avg.item()), float(self.adv_top5.avg), float(self.entropy_losses.avg)))
return self.adv_losses.avg, self.adv_top1.avg, self.adv_top5.avg, self.entropy_losses.avg
def train_adversary(self, model_filename=None, total_epoch=100): self.adv_logger = Logger(os.path.join(’checkpoint/’, self.adv_log_file_name), title=’Problem’) self.adv_logger.set_names([’LR’, ’Train-Loss’, ’Test-Loss’, ’Train Acc.’, ’Train Acc5.’, ’Test Acc.’, ’Test Acc5.’,
’Train Entropy’,’Test Entropy’])
self.adv_best_acc = 0 scheduler = CosineAnnealingLR(self.adv_optimizer, T_max=total_epoch, eta_min=1e-7) if model_filename is not None:
checkpoint = torch.load(os.path.join(’checkpoint/’, model_filename)) self.net = checkpoint[’net’] self.net.eval()
for epoch in range(total_epoch): print(’\nEpoch: %d’ % epoch) scheduler.step() train_loss, train_acc, train_acc5, train_entropy =
self.perform_epoch_adversary(epoch=epoch, test_flag=False) with torch.no_grad():
test_loss, test_acc, test_acc5, test_entropy = self.perform_epoch_adversary(epoch=epoch, test_flag=True)
self.adv_logger.append([self.adv_optimizer.param_groups[0][’lr’], float(train_loss), float(test_loss), float(train_acc),
float(train_acc5), float(test_acc), float(test_acc5), float(train_entropy), float(test_entropy)])
# Save checkpoint. if test_acc > self.adv_best_acc:
print(’Saving..’) state = {
’net’: self.adv_net.module if self.use_cuda else self.adv_net, ’state_dict’: self.adv_net.state_dict(), ’acc’: test_acc, ’epoch’: epoch,
’optimizer’: self.adv_optimizer.state_dict() } if not os.path.isdir(’checkpoint’):
os.mkdir(’checkpoint’) torch.save(state, ’checkpoint/’ + self.adv_checkpoint_filename) self.adv_best_acc = test_acc
self.adv_logger.close() print("Adversary Done.")
def perform_epoch_target(self, epoch, test_flag=False): if test_flag:
self.target_net.eval() loader = self.test_loader str = "Test" print_interval = self.print_interval_test
else: self.target_net.train() loader = self.train_loader str = "Train" print_interval = self.print_interval_train
self.net.eval() iteration = 0 self.target_losses.reset() self.target_top1.reset() self.target_top5.reset() self.target_entropy_losses.reset()
for batch_idx, (inputs, target_label, sensitive_label) in enumerate(loader): batch_size = inputs.size(0) iteration += 1 if self.data.name == ’mnist’:
inputs = torch.unsqueeze(inputs, 1).float()
self.inputs.resize_(inputs.size()).copy_(inputs) self.target_label.resize_(target_label.size()).copy_(target_label) self.sensitive_label.resize_( sensitive_label.size()).copy_(sensitive_label)
with torch.no_grad(): outputs, z, _ = self.net(self.inputs)
d_outputs, _, prob = self.target_net(z.detach()) d_loss = self.cross_entropy_loss(d_outputs, self.target_label)
with torch.no_grad(): entropy_loss = -self.entropy_loss(prob)
if not test_flag: self.target_optimizer.zero_grad() d_loss.backward() self.target_optimizer.step()
d_prec1 = accuracy(prob.data, self.target_label.data) d_prec5 = accuracy(prob.data, self.target_label.data, topk=(int(np.min([5, self.n_target_class])),)) self.target_losses.update(d_loss.data.item(), batch_size) self.target_top1.update(d_prec1[0], batch_size)
self.target_top5.update(d_prec5[0], batch_size) self.target_entropy_losses.update(entropy_loss.data.item(), batch_size)
if iteration % print_interval == 0: print(str + ’ Epoch:[{0}][{1}/{2}] |’
’ T_Loss: {3:.5f} |’ ’ T_Prec: {4:.2} |’ ’ T5_Prec: {5:.2f} |’ ’ Entropy: {6:.3f} |’
.format( epoch, batch_idx, len(loader), float(self.target_losses.avg), float(self.target_top1.avg.item()), float(self.target_top5.avg), float(self.target_entropy_losses.avg)))
return self.target_losses.avg, self.target_top1.avg, self.target_top5.avg, self.target_entropy_losses.avg
def train_target(self, model_filename=None, total_epoch=100): self.target_logger = Logger(os.path.join(’checkpoint/’, self.target_log_file_name), title=’Problem’) self.target_logger.set_names([’LR’, ’Train-Loss’, ’Test-Loss’, ’Train Acc.’, ’Train Acc5.’, ’Test Acc.’, ’Test Acc5.’,
’Train Entropy’,’Test Entropy’])
self.target_best_acc = 0 scheduler = CosineAnnealingLR(self.target_optimizer, T_max=total_epoch, eta_min=1e-7) if model_filename is not None:
assert os.path.isdir(’checkpoint’), ’Error: no checkpoint directory found!’ checkpoint = torch.load(os.path.join(’checkpoint/’, model_filename)) self.net = checkpoint[’net’] self.net.eval()
for epoch in range(total_epoch): print(’\nEpoch: %d’ % epoch) scheduler.step() train_loss, train_acc, train_acc5, train_entropy = self.perform_epoch_target(epoch=epoch, test_flag=False) with torch.no_grad():
test_loss, test_acc, test_acc5, test_entropy = self.perform_epoch_target(epoch=epoch, test_flag=True)
self.target_logger.append([self.target_optimizer.param_groups[0][’lr’], float(train_loss), float(test_loss), float(train_acc),
float(train_acc5), float(test_acc), float(test_acc5), float(train_entropy), float(test_entropy)])
if test_acc > self.target_best_acc: print(’Saving..’) # Save checkpoint. state = {
’net’: self.target_net.module if self.use_cuda else self.target_net, ’state_dict’: self.target_net.state_dict(), ’acc’: test_acc,
os.mkdir(’checkpoint’) torch.save(state, ’checkpoint/’ + self.target_checkpoint_filename) self.target_best_acc = test_acc
self.target_logger.close() print("Target Done")

heading: E DG SUPPLEMENTS
text: Compute Resources We run the 10, 560 jobs on NVIDIA GEFORCE RTX 2080 TI GPUs as well as NVIDIA TITAN RTX GPUs for the more resource intensive jobs.
Baseline models We compare our algorithms to the following existing algorithms:

heading: Dataset Name Support of D Number of Samples Image Dimensions Number of Classes
text: Implementation To be more concrete, we change the code that computes the distributional distance penalty from this:
for i in range(nmb): for j in range(i + 1, nmb):
penalty += self.dist_loss(features[i], features[j])
if nmb > 1: penalty /= (nmb * (nmb - 1) / 2)
to this:
first = None second = None
for i in range(nmb): slice = random.randint(0, len(features[i]))
if first is None: first = features[i][:slice] second = features[i][slice:]
else: first = torch.cat((first, features[i][:slice]), 0) second = torch.cat((second, features[i][slice:]), 0)
penalty = self.dist_loss(first, second)
.
Here is the concrete full class of our CausIRL with MMD model:
class CausIRL_MMD(ERM): def __init__(self, input_shape, num_classes, num_domains, hparams):
super(CausIRL_MMD, self).__init__(input_shape, num_classes, num_domains, hparams)
self.kernel_type = "gaussian"
def my_cdist(self, x1, x2): x1_norm = x1.pow(2).sum(dim=-1, keepdim=True) x2_norm = x2.pow(2).sum(dim=-1, keepdim=True) res = torch.addmm(x2_norm.transpose(-2, -1),
x1, x2.transpose(-2, -1), alpha=-2).add_(x1_norm)
return res.clamp_min_(1e-30)
def gaussian_kernel(self, x, y, gamma=[0.001, 0.01, 0.1, 1, 10, 100, 1000]):
D = self.my_cdist(x, y) K = torch.zeros_like(D)
for g in gamma: K.add_(torch.exp(D.mul(-g)))
return K
def mmd(self, x, y): Kxx = self.gaussian_kernel(x, x).mean() Kyy = self.gaussian_kernel(y, y).mean() Kxy = self.gaussian_kernel(x, y).mean() return Kxx + Kyy - 2 * Kxy
def update(self, minibatches, unlabeled=None): objective = 0 penalty = 0 nmb = len(minibatches)
features = [self.featurizer(xi) for xi, _ in minibatches] classifs = [self.classifier(fi) for fi in features] targets = [yi for _, yi in minibatches]
first = None second = None

heading: Algorithm ColoredMNIST RotatedMNIST VLCS PACS OfficeHome TerraIncognita DomainNet Avg

heading: E.0.1 MODEL SELECTION: TRAINING-DOMAIN VALIDATION SET
text: We present here the results of our DG experiments for the training-domain validation model selection method. Result are summarized in Table 10. For CausIRL with CORAL, the overall performance is slightly below vanilla CORAL. CausIRL with CORAL especially underperforms CORAL on the PACS dataset. On the other hand, CausIRL with CORAL performs better than CORAL on DomainNet. For CausIRL with MMD, the overall performance is significantly better than MMD. This overperformance is mainly driven by the results on TerraIncognita and DomainNet, where for the latter we observe a leap in accuracy from 23.4% to 40.3%.

heading: E.1 MODEL SELECTION: LEAVE-ONE-DOMAIN-OUT CROSS-VALIDATION
text: We here present the complete results for the leave-one-domain-out cross-validation model selection method in Table 11.

heading: Algorithm ColoredMNIST RotatedMNIST VLCS PACS OfficeHome TerraIncognita DomainNet Avg

heading: Algorithm ColoredMNIST RotatedMNIST VLCS PACS OfficeHome TerraIncognita DomainNet Avg

heading: E.1.1 MODEL SELECTION: TEST-DOMAIN VALIDATION SET (ORACLE)
text: Finally, we here look at the DG experiment results for the test-domain validation set model selection. The results are summarized in Table 12. This setting is less realistic as we have access to test samples during training, but it is still useful as it shows the best possible model for each algorithm. It allows us to evaluate whether there is headroom for improvement for each algorithm and to see which algorithm has the inductive bias that more closely fit the task.
For both CausIRL with CORAL and CausIRL with MMD, we observe a better overall performance compared to their vanilla counterparts. We even have that CausIRL with CORAL is the best overall performing algorithm among the evaluated algorithms. Once again, we observe a large difference in performance on DomainNet between MMD and CausIRL with MMD, going from an average accuracy of 23.5% to 40.6%. We also again have that CausIRL with CORAL is the best algorithm for DomainNet compared to all the other algorithms, which explains why we perform well on this dataset.

Review summaries:
- The paper claims to provide a unifying causal framework for invariance-based algorithms. - It uses the graphical model in figure 1 to derive some conditions for achieving invariance. - It also develops a new algorithm for invariance-based representation learning. The idea is based on conjecture 4.3, which hypothesizes that one can create new domains using mixtures of existing ones. 

This paper considers a causal framework for invariant representation learning, with some analysis about the invariance and independences based on the proposed causal graph. An algorithm is then proposed based on the above analysis and empirical result validates the performance of the proposed algorithm. The paper is well written.

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.