Causal Balancing for Domain Generalization
Authors: Xinyi Wang
Michael Saxon
Jiachen Li
Hongyang Zhang
Kun Zhang
William Yang Wang
Venue: ICLR 2023 poster (2023)
Keywords: domain generalization
causality
latent variable model

Abstract: While machine learning models rapidly advance the state-of-the-art on various real-world tasks, out-of-domain (OOD) generalization remains a challenging problem given the vulnerability of these models to spurious correlations. We propose a balanced mini-batch sampling strategy to transform a biased data distribution into a spurious-free balanced distribution, based on the invariance of the underlying causal mechanisms for the data generation process. We argue that the Bayes optimal classifiers trained on such balanced distribution are minimax optimal across a diverse enough environment space. We also provide an identifiability guarantee of the latent variable model of the proposed data generation process, when utilizing enough train environments. Experiments are conducted on DomainBed, demonstrating empirically that our method obtains the best performance across 20 baselines reported on the benchmark.

Published as a conference paper at ICLR 2023
CAUSAL BALANCING FOR DOMAIN GENERALIZATION
Xinyi Wang', Michael Saxon', Jiachen Li", Hongyang Zhang', Kun Zhang
William Yang Wang
DDeaarmeett of Computer Science, University of California, Santa Barbara, USA
2 David R. Cheriton School of Computer Science, University of Waterloo, Canada "Department of Philosophy, Cameeii Mellon University, USA
Maachnne Learning Department, Mohamed bin Zayed University of Artificial Intelligence, UAE xinyi wang@uccea Sungguucbb... edu, saxon@ucets ucsb. edu, j jiachen 11 LiBucsb. edu, hongyang. zhang@uwaterloo kunz @cmu edu, wi aalllimm aimmmaaee ucsb.edu
ABSTRACT
While machine learning models rapidly advance the state-of-the-art on vari- ous real-world tasks, out-of-domain (OOD) generalization remains a challeng- ing problem given the vulnerability of these models to spurious correlations. We propose a balanced mini-batch sampling strategy to transform biased data dis- tribution into spurious-free balanced distribution, based on the invariance of the underlying causal mechanisms for the data generation process. We argue that the Bayes optimal classifiers trained on such balanced distribution are minimax optimal across diverse enough environment space. We also provide identi- fiability guarantee of the latent variable model of the proposed data generation conducted
process, when utilizing enough train environments. Experiments on DomainBed, demonstrating empirically that our method obtains the best per- formance across 20 baselines reported on the benchmark.
INTRODUCTION
Machine learning is achieving tremendous success in many fields with useful real-world applica- tions (Silver et al., 2016; Devlin et al., 2019; Jumper et al., 2021). While machine learning models perform well on in-domain data sampled from seen environments, they often fail to generalize to out-of-domain (OOD) data sampled from unseen environments (Quiñonnero (Quiñonero-Candlla et al., 2009; Szegedy et al., 2014). One explanation is that machine learning models prone to learning spu- rious correlations that change between environments. For example, in image classification, instead of relying on the object of interest, machine learning models easily rely on surface-level textures (Jo & Bengio, 2017; Geirhos et al., 2019) or background environments (Beery et al., 2018; Zhang al., 2020). This vulnerability changes in environments can cause serious problems for machine learning systems deployed in the real world, calling into question their reliability over time.
Various methods have been proposed improve the OOD generalizability by considering the in- variance of causal features or the underlying causal mechanism (Pearl, 2009) through which data is generated. Such methods often aim find invariant data representations using new loss function de- signs that incorporate some invariance conditions across different domains into the training process (Arjovsky al., 2020; Mahajan et al., 2021; Liu et al., 2021a; Lu et al., 2022; Wald et al., 2021). Unfortunately, these approaches have to contend with trade-off's between weak linear models ap- proaches without theoretical guarantees (Arjovsky et al., 2020; Wald et al., 2021), and empirical studies have shown their utility in the real world to be questionable (Gulrajani & Lopez-Paz, 2020).
In this paper, we consider the setting that multiple train domains/environmentt are available. We the- oretically show that the Bayes optimal classifier trained on a balanced (spurious-five free)) distribution is minimax optimal across all environments. Then we propose principled two-step method to sam- ple balanced mini-batches from such balanced distribution: (1) learn the observed data distribution using a variational autoencoder (VAB) and identify the latent covariate; (2) match train examples
We publicly release mansal-balanhe foreiinn--oooddoooo for --neeerrlll zation.
code at https://g http:/ttt...c......oo.toom

---

