Causal Balancing for Domain Generalization
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 causally-motivated balanced mini-batch sampling strategy to train robust classifiers that is minimax optimal across a diverse enough environment space, by utilizing multiple training sets from different environments.
We provide an identifiability guarantee of the latent covariates in the proposed causal graph and show that our proposed approach samples train data from a balanced, spurious-free distribution under ideal scenario.
Experiments are conducted on three domain generalization datasets, demonstrating empirically that our balanced mini-batch sampling strategy improves the performance of four different established domain generalization model baselines compared to the random mini-batch sampling strategy.
1. Introduction
Machine learning is achieving tremendous success in many fields with useful real-world applications (Silver et al., 2016;
Devlin et al., 2019; Jumper et al., 2021).
However, machine learning models often fail to generalize to out-of-domain (OOD) data sampled from unseen environments (Qui˜noneroCandela et al., 2009; Szegedy et al., 2014).
One explanation for such a failure is that the models are prone to learning spurious correlations that change between environments.
Recently, various causal-inspired (Pearl, 2009) methods have been proposed to improve the OOD generalizability by considering the invariance of causal features or the underlying causal mechanism through which data is generated.
Such methods often aim to find invariant data representations using new loss function designs that incorporate the invariance conditions across different domains into the training process (Arjovsky et al., 2020; Mahajan et al., 2021;
Liu et al., 2021; Lu et al., 2022; Wald et al., 2021).
Unfortunately, these approaches have to contend with trade-offs between weak linear models or approaches 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 first demonstrate that the Bayes optimal classifier trained on a balanced (spurious-free) distribution is minmax optimal across all environments. Then we propose a two-step method to create such balanced distribution from multiple train datasets collected form different environments: (1) learn the observed data distribution using a variational autoencoder (VAE) by latent covariate learning , and (2) use the learned latent covariate to create balanced mini-batches that follow a balanced distribution.
Because the only modification at train time is a resampling of train examples, it makes our method lightweight and highly flexible, enabling seamless incorporation with off-the-shelf domain generalization methods (Gulrajani & Lopez-Paz, 2020; Sagawa et al., 2019; Sun & Saenko, 2016).
Our contributions are as follows: (1) We propose the balanced distribution without spurious correlation and prove that it can produce mimmax optimal classifiers for OOD generalization; (2) We demonstrate that the source of spurious correlation, as a latent variable, can be identified given a large enough set of training environments under mild conditions in a nonlinear setting; (3) We propose a novel balanced mini-batch sampling algorithm that, in an ideal scenario with exact matches of the true source of spurious correlation, can remove the spurious correlations in the observed data distribution; (4) Our empirical results show that our two-phased method obtains significant performance gain on three domain generalization datasets, ColoredMNIST (Arjovsky et al., 2020), PACS (Li et al., 2017) and TerraIncognita (Beery et al., 2018), across four different domain generalization methods over a random sampling strategy. [Figure]
2. Preliminaries 2.1. Problem Setting
We consider a standard domain generalization setting with a potentially high-dimensional variable X (e.g. an image), a label variable Y and a discrete environment (or domain) variable E in the sample spaces X , Y , E , respectively.
Here we focus on the classification problems with Y = { 1 , 2 , ..., m } and X ∈ R d .
We assume that the training data are collected from a finite subset of training environments E train ⊂ E .
The training data D e = { ( x ei , y ei ) } N e i =1 is then sampled from the distribution p e ( X, Y ) = p ( X, Y | E = e ) for all e ∈ E train .
Our goal is to learn a classifier C ψ : X → Y that performs well in a new, unseen environment e test ̸∈ E train .
We assume that there is a data generation process of the observed data distribution p e ( X, Y ) represented by an underlying structural causal model (SCM) shown in Figure 1a.
More specifically, we assume that X is caused by label Y , an unobserved latent variable Z (with sample space Z ∈ R n ) and an independent noise variable ϵ with the following for- mulation:
[Equation]
Here, we assume the causal mechanism is invariant across all environments e ∈ E and we further characterize f with the following assumption: Assumption 2.1. f : { 1 , 2 , ..., m } × Z → X is injective.
f − 1 : X → { 1 , 2 , ..., m } ×
Z is the left inverse of 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
[Equation]
While the causal mechanism is invariant across environments, we assume that the correlation between 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 a function of Z , the generation process of X can completely ignore Y and f would not be injective.
Z can be understood as a potentially massive set of latent features that do not determine the “ Y -ness” and that can be spuriously correlated with environment E .
Let e ∈ E index a family of distributions F = { p e ( X, Y, Z )
=
p f ( X | Z, Y )
p e ( Z | Y ) p e ( Y ) } e , where p e ( Z | Y ) > 0 and p e ( Y ) > 0 .
Note that any mixture of distributions from F would also be a member of F .
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 ̸∈ E train as it involves E and we only want to learn the stable causal relation Y → X .
However, it is inevitable that the learned predictor may absorb the unstable relation between X and Y if we simply train it on the observed train distribution p e ( X, Y ) with empirical risk minimization.
The causal graphs of two examples of the realization of our data generation model are shown in Appendix A.
To avoid learning the unstable relations, we propose to consider a balanced distribution p B ( X, Y, Z ) such that
Y ⊥⊥ B Z while the causal mechanism Z → X ← Y unchanged, as shown in Figure 1b, which is defined below:
Definition 2.2.
A balanced distribution can be written as p B ( X, Y, Z ) =
p f
( X | Y, Z ) p B ( Z ) p B ( Y ) , where p B ( Y ) =
U { 1 , 2 , ..., m } and Y ⊥⊥ B Z .
In this new distribution, X and Y is only correlated through the stable causal relation Y → X . Here we do not specify p B ( Z ) . Note that p B ( X | Y, Z ) = p f
( X | Y, Z ) is a result of the unchanged causal mechanism Z → X ← Y , and that p B ( X, Y, X ) ∈ F can also be regarded as from an environment B ∈ E .
Under an additional conditional independence assumption Y ⊥⊥ B Z | X , we can prove that p B ( Y | X ) is invariant for any choice of p B ( Z ) .
Then we have the following theorem 1 : Theorem 2.3. Consider a classifier C ψ ( X ) = argmax Y p ψ ( Y | X ) with parameter ψ .
We denote the cross entropy loss of such a classifier on environment e by L e ( p ψ ( Y | X ))
= − E p e ( X,Y ) log p ψ ( Y | X ) .
Assume that
(1) Y ⊥⊥ B Z | X , and (2) E satisfies:
[Equation]
Then the Bayes optimal classifier trained on any balanced distribution p B ( X, Y ) is a minimax optimal classifier with respect to cross entropy loss across all environments in E :
[Equation]
The first assumption implies the noise variable ϵ can be disentangled into ( ϵ Y , ϵ Z ) , such that there exist functions g Y , g Z with ( Y, Z )
=
f − 1 ( X − ϵ )
= ( g Y ( X − ϵ Y ) , g Z ( X − ϵ Z )) .
The second assumption implies that the environment space E is large and diverse enough such that a perfect classifier on one environment will always perform worse than random guessing on some other environment.
Under these two assumptions, no other Byes optimal classifier produced by an environment in E would have a better worst case OOD performance than the balanced distribution.
3. Method
We propose a two-phased method that first use an VAE to learn the underlying data distribution p e ( X, Y, Z ) with latent covariate Z for each e ∈ E train , and then use the learned distribution to calculate a balancing score to create a balanced distribution based on the training data.
3.1. Latent Covariate Learning
We argue that the underlying joint distribution of p e ( X, Y, Z ) can be learned and identified by a VAE, given a sufficiently large set of train environments E train .
To specify the correlation between Z and Y , we assume that the conditional distribution p e ( Z | Y ) is conditional factorial with an exponential family distribution:
Assumption 3.1.
The correlation between Y and Z in environment e is characterized by p eT,λ
( Z | Y ) as follows:
[Equation]
where Z i is the i -th element of Z , Q =
[ Q i ] i : Z → R n is the base measure, W e =
[ W ei ] i : Y → R n is the normalizing constant, T
=
[ T ij ] ij :
Z → R nk is the sufficient statistics, and λ e = [ λ eij ] ij :
Y → R nk are the Y dependent parameters.
Here n is the dimension of the latent variable Z , and k is the dimension of each sufficient statistic determined by the type of chosen exponential family distribution.
Note that the exponential family assumption is not very restrictive as it has universal approximation capabilities (Sriperumbudur et al., 2017).
We then consider the following conditional generative model in each environment e ∈ E train , with parameters θ = ( f , T , λ ) :
[Equation]
We use a VAE to estimate the above generative model with a variational approximation q eϕ ( Z | X, Y ) of the prior probability of latent variable p eθ
( Z | X, Y ) .
We denote the empirical data distribution given by dataset D e = { (
[Equation]
ei ,
[Equation]
ei
[Equation]
[Equation]
e
[Equation]
= collected from environment e .
The evidence lower bound (ELBO) of the data log-likelihood in each environment e ∈ E train is then defined as follows:
[Equation]
The KL-divergence term can be calculated analytically.
To sample from the variational distribution q eϕ ( Z | X, Y ) , we use reparameterization trick (Kingma & Welling, 2013).
We then maximize the above ELBO
[Equation]
over all training environments to obtain model parameters ( θ, ϕ ) .
To show that we can uniquely recover the latent variable Z up to some simple transformations, we want to show that the model parameter θ is identifiable up to some simple transformations.
That is, for any { θ = ( f
[Equation] [Equation] where Θ is the parameter space and ∼ represents an equivalent relation. Specifically, we consider the following equivalence relation from (Motiian et al., 2017):
Definition 3.2.
If ( f , T , λ ) ∼ A ( f ′ , T ′ , λ ′ ) , then there exists an invertible matrix A ∈ R nk × nk and a vector c ∈ R nk , such that T ( f − 1 ( x ))
=
A T ′ ( f ′− 1 ( x ))
+ c , ∀ x ∈ X .
When the underlying model parameter θ ∗ can be recovered by perfectly fitting the data distribution p eθ ∗ ( X | Y ) for all e ∈ E train , the joint distribution p eθ ∗ ( X, Z | Y ) is also recovered.
This further implies the recovery of the prior p eθ ∗ ( Z | Y ) and the true latent variable Z ∗ .
The identifiablity 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 (6) , with parameters θ = ( f , T , λ ) .
In addition to Assumption 2.1 and Assumption 3.1, we assume the following conditions holds
: (1) The set { x ∈ X|
ϕ
ϵ
( x ) = 0 } has measure zero, where ϕ ϵ is the characteristic function of the density p ϵ .
The sufficient statistics T ij are differentiable almost everywhere, and ( T ij ) 1 ≤ j ≤ k are linearly independent on any subset of X of measure greater than zero.
(3) There exist nk + 1 distinct points ( y 0 , e 0 ) , . . .
, ( y nk , e nk ) such that the nk × nk matrix
[Equation]
[Equation]
is invertible.
Then we have the parameters θ =
( f , T , λ ) are ∼ A -identifiable.
Note that the last assumption in Theorem 3.3 implies that the product space Y × E train has to be large enough to ensure the identifiability of θ when perfectly fitting the given training data distribution.
i.e. We need m |E train | > nk .
The invertibility of L implies that λ e i ( y i ) − λ e 0 ( y 0 ) need to be orthogonal to each other which further implies the diversity of environment space E .
3.2. Balanced mini-batch sampling
We consider a classic method that has been widely used in the average treatment effect (ATE) estimation — balancing score matching — to sample balanced mini-batches that mimic a balanced distribution shown in Figure 1b.
Causal effect estimation studies the effect a treatment would have had on a unit which in reality received another treatment.
A causal graph similar to Figure 1a is usually considered in a causal effect estimation problem, where Z is called the covariate (e.g. a patient profile), which is observed before treatment Y ∈ { 0 , 1 } (e.g. taking drug or placebo) is applied.
We denote the effect of receiving a specific treatment Y =
y as X y .
Note that this causal graph implies that we make the Strong Ignorability assumption.
i.e. Z includes all variables that are related to both X and Y .
In the case of a binary treatment, the ATE is defined as the expected difference of effect after receiving different treatments: E [ X 1 − X 0 ] (e.g. difference in blood pressure).
For a randomized controlled trial, we can directly estimate the difference between E [ X | Y = 1] and E [ X | Y = 0] from the observed data as the true treatment effect, as in this case we force Z ⊥⊥ Y and there would not be systematical difference between units exposed to one treatments and units exposed to another.
However, in most observed datasets, Z is correlated with
Y .
Thus E [ X | Y = 1] and E [ X | Y = 0] are not directly comparable.
We can then use balancing score b ( Z ) (Dawid, 1979) to de-correlate Z and Y :
Definition 3.4.
A balancing score b ( Z ) is a function of covariate Z s.t.
Z ⊥⊥ Y | b ( Z ) .
The ATE can then be estimated by matching units with same balancing score but different treatments:
E b ( Z ) [ E [ X | Y = 1 , b ( Z )] − E
[ X | Y = 0 , b ( Z )]] .
There is a wide range of functions of Z that can be used as a balancing score, where the scalar propensity score p ( Y = 1 | Z ) is the coarsest one and the covariate Z itself is the finest one (Rosenbaum & Rubin, 1983).
To extend this statement to non-binary treatments, we first define propensity score s ( Z ) for Y ∈ Y = { 1 , 2 , ..., m } as a vector:
Definition 3.5.
The propensity score for Y ∈ { 1 , 2 , ..., m } is s ( Z ) =
[ p ( Y = y | Z )] my =1 .
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 b ( Z ) is a balancing score, if and only if b ( Z ) is finer than s ( Z ) . i
[Table] e. exists a function g such that s ( Z )
=
g ( b ( Z )) .
We use b e ( Z ) to denote the balancing score for a specific environment e . The propensity score in a training environ- ment e would then be s e ( Z ) = [ p eθ ( Y = y | Z
)] my =1 , which [Table] O
[Equation] [Table] CORAL [Table]
can be derived from the learned conditional prior p eθ ( Z | Y ) :
[Equation]
where p e ( Y = i ) can be directly estimated from the training data D e .
We propose to construct balanced mini-batches by matching 1 ≤ a ≤ m − 1 examples with a different label Y but the same/closest balancing score b e ( Z ) for each example sampled from the training environment e .
(The detailed sampling algorithm is shown in Appendix A.)
With perfect match at every step (i.e., b e ( z j ) = b e ( z ) ) and a = m − 1 , we can obtain a completely balanced minibatch sampled from the balanced distribution with Y ⊥⊥ Z .
However, an exact match of balancing score is unlikely in reality, so the quality of matched data point would likely be lower than the referencing data point in terms of having the same balancing score.
This can be mitigated by choosing a smaller a .
However, this would make Y and Z not completely independent.
In fact, if we have exact match of balancing score, the larger a is, the weaker the correlation between Y and Z would be.
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 .
The above arguments can be summarized as below:
Theorem 3.7.
A balanced mini-batch with exact matches of balancing score and a = m − 1 can be regarded as sampling from the balanced distribution over all training environments p B ( X, Y, Z ) .
In general, the balanced mini-batch can be regarded as sampling from a semi-balanced distribution with ˆ p B ( Y | Z, E )
= 1 a +1 ( am − 1 + m − a − 1 m − 1 p
( Y | Z, E )) .
4. Experiments
To verify the effectiveness of our proposed mini-batch balancing method, we conduct experiments on three domain generalization datasets: ColoredMNIST (Arjovsky et al.,
, PACS (Li et al., 2017) and TerraIncognita (Beery et al., 2018).
We use the training-domain validation as defined in (Gulrajani & Lopez-Paz, 2020) for model selection.
As our method only modifies the mini-batch sampling strategy, we applied our proposed method along with four widely-used domain generalization baselines:
ERM (Vapnik, 1998), IRM (Gulrajani & Lopez-Paz, 2020),
DRO (Sagawa et al., 2019) and CORAL (Sun & Saenko, 2016), and compare the performance of using our method with using the usual random mini-batch sampling strategy.
We use both b ( Z ) =
Z and b ( Z )
=
s e ( Z ) to construct balanced mini-batches.
For details, see Appendix C.
ColoredMNIST:
In Table 1, b ( Z ) =
s e ( Z ) significantly outperforms b ( Z )
=
Z , and outperforms the random minibatch sampling baseline by more than 5% (absolute) with all four classifiers.
This is likely due to the relatively large noise (25%) in the label assignment: Z captures all other features except the ones directly determines the label, while s e ( Z ) only capture the features that has strong correlation with the label Y .
i.e. the color information.
As Z is finer and contains more information than s e ( Z ) , it is likely that the closest match by b ( Z ) =
Z would be an image of the same digit with the same color.
i.e. The matched example is likely to be an image labeled with the “wrong” class.
This will hurt the performance as we up-sample the noise data.
PACS and TerraIncognita:
Table 1 shows that our balanced mini-batch sampling method outperforms the random mini-batch sampling baseline on two real-world datasets.
We observe that the performance difference between b ( Z ) =
s e ( Z ) and b ( Z )
=
Z is not as large as that on the ColoredMNIST dataset.
It is partially because these two datasets have less label noise than ColoredMNIST.
5. Related Work
A growing body of work has investigated the out-of-domain (OOD) generalization problem with causal modeling.
One prominent idea is to learn invariant causal features that describe the true causal mechanism of interest across domains.
When multiple training domains are available, this can be approximated by enforcing some invariance conditions across the observed 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).
However, recent work claims that many of these approaches still fail to achieve the intended invariance property (Kamath et 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).
Some works propose to use an auxiliary variable different from the label instead of datasets from multiple domains to solve the OOD problem (Makar et al., 2022;
Puli et al., 2022).Their methods are two-phased that first balance the train data distribution with the help of auxiliary variable and then add invariance regularizations on the training objective.
With appropriate assumptions of the train data distribution, some other OOD works propose to use variational autoencoder (VAE) to learn latent variables in the assumed causal graph (Liu et al., 2021;
Lu et al., 2022), instead of using an observed auxiliary variable.
The identifiability guarantee is usually based on the pioneer work on identifiable VAE by (Motiian et al., 2017).
Our method is based on the idea of both distribution balancing and latent variable learning.
To better utilize the learned latent variable, we use a classic method for average treatment effect (ATE) estimation (Holland, 1986) – balancing score matching (Rosenbaum & Rubin, 1983).
Recently, perfect match (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
We propose a causality-based domain generalization method that samples balanced mini-batches to reduce spurious correlations.
We show that our assumed data generation model with invariant causal mechanism can be identified up to sample transformations.
We demonstrate theoretically that the balanced mini-batch is approximately sampled from a spurious free data distribution with the same causal mechanism under idea scenarios.
Our experiments empirically show the effectiveness of our method on both semi-synthetic dataset and real-world datasets.
A. Algorithm details
ϵ
[Figure]
E
[Figure]
environment Number color: red Appearance A small number (<5)
Noise (a) As realized by ColoredMNIST ϵ Y X Z E
Sketch style class (rather than photo) Artist style Object class: Giraffe Noise
C Background, Drawing-specific choices Appearance (b) As realized by PACS Input: |E train | training datasets D e = { ( x ei , y ei ) } N e i =1 sampled from distribution P ( X, Y | E = e ) for all e ∈ E train , a balancing score b e ( z i ) calculated for each training data point ( x ei , y ei ) , and a distance metrics d ( · , · ) that calculates the distance between two balancing scores
D balanced ← Empty for e ∈ E train do
Randomly sample B data points D erandom from D e
Add D erandom to D balanced for ( x e , y e ) ∈
D erandom do
Uniformly sample α different labels Y alt from Y = { 1 , 2 , ..., m } such that y ̸
= y e for all y ∈ Y alt
Suppose the balancing score of ( x e , y e ) is b e ( z ) for y ∈ Y alt do
Search across D e for the data point ( x ej , y ej ) such that y ej =
y and has the smallest d ( b e ( z j ) , b e ( z ))
Add ( x ej , y ej ) to D balanced end for end for end for
B. Proofs
In this section, we give full proofs of the main theorems in the paper.
B.1. Balanced mini-batch sampling B.1.1. P ROOF FOR T HEOREM 3.
Our proof of all possible balancing scores is an extension of the proof of Theorem 2 from (Rosenbaum & 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 (Definition 3.4) and Bayes’ rule, we have:
[Equation]
On the other hand, since b ( Z ) is a function of Z , we have:
[Equation]
Equation (11) and Equation (12) give us p ( Y | b ( Z ))
=
p ( Y | Z ) .
So to show b ( Z ) is a balancing score, it is sufficient to show p ( Y | b ( Z ))
=
p ( Y | Z ) .
Let the y -th entry of S ( Z ) be s y ( Z )
=
p ( Y = y
| Z
) , then:
[Equation]
But since b ( Z ) is finer than s ( Z ) , b ( Z ) is also finer than s y ( Z ) , then
[Equation]
Then by Equation (13) and Equation (14) we have P ( Y = y | Z ) =
P ( Y = y
| b ( Z )) as required.
So b ( Z ) is a balancing score.
For the converse, suppose b ( Z ) is a balancing score, but that b ( Z ) is not finer than s ( Z ) . Then there exists z 1 and z 2 such that s ( z 1 )
̸
= s ( z 2 ) , but b ( z 1 ) =
b ( z 2 ) .
By the definition of s ( · ) , there exists y such that P ( Y = y | z 1 ) ̸
= P ( Y = y
| z 2 ) .
This means, Y and Z are not conditionally independent given b ( Z ) , thus b ( Z ) is not a balancing score.
Therefore, to be a balancing score, b ( Z ) must be finer than s ( Z ) .
Note that s ( Z ) is also a balancing score, since s ( Z ) is also a function of itself.
B.1.2. P ROOF FOR T HEOREM 3.
We provide a proof for Theorem 3.7, demonstrating the feasibility of balanced mini-batch sampling.
Proof.
In Algorithm 1, by uniformly sampling a different labels such that y ̸ = y e , we mean sample Y alt = { y 1 , y 2 , ..., y a } by the following procedure:
[Equation]
Where U denotes the uniform distribution.
Suppose D balanced ∼ ˆ p B ( X, Y ) , and data distribution D e ∼ p
( X, Y | E = e ) , ∀ e ∈ E train .
Suppose we have an exact match every time we match a balancing score, then for all e ∈ E train , we have ˆ p
[Equation]
Z ) ,
[Equation]
By the definition of balancing score, p ( Y | Z, E = e ) =
p ( Y | b e ( Z ) , E = e ) and ˆ p B ( Y | Z, E = e ) =
ˆ p B ( Y | b e ( Z ) , E = e ) , then we have [Equation]
[Equation]
When a = m − 1 , we have ˆ p B ( Y | Z, E ) = 1 m = U { 1 , 2 , ..., m } , which means ˆ p B ( X, Y, Z ) = p B ( X, Y, Z ) .
i.e. D balanced can be regarded as sampled from the balanced distribution p B as defined in Definition 2.2. B.1.3. P ROOF FOR T HEOREM 2.  Here we give a proof of the minimax optimality of the Bayes optimal classifier trained on a balanced distribution. Proof.
The Bayes optimal classifier trained on a balanced distribution p B ( X, Y ) has p ψ ( Y | X ) = p B ( Y | X ) . Then consider the expected cross entropy loss of such classifier on an unseen test distribution p e : [Equation]
[Equation]
• Equation (16) is the definition of cross entropy loss.
• Equation (18) is obtained by Y ⊥⊥ B Z and Y ⊥⊥ B Z | X .
Thus we have the cross entropy loss of p B ( X, Y ) in any environment e is smaller than that of p B ( Y )
= 1 m (random guess):
[Equation]
Which means:
[Equation]
That is, the performance of p B ( X, Y ) is at least as good as random guess in any environment.
Since we make an assumption of the environment diversity, that is for any p e with Y ̸⊥⊥ e Z , there exist a environment e ′
such that p e
( Y | X ) performs worse than random guess.
So we have:
[Equation]
Now we want to prove that ∀ e ∈ E , Y ⊥⊥ e Z, Y ⊥⊥ e Z | X, p e ( Y )
= 1 m = ⇒
p e ( Y | X ) =
p B ( Y | X ) .
For any Z ∈ Z , we have:
[Equation]
Thus we have the following minimax optimality:
[Equation]
B.2. Latent Covariate Learning B.2.1. P ROOF FOR T HEOREM 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 of this Theorem is based on the proof of Theorem 1 in (Motiian et al., 2017), with the following modifications:
We use both E and Y as auxiliary variables.
We include Y in the causal mechanism of generating X by X [Equation]
Y [Equation] Z ) [Equation]
Proof.
Step I.
In this step, we transform the equality of the marginal distributions over observed data into the equality of a noise-free distribution.
Suppose we have two sets of parameters θ =
( f , T , λ ) and θ ′ =
( f ′ , T ′ , λ ′ )
such that p θ ( X | Y, E = e ) = p θ ′ ( X | Y, E = e ) , ∀ e ∈ E train , then:
[Equation]
[Equation]
[Equation]
[Equation]
⇒
[Equation]
⇒
[Equation]
⇒
[Equation]
⇒
[Equation]
• In Equation (24), we denote the volume of a matrix A as vol A := √ det A T A .
J denotes the Jacobian.
We made the change of variable ¯ X = f Y ( Z ) on the left hand side and ¯ X
= ¯ f Y ( Z ) on the right hand side.
Since f is injective, we have f − 1 ( ¯ X ) =
( Y, Z ) .
Here we abuse f − 1 ( ¯ X ) to specifically denote the recovery of Z ,
i
[Equation]
e. f − 1 ( ¯ X )
=
Z . • In Equation (25), we introduce
[Equation]
on the left hand side, and similarly on the right hand side
[Equation] • In Equation (26), we use ∗ for the convolution operator.
• In Equation (27), we use F [ · ] to designate the Fourier transform, and the characteristic function of ϵ is
ϕ ϵ = F [ p ϵ ] .
• In Equation (28), we dropped ϕ ϵ ( ω ) from both sides as it is non-zero almost everywhere (by assumption (1) of the Theorem).
Step II.
In this step, we remove all terms that are either a function of X or Y or e .
By taking logarithm on both sides of
Equation (29) and replacing P T ,λ by its expression from Equation (3) we get:
[Equation]
Let ( e 0 , y 0 ) , ( e 1 , y 1 ) , ..., ( e nk , y nk ) be the points provided by assumption (3) of the Theorem.
We evaluate the above equations at these points to obtain k + 1 equations, and subtract the first equation from the remaining k equations to obtain:
[Equation]
Let L be the matrix defined in assumption (3) and L ′ similarly defined for λ ′ ( L ′ is not necessarily invertible).
Define[-b l = (cid:80) ni =1 log W e ′ 
 i ( y 0 )-]
[-W eli ( y l ) W e 
 i ( y 0 ) W e ′ L i ( y l ) and b = [ b l ] nkl =1 .-][Equation]
Then Equation (32) can be rewritten in the matrix form:
[Equation]
We multiply both sides of Equation (33) by L − T to get:
[Equation]
Where A = L − T L ′ and c =
L − T b .
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 × n matrix of rank n .
This implies that the Jacobian of T ′ ◦
f ′− 1 exists and is of rank n and so is A .
We distinguish two cases:
If k = 1 , then A is invertible as A ∈ R n × n .
If k > 1 , define ¯ x = f − 1 ( x ) and T i (¯ x i )
=
( T i, 1 (¯ x i ) , ..., T i,k (¯ x i )) .
Suppose for any choice of ¯ x 1 i , ¯ x 2 i , ..., ¯ x ki , the family ( d T i (¯ x 1 i )
d ¯ x 1 i , ..., d T i (¯ x ki )
d ¯ x ki ) is never linearly independent.
This means that T i ( R ) is included in a subspace of R k of dimension of most k − 1 .
Let h be a non-zero vector that is orthogonal to T
i ( R ) .
Then for all x ∈ R , we have ⟨ d T i ( x ) dx , h ⟩ = 0 .
By integrating we find that ⟨ T i ( x ) , h ⟩ = const.
Since this is true for all x ∈ R and for a h ̸ = 0 , we conclude that the distribution is not strongly exponential.
So by contradiction, we conclude that there exist k points ¯ x 1 i , ¯ x 2 i , ...
¯ x ki such that ( d T i (¯ x 1 i )
d ¯ x 1 i , ..., d T i (¯ x ki )
d ¯ x ki ) are linearly independent.
Collect these points into k vectors (¯ x 1 , ..., ¯ x k ) and concatenate the k Jacobians J T (¯ x l ) evaluated at each of those vectors horizontally into the matrix Q =
( J T (¯ x 1 ) , ..., J T (¯ x k )) and similarly define Q ′ as the concatentation of the
Jacobians of T ′ ( f ′− 1 ◦ f (¯ x )) evaluated at those points.
Then the matrix Q is invertible.
By differentiating Equation (34) for each x l , we get:
[Equation]
The invertibility of Q implies the invertibility of A and Q ′ .
This completes the proof.
C. Experiment Details
Baselines : 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 minimization (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) performs ERM while increasing the weight of the environments with larger errors.
Deep CORAL matches the mean and covariance of feature distributions across training domains.
We use a multi-layer perceptron based VAE (Kingma & Welling, 2013) to learn the latent covariate
Z
and we choose the conditional prior p θ ( Z | Y, E = 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 fixed variance.
For the architecture of the image classifiers (ERM (Vapnik, 1998), IRM (Gulrajani & Lopez-Paz, 2020), GroupDRO (Sagawa et al., 2019), CORAL (Sun & Saenko, 2016)), following the setting of DomainBed (Gulrajani & Lopez-Paz, 2020), we train a convolutional neural network from scratch for ColoredMNIST (Arjovsky et al., 2020) dataset, and use a pretrained ResNet50 (He et al., 2016) for PACS (Li et al., 2017) and TerraIncognita (Beery et al., 2018).
Each experiment is repeated with 3 different random seeds.
We perform our experiments on the DomainBed codebase 2 and follow its default settings and hyperparameters.
For all datasets, we set the number of matched examples to be a = 1 .
[Table]
ColoredMNIST: The ColoredMNIST (Arjovsky et al., 2020) dataset is generated from the regular MNIST hand-written digits dataset (Lecun et al., 1998) by first assigning a binary label to indicate whether the digit is smaller than 5 to each MNIST image with some random noise [Table] and then assigning a binary color (green or red) based on different correlations with the label across environments.
Because of the injected noise in label assignment, the maximum accuracy a classifier can achieve is 75%.
As shown in Table 2, there are in total 3 environments in the dataset, each of which has p ( Color
=
Red | Y = 1)
= 0 .
9 , with 70,000 examples in each environment.
In this case we have |E train | = 2 and m = 2 .
We also set k = 1 by fixing the variance of the conditional prior p θ ( Z | Y, E = e ) to be 1.
Then we take the maximum possible dimension of the latent Z , n = 3 , according to the identifiability condition in Theorem 3.3.
Note that as m = 2 , the balanced mini-batches can be regarded as sampling from a balanced distribution with perfect match in balancing score and a = 1 as stated in Theorem 3.7.
Figure 3 shows three sets of reconstructed images with the same latent variable 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 with corresponding label
Y .
As shown in Table
[Table]
our proposed balanced mini-batch sampling method significantly outperforms random mini-batch sampling in the worst test environment (0.9) by 20% (absolute), while achieving comparable performance on the other two environments.
Note that 0.9 is the hardest test environment as the train environments 0.1 and 0.2 are significantly different from it.
While the accuracy on 0.9 is low, we are able to get around 75% (maximum) accuracy on the train domain validation sets.
[Table] [Table]
PACS and TerraIncognita:
We also conduct experiments on two more realistic datasets, PACS (Li et al., 2017) and
TerraIncognita (Beery et al., 2018).
The PACS has four environments, image style: art, cartoons, photos and sketches, with 9,991 examples in total.
There are 7 classes in the dataset, indicating the object in the image.
The
TerraIncognita dataset contains photographs of wild animals taken by camera traps at four different locations: L100, L38,
L43 and L46, with 24,788 examples in total.
There are 10 classes in the dataset, indicating the animal appears in the image.
In these two datasets, we do not require |E train | , m , n and k to explicitly satisfy the identifiability condition in Theorem 3.3, as the latent variable could require a larger dimensionality n to capture than the dataset allows.
Also, as m > 2 , a =means we cannot completely eliminate the spurious correlation in the training dataset in the ideal scenario.
All experiments were conducted on NVidia A-100 and Titan X GPUs.
The DomainBed codebase is released under an MIT license at https://github.com/facebookresearch/DomainBed.
We utilize versions of the datasets ColoredMNIST, PACS, and
TerraIncognita (Caltech Camera Traps) that are distributed within DomainBed, under the Creative Commons AttributionNonCommercial 4.0, CopyLeft/No Rights Reserved, and Community Data License Agreement (CDLA) licenses respectively.
Our code and corresponding instructions are included in the supplementary materials.