Published as a conference paper at ICLR 2023
Z
2
(a) Observed distribution p(X, =
(b) Balanced distribution p(X; =
Figure 1: The causal graphical model assumed for data generation process in environment E. Shaded nodes mean being observed and white nodes mean not being observed. Black arrows mean causal relations invariant across different environments. The Red dashed line means correlation varies across different environments.
with the closest latent covariate to create balanced mini-batches. By only modifying the mini-batch sampling strategy, our method is lightweight and highly flexible, enabling seamless incorporation with complex classification models or improvement upon other domain generalization methods.
Our contributions are as follows: (1) We propose a general non-linear causality- framework for the domain generalization problem of classification tasks; (2) We prove that a spurious- free bal- anced distribution can produce minimax optimal classifiers for OOD generalization; (3) We rigor- ously demonstrate that the source of spurious correlation,
latent variable, can be identified given large enough set of training environments in a nonlinear setting; (4) We propose novel and prin- cipled balanced mini-batch sampling algorithm that, an ideal scenario, can remove the spurious correlations in the observed data distribution; (5) Our empirical results show that our method obtains significant performance gain compared to 20 baselines on DomainBed (Arjovsky et al., 2020).
PRELIMINARIES
Problem Setting. We consider standard domain generalization setting with a potentially high- dimensional variable X (e.g. an image), a label variable Y and discrete environment (or domain) variable E' in the sample spaces X, X.Y, E, respectively. Here we focus on the classification problems with y = (1,2, ..., m and X' C Rd. We assume that the training data are collected from a finite subset of training environments Enain C E. The training data D = {(r] ) from the distribution p" Y) p( E e) for all Erain- Our goal is to learn classifier Cg:: X + )' that performs well in a new, unseen environment exest & Emain-
No is then sampled
We assume that there is data generation process of the observed data distribution (X, repre- sented by an underlying structural causal model (SCM) shown in Figure la. More specifically, we assume that X is caused by label Y, an unobserved latent variable (with sample space Z E R") and independent noise variable with the following formulation:
X f(Y, Z) + fy (Z)
Here, assume the causal mechanism is invariant across all environments e E E and we further characterize with the following assumption:
Assumption 2.1. f 21,2, m} x X injective. f-1 X {1, 21.2..., m] x Z the left inverse f.
Note that this assumption forces the generation process of X to consider both Z and Y instead of only one of them. Suppose has a known probability density function P. > 0. Then we have
Pr(X|Z,Y) = (x (Z)).
While the causal mechanism is invariant across environments, we assume that the correlation be- tween label Y and latent Z is environment-variant and Z should exclude Y information. i.e., Y cannot be recovered as a function of Z. If Y is function of Z, the generation process of X can completely ignore and f would not be injective. We consider the following family of distributions:
(1) Then the environment space we consider would be all the index of F: E = {el E Note that any mixture of distributions from F would also be member of F. i.e. Any combination of the environments from would also be an environment in
F (X,,, Z) = px Y (Z(Y)p² (P) (Z(Y), pf (V) > 20e.
To better understand our setting, consider the following example: image X of object in class Y has an appearance driven by the fundamental shared properties of Y as well as other meaning- ful latent features Z that do not determine "e-nsss". but can be spuriously correlated with Y. In
2

---

Published as a conference paper at ICLR 2023
Rackgrand Poowegg-ppcllccceee
Noise
Number color: red
Z
"Small numbors rod' color environment
Noise
Sketch class
(rather photo)
Appesrance
Appesrance
Object cless Giraffe
A small number (x5)
4
(a) realized by ColoredMNIST
(b) As realized by PACS
Figure 2: Annotated example causal graphs of two realizations of the joint distribution p(X,Y,
Figure 2, we plot causal diagrams for the joint distributions p(X, Y. E) of two example domain generalization datasets, ColoredMNIST (Arjovsky al., 2020) and PACS (Li al., 2017). In ColoredMNIST, Z indicates the assigned color, which is determined by the digit label and the environment E' p(Z] Y). In PACS, images of the same objects in different styles (e.g. sketches and photographs) occur in different environments, with Z containing this stylistic information.
In this setting, we can see that the correlation between X and Y would vary for different values of e. We argue that the correlation Y ++ Z + X is not stable in an unseen environment E &uain as involves E and we only want to learn the stable causal relation Y X. However, the learned predictor may inevitably absorb the unstable relation between X and Y if we simply train it on the observed train distribution yf (X, Y) with empirical risk minimization.
Balanced Distribution. To avoid learning the unstable relations, we propose to consider balanced distribution:
Definition 2.2. A balanced distribution can be written as p (X, - (2)) (Y),
where pd (Y)= U(1,2,... m) and LB B Z.
Here we do not specify (Z). Note that 2) F can also be regarded constructing
a result of the unchanged
causal mechanism % - X Y, and that (X,Y, an new environment B E E. In this new distribution, X and Y are only correlated through the stable causal relation + X. We want argue that the Bayesian optimal classifier trained on such a balanced distribution would have the lowest worst-case risk, compared to Bayesian optimal classifiers trained on other environments in & as defined in Equation (1). To support this statement, we further assume some degree of disentanglement of the causal mechanism:
Assumption 2.3. There exist functions gy, gz and noise variables ey, ez, such that (Y,2) f- = (gy (X- - cy),8z(X - -zz),, and ey 1L ez.
The above assumption implies that Y Le Z We can then have the following theorem?:
Theorem 2.4. Consider classifier (X) = arg maxy PY (X) with parameter th. The risk of such classifier
environment - (X.) lYgp,(Y(X)) Assume that ε satisfies:
E is its cross entropy: (Y
Ve E,Y 14. pz - 3e' E E s.l. L' (Vx ()" >0.
Then the Bayes optimal classifier trained on any balanced distribution DB (x is minimax opti- mal across all environments in E:
B (YIX)
arg min max L (x(p+ eEE
P.ET
The assumption implies that the environment space & is large and diverse enough such that a perfect classifier on one environment will always perform worse than random guessing on some other envi- ronment. Under such assumption, no other Byes optimal classifier produced by an environment in & would have a better worst case OOD performance than the balanced distribution.
3 METHOD
We propose a two-phased method that first use VAE to learn the underlying data distribution px(X, Y, Z) with latent covariate Z for each e E Emains and then use the learned distribution to calculate balancing score to create a balanced distribution based on the training data.
See Appendix A for proofs of all theorems.
3

---

Published as a conference paper at ICLR 2023
3.1 LATENT COVARIATE LEARNING
We argue that the underlying joint distribution of p² (X) 2) can be learned and identified by a VAE, given sufficiently large of train environments Erainn To specify the correlation between Z and Y, we assume that the conditional distribution y (Z is conditional factorial with an expo- nential family distribution:
Assumption 3.1. The correlation between Y and Z in environment e is characterized by:
II (Z)).) 2=1 (YY
PTA (Z (Z
Σ J=1
exp
where Ze is the i-th element of Z. Q [Q1]: Z +
RR is the base measure, W [W]] :
y + Th is the normalizing constant, T = [Ty]ls - Tank is the sufficient statistics, and
x [igl : RRn are the Y dependent parameters.
Here n is the dimension of the latent variable Z. and k is the dimension of each sufficient statis- tic. Note that k. Q, and T determined by the type of chosen exponential family distribution thus independent of the environment. The simplified conditional factorial prior assumption is from the mean-field approximation, which can be expressed as a closed form of the true prior (Blei et al., 2017). Note that the exponential family assumption is not very restrictive as it has universal ap- proximation capabilities (Sriperumbudur et al., 2017). We then consider the following conditional generative model in each environment Exrinn with parameters @ (f, A))):
XX, Z pr( T) (Z
We use a VAE to estimate the above generative model with the following evidence lower bound (ELBO) in each environment Eurinn
Ep [Xox)((XI)) 2 L5 = Ep [Bag(2)x,y)
DDr. (xt.
The KL.-divergence term can be calculated analytically. To sample from the variational distribution & (X), we use reparameterization trick (Kingma & Welling, 2013).
We then maximize the above ELBO Enain EeEEemn Ce over all training environments to obtain model parameters (0, d). To show that can uniquely recover the latent variable Z up to some simple transformations, we want to show that the model parameter 0 identifiable up to some simple transformations. That is, for any {0- (f, T, A), 0 (f', (')) E O
() = pe. (X Ve E Egain = 0 0'
where θ is the parameter space and ~ represents an equivalent relation. Specifically, we consider the following equivalence relation from Motian al. (2017):
Definition 3.2. If (f, T, A) ~ (f',T', T' X'), then there exists an invertible matrix A E Bunxxnk and vector n such that T(f-1()) (f- (x)) + c, Vr X.
When the underlying model parameter 6" can be recovered by perfectly fitting the data distribution ph- (X for all Errain the joint distribution pf. (X, Z|Y) is also recovered. This further implies the recovery of the prior Pf- (2lY) and the true latent variable 2 The identifiability of our proposed latent covariate learning model can then be summarized as follows:
Theorem 3.3. Suppose we observe data sampled from the generative model defined according to Equation (2), with parameters 0 (f,T, >). addition Assumption 2.1 and Assumption 3.1, we assume the following conditions holds: (1) The set {i E X (x)x(x)= = 0} has measure zero, where de the characteristic function of the density Pe. (2) The sufficient statistics Tiş are differentiable almost everywhere, and (Tú) (iijk are linearly independent on any subset of X of measure greater than zero. (3) There exist nk + distinct pairs (yo, eo)..... (ynnk Crn) such that the nk x nkk matrix
L L= (X c0 (yoo..... A Crck (yrr)) xx0 invertible. Then we have the parameters 0 (f, T, are ~ --deentifiable
Note that in the last assumption in Theorem 3.3, since there exists nk + 1 distinct points (yi, ez), the product space y' Evain has to be large enough. i.c. We need m m|Euaia > nk. The invertibility of L implies that (y:)) x (y)) need to be orthogonal to each other which further implics the diversity of environment space E.
4

---

Published as a conference paper at ICLR 2023
3.2 BALANCED BANIH SAMPLING
We consider using a classic method that has been widely used in the average treatment effect (ATE) estimation balancing score matching (Rosenbaum & Rubin, 1983) to sample balanced mini- batches that mimic a balanced distribution shown in Figure lb. A balancing score is used to balance the systematical difference between the treated unites and the controlled units, and to reveal the true causal effect from the observed data, which is defined as below:
Definition 3.4. A balancing score b( b() is a function of covariate Z s.I. Z TL Y|b(Z).
There is a wide range functions of Z that can be used as a balancing score, where the propensity score pYY 1[2) is the coarsest one and the covariate Z itself is the finest one (Rosenbaum & Rubin, 1983). To extend this statement to non-binan non-binary treatments, we first define propensity score zZ) for Y E y {1.2, m) as vector: Definition 3.5. The propensity score for Y E {1,,2,... m) ((Z) = lp -
We then have the following theorem that applies to the vector version of propensity score s(Z): Theorem 3.6. Let b(Z) be a function of Z. Then (( is a balancing score, if and only if ((Z) is finer than s()). i.e. exists function g such that s(Z) g(b(2)).
We use (Z) denote the balancing score for a specific environment e. The propensity score would then be (Z) [p=(Y y y/ y|>)|}== which can be derived from the VAE's conditional prior
PT, (2)Y) as defined in Equation (2):
p² (Y = v1 = P (Z = v) =
(3)
Σ 12)
where p² (Y 1) can be directly estimated from the training data D
In practice, we adopt the propensity score computed from Equation (3) as our balancing score (b(Z) s" (Z)) and propose to construct balanced mini-batcher by matching 1 S m different examples with different labels but the same/closest balancing score, If (Z) E B, with each train example. The detailed sampling algorithm is shown in Algorithm
Algorithm 1: Balanced Mini-batch sampling.
Input: |rrii training datasets 2 {(1f of ) for all e Euain a balancing score (z(z) inferred from each training data point (x) of ). and a distance metrics d: B x B + R; Output: balanced batch of data Dbalenced consisting of B x |Euaial (a 1) examples;
D Dbalaweed t- Empty:
for e E Exnin do
Randomly sample B data points D Frandom from D"; Add D Drandom to Dbatancedi
for (x.,v") Drandom, do
Yau = {y: ~ 2,,,,,..., m) \ fv" yi..., y1 1i E [1,a]]: Compute balancing score be (2¹) from
for yi E Yatt do
j = arg min d(b (2z), (e (x*)) such that yf yi and (f)) E Dc:
Add vf) to Dbalanced-
We denote the data distribution obtained from Algorithm by (( Y, Z, F),, then we have:
Theorem 3.7. If ((b)),)), 5 (x²)) = O in Algorithm 1, the balanced mini-matee can be regarded as sampling from a bami--alanced distribution with & (Y(()Z,, E) = a+I - + m- When m 8 (Y E)
B (Y).
With perfect match at every step (i.c., (2,) (2)) and a = m - 1, we can obtain a completely balanced mini-batch sampled from the balanced distribution. However, an exact match of balancing score unlikely in reality, so larger a will introduce more noises. This can be mitigated by choosing smaller a, which the other hand will increase the dependency between Y and Z. So in practice, the choice of a reflects a trade-off between the balancing score matching quality and the degree of dependency between Y and Z.
5

---

Published as a conference paper at ICLR 2023
12
22 7
24 8
/
5
6
9
6
0
3
8
4
5
6
9
PM
yd
P2
PM
n
H1
Hd
Pd
12
15
95
m7
Mt 8
6
4
4
2
6
9
5
3
8
O
2
3
4
5
6
7
9
3
1
2
8
8
9
/
2
3
4
1
6
7
8
9
(a) A random mini-batch.
(b) A balanced mini-batch (obtained by our method).
Figure 3: A random mini-batch and a balanced mini-batch from the ColoredMNISTT dataset. Note that there is 25% label noise so mismatches of label y and image are expected.
4 EXPERIMENTS
Datasets: To verify the effectiveness of our proposed balancing mini-batch method, we conduct experiments on DomainBed standard domain generalization benchmark, which contains seven different datasets: ColoredMNIST (Arjovsky et al., 2020), RotatedMNIST (Ghifary et al., 2015), VLCS (Fang et al., 2013), PACS (Li et al., 2017), OfficeHome (Venkateswara et al., 2017), Ter- raIncognita (Beery al., 2018) and DomainNet (Peng et al., 2019). We also report results slightly modified version of ColoredMNIST dataset, ColoredMNISTT (Bao et al., 2021), which classify digits into 10 classes instead of binary classes.
Baselines: We apply our proposed balanced mini-batch sampling method along with four represen- tative widely-used domain generalization algorithms: empirical risk minimization (ERM) (Vapnik, 1998), invariant risk minimization (IRM) (Arjovsky et al., 2020), GroupDRO (Sagawa et al., 2019) and deep CORAL (Sun& & Saenko, 2016), and compare the performance of using our balanced mini- batch sampling strategy with using the usual random mini-batch sampling strategy. We compare our method with 20 baselines
total (Xu et al., 2020; Li et al., 201 Sa; Ganin et al., 2016; Li et al., 201 1c;;;; Krueger et al., 2021; Blanchard et al., 2021; Zhang et al., 2021; Nam et al., 2021; Huang et al., 2020; Shi et al., 2022: Parascandolo et al., 2021; Shahtalebi et al., 2021; Rame al., 2022; Kim et al., 2021) reported on DomainBed, including recent causality based baseline L and CausTRL (Chevalley et al., 2022) that also utilize the invariance of causal mechanisms. We also compare with group- based method PI (Bao et al., 2021) that interpolates the distributions of the correct predictions and the wrong predictions on ColoredMNNIST"".
To control the effect of the base algorithms, we use the same set of hyperparameters for both the random sampling baselines and our methods. We primarily consider train domain validation for model selection, it is the most practical validation method. A detailed description of datasets and baselines, and hyperparameter tuning and selection can be found in Appendix B.
ColoredMNIST: We use the ColoredMNIST dataset proof of concept scenario, as we already know color is dominant latent covariate that exhibits spurious correlation with the digit label.
For ColoredMNIST". of the original ColoredMNIST dataset (Arjovsky et al., 2020). The label is assigned according to the numeric digit of the MNIST image with 25% random noise. Then we assign one of set of 10 colors (cach indicated by a separate color channel) to the image according to the label with probability that we assign the corresponding color and probability 1 - e we randomly choose another color. Here e E {0.1, 0.2} for two train environments and e = 0.9 for the test environment. For ColoredMNIITT we adopt the original setting from (Arjovsky et al., 2020), which only has two classes (digit smaller/larger than 5) and two colors, with three environments E (0.1, 0.2, 0.9).
adopt the setting from (Bao et al., 2021), which is multiclass version
Balanced mini-batch example. An example of a balanced mini- batc- created by our method from digit 4, 5 and 7 in ColoredMNISTT is illustrated in Figure 3. In the random mini-batch, labels are spuriously correlated with color. e.g. most are blue, most 1 are red and most are yellow. In the balanced mini-batch, example with an example with a different label but the same color. Here, the color information is implicitly learned by latent covariate learning.
force each label to have uniform color distribution by matching each
ColOredMNIST main results. Table 1 shows the out-of-domain accuracy of our method combined with various base algorithms on ColoredMNIST"0 and ColoredMNIST dataset. Our balanced mini-
https: com http:////e..cccooeeeeeeeeeeeeeee Doma niedd

---

Published as a conference paper at ICLR 2023
Table 1: Out-of-domain accuracy on ColoredMNIST and ColoredMNIST with two train environ- ments [0.1, 0.2] and one test environment [0.9].
Validation
Dataset
Sampling
ERM
IRM
GroupDRO
CORAL
CausTRI.
PT
CMNIST
Random
14.25
13.13
21.06
13.1 0.3
12.5 + 0.1
69.68
Ours
69.8 03
63.8 us
69.3 20
70.1 1:02
69.6 1:
Train
CMNIST
Random
10.0 au
10.2 + 0.
10,0 =02
9.9 0.1
10.0
Ours
37.6 2.
31.1 1
17.0 33
57.2 4:44
43.7 299
CMNIST
Random
26.15
45.41
32.51
21.1 +: 0.1
20.8
69.44
Ours
70.5 1.
63.8
69.4 0.3
70.1
69.6 03
Test
Random
28.7 a
58.5 + sa
36.8 #2
31.1 + 1.6
27.4 + 0.3
CMNIST
Ours
38.4 1
69.7 16.
44.8 1100
60.5 4.1
43.3 1 192
70
80
Accuracy
Accuracy
Pccurasy
60
40
20
train
train val
random
test val
40
test vel
balanced (ours)
0. 1.0
4
0.2
0.8
(a) Degree of balancing
(b) number of matched examples
(c) Test env
Figure 4: The out-of-domain accuracy versus (a) degree of balancing, (b) number of matched exam-
ples a, and (c) test environment, on ColoredMNIST dataset with ERM base algorithm.
batch sampling can increase the accuracy of all base algorithms by large margin, with CORAL improving the most (57% and 47.3%). Note that the highest possible accuracy without relying on the color feature is 75%.
In Figure 4, we study important factors in our proposed method by ablating on the ColoredMNIST dataset with ERM.
The effectiveness of balancing. We construct oracle balanced mini-batches with b(Z) Color, and then control the degree of balancing by varying the fraction of balanced examples in a mini-batch for each randomly sampled example, with probability B. we match with 9 examples with the same color but different labels to balance the mini-batch; otherwise, we match it with 9 examples with the same color and label to maintain the original distribution. Figure 4a shows that increasing the balancing fraction would increase the OOD performance.
The effect of the number of matched examples a. Figure 4b shows that when a increases, the OOD performance first increases, then becomes stable with slightly decreasing trend. This result consistent with our analysis in Section 3.2, that a large a will increase balancing in theory, but due imperfection of the learning of latent covariate Z, large a will eventually introduce more low-quality matches, which may hurt the performance. It can also be observed that we do not need very large a to reach the maximum performance.
The effect of different test environments. In Figure 4c, we fix the train environments as [0.1, 0.2] and test on different test environments. We report the results chosen by train domain validation, as the results with test domain validation are almost the same as the training domain validation results. The accuracy of the model trained with random mini-batches drops linearly when the test environment changes from 0.1 to 0.9, indicating that the model learns to use the color feature as the main predictive evidence. On the other hand, the accuracy of the model trained with balanced mini-batches produce by our method almost stays the same across all test domains, indicating that the model learns to use domain- invariant features.
DomainBed: We investigate the effectiivness of our method under different situations.
DomainBed main results. In Table 2, we consider combining our method with four representative base algorithms: ERM, IRM, GroupDRO, and CORAL. IRM represents wide range of invariant representation learning baselines. GroupDRO represents group-based methods that minimize the worst group errors. CORAL represents the distribution matching algorithms that match the feature distribution across train domains. In general, our method can improve the average performance of all the base algorithms by one to two points (1.6% for ERM, IRM and GroupDRO), while CORAL
7

---

Published as a conference paper at ICLR 2023
Table 2: Out-of-domain accuracy on DomainBed benchmark. Numbers are averaged over all test environments with standard deviation over 3 runs. The training domain validation scheme is used. Full results on each test environment can be found in Appendix B.4.
Algorithm
CMNIST
RMNIST
VLCS
PACS
Offece-Home
Terrainc
DomainNet
Avg
ERM
51.5
0.0
77.5
85.5
0.2
66.5 64.3
46.1
1.8
40.9
66.6
IRM
52.0 0.1
97.7
78.5 0.5
83,5 0.8
47.6 = 0.8
33,9 + 2.8
65,4
GroupDRO
0.0
9.00
76.7 0.6
8.84
66.0
43.2 1.1
33.3 0.2
64.8
Mixup
5.21
98.0
0..4
0.66
6,
47.9 = 0.8
0.1
66.7
MLDG
51.
97
7724044
1.9
6..8
47.7
4..
CORAL
51.5
98.000.11
7.
8..2
68.7
47.6 1.0
4.1
67.5
MMD
5.2
97.9
7..
0.6
66.
4.
2..4
63.3
DANN
0.3
97.
78.6
83.6
6.
46.7
3..
66.1
CDANN
5.17
9..
7..
8,66
63
458
0.3
MTL
5.4
9.
84.6
6.
45.6 1.2
0.1
66.2
SagNet
5.0
9..0
0.5
8.2
68.1
4..
67.2
ARM
5.2
98.1
7.
85
6,
45.5
3,25
66.1
VREx
518
9..
4.22
6.44066
33.6
RSC
5.2
97.6
0.5
0,9
0.
46.6 1.0
3.
66.1
Fish
5.66
9.00
0.3
0.3
68.6
4..1
42
67.1
Fishr
52.0
9.8
77.8
0.01
85.5 0.4
6.18
474
4..7
67.1
AND-mask
5.2
9..6
7..9
8.
65.6
0.
3.6
SAND-mask
5.2
9..
7.2
0.96
65
4.7
0.61
64.6
SelfReg
5.21
98.0
7.99
8..6
6.7
47.0 0.3
4.8
67.3
CausIRLcORAL
5.17
97.
7.6
8.1
68
47.3 0.8
4.1
67.3
Cays|RLy
51.6
9.09
7..6
0.8
6.67
46.3 0.9
40.
66.2
Ours+ERM
60.1 1.0
9.07
7..1
86.1 66.10..
67.1
48.0 1.7
42.6 1.0
68.2
Ours+IRM
85.2
646
40.5
Ours+DRO
53,9 =
9.16
7..0
84.9
0.5
45.4 0.4
40.8 + 0.6
66.4
Ours+CORAL
66.6 + 1.2
0.1
7.54
86.7 0.1
69.6 0.2
47.0 = 1.2
43.9 + 0.1
69.7
improves the most (2.2%). The reason why CORAL works the best with our method, and achieves the state-of-the-art O01) accuracy not only on average but also on ColoredMNIST, PACS, Office- Home and DomainNet dataset, is likely because our method aims to balance the data distribution and close the distribution gap between domains, which line with the objective of distribution matching algorithms.
Our proposed method improves most on ColoreMNIST, OfficeHome, and DomainNet, while our method is not very effective RotatedMNIST and VLCS.
Reason for significant improvements. The large improvement on ColoredMNIST (8.6% for ERM, 7.2% for IRM, 1.8% for GroupDRO and 15.1 1 % for CORAL) is likely because the dominant latent covariate, color, is relatively easy to learn with a low dimensional VAE. The good performance on OfficeHome and DomainNet (1.7% for ERM, 6.6% for IRM, 7.5% for GroupDRO and 2.4% for CORAL) is likely because of the large number of classes. OfficeHome has 65 classes, and Domain- Net has 345 classes, while all the other datasets have less equal to 10 classes. According the conclusion of Theorem 3.3, larger number of labels or environments will enable the identification of a higher dimensional latent covariate, which is more likely capture the complex underlying data distribution.
Reason for insignificant improvements. The lower performance on RotatedMNIST's because the digits in each domain are all rotated by the same degree. Since classes are balanced, images in each domain are already balanced for rotation, the dominant latent covariate. As the performance with random mini-batches is already very high, the noise introduced by the matching procedure may hurt performance. VLCS on the one hand has a pretty complex data distribution as the images from each domain are very different realistic photos collected in different ways. However, VLCS only has classes and 4 domains, which only enables the identification of a very low dimensional latent covariate, which insufficient to capture the complexity of each domain.
In practice we suggest using our method when there is a large number of classes or domains, and preferably combined with distribution matching algorithms for domain generalization.
RELATED WORK
A growing body of work has investigated the oft-oi-dommin (OOD) generalization problem with causal modeling. One prominent idea leamn invariant features. When multiple training domains
8

---

Published as a conference paper at ICLR 2023
available, this can be approximated by enforcing some invariance conditions across training domains by adding a regularization term to the usual empirical risk minimization (Arjovsky et al., 2020: Krueger et al., 2021: Bellot & van der Schaar, 2020; Wald et al., 2021; Chevalley al., 2022). There are also some group-based works (Sagawa et al., 2019; Bao et al., 2021; Liu et al., 2021b; Sanh et al., 2021; Piratla et al., 2021; Zhou et al., 2021) that improve worst group performance and can be applied to domain generalization problem. However, recent work claims that many of these approaches still fail to achieve the intended invariance property (Kamath al., 2021; Rosenfeld et al., 2020; Guo et al., 2021), and thorough empirical study questions the true effectiveness of these domain generalization methods (Gulrajani & Lopez-Paz, 2020).
Instead of using datasets from multiple domains, Makar et al. (2022) and Puli et al. (2022) propose to utilize an additional auxiliary variable different from the label to solve the OOD problem, using single train domain. Their methods are two-phased: (1) reweight the train data with respect to the auxiliary variable: (2) add invariance regularizations to the training objective. The limitation of such methods is that they can only handle distribution shifts induced by the chosen auxiliary variable. Little & Badawy (2019) also propose a bootstrapping method to resample train data by reweighting to mimic randomized controlled trial. There is also which proposes new training objectives to reduce spurious correlations.
methods like Wang et al. (2021)
Some other OOD works aim to improve OOD accuracy without any additional information. Liu et al. (2021a) and Lu et al. (2022) propose to use VAE to learn latent variables in the assumed causal graph, with appropriate assumptions of the train data distribution in single train domain. The identifiability of such latent variables is usually based on Khemakhem et al. (2020), which assumes that the latent variable has factorial exponential family distribution given an auxiliary variable. Our identifiability result is also an extension of Khemakhem et al. (2020), where we use both label Y and training domain E as the auxiliary variable and include the label Y in the causal mechanism of generating X instead of only using the latent variable Z to generate X. Christiansen et al. (2021) use interventions on a different structural causal model to model the OOD test distributions and show a similar minimax optimal result.
To sample from the balanced distribution, we use a classic method for average treatment effect (ATE) estimation (Holland, 1986) - balancing score matching (Rosenbaum & Rubin, 1983). Causal effect estimation studies the effect a treatment would have had on a unit that in reality received another treatment. A causal graph (Pearl, 2009) similar to Figure la is usually considered in causal effect estimation problem, where Z is called the covariate (c.g. a patient profile), which is observed before treatment Y {0, (e.g. taking placebo or drug) is applied. We denote the effect of receiving specific treatment Y = y as X, (e.g. blood pressure). Note that the causal graph implies the Strong Ignorability assumption (Rubin, 1978). i.e. Z includes all variables related to both X and Y. In the case of a binary treatment, the ATE is defined as E[X1 - -Xo]-
For randomized controlled trial, ATE can be directly estimated by E[X Y = -
0].
as in this case Z JLI Y and there would not be systematic differences between units exposed to one treatment and units exposed to another. However, in most observed datasets, Z is correlated with Y. Thus [[X,] and [[[ are not directly comparable. We can then use balancing score b( (Z) (Dawid, 1979) to de-correlate Z and Y, and ATE can then be estimated by matching units with same balanc- ing score but different treatments: T[X; Xo) Es(z) [ 1, b.)()) E[X|Y O, b( b())]] Recently, Schwab et al. (2018) extends this method to individual treatment effect (ITE) estimation (Holland, 1986) by constructinng virtually randomized mini-batches with balancing score.
6 CONCLUSION
Our novel causality-based domain generalization method for classification task samples balanced mini-batches to reduce the presentation of spurious correlations in the dataset. We propose a spurious-free balanced distribution and show that the Bayes optimal classifier trained on such distri- bution is minimax optimal over all environments. We show that our assumed data generation model with an invariant causal mechanism can be identified up to sample transformations. We demonstrate theoretically that the balanced mini-batch is approximately sampled from spurious courious-free balanced distribution with the same causal mechanism under ideal scenarios. Our experiments empirically show the effectiveness of our method in both semi-synthetic settings and real-world settings.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
This work was supported by the National Science Foundation award #2048122. The views expressed are those of the author and do not reflect the official policy or position of the US government. We thank Google and the Robert N. Noyce Trust for their generous gift the University of Califor- nia. This work was also supported in part by the National Science Foundation Graduate Research Fellowship under Grant No. 16501 14. This work was also partially supported by the National Insti- tutes of Health (NIH) under Contract RO1HL.159803 by the NSF-Convergence Accelerator Track-I) award #21 134901, by grant from Apple Inc., a grant from KDDI Research Inc, and generous gifts from Salesforce Inc., Microsoft Research, and Amazon Research. This work was also supported by NSERC Discovery Grant
DGECR-2022-0033 2022-00377.
REFERENCES
Kartik Ahuja, Ethan Caballero, Dinghuai Zhang, -ha-hhhrsttpphe Gagnon-Audet, Yoshua Bengio, loannis Mitliagkas, and Irina Rish. Invariance principle meets information bottleneck for out-of- distribution generalization. Advances in Neural Information Processing Systems, 34:3438 -3500, 2021.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization, 2020.
Yujia Bao, Shiyu Chang, and Regina Barzilay. Predict then interpolate: A simple algorithm to learn stable classifiers. In International Conference on Machine Learning, pp. 440-650. PMLR, 2021.
Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In Proceedings of the European conference computer vision (ECCV), pp. 456-473, 2018.
Alexis Bellot and Mihaela van der Schaar. Accounting for unobserved confounding in domain generalization. arXiv preprint arXiv:2007.1065: 2020.
Gilles Blanchard, Aniket Anand Deshmukh, Urun Dogan, Gyemin Lee, and Clayton Scott. Domain generalization by marginal transfer learning. The Journal of Machine Learning Research, 22(1): 46-100, 2021.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisti- cians. Journal of the American statistical Association, 112(518):859-8777 877, 2017.
Yang Chen, Yu Wang, Yingwei Pan, Ting Yao, Xinmei Tian, and Tao Mei. A style and semantic memory mechanism for domain generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 9164-9173, 2021.
Mathieu Chevalley, Charlotte Bunne, Andreas Krause, and Stefan Bauer. Invariant causal mecha- nisms through distribution matching. arXiv preprint arXiv:2206. 1666, 2022.
Rune Christiansen, Niklas Pfister, Martin Emil Jakobsen, Nicola Gneeco, and Jonas Peters. A causal framework for distribution generalization. IEEE Transactions Pattern Analysis and Machine Intelligence, 2021.
A. P. Dawid. Conditional independence in statistical theory. Journal of the Royal Statistical Society, Series B (Methodological), 41(1):1-31, 1979. ISSN 00359246. URL http:/t////. /wwww jstor. org/ Dataalee2984720
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Pupers), pp. 4171-4186, 2019.
Chen Fang, Ye Xu, and Daniel N Rockmore. Unbiased metric learning: On the utilization of multiple datasets and web images for softening bias. ICCV, 2013.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. works. The journal of machine learning research, 17(1):29662008 2016.
training of neural net-
10

---

Published as a conference paper at ICLR 2023
Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A. Wichmann, and Wieland Brendel. Imagenet-trained cnns are biased towards texture; increasing shape bias im- proves accuracy and robustness, 2019.
Muhammad Ghifary, W Bastiaan Kleijn, Mengjie Zhang, and David Balduzzi. Domain generaliza- tion for object recognition with multi-task autoencoders. ICCV, 2015.
Ishaan Guirajani and David Lopez-Paz. In search of lost domain generalization. In International Conference on Learning Representations, 2020.
Ruocheng Guo, Pengchuan Zhang, Hao Liu, and Emre Kiciman. Out-of-distribution prediction with invariant risk minimization: The limitation and an effective fix. arXiv preprint arXiv:2101.077:2, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp- 770-778, 2016.
Paul W. Holland. Statistics and causal inference. Jourmal of the American Statistical Associa- fion, 816396),445-96, 1986. doi: 00.100000622459.1000 10478354, URL https: /www. handfonline. htom/doi//e atoii/abs//e / wtw.://///1111..11111 10478354.
Kevin D. Hoover. The logic of causal inference: Econometries and the conditional analysis of
causation. Economics and Philosophy, 6(2):207-234, 1990. doi: 10.1017/S9222661100012XX
Zeyi Huang, Haohan Wang, Eric Xing, and Dong Huang. Self-challenging improves cross-domain generalization. In European Conference Computer Vision, pp. 124-140. Springer, 2020.
Jason Jo and Yoshua Bengio. Mcasuring the tendency of CNNs to learn surface statistical regulari- ties, 2017.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-589) 2021.
Pritish Kamath, Akilesh Tangella, Danica Sutherland, and Nathan Srebro. Does invariant risk minimization capture invariance? In AISTATS, 2021.
Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoen- coders and nonlinear ica: A unifying framework. International Conference on Artificial Intel- ligence and Statistics, pp. 2207-2211 PMLR, 2020.
Dachee Kim, Youngjun Yoo, Seunghyun Park, Jinkyu Kim, and Jaekoo Lee. Selfreg: Self- supervised contrastive regularization for domain generalization. In Proceedings the International Conference Computer Vision, pp. 9619-9628, 2021.
the IEEE/CVF
Diederik P Kingma and Max Welling, arXiv:1312.61 14, 2013.
Auto-encoding variational bayes. arXiv preprint
David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapo- lation (rex). In International Conference on Machine Learning, pp. 5815-5826. PMLR, 2021.
Bo Li, Yifei Shen, Yezhen Wang, Wenzhen Zhu, Dongsheng Li, Kurt Keutzer, and Han Zhao. In- variant information bottleneck for domain generalization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 7399-7407, 2022.
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain generalization. Proceedings of the IEEE international conference on computer vision, pp- 5542-5550, 2017.
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy Hospedales. Learning to generalize: Meta-learning for domain generalization. Proceedings of the AAAI conference on artificial intelligence, vol- ume 32, 2018a.
11

---

Published as a conference paper at ICLR 2023
Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex C. Kot. Domain generalization with adver- sarial feature learning. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2018b.
Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao. Deep domain generalization via conditional invariant adversarial networks. In Proceedings the European Conference on Computer Vision (ECCV), September 2018c.
Max A Little and Reham Badawy. Causal bootstrapping. arXiv preprint
2019.
Chang Liu, Xinwei Sun, Jindong Wang, Haoyue Tang, Tao Li, Tao Qin, Wei Chen, and Tie-Yan Liu. Learning causal semantic representation for -uttt--ddittibutiinn prediction. Advances in Neural Information Processing Systems, 34, 2021a.
Evan Z Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In International Conference on Machine Learning, pp. 6781-6792. PMLR, 2021b.
Chaochao Lu, Yuhuai Wu, José Miguel Hemáández-Loato and Bernhard Schölkopf. Invariant causal representation learning for out-of-distribution generalization. International Confer- Learning Representations, 2022. URL https: / https:://penreeeeeeeee /owenrvveewe htt.:/ewwwwwweeeeeeeeeeeeeeeeeee forum?id= -e4EXDWXnSn.
Divyat Mahajan, Shruti Tople, and Amit Sharma. Domain generalization using causal matching, In International Conference on Machine Learning, pp. 7313-7324. PMLR, 2021.
Maggie Makar, Ben Packer, Dan Moldovan, Davis Blalock, Yoni Halpern, and Alexander D' A'mour Causally motivated shortcut removal using auxiliary labels. In Gustau Camps- Valls, Francisco J. R. Ruiz, and Isabel Valera (eds.), Proceedings of The 25th International Conference Arti- ficial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pp. 739-766. PMLR, 28-30 Mar 2022. makar22a. html.
URL https: https://proceeetteeeeeeeeeeeeeee http///poocc/ceee// ngs mlr press/v1511
Saeid Motiian, Marco Piccirilli, Donald A Adjeroh, and Gianfranco Doretto. Unified deep super- vised domain adaptation and generalization. In Proceedings of the IEEE international conference on computer vision, pp. 5715-5725, 2017.
Hyeonseob Nam, HyunJae Lee, Jongchan Park, Wonjun Yoon, and Donggeun Yoo. Reducing do- main gap by reducing style bias. In Proceedings the IEEE/CVF Conference Computer Vision and Puttern Recognition, pp. 8690-8699, 2021.
Giambattista Parascandolo, Alexander Neitz, ANTONIO ORVIETO, Luigi Gresele, and Bernhard Schölkopf. Learning explanations that are hard to vary. In International Conference on Learning Representattons, 2021. URL https: wtwp///pppeeeee.eeeeeeeee net forum? hhl
Judea Pearl. Causality. Cambridge university press, 2009.
Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In Proceedings of the IEEE International Conference on Computer Vision, pp. 1406-1415, 2019.
Vihari Piratla, Praneeth Netrapalli, and Sunita Sarawagi. Focus on the common good: Group distri- butional robustness follows. In International Conference Learning Representations, 2021.
Aahlad Manas Puli, Lily H Zhang, Eric Karl Oermann, and Rajesh Ranganath. distif-dist generalization in the presence of muisance-induced spurious correlations. In International Confer- ence on Learning Representations, 2022. URL https:/ /wwpoeeeevvveeeeeeee iew. net / forum? id 12.0R2032T.
Joaquin Quiñonero Quinonero-Candela. Masashi Sugiyama, Neil D Lawrence, and Anton Schwaighofer.
Dataset shift in machine learning. Mit Press, 2009.
12

---

Published as a conference paper at ICLR 2023
Alexandre Rame, Corentin Dancette, and Matthieu Cord. Fishr: Invariant gradient variances for out-of-distribution generalization. International Conference on Machine Learning, pp. 18347- 18377. PMLR, 2022.
Paul Rosenbaum and Donald B. Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1):41-55, 04 1983. ISSN 0006-3444. doi: 10.1093/ biomet/70.1.44 URL wwwpt:ttp///to
/00. http/////ttt///////ee//e//eee/ee
1.41.
Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. The risks of invariant risk minimization. arXiv preprini arXiv:2010.05761, 2020.
Donald B. Rubin. Bayesian Inference for Causal Effects: The Role of Randomization. The Annals of Statistics, 6(1):34 - 58, 1978. doi: 10. 176344464, URL ntttps://doi. 10. 10l11214/osss / 1/1//11133388 63444664
Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generaliza- tion. arXiv preprint wrXiv:/91.0055
2019.
Victor Sanh, Thomas Wolf, Yonatan Belinkov, and Alexander M Rush. Learning from others' mis- takes: Avoiding dataset biases without modeling them. In International Conference on Learning Representations, 2021. URL https: : http///ppeeee.eeeeeeeeeee htte://fouumnii-ffffoooooooo
Patrick Schwab, Lorenz Linhardt, and Walter Karlen. Perfect Match: A Simple Method for Learning Representations For Counterfactual Inference With Neural Networks. arXiv preprint arXiv:1810.00650, 2018.
Soroosh Shahtalebi, Jean- -ean-Christophe Gagnon-Audet, Touraj Laleh, Mojtaba Faramarzi, Kartik Ahuja, and Irina Rish. Sand-mask: enhanced gradient masking strategy for the discovery of invariances in domain generalization. arXiv preprint arXiv:2106.022 02266, 2021.
Yuge Shi, Jeffrey Seely, Philip Torr, Siddharth N, Awni Hannun, Nicolas Usunier, and Gabriel Synnaeve. Gradient matching for domain generalization. In International Conference on Learning Representations, 2022. URL https: wwwp//wwwww..p...opppppp net www.mmii-dooor Fid-vDWWWW 9HmO.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game go with deep neural networks and tree search. nature, 529(75877:4844488, 2016.
Bharath Sriperumbudur, Kenji Fukumizu, Arthur Gretton, Aapo Hyvärinen, and Revant Kumar. Density estimation in infinite dimensional exponential families. Journal of Machine Learning Research, 18(57):1 1-,9,,959229901 2017. URL http: jjmlr. org/ papers/pp/paae/eppp /vl.es/vlll. www./w110000 html.
Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In European conference on computer vision, pp. 443-450, Springer, 2016.
Xinwei Sun, Botong Wu, Xiangyu Zheng, Chang Liu, Wei Chen, Tao Qin, and Tie- Yan Liu. Recovering latent causal factor for generalization to distributional shifts. In M. Ran- zato, A. Beygelzimer, Y. Dauphin, P.S. Liang and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 16846-16859, Curran Asso- cc/
ciates, Inc., 2021. URL 8cC744242222238888533374444440
wttparipss.cppppp///2/2... ctcepppaer/22////1ee. pdf.
Christian Szegedy, Wojciech Zaremba, Ilya Sutskever, Joan Bruna, Dumitru Erhan, Ian Goodfellow, and Rob Fergus. Intriguing properties of neural networks, 2014.
Vladimir Vapnik. Statistical learning theory wiley. 1998.
Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan. Deep hashing network for unsupervised domain adaptation. CVPR, 2017.
Matthew Vowels, Necati Cihan Camgoz, and Richard Bowden. D'ya like dags? a survey on structure learning and causal discovery. ACM Computing Surveys (CSUR), 2011.
13

---

Published as a conference paper at ICLR 2023
Yoav Wald, Amir Feder, Daniel Greenfeld, and Uni Shalit. On calibration and out-of-domain gener- alization. In A. Beygelzimer, Y. Dauphin, Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https: http:////pppeeeeeeeeeemee net forum?
[d=xWYJ25-yTRR.
Xinyi Wang, Wenhu Chen, Michael Saxon, and William Yang Wang. Counterfac- tual maximum likelihood estimation for training deep networks. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 25072-25085, Curran Asso- ciates, Inc., 2021. URL d30d0f522a86b33666655 35500122886336666333351111222222 ottp:aaaeeeeaaaa11aae-aa-areerre Paper.pdf.
http://ceeeinnn..u........e..eee cc scc/papper2221//11ee
Minghao Xu, Jian Zhang, Bingbing Ni, Teng Li, Chengjie Wang. Qi Tian, and Wenjun Zhang. Adversarial domain adaptation with domain mixup. In Proceedings the AAAI Conference Artificial Intelligence, volume 34, pp. 6502-6509, 2020.
Junkun Yuan, Xu Ma, Kun Kuang, Ruoxuan Xiong, Mingming Gong, and Lanfen Lin. Learn- relationship with instrumental variable for domain generalization. arXiv preprint arXiv:2 10.01438, 2021.
ing
Marvin Mengxin Zhang, Henrik Marklund, Nikita Dhawan, Abhishek Gupta, Sergey Levine, and Chelsea Finn. Adaptive risk minimization: Learning to adapt to domain shift. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https::/p/// /epprrevv/vevvveeeeeeeee ew net forum? i d== wgw.vvvV
Yubo Zhang. Hao Tan, and Mohit Bansal. Diagnosing the environment bias in vision-and- Christian Bessiere (ed.). Proceedings of the Twenty-Ninth Interna- tional Joint Conference on Artificial Intelligence, IJCAI 2020, pp. 890-897, ijcai.org, 2020. doi: 20.44963/jjcai.2222/144 URL https: htt.:/iioooo/ooooooooo/ooo /00 wtt.a///....299//11100 249.2/ 2020 124.
language navigation.
10.
Chunting Zhou, Xuezhe Ma, Paul Michel, and Graham Neubig. Examining and combating spurious features under distribution shift. In International Conference Machine Learning, pp. 12857- 12867. PMLR, 2021.
14

---

Published as a conference paper at ICLR 2023
A
PROOFS
In this section, we give full proofs of the main theorems in the paper.
A.1 BALANCED DISTRIBUTION
A.1.1 PROOF FOR THEOREM 2.4
Here we give a proof of the minimax optimality of the Bayes optimal classifier trained on balanced distribution.
Proof. The Bayes optimal classifier trained on a balanced distribution (x, Y) has (Y = p (X|X) Then consider the expected cross entropy loss such classifier on an unseen test distri- bution p":
L"(p" (x,) log logp"(Y (X
(4)
Ex log (Y) (V) + Ex log PV
(Y(()x)
(V)
=L"(p"(Y)) (x
log (X(()))
log (Y
-)((B(()) E (X,2)
(Y
(Y(1)))
-L"("(Y)) + E (X
BXp(XIV:2)
log PB 7)
(5)
p²(XX))
-'"(D(()) (x)))
log (X Z)
- () (X 1D
Equation (4) is the definition of cross entropy loss.
Equation (5) is obtained by Y 1L Z and Y 1L Z
Thus we have the cross entropy loss of (X, in any environment is smaller than that of (Y (random guess):
-y(p((()x))- - L " 2
s0.
which means:
L (g" L L - (+ s
max e3
That is, the performance of (X, Y) is at least as good as random guess in any environment.
Since we assume the environment diversity, that is for any pe with _ Z, there exists an environ- ment el such that pY performs worse than random guess. So we have:
I (Y(x)))) L (n (r)) 5 max [ - (x))) e328
max e'38
15

---

Published as a conference paper at ICLR 2023
Now want to prove that ve E, Y lhe Z, Y de Z|X, P = + p((YIx)
p (YIX). For any Z E 2, we have:
p'(XIx) 1f (VIX, 2)
( 2) =V (Y1(( 1 (X]
= -Y p (XlY, Z)
= =p"(Y|X,Z) = pp (Y|X).
(Y(x))
Thus we have the following minimax optimality:
(Y|x)) arg min max L (p (YIX)) D.EF
A.2 LATENT COVARIATE LEARNING
A.2.1 PROOF FOR THEOREM 3.3
We now prove Theorem 3.3 setting up the identifiability of the necessary parameters that capture the spuriously correlated covariate features in the VAE. The proof is based on the proof of Theorem (Motiian et al., 2017), with the following modifications:
1. We use both E and Y as auxiliary variables.
2. We include in the causal mechanism of generating X by X (c=
Proof. Step I. In this step, we transform the equality of the marginal distributions over observed data into the equality of noise-free distribution. Suppose have two sets of parameters O (f, A) and 0' = (f', T', X') such that (X E e) po (X F = e). Ve E Eirin. then:
PT. PTAA = e)p /
Z
/ PTA F e)p Fy (Z))d7
FTrx(ZYYF=eehnd - (Z))d)
pf ( (X)JY, =))ool - x)dX=
(P = - - 8)d8
(6)
J (X) - 8))X 1 PT P - 8)dX)
(7)
PE)(X) FlPrARYd(W)oo. (w) P FiTAARYO F F PTAAEY (..
(PTANY. PXX))
(8)
(10)
(11)
In Equation (6), denote the volume of matrix A volA :- det A A. J denotes the Jacobian. We made the change of variable X = fy (Z) on the left hand side and X fy (Z) the right hand side. Since f is injective, we have f f (Y, Z). Here we abuse f-1 (X) to specifically denote the recovery of 2, i.e. (x) = Z.
In Equation (7), we introduce
(X) = - (X)xxX)
e)volJ
the left hand side, and similarly on the right hand side.
16

---

Published as a conference paper at ICLR 2023
In Equation (8), we use for the convolution operator.
In Equation (9), of is then &e = # [pe]-
[-] designate the Fourier transform. The characteristic function
In Equation (10), we dropped (w)) from both sides as is non-zero almost everywhere (by assumption (1) of the Theorem).
Step II. In this step, we remove all terms that are either function of X or Y By taking logarithm on both sides of Equation (11) and replacing Pr.x by its expression from Equation (3) we get:
k
log vol.4 (X) +
Σ
(log Q:(f (X)) - log (Y() +
Σ TT (X))AE(()))
1=1
Σ = -log (X)
(log Q2
log (Y)
T (C) (X))
Let (Co: Vo): (e) 81 B1).... (enk: Yok) be the points provided by assumption (3) of the Theorem. We evaluate the above equations at these points to obtain & equations, and subtract the first equation from the remaining &t equations to obtain:
W (yo)
(X), (30 +
log
W. (3n)
W
(X)), ( (vo))
log
(12)
W (y1)
Let L be the matrix defined in assumption (3) and L' similarly defined for A (L' is not necessarily w (wo """" and [00]
invertible). Define bj Er= log w (wo) (3n)
Then Equation (12) can be rewritten in the matrix form:
LTT(F-'(XX) = LF + b.
(13)
We multiply both sides of Equation (13) by - to get:
T(f (X))
AT +c.
(14)
Where A = L- and cnd
Step III. To complete the proof, we need to show that A is invertible. By definition of T and according to Assumption (2), its Jacobian exists and is an nk x n matrix of rank n. This implies that the Jacobian of T' o f' exists and is of rank and so is A.
We distinguish two cases:
1. kf 1, then A is invertible A Ruxn
2. I 1, define x f (x) and TT = (T), (3.), Tex(E)))
Suppose for any choice of F.. ,?,,..,, J. the family ( (T.
dT. is never linearly
de
dan
independent. This means that (R() is included in subspace of g of the dimension of
most & - 1. Let h be non-zero vector that is orthogonal to (R). Then for all x E IR we have (x) h) 0. By integrating find that (Tx( h) const.
17

---

Published as a conference paper at ICLR 2023
Since this is true for all I E R and a h O, conclude that the distribution is not strongly exponential. So by contradiction, we conclude that there exist & points E2....... such
that di
d'T (TE
(TT) are linearly independent.
d2
Collect these points into vectors (x',..., x*) and concatenate the Jacobians JT JT(x') evaluated at each of those vectors horizontally into the matrix Q (JT(x') (...)..... Jq...JJ JT (x*))
and similarly define Q' as the concatenation of the Jacobians of T'' o f(x)) evaluated at those points. Then the matrix Q is invertible. By differentiating Equation (14) for each x", we get:
=AQ'.
The invertibility Q implies the invertibility A and Q'. This completes the proof.
A.3 BALANCED MINI-BATCH SAMPLING
A.3.1 PROOF FOR THEOREM 3.6
Our proof of all possible balancing scores is an extension the proof of Theorem 2 from (Rosen- baum & Rubin, 1983). by generalizing the binary treatment to multiple treatments.
Proof. First, suppose the balancing score b(Z) is finer than the propensity score s(Z). By the definition of a balancing score (Theorem 3.4) and Bayes' rule, we have:
p pY|Z,,(Z) p(Y(b(Z))
(15)
On the other hand, since b(Z) is function of Z, we have:
p(Y 2(
(16)
Equation (15) and Equation (16) give us p(Y]b(2)) =>(YIZ). So to show ((2) is a balancing
score, sufficient show ((Y|b(()) = 12)
Let the y-th entry of s(Z) be sy (2) p(Y = 2), then:
E| E[s,(Z)(b(2)) = / p(Y = y Z = =)p(Z = =(bb))
(17)
But since b(Z) is finer than ((2),62)) is also finer than (2), then
((s((Z) = sy(
(18)
Then by Equation (17) and Equation (18) have P(Y y/ = P(Y ylb(Z)) required. So b(Z) a balancing score.
For the converse, suppose b(2) is a balancing score. but that b(Z) ) is not finer than (Z). Then there exists 21 and 2n such that s(2) 전 8(22). but b(21) = b(z2). By the definition of s(-). there exists such that P(Y y v|2 # P(Y y/22). This means, Y and Z are not conditionally independent given b(Z), thus b(Z) is not balancing score. Therefore, to be a balancing score, b(Z) must be finer than s(Z).
Note that s s(2) is also a balancing score, since s(2) is also function of itself.
18

---

Published as a conference paper at ICLR 2023
A.3.2 PROOF FOR THEOREM 3.7
We provide proof for Theorem 3.7, demonstrating the feasibility of balanced mini-batch sampling.
Proof. In Algorithm 1, by uniformly sampling different labels such that y * yf, mean sample Yale = {V1. 1/, ya by the following procedure:
V1 Uf1,2, .... mJ \ {ye}
1/2 ~ Uf1,2,... ..., m} \ file 8/1
: Je 21., m} \ {yer U1s U2.--Va-1)
where U denotes the uniform distribution. Suppose Dralancod De '(X, Y = e), Ve Erain-
(X, and data distribution
Suppose we have an exact match every time we match a balancing score, then for all e E Erains we have
(2)
+ p (2),
a+1 (1 p( (2), E ==
+
m
1
a+1 (1 (()
+...
a
m
m - 2 (1
a+1 (1 P( (2), e)
m
m 2
m a +1 1
m -
a
m. - a
+
a
a
m
m
By the definition of balancing score, p PY E 0) = pY (Z). E =) and Bd (Y Z, E =e) B ( (Z), E = e), then we have
B OY
1 + m m p Y E)).
m 1
When a 1, we have j (Y| F)) = U11,2, m). which means ( Y. (X,
(X, Y, 2). i.e. Doalanced can be regarded sampled from the balanced distribution n as defined in Definition 2.2.
B EXPERIMENT DETAILS
In this section, give more details of experiments. We perform experiments the Do- mainBed codebase* (Gulrajani & Lopez-Paz, 2020).
B. DATASETS
ColoredMNIIST is a variant of the MNIST handwritten digit classification dataset. Each domain in 10.1, 0.3, 0.9] is constructed by digits spuriously correlated with their color. This dataset contains 70, 000 examples of dimensions (2, 28, 28) and 2 classes, where the class indicates the digit is less than 5, with 25% noise. RotatedMNIST is another variant of MNIST where each domain
https:://gi thub. com http///aceoookeeeeeeree
http://hh..c.c.ccccccoeeeeeeeeee comaa Doma nHed
19

---

Published as a conference paper at ICLR 2023
contains digits rotated by degrees, where E {0, 30, 45, 75). This dataset contains 70, 000 examples of dimensions (1, 28, 28) and 10 classes, where the class indicates the digit. PACS comprises four domains: art, cartoons, photos, and sketches. This dataset contains 9, 991 examples of dimensions (3, 224, 224) and 7 classes, where the class indicates the object in the image. VLCS comprises four photographic domains: Caltech LabelMe, SUN09, and VOC2007. This dataset contains 10, 729 examples of dimensions (3, 224, 224) and 5 classes, where the class indicates the main object in the photo. OfficeHome includes four domains: art, clipart, product, and real. This dataset contains 15, 588 examples of dimension (3, 224, 224) and 65 classes, where the class indicates the object in the image. Terralncognita contains photographs of wild animals taken by camera traps at four different locations: L100, L38, L43, and L46. This dataset contains 24, 788 examples of dimensions (3, 224, 224) and 10 classes, where the class indicates the animal in the image. DomainNet has six domains: clipart, infographics, painting, quickdraw, real, and sketch.
This dataset contains 586, 575 examples of size (3, 224, 224) and 345 classes.
B.2 BASELINES
We choose ERM. IRM, GroupDRO and CORAL as base algorithms to apply our method be- cause they are representative methods for domain generalization, and they serve strong baselines when compared to wide range of domain generalization methods. Empirical risk minimization (ERM) is a default training scheme for most machine learning problems, merging all training data into one dataset and minimizing the training errors across all training domains. Invariant risk mini- mization (IRM) represents a wide range of invariant representation learning baselines. IRM learns a data representation such that the optimal linear classifier on top of it is invariant across training domains. Group distributionally robust optimization (GroupDRO) represents group- sasud methods that minimize the worst group errors. GroupDRO performs ERM while increasing the weight of the environments with larger errors. Deep CORAL represents the distribution matching algorithms. CORAL matches the mean and covariance of feature distributions across training domains, Ac- cording (Gulrajani & Lopez-Paz, 2020), CORAL is the best performing domain generalization algorithm averaged across 7 datasets, compared to other 13 baselines.
B.3 HYPERPARAMETE SELECTION
Base algorithms: For the architecture of image classifiers, following the DomainBed setting, we train a convolutional neural network from scratch for ColoredMNIST and RotatedMNIST datasets, and use a pre-trained ResNet50 (He et al., 2016) for all other datasets. Each experiment is repeated with 3 different random seeds. We choose the hyperparameters of base algorithms based on the default hyperparameter search with random mini-batch sampling. More specifically, we extract the hyperparameters from the official experimental logs provided in the DomainBed GitHub repository. 5 To retrieve hyperparameters, we ran the script collect results deta metailes ailed. . py, modified from the provided co eect resul 1 t.s.. . py script, to collect the hyperparameters that are used to produce the DomainBed results table with train domain validation.
Balanced mini-batch construction: We use multi-layer perceptron (MLP) based VAE (Kingma & Welling, 2013) to learn the latent covariate Z. For ColoredMNIST. ColoredMNIST" and Rotat- edMNIST, we use a 2-layer MLP with 512 neurons in each layer. For all other datasets, we use a 3-layer MLP with 1024 neurons in each layer. We choose the conditional prior Pt (z = e) to be a Gaussian distribution with diagonal covariance matrix. We also choose the noise distribution P. to be a Gaussian distribution with zero mean and identity variance matrix. We choose the largest possible latent dimension n according to Theorem 3.3 up to 64. We choose KI. divergence as our distance metric d on DomainBed.
The hyperparameters we use are shown in Table 3. We control is by choosing different distributions model the latent covariate: for A = 2, we choose Normal distribution, and for k = 1, we choose Normal distribution with a fixed variance equal to the identity matrix. When choosing the latent dimension n, we follow the identifiability requirement m Erainnl nk Section 3.1. and we chose the maximum allowed up to À = 64 for large images (224 x 224) and up to À = 16 for small images (28 x 28). i.e. n = min [m]mmni [kk], {}. For the distance metric d, we choose the KI. divergence on all datasets except on ColoredMNiST we choose the Loo distance. Different choice
'https://drivee wwwglo.oooooooooooooo com/
20

---

Published as a conference paper at ICLR 2023
of distance metric usually does not affect the final results too much, shown in Table 4. We tune the number of matching examples a for each base algorithm with a train domain validation, and the best for each base algorithm is shown in the order of ERR//R///GouDDRO/CORAA in the last column of Table Typically, the best a for a dataset across different base algorithms is similar.
Choice of hyperparameters for constructing balanced mini-batches, including training the VAE model for latent covariate learning (n, Ir, batch size) and the balancing score matching (a. d).
Table 3:
Evain
m
T2
Ir
batch size
d
a
ColoredMNIST
2
10
16
le-3
64
Loc
4/4/4/4
1olOredMNNIST
2
2
1
3
le-3
64
KLD
1/1/1/1
RotatedMINIST
5
10
16
le-3
64
KLD
1/2/1/1
VLCS
le-4
KLD
2/1/1/2
PACS
3
7
2
10
le-4
32
KLD
3/2/1/2
OfficeHome
3
65
2
64
le-4
32
KI.D
2/2/2/2
Terralncognita
3
10
2
14
le-4
32
KI.D
2/1/1/2
DomainNet
345
64
le-4
32
KLD
5/5/5/5
Table 4: Out-of- accuracy ColoredMNIST" when using different distance metrics.
L1
L2
Loc
KLD
Train Val
69.3 = a.1
69.5 1
69.8 = 0.1
69.2 0.
Test Val
70.2 a.s
70.3 n4
70.5 =
69.9 2.03
Figure 5 shows three of reconstructed images with the latent covariate Z and different label Y using our VAE model. We can see that Z keeps the color feature and some style features, while the digit shape is changed the closest digits belongs to class Y.
3
8
3
4
9
4
S
S
3
(a)
(b)
(c)
Figure 5: Reconstructed ColoredMNIST images from our VAE model. In each sub-figure, we infer Z from the leftmost image, then generate images with labels Y (middle) and Y (right).
B.4 DETAILED RESULTS
All experiments conducted on NVidia A100, Titan RTX and RTX A6000 GPUs. Here report detailed results on each domain of all seven datasets on DomainBed, with base algorithms ERM, IRM, GroupDRO, and CORAL. We training domain validation.
Table 5: ColOredMNIST
Algorithm
+90
+80 %
-90 %
Avg
ERM
71.7 + 0.1
72.9 + 0.2
10.0 + 0.1
$1.5
IRM
72.5 +
73.3 0.5
10.2 0.3
$2.0
GroupDRO
0.3
10.0 0.2
52.1
CORAL
0.3
0.1
9.9
51.5
Ours+ERM
0.3
71.2
37.6
60.1
Ours+IRM
75.4
71.0
8.61
59.2
Ours+GroupDRO
72.0 0.6
72.8 + 0.2
3.5
53.9
Ours+CORAL
70.5
72.0 - 0.2
57.2
66.6
21

---

Published as a conference paper at ICLR 2023
Table 6: RotatedMNIST
Algorithm
15
30
45
60
75
Avg
ERM
95.9
98.9
98.9
98.8 0.0
98.9 +
98.9
96.4
IRM
9..
98.8
98.7
98.6
9.07
0.2
97.7
GroupDRO
9.16
9.1
98.9
99.0
98.9
9..5
98.0
CORAL
95.8
98,8
9.0
9..0
9..
9.4
98.0
Ours+ERM
94.8
0.1
9..7
9.0
9..8
9.44
Ours+IRM
93.0 0.5
0.1
9.1
9.2
98.6
9.
96.8
Gurs+GroupDRO
94.8 + 0.2
0.1
9.09
9.0
98.9
95.9
97.6
Ours+CORAL
94.5
98.7
9.1
99.0
98.9
9.2
97.7
Table 7: VLCS
Algorithm
C
S
Avg
ERM
97.7 + 0,4
64.3 + 0.9
73.4 + 0.5
74.6 +
1.3
77.5
IRM
98.6 0.1
64.9 0.9
73.4
0.6
77.3
78.5
GroupDRO
97.3
0.9
69.5
7.77
76.7
CORAL
9..
66.1 1.2
0.3
7.2
78.8
OurS+ERM
96.9 + 0.4
64.8
0.8
72.6
76.1
Ours+IRM
775
61.6
1.2
7.25
76.5
GrosoGrouDDOO
0
0.9
69.2 0.8
72.6 0.6
76.0
Ours+CORAL
98.3 + 0.1
63.9 + 0.2
1.1
73.7 + 1.3
76.4
Table 8: PACS
Algorithm
A
C
P
S
Avg
ERM
84.7 : 0.4
80.8 + 0.6
0.3
79.3
1.0
85.5
IRM
84.8 + 1.3
1.1
0.6
76.1 1
83.5
GroupDRO
83.5 0.9
7. 9..7 79.1 0.6
84.4
CORAL
88.3 + 0.2
0.5
9.3
78.8
86.2
Ours+ERM
87.9
1.0
0.3
79.1 1.2
86.1
Ours+IRM
84.6 1.1
0.1
9.44
80.0 1.2
85.2
Ours+GroupDRO
86.3 + 0.6
79.2 0.8
96.5 0.2
77.7 +
84.9
Ours+CORAL.
87.8 + 0.8
81.0 +.1
97.1 0.4
81.1 + 0.8
86.7
Table 9: OfficeHome
Algorithm
A
R
Avg
ERM
61.3 + 0.7
52.4 + 0.3
75.8 + 0.1
76.6 + 0.3
66.5
IRM
58.9
52.2 + 1.6
72.1 2.9
7..0
64.3
GroupDRO
0.74
52.7
7.70 7.70
66.0
CORAL
6.4
0.5
7.15
7.54
68.7
OurS+ERM
61.5 0.4
53.8 0.5
7..9
77.4
67.1
Ours+IRM
59.2
49.8 0.9
2.3
75.5
64.6
Ours+GroupDRO
61.7 11.7
52.5 0.8
0.8
76.9 + 0.6
66.5
Ours+CORAL
6..6
56.5 + 0.6
0.3
78.8 0.5
69.6
22

---

Published as a conference paper at ICLR 2023
Table 10: TerraIncognita
Algorithm
L,000
L38
1.43
1,46
Avg
ERM
49.8 4.4
42.1
1.4
56.9 + 1.8
35.7
46.1
IRM
54.6 + 1.3
39.8 + 1.9
56.2 + 1.8
39.6
47.6
GroupDRO
41.2 0.7
38.6
56.7
0.9
3.14
43.2
CORAL
2.4
42.2 1.0
57.0
3.9
47.6
Ours+ERM
53.
47.2
+.7
3.0
48.0
Ours+IRM
50.0
41.3 1.
54.0 + 2.7
40.5 0.6
46.5
Ours+GroupDRO
51.2
35.4 2.5
56.0 + 1.0
38.9 1.4
45.4
Ours+CORAL
55.2
42.3 + 3.6
54.7
+ 0.4
36.0 - 1.0
47.0
Table DomainNet
Algorithm
clip
info
paint
quick
real
sketch
Avg
ERM
58.1 0.3
0.3
46.7 0.3
12.2 +
59.6 0.1
49.8 + 0.4
40.9
IRM
4
15.0 1.5
4.3
10.9 1: 0.5
48.2
5.2
42.3 3.1
33,9
GroupDRO
4
17.5
0.4
33.8
9.3
51.6 0.4
40.1 0.6
33.3
CORAL
59.2
19.7 0.2
46.6
1..4
59.8 0.2
50.1 0.6
41.5
Ours+ERM
61.2
48.6 0.3
13.0
6..0
51.9
Ours+IRM
$7.9
18.2 1.3
46.0 1.5
13.2
57.2
1.
40.5
Grou+
59.3 0.3
18.4 0.2
45.3 + 0.3
12.2 0.4
60.5 1.5+00.4
48.9 + 0.2
40.8
Ours+CORAL
0.1
20.7 - 0.2
50.4 0.1
13.6 1 0.4
62.7 4 0.1
52.8
0.3
43.9
C DISCUSSIONS AND LIMITATIONS
The experiments show that our balanced mini-batch sampling method outperforms the random mini- sampling baseline when applied multiple domain generalization methods, both semi-synthetic datasets and real-world datasets. While our method can be easily incorporated into other domain generalization methods with good performance, there are some potential drawbacks of our method. First, the computation complexity of our method grows quadratically with the dataset size, as for each training example, our method requires searching across the dataset to find the closest match in balancing score, which could become computation bottleneck on large datasets. However, this could solved by matching examples offline before training, or with more efficient searching methods. The second caveat is that we do not provide optimized model selection method to complement our method. While it is possible to balance the held-out validation set with our method and choose the best model based the accuracy of the balanced validation set, the quality of such a balanced validation set is questionable given the small size of a typical validation set. For now, we recommend the training-domain validation scheme in practice.
D IN-DEPTH COMPARISON WITH RELATED WORK
D.1 COMPARISON OF ASSUMPTIONS
Certain assumptions are needed for our paper, as in other works on domain generalization. Our assumptions are not stronger than other domain generalization works that give similar generalization guarantees. Arguably, ours are weaker than most them.
We provide the identifiability of the balanced distribution given finite set of train environments and prove that the Bayesian optimal classifier trained on the balanced distribution would be minimax optimal across all environments. Our main assumptions are the factorial exponential distribution of the latent covariate given the label, the invertible causal function f. and the additive noise. Similar assumptions have been made in Sun et al. (2021).
Works without constraints on environments usually can only provide generalization guarantee when optimizing overall environments (Mahajan et al., 2021) or do not provide any such guarantees
23

---

Published as a conference paper at ICLR 2023
(Chen et al., 2021; Li et al., 2022). To provide generalization guarantee with single small number of train environments, Yuan et al. (2021); Wald et al. (2021); Ahuja et al. (2021) use a more restrictive linear causal model, Arjovsky et al. (2020) only provide full solution for linear classifiers, Christiansen et al. (2021) assume additive confounders, Yuan et al. (2021); Makar et al. (2022); Puli al. (2022) need to utilize the observation of the variable spurious correlated with the label Y.
In practice, the model built with our assumptions works well on real-world datasets that do not exactly fit our assumptions, which empirically demonstrates that our method robust against viola- tions of our assumptions.
D.2 COMPARISON OF CAUSAL MODEL
In general, the assumption of the underlying Structural Causal Model (SCM) is determined by the nature of the task. Sometimes, such SCM can be designed by human expert who knows the data generation process of the task. In our paper, we propose to adopt coarse baase-grained SCM for general image classification tasks with only three variables: image X, label Y, and latent variable Z.
Our high-level philosophy is that the image itself is merely a record of what has been done, and driving force of the recorded event. When one intervenes on
the label can usually be regarded image X, the label Y of the image does not necessarily change. However, if the intervention is on the class label Y. the image X changes almost for sure for well-defined image classification task. For example, in the medical domain, a disease (Y) would cause some lesions, further driving the different appearance of MRI images (X). Another example is when Y is the object class of the item appearing in the image X, which is usually the case for the most widely used image classification benchmarks like ImageNet. We have also discussed this in Section 2.1 of our paper.
However, there could be exceptions. For example, if we are asked to classify whether we feel happy or sad after seeing picture, picture would become the cause of the sentiment label Y. Such a scenario is less likely to happen in real-world image classification tasks. To resolve the issue of different SCM for different tasks, Christiansen et al. (2021) consider all SCMs that can be transformed into a specific linear form with plausible interventions. Wald et al. (2021) assume X can be disentangled into features causing and features caused by Y, and derive their theoretical results with linear SCM. We assume more general nonlinear SCM with Y + X, which is suitable for most of the image classification tasks we consider. On the other hand, Yuan et al. (2021) directly assumes SCM with X + Y. Empirically, they obtained worse results PACS (84.4 v.s. 86.7) and OfficeHome (64.2 v.s. 69.6) datasets, which confirms that our SCM is more suitable.
A principled way of identifying the causal relationship (if there is any) between X and Y is causal discovery. However, current causal discovery techniques cannot handle the complex high- dimensional image data we consider in the paper (Vowels et al., 2021). A slightly related work is Hoover (1990), which proposes that decomposing joint distribution following the causal graph more stable for interventions than a random decomposition. Our paper uses the invariant of P(X Z), where Z represents domain-dependent features like camera positions and picture style. is hard to find such invariance in other ways of decomposition.
On the other hand, quite few works assume no direct causal relationship between X and Y (Chen et al., 2021; Mahajan al., 2021; Liu al., 2021a; Sun al., 2021: Ahuja et al., 2021: Li al., 2022). Instead, they assume there is causal feature Zeausal directly causing X, together with another non-causal feature Znon-cussal- Y caused by Zemusal, which implies that Zemusal may contain more information than Y. Such causal model can be viewed as noisy version of ours, the same as the causal feature Zeausal, and Z the same as the non-causal feature Znoo-causal- Different paper model the spurious correlation between Zcausl and 7mmm causal in different way in the SCM, while we just ensure Zegusal and Zoun-causal are correlated, without specifying how they are correlated.
consider Y
24


Reviews:
[{"review_id": "RvjmuUP27X", "reviewer": "Reviewer_KTyH", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper introduces a balanced sampling strategy to make data distribution spurious-free. It is based on the causal assumption of invariant data generating process. The model trained on the unbiased data is assumed to be minimax optimal across different environments. The authors also provide identifiability guarantee of the causal model. Experiments on the DomainBed dataset show its favorable out-of-distribution generalization performance. The method is basically sound and the writing is good.", "summary_of_the_review": "This paper introduces a causality-inspired balanced mini-batch sampling method for domain generalization. The details of the proposed method and the theoretical insights are introduced clearly. The experiments on seven popular datasets and the DomainBed codebase are fair and good. My concerns about this paper are its novelty in causal graph, its sampling efficiency, and its performance improvement.", "strength_and_weaknesses": "Strength:\n1.\tThis paper proposes to debias the data distribution through a balanced mini-batch sampling, and then it uses the balanced data to train the model. It is technically sound.\n2.\tThe authors give theoretical analyses for the insights of this sampling strategy. \n3.\tThe authors implement the experiments on the public DomainBed codebase with seven datasets. It is considered fair and reproducible.\n4.\tThe writing is good and easy to follow.\nWeakness:\n1.\tSince there are numerous causality-based domain generalization methods, e.g., [1-9], and many of them are also based on the invariant causal structure assumption like this paper. I would like to see more in-depth discussions about the difference between them and this work. For example, some papers [1, 4] may assume that X is the cause of Y, some papers [2, 3, 5, 6, 9] may assume that there is not a direct causal relation between X and Y, but this paper supposes that Y causes X. What are the differences between them and this work? And what is the novelty of this paper in designing the causal graph in Fig. 1?\n2.\tAs stated in Algorithm 1, the sampling strategy should be performed for each sample in the dataset. So, I guess it may not be efficient. The authors may discuss more about this limitation. For example, the total running time before/after introducing balanced sampling process.\n3.\tI see the main improvement of average accuracy comes from the CMNIST dataset. For example, it improves the accuracy of ERM from 51.5% to 60.1%. But the improvement on other datasets is not such significant, or even worse on RMNIST and VLCS dataset. I suggest the authors to give more discussions about this.\n4.\tFor Fig.4 (a), if I understand correctly, training data is from the source domain and test data is from the target domain. So, I don’t see why the accuracy of training data is lower than the accuracy of test data. Could the authors explain it?\n5.\tWhy does the results in Table 4 with KLD metrics are not consistent with the results of CMNIST in Table 2?\n\n[1] A Causal Framework for Distribution Generalization [TPAMI 2021]\n[2] A Style and Semantic Memory Mechanism for Domain Generalization [ICCV 2021]\n[3] Domain Generalization using Causal Matching [ICML 2021]\n[4] Learning Domain-Invariant Relationship with Instrumental Variable for Domain Generalization [arXiv 2021]\n[5] Learning Causal Semantic Representation for Out-of-Distribution Prediction [NeurIPS 2021]\n[6] Recovering Latent Causal Factor for Generalization to Distributional Shifts [NeurIPS 2021]\n[7] On Calibration and Out-of-domain Generalization [NeurIPS 2021]\n[8] Invariance Principle Meets Information Bottleneck for Out-Of-Distribution Generalization [NeurIPS 2021]\n[9] Invariant Information Bottleneck for Domain Generalization [AAAI 2022]\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "PVfgOzsnvm", "reviewer": "Reviewer_ypmL", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper develops a method for using multiple environments during training to obtain a model that performs well on an unseen test environment. The method has both theoretical and experimental support.", "summary_of_the_review": "The paper develops an interesting method for attempting to generalize well to unseen test environments. The theoretical support is interesting, with some strong assumptions, and the method shows significant improvements on some datasets (with much more modest improvements on some others).", "strength_and_weaknesses": "Strengths:\n\nThe method developed in the paper is based on an interesting way to model different environments, and the idea of selecting the minimax model across the training environments is an interesting idea.\n\nThe results are very promising on ColoredMNIST and the 10-class version.\n\nWeaknesses:\n\nThe improvements on the datasets other than ColoredMNIST are much more modest.\n\nSome of the assumptions about how the environments are related are relatively strong (for example, it is quite possible that some of the datasets have more noise than the others).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "90LAY3ApEz_", "reviewer": "Reviewer_AxQ6", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper aims to design a domain generalization model based on causal balancing. To this end, the authors assume that the observation is causally determined by the label and some latent factors and a random variable. Based on such assumption, the authors have provided many theories to demonstrate that the causal model is identifiable. In the experiments, a lot of experiments have been conducted to demonstrate the effectiveness of the proposed model.", "summary_of_the_review": "See the above comments.", "strength_and_weaknesses": "\nIn general, the paper is well written, and the assumptions are well described and the theories are solid. The balancing problem is important for domain generation, and the authors have provided a valid solution to this problem. \nMy concerns are mainly about the assumption made in this paper. I would like to see how strong of these assumptions, for example, in assumption 3.1, whether is it possible to drop the exponential assumption. If the noise variable can not be added to the function f, how about the conclusion of this paper. \nIn addition, shall we have more advanced methods for addressing the balancing problem, for example representation learning and so on.\n\nOverall, I think this is a good paper. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "-PxqXXAuo-M", "reviewer": "Reviewer_2Bti", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper aims to address OOD generalisation by providing a two step method that aims to generate samples from a 'balanced distribution' that can be used to train subsequent off the shelf classifiers that generalise to new environments well. This involves a two step procedure, the first involves learning p(V) using a Variational Autoencoder, and then using a matching method to create balanced mini batches. Finally, experimental validation is performed on DomainBed, where an improvement is demonstrated. ", "summary_of_the_review": "Overall, this paper provides a novel framework to first learn the balanced distribution that has invariance properties across different environments, and then train a classifier based on samples generated using the proposed sampling algorithm. Theorems regarding identifiability and minimax results are proved, and finally experimental validation on the DomainBed datasets show improvement in performance. ", "strength_and_weaknesses": "Strength: This paper provides identifiably results regarding the auto-encoder and its ability to recover the parameters of true distributions. The balanced mini-batch sampling strategy allows for the use of any ML classifier after, which is very nice. The experimental results are good.\n\nWeakness: I would like to see a more rigorous discussion about why the DAG that models the data generating process is suitable for this problem as this would benefit readers. Additionally, the condition regarding the Cartesian product of support of Y and the support of E-train I would like to see discussed more and it's relation to dimensionality and the amount of data points across various environments needs. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster