A SYMMETRY L EARNING FOR C OUNTERFACTUALI NVARIANT C LASSIFICATION IN OOD T ASKS
A BSTRACT
Generalizing from observed to new related environments (out-of-distribution) is central to the reliability of classiﬁers.
However, most classiﬁers fail to predict label Y from input X when the change in environment is due a (stochastic) input transformation T te ˝ X 1 not observed in training, as in training we observe T tr ˝ X 1 , where X 1 is a hidden variable.
This work argues that when the transformations in train T tr and test T te are (arbitrary) symmetry transformations induced by a collection of known m equivalence relations, the task of ﬁnding a robust OOD classiﬁer can be deﬁned as ﬁnding the simplest causal model that deﬁnes a causal connection between the target labels and the symmetry transformations that are associated with label changes.
We then propose a new learning paradigm, asymmetry learning , that identiﬁes which symmetries the classiﬁer must break in order to correctly predict Y in both train and test.
Asymmetry learning performs a causal model search that, under certain identiﬁability conditions, ﬁnds classiﬁers that perform equally well in-distribution and out-of-distribution.
Finally, we show how to learn counterfactually-invariant representations with asymmetry learning in two simulated physics tasks and six image classiﬁcation tasks
I NTRODUCTION
A signiﬁcant challenge in classiﬁcation tasks happens when the test distribution differs from the training distribution (i.e., the task requires out-of-distribution (OOD) generalization), since not accounting for the distribution shift can lead to poor generalization accuracy (Geirhos et al., 2020;
Hu et al., 2020; Koh et al., 2020; D’Amour et al., 2020).
If the learner sees examples from the test distribution, ﬁnding a classiﬁer invariant to the distribution shift can still be a data-driven task (e.g., classical domain adaptation Ben-David et al.
; Muandet et al.
; Zhao et al.
This includes cases such as invariant risk minimization (Arjovsky et al., 2019) and its generalizations (Bellot & van der Schaar, 2020), where the training data and the test data distributions overlap in a way that can be exploited by data-driven algorithms (Creager et al., 2021; Krueger et al., 2021; Rosenfeld et al., 2020).
However, if the learner sees no examples from the test distribution, the task is not purely data-driven and requires assumptions about the data generation process.
More formally, our work considers general OOD tasks with training distribution P p Y tr , X tr q , where X tr : “ T tr ˝ X : , with X : as a hidden variable with distribution P p X : q and T tr P T is a random input transformation in training
T tr : X Ñ X , where t ˝ x is the application of transformation t P T on x P X .
The difference between train and test is a change in input transformation with Y te : “ Y tr and X te : “ T te ˝ X : , where P p T tr q ‰ P p T te q .
We are interested in learning an invariant classiﬁer that generalizes well in held out examples from the training and test distributions.
The deﬁnition of transformation matters in this task.
We ﬁrst seek to generalize the existing literature on transformation invariances, e.g. (Shawe-Taylor, 1993; Kondor & Trivedi, 2018; Finzi et al., 2021;
Maron et al., 2018; Murphy et al., 2019b; Mouli & Ribeiro, 2021; Bronstein et al., 2017).
Our transformations are tied to equivalence relations rather than transformation groups, which frees them from the need to have inverses (in order to form a transformation group).
Our transformations may not have inverses.
We also explain why the task of learning an invariant OOD classiﬁer is not, in general, solvable via traditional data augmentation.
Before we continue describing our OOD learning task, it is important to clarify the connection between Pearl’s causal hierarchy and invariant representation learning.
Pearl’s causal hierarchy and invariant representation learning.
Pearl’s causal hierarchy (Pearl & Mackenzie, 2018; Bareinboim et al., 2020)) has three layers: Observational (Layer 1), interventional (Layer 2), and counterfactual (Layer 3).
Upper layers can perform lower layer tasks, but not vice-versa (see Bareinboim et al.
Tasks should be described using the lowest layer that can solve them.
Layer 1: Any task that can be performed without constraints on the causal model, i.e., by data alone, is observational (Layer 1).
Traditional domain adaptation is a Layer 1 task.
Note that a classiﬁer that performs well OOD is itself a Layer 1 classiﬁer, since it tries to predict P p Y te | X te q .
Layer 2: Without observations from P p X te q and/or P p Y te | X te q , learning an OOD classiﬁer requires some assumptions about the data generation process (Layers 2 or 3 assumptions).
Data augmentation is traditionally an interventional task (Layer 2), with new interesting methods increasingly using causal language (Ilse et al., 2021; Teney et al., 2020).
For instance, in a task predicting an image’s foreground, knowing how to act on an image in training X tr to change the background seen in training to the backgrounds seen in test X te “ T ˝ X tr with a transformation T , implies we know how to predict P p Y | X, do p T qq .
Layer 3: Counterfactuals are the most challenging task.
We start our description with an example.
Consider a random continuous transformation T tr 2 (in training) which changes to random transformation T te 2 (in test).
Let X : describe a hidden variable such that X tr : “ T 1 ˝ T tr 2 ˝ T 3 ˝
X : and
X te : “ T 1 ˝ T te 2 ˝ T 3 ˝
X : , where T 1 and T 3 are independent continuous random transformations and
P p T tr 2 q ‰ P p T te 2 q .
Assume the target variable Y depends only on X : , T 1 , and T 3 .
To counterfactually ask what would have happened to the observed input x if we had forced do p T tr 2 “ ˜ t 2 q , we are inquiring about X p T tr 2 “ ˜ t 2 q| X tr “ x .
Note that do p T tr 2 “ ˜ t 2 q does not change Y .
Also note that the knowledge of X tr “ x is an indirect statement about T tr 2 since P p T tr 2 | X tr “ x q ‰ P p T tr 2 q .
That is, for x, x 1 P X ,
[Equation]
Equation (1) and the difference between the causal hierarchy layers will be relevant for our results.
Contributions.
Our contributions can be described as follows:
We introduce a generalization of transformation groups via symmetry transformations tied to equivalence classes that removes the requirement of invertible transformations common in deﬁnitions using transformation groups.
We introduce the concept of counterfactual invariant representations for symmetry transformations and show how it can be described as a counterfactual task for causal structure discovery.
Finally, we introduce asymmetry learning , which describes a representation regularization that, under a set of assumptions, learns the correct counterfactual invariant OOD classiﬁer.
S YMMETRIES AND TRANSFORMATIONS
Geometrically, an object is called symmetric if there is a transformation on the object that does not change its shape (in some deﬁnition of shape).
For example, a square is symmetric with respect to rotations.
The notion of symmetry however is not restricted to geometric notions.
In general, we can deﬁne a mathematical object as symmetric if there is a transformation on the object that returns another object equivalent to the ﬁrst (Rosen, 2008, Chapter 10).
It is clear from this deﬁnition of symmetry that we ﬁrst need to deﬁne what we mean by equivalent objects.
For instance, we say two geometrical objects are equivalent if they have the same shape, but we need a more general deﬁnition.
We deﬁne an input symmetry in a space X with at least two elements as an equivalence relation „ .
An equivalence relation in X is a binary relation „ such that for all a , b , c P X , we have (i) a „ a , (ii) a „ b ðñ b „ a , and (iii) ( a „ b and b „ c q ùñ a „ c .
Equivalence relations allow us to deﬁne equivalent objects in X : a „ b means a is equivalent to b .
The set of all objects equivalent to some a P X is called the equivalence class of a , deﬁned as r a s : “ t x P X : x „ a u .
Note that one can deﬁne m ě 2 equivalence relations on the same input space.
The equivalence class of x with respect to equivalence relation k is denoted r x s p k q , k “ 1 , . . .
, m .
Two inputs a , b P X might be equivalent under one equivalence relation „ 1 , but not equivalent under a different equivalence relation „ 2 , that is, we can have both b P r a s p 1 q and b R r a s p 2 q .
Still, even in this last case it is possible that a is equivalent to some other input c ‰ b in both equivalence relations, i.e., it is possible
D c P X , c ‰ a , s.t.
c P r a s p 1 q X r a s p 2 q .
We denote the collection of equivalence classes of X under the equivalence relation „ k as the quotient space X { „ k : “ tr x s p k q | x P X u .
Transformation group example.
Consider the bijective transformations t : X Ñ X of a transformation group G , t P G .
We now deﬁne an equivalence relation over G as t ˝ x „ G x for all t P G .
The equivalence class r x s p G q is x ’s orbit deﬁned as r x s p G q : “ t x 1 : D t P G, x 1 “ t ˝ x u .
For example, if
G is the group that permutes the elements of vectors in R 3 , then p 1 , 2 , 3 q „ G p 2 , 1 , 3 q .
Property functions example.
Another way of deriving an equivalence relation is via functions of the input space z : X Ñ R p , where the output z p x q is a particular property of the vector x P X .
For example, given an observation of length T from a dynamical system, x P R d ˆ T , a possible property function could be z energy p¨q that computes the energy of the dynamical system.
Assuming there are m known properties z 1 , . .
. , z m with z i : X Ñ R p
i , we can construct corresponding equivalence relations „ 1 , . . .
, „ m such that for any x , x 1 P X , x „ i x 1 if
z j p x q “ z j p x 1 q , @
j ‰ i .
In words, two inputs are equivalent under „ i if they have the same properties for all z j , j ‰ i .
Symmetry transformations.
As seen above, symmetries can be deﬁned without deﬁning how the input is transformed to create the equivalence classes, although deﬁning a set of transformations is useful when describing the equivalence class.
Given an equivalence relation „ , we can deﬁne a set of transformations T that respect the equivalence relation such that @ t P T , @
x P X , t ˝ x „ x .
We call
T the set of symmetry transformations of „ .
Similar to transformations groups, T always has the identity transformation t id ˝ x “ x , but in contrast, all the transformations in T need not be bijective.
Join of equivalence relations.
Similar to how two groups can be joined to form a larger group, two equivalence relations can be joined to form a coarser equivalence relation.
Given two equivalence relations, „ 1 and „ 2 , their join „ 1 _ „ 2 is deﬁned as: for all x , x 1 , x p„ 1 _ „ 2 q x 1 if and only if there exists a chain of equivalence relations x „ k 1 x 1 , . . .
, x h ´ 1 „ k h x 1 with all k j P t 1 , 2 u .
It is easy to check that „ 1 _ „ 2 is an equivalence relation.
We are now ready to deﬁne a general causal model that deﬁnes the training and test distributions in our setting.
SCM FOR SYMMETRY - BASED OOD TASKS
Let X , Y denote the input and output spaces respectively.
We deﬁne our general structural causal model (SCM) as follows.
We deﬁne X : P X as the unobserved canonically ordered input
[Equation]
with U u a background random variable and g : U u Ñ X is a measurable map.
This deﬁnition is general enough to deﬁne any task.
There are m possible symmetries given in the form of equivalence relations „ 1 , . . .
, „ m over the input space X .
Let T p
k
q denote a set of symmetric transformations t on X corresponding to the equivalence relation „ k , 1 ď k ď m .
In other words, for all x P X and t P T p k q , we have p t ˝ x q „ k x .
Similarly, let T be the set of all symmetric transformations with respect to the join equivalence relation „ 1 ,...,m ”„ 1 _ . . .
_ „ m .
We can think of transformation t P T as a path x
t p k 1 q ÝÝÝÑ x 1 ¨ ¨ ¨ x h ´ 1 t p kh q ÝÝÝÑ x h that starts at x , applies a transformation t p k 1 q P T p k 1 q to get x 1 P r x s p k 1 q , and so on until it stops and outputs a value x h , h ě 1 .
Let U 1 , . .
, U m be independent background variables associated with the m symmetries, where
U i P U
i , i “ 1 , . . .
, m .
These background variables together select a function t p U 1 , . .
, U m q from the set T as follows.
Each U k independently selects a countable sequence of transformations t p k q 1 ,U k , t p k q 2 ,U k , . . .
P T p k q .
Then, t p U 1 , . .
, U m q is deﬁned by interleaving these transformations
[Figure] [Figure]
C oun t e r f a c t ua l i npu t s w e m u s t be i n v a r i an t t o ...
[Equation]
¨ to construct the path described above.
Since T p 1 q , T p 2 q , . . . contain the identity transformation, t p U 1 , . .
, U m q can be described by a ﬁnite sequences of transformations.
The observed X is the result of a transformation of X :
[Equation]
Finally, the label Y is deﬁned as a function of the untransformed canonical input X : as
[Equation]
where D Ď t 1 , . .
, m u is unknown.
This means that Y is not invariant with respect to equivalence relations „
i , i P D , i.e., examples x and x 1 P r x s p
i q can have different labels.
A distribution over the variables U u , t U i u mi “ 1 , U Y entails a joint distribution P p X, Y q over the observed variables.
Illustrative SCM example.
Figure 1 illustrates our data generation process.
The training data
Figure 1(a) has X : deﬁned as a centered upright brown rod (i.e., X : is deterministic).
The label Y is deﬁned by the rotation transformations T rot “ t T rot 0 ˝ , T rot 90 ˝ u .
The image can also be horizontally translated by t´ 5 , 0 , 5 u units via transformations T trans “ t T trans ´ 5 , T trans, T trans ` 5 u (only 0 and `translations are depicted), but Y does not depend on these horizontal translations.
The transformations applied to X : are randomly chosen via U rot and U trans , which are two bidimensional vectors indexing a sequence four transformations that interleave rotations and translations (see Figure 1).
A representation that counts the number of brown pixels in the green shaded area of X tr is enough to achieve 100% accuracy in the training distribution.
We formally deﬁne OOD distribution shifts next using Figure 1 for illustration.
OOD distribution shift.
Let s D “ t 1 , . . .
, m uz D be the complement of the set of symmetry relations D that Y depends on.
We deﬁne the OOD distribution shift between train and test as a shift in the distribution of P pp U
i q i P s D q , inﬂuencing the distribution of input transformations in Equation (3), which in turn can shift the distributions P p X tr q , P p Y tr | X tr q , P p Y tr , X tr q to
P p
X te q , P p Y te
| X te q , P p Y te , X te q respectively.
Since X does not causally affect Y in our structural causal model (Equation (4)), changes in input transformations are able to shift P p Y | X q .
For example, in Figure 1(b) the test data (only a single example shown) could suffer an OOD shift due to a different distribution over P p U trans q that introduces non-zero translations before the second rotation.
Note that the representation that counted the number of brown pixels in the green shaded area, which was perfect for the training inputs X tr , will achieve poor accuracy in the test inputs X te .
Learning OOD classiﬁers.
Equation (4) shows that the label Y is invariant to changes in the distribution of p U
i q
i P s D in the test distribution, but we do not know s D .
Hence, if our representation of
X is invariant to changes in the distribution of p U
i q
i P s D , we will be able to perform the OOD task.
A SYMMETRY LEARNING & FINDING THE RIGHT REPRESENTATION SYMMETRY FOR THE OOD TASK 4. F INDING OOD- INVARIANT REPRESENTATIONS AS CAUSAL STRUCTURE DISCOVERY
We ﬁrst deﬁne the process of ﬁnding an OOD invariant representations for the symmetries t„ i u i
P
s D our classiﬁer should be invariant to in the test data.
Since Y does not depend on t U i u i P s D , we will make a representation of X that is invariant to transformations driven by t U i u
i P s D .
Deﬁnition 1 introduces the concept of counterfactual invariance for symmetry transformations.
We note that this deﬁnition is less restrictive than the parallel work of Veitch et al.
(2021, Deﬁnition 1.1): whereas Veitch et al.
(2021, Deﬁnition 1.1) require invariance over the entire sample space, we only require invariance over the test support of transformation variable
U i .
The deﬁnitions are equivalent if the test support is the entire sample space of U i .
Deﬁnition 1 (Counterfactual-invariant representations for symmetric transformations) .
Assume the
SCM deﬁned in Equations (2) to (4) .
A representation Γ i : X Ñ R d , d ě 1 , is counterfactualinvariant to the transformations T 1 ,U i , T 2 ,U i , . . .
of equivalence relation „ i , 1 ď i ď m , if Γ i p x q “ Γ i p X p U i
“ ˜ u i q| X “ x q almost everywhere, @
˜ u
i P supp p
U te i q , @
x P supp p X tr q , where supp p
A q is the support of random variable A .
A representation Γ S : X Ñ R d , d ě 1 , is counterfactual-invariant to a subset S Ď t 1 , . . .
, m u if it is jointly counterfactual-invariant to the transformation indices
t U j u j P S of equivalence relations t„
j u j P S .
We refer the reader to Equation (1) for the relationship between the counterfactual variables
X p
U i
“ ˜ u q| U i “ u and X p U i “ ˜ u q| X “ x . Figure 1(d) illustrates why counterfactual language is important for our task: It states that given an input X tr “ x we need to know how it would have been different if we had chosen a different distribution P p
U trans q resulting in a different sequence of transformations
T 1 ,U trans , T 2 ,U trans .
From Figure 1(c) it is clear that we cannot simply data-augment our training data with translations, since we would think that counting brown pixels in the green shaded area is an invariant representation for U trans .
Up until now we have not imposed restrictions on the types of transformations T p i q , i “ 1 , . . .
, m , we consider in this work.
Our next results require imposing conditions on these transformations.
Deﬁnition 2 (Equivalence class lumpability) .
The quotient space X { „ i is the set of equivalence classes of X with respect to equivalence relation „ i , i “ 1 , . . .
, m .
Let r x s p
i
q P X { „ i be the equivalence class of x P X with respect to equivalence relation „ i .
Then, X { „ i is said to be lumpable with respect to a transformation set T if
[Equation]
In words, if the lumpability condition in Deﬁnition 2 holds for an equivalence relation „ i with respect to a set of transformations T , then every transformation in T maps all points within an equivalence class r x s p i q P X { „ to points in a another equivalence class r x 1 s p
i q P p X { „q .
To illustrate the lumpability condition, consider two transformation groups G 1 and G 2 whose transformations commute, i.e., @p t 1 , t 2 q P G 1 ˆ G 2 , t 1 ˝ t 2 “ t 2 ˝ t 1 .
Then the equivalence classes imposed by G i , i.e., the orbits r x s p i q “ t t i ˝ x : @ t i P G i u , are lumpable with respect to the transformations G j , for i, j P t 1 , 2 u and j ‰ i .
[Figure]
[Figure] [Figure]
[Figure] [Figure]
Figure 2a(i) shows our structural causal graph where an edge U i Ñ Y exists only if i P D .
Then, we use the deﬁnition of lumpability to prove that, under certain conditions, a most-expressive representation Γ i invariant with respect to „ i allows us to identify if there is no edge U i Ñ Y in the causal DAG.
Theorem 1 (Counterfactual invariance & causal DAG identiﬁcation) .
Let X { „ i be lumpable given every T p j q , j ‰ i as in Deﬁnition 2.
Then, the structural causal DAG implied by Equations (2) to (4) (depicted in Figure 2a(i)) does not contain the edge U
i Ñ Y
iff
[Equation]
P
[Equation]
X : q
[Equation]
P p U
[Equation]
q ,
[Equation]
P p U m q , where Γ i is a most-expressive representation that is invariant with respect to „ i .
The proof is in the Appendix.
With the lumpability assumption of X { „ i , Γ i in Theorem 1 is a counterfactual-invariant representation.
We now use Figure 2a(ii) to describe the result in Theorem 1.
We ﬁrst note that the representation Γ s D depicted in the ﬁgure is counterfactual invariant to s D , and hence also counterfactual invariant to k P s D .
Next we see that since the representation Γ s D is counterfactual invariant to U k , there is no arrow U k Ñ Γ s D p X q in Figure 2a(ii).
If there is no arrow
U k Ñ Y , the missing arrows from U k to Γ s D p X q will have no inﬂuence in the ability of Γ s D p X q to predict Y , assuming Γ s D is most-expressive.
If there is an arrow U k Ñ Y , cutting the arrow
U k Ñ Γ s D p X q creates a loss in predictive performance from Γ s D p X q to Y for some distribution of the background and observable variables.
If Γ s D p X q never loses any predictive power over Y for any distribution of the background and observable variables, then there is no arrow U k Ñ Y .
Assumption 1 (Asymmetry learning training data) .
In asymmetry learning we assume that every
X { „ i , i P t 1 , . . .
, m u is lumpable given T p
j q , j ‰ i , and that in a large training dataset sampled from p Y tr , X tr q , an arrow U j Ñ Y in the causal DAG of Figure 2a(i), j P t 1 , . . .
, m u , contains observations of t U j u j P D that violate Equation (5) .
Hence, if Equation (5) holds for some i P t 1 , . .
, m u in this dataset, we can conclude that there is no arrow U
i Ñ Y in the true causal DAG.
See Appendix A for a justiﬁcation of this assumption.
Next we use Assumption 1 and the previous results to search for the right OOD invariance.
4. C AUSAL STRUCTURE DISCOVERY OF RELEVANT SYMMETRIES
We need a general procedure for obtaining the unknown set D , which is equivalent to ﬁnding all transformations indices
t U i u
i P D Ď t U 1 , . .
, U m u that act as confounders between Y and X in the causal DAG in Figure 2a(i).
Finding whether an edge exists or not in the causal DAG is known as the causal structure discovery problem (e.g., Heinze-Deml et al.
The principle of our search is learning the causal structure with the fewest possible edges into Y (i.e., where Y is invariant to most U i , i “ 1 , . . .
, m ) while also maximizing the likelihood of the observed data.
Accordingly, we take the score-based causal discovery approach (Chickering (2002); Huang et al.
that assigns scores to each allowed DAG based on the training data and the complexity of the DAG to ﬁnd a minimal causal structure that ﬁts the training data.
This idea is visualized in Figure 2a(iii) where causal graphs with more edges between the transformation indices into Y are deﬁned to have higher complexity and are higher up in the partial ordering.
Our search space is simpler than typical structure discovery tasks: The DAGs in our search space have the same structure for X and only differ in edges of the form U
i Ñ Y, i P t 1 , . .
, m u .
Next, we describe a scoring criterion that uses
Theorem 1 and counterfactual-invariant representations to assign scores to the corresponding causal structures.
Proposed DAG scoring criterion.
For each DAG in the search space, we wish to assign a score based on the training data D “ tp x p
i
q , y p
i q u n tr i “ 1 under Assumption 1 for a classiﬁcation task with
C classes.
Theorem 1 shows that there is a correspondence between a causal structure without the edge
U i Ñ Y and a predictive probability gap between the original input and a most-expressive representation Γ i that is counterfactually-invariant to U i .
Thus, under Assumption 1, we can represent the causal search from Figure 2a(iii) in terms of a search over counterfactually-invariant representation function classes as shown in Figures 2a(iii) and 2b.
Formally, we are given a collection of function classes F : “ t F S : S Ď t 1 , . . .
, m uu , where F S is a family of functions Γ S that are counterfactuallyinvariant to all
U
i , i P S (Deﬁnition 1).
We wish to score each of the function classes F S P F to indirectly learn the correct causal structure.
The minimum description length (MDL) principle (Schwarz, 1978) is commonly used for causal structure discovery (Budhathoki & Vreeken, 2016; 2017) and comes with the key insight that learning from data can be viewed as compressing it.
Given the collection F and the training dataset D ,
MDL ﬁnds the function class F S P F that compresses D the most.
While there are several ways of encoding a dataset given the function class, normalized maximum likelihood (NML) code is known to be optimal (Shtarkov, 1987).
NML code is computed as follows
[Equation]
where L p F S | D q “ sup Γ S P F S ř n tr i
[Equation]
[Equation]
P p
y p
i q
[Equation]
S p x p
i q qq is the maximum log-likelihood of F S given the data and
[Equation]
measures the complexity of the function class F S by computing how well it can represent different label distributions for the given inputs t x p i q u n tr i “ 1 in training.
We can estimate the combinatorial sum in Equation (7) by uniformly sampling random labels for all the training examples.
Since COMP p F S , D q is computed using the training data, it may underestimate the complexity of function classes if, for instance, all the training examples are generated with U
i “ u i .
Then, F t i u and
F H are given the same score even though F t i u is clearly more invariant and thus, a simpler function class.
This can happen in practice if, say, all images are upright in training with no rotations applied; both rotation-invariant and rotation-sensitive function classes get the same complexity score.
In order to break the above ties of our COMP score, asymmetry learning adds an additional term to the
NML score that chooses models that have higher invariance based on the partial order (see Figure 2b).
We extend the penalty proposed by Mouli & Ribeiro (2021) and use R p F S q : “ |t F 1 : F 1 P F , F 1 ą
F S u| , the number of function classes that are higher in the partial order than F S , as the tie-breaking term.
For example, in ﬁgure R p F t 1 u q “ |t F t 1 , 2 u , F t 1 , 3 u , F t 1 , 2 , 3 u u| “ 3 .
We deﬁne the ﬁnal score of each function class F S P F as
[Equation]
The score in Equation (8) can be minimized by a score-based causal discovery algorithm to obtain the ﬁnal DAG.
We use Greedy Equivalence Search (Chickering, 2002) to showcase a concrete instantiation of asymmetry learning.
Other score-based structure discovery algorithms could also be used.
Greedy Equivalence Search.
Greedy Equivalence Search (GES) is a greedy search algorithm that optimizes a given scoring function over DAGs.
In our setting, the search begins with a DAG with no edges of the form U i
Ñ Y, i P t 1 , . . .
, m u .
In the ﬁrst phase, GES adds these edges one at a time
[Table]
that maximally improve the score in Equation (8) until there is no improvement.
In the second phase,
GES begins from the DAG obtained at the end of ﬁrst phase and deletes edges one at a time until such deletions do not improve the score.
The DAG obtained at the end of the second phase is the ﬁnal output of the algorithm.
Under the causal Markov and faithfulness assumptions, Chickering (2002) showed that GES is optimal in the large sample limit if the scoring function is locally consistent.
R ESULTS
Pendulum task description.
We evaluate the proposed method in a simulated classiﬁcation task.
Our input x is a motion vector over time p θ t , dθ t dt q Tt “ 1 of a simple pendulum of an unknown length l after it is dropped from some initial angle θ 0 with dθ 0 dt “ 0 .
After an initial τ seconds of uninterrupted motion, we simulate an elastic collision by placing another object of same mass at the bottom.
The classiﬁcation task is to predict whether the kinetic energy imparted by the pendulum is enough to move the second object beyond a certain threshold.
Physical properties and equivalence relations.
We consider the following two properties of the dynamical system described above: z 1 : X Ñ R which computes the initial potential energy of the system and z 2 : X Ñ R which returns the time of collision.
The equivalence relations „ 1 and „ 2 are deﬁned using these properties as deﬁned in Section 2.
For instance, two pendulum motion curves x , x 1 are equivalent with respect to „ 1 , i.e., x „ 1 x 1 , if they have the same time of collision, z 2 p x q “ z 2 p x 1 q .
Then T p 1 q consists of transformations that change the initial potential energy of the system (for example, by changing the length of the pendulum or the initial dropping angle θ 0 ) while keeping the time of collision same.
Similarly, x „ 2 x 1 if their respective potential energies are the same and transformations in T p 2 q change the time of collision while keeping the same initial potential energies.
Note that the space of equivalence classes X { „ 1 is lumpable with respect to
T p 2 q and vice versa (Deﬁnition 2).
Thus, by Theorem 1, we can use predictive performance of counterfactual-invariant representations for scoring the causal DAGs.
Unknown D and OOD classiﬁcation.
We consider two scenarios for the label Y given X .
First, if the motion of the pendulum is not damped by friction, then Y depends only on z 1 p x q , i.e, D “ t 1 u .
Second, if the motion of the pendulum is damped, then Y depends on both z 1 p x q and z 2 p x q , i.e., D “ t 1 , 2 u .
The extrapolation test data is generated by shifting the distribution of the background variables
t U i u
i P s D .
The task of a structure discovery algorithm is to correctly identify D .
Results.
We use the greedy equivalence search (GES, Section 4.2) algorithm to search over the different causal graphs with the proposed scoring criterion deﬁned in Equation (8).
We build classes of counterfactual-invariant representations F S corresponding to each possible value of S Ĺ t 1 , 2 u , where every Γ S P F S is invariant to t U i u i P S .
For example, F t 1 u is a family of feedforward neural networks that only take z 2 p x q as input, i.e., invariant to z 1 p x q , whereas F H is a sequence model (e.g., LSTM) with no invariance.
Table 1 reports the estimated complexity { COMP p F , D q and the ﬁnal scores S p F , D q for the different function classes for the two tasks.
The bold values indicate the function class chosen by the GES algorithm.
When D “ t 1 u , the greedy search stops after adding the edge U 1 Ñ Y as adding the second edge U 2 Ñ Y only worsens (increases) the score.
When D “ t 1 , 2 u , the greedy search is able to improve the score by adding both edges, ﬁrst U 1 Ñ Y and then U 2 Ñ Y .
In both the cases, the extrapolation test accuracy achieved by the chosen model class is the highest.
Image classiﬁcation task.
Appendices A.4 and A.5 also offers an application to image classiﬁcation using image transformation sets (groups and nongroups).
R ELATED W ORK
Counterfactual inference and invariances.
Recent efforts have brought causal inference to machine learning (extensively reviewed in Schölkopf et al. (2021); Schölkopf (2022)). Invariant Causal Prediction (Peters et al., 2015; Heinze-Deml et al., 2018) and Invariant Risk Minimization methods (Arjovsky et al., 2019; Bellot & van der Schaar, 2020) learn representations that are invariant across multiple environments but have been shown to be insufﬁcient for OOD generalization in classiﬁcation tasks without additional assumptions
Ahuja et al.
Wang & Jordan (2021) use counterfactual language to formally deﬁne and learn non-spurious representations from a single environment that can extrapolate to new environments.
Veitch et al.
(2021) deﬁne counterfactual invariant predictors f p X q when X has a single parent Z and provide conditions such predictors must satisfy over the observed distribution (given an SCM).
Kaushik et al.
(2020; 2021) propose counterfactual data augmentation for text datasets but they either require a fully-speciﬁed toy SCM or rely on humans-in-the-loop to generate the counterfactual data.
Other counterfactual methods (Johansson et al., 2016; Shalit et al., 2017; Qidong et al., 2020) learn representations to predict counterfactual change in some observed variables whereas in our setting, the transformation variables U i that generate the observed X are unobserved.
In-depth comparison of our work with the existing counterfactual methods is presented in Appendix A.3.
Domain adaptation and domain generalization.
Domain adaptation and domain generalization (e.g. (Long et al., 2017; Muandet et al., 2013; Quionero-Candela et al., 2009; Rojas-Carulla et al., 2018; Shimodaira, 2000; Zhang et al., 2015) and others) consider observed or known shifts in the data distribution, for instance, given the test distribution P
p X te q , rather than
counterfactual questions.
Causal structure discovery.
The methods for causal structure discovery can be broadly classiﬁed into two categories.
Constraint-based approaches (e.g., Spirtes et al.
; Sun et al.
(2007)) use conditional independence tests and reject causal graphs that impose more independence than what is observed in data.
On the other hand, score-based causal discovery approaches (e.g., Chickering (2002); Huang et al.
; Ding et al.
Zhu et al.
assign scores to each allowed causal graph based on the data and ﬁnd the one with best score.
While there are several works (Budhathoki &
Vreeken, 2016; 2017; Bornschein et al., 2021) that use minimum description length (MDL) (Schwarz, 1978) as a scoring criterion, we show why it is insufﬁcient for out-of-distribution tasks and use an additional term for tie-breaking.
Goudet et al.
(2017) minimize the divergence between a distribution generated by a learnt causal DAG and the observed data distribution; however the method is limited to orienting edges over observed variables, whereas our transformation variables U i are unobserved.
Recently, GFlowNets Bengio et al.
(2021a;b) have been used to sample DAGs proportional to a score function for Bayesian structure learning Deleu et al.
(2022), however we are interested in ﬁnding the best DAG with the minimum score.
Group-invariant representations.
Majority of the works strictly enforce G-invariances either within the architecture (e.g., Zaheer et al.
; Cohen et al.
Lyle et al.
; Murphy et al.
(2019a)) or via data-augmentation (Chen et al., 2020)
and do not handle the case when the target is actually inﬂuenced by the transformation of the input.
Other works (Benton et al., 2020;
Zhou et al., 2020; van der Wilk et al., 2018;
Anselmi et al., 2019) consider learning symmetries from the training data but do not consider the extrapolation task that we show can be solved only under certain conditions.
Mouli & Ribeiro (2021) consider the special case where the transformations are from normal subgroups and do not formally describe the causal task.
These works rely on invertible transformations while we deﬁne symmetries more generally via equivalence relations.
Dubois et al.
(2021) also deﬁne invariances via equivalence relations and, under the assumption that all such invariances hold in the data, the authors design methods for data compression.
Our goal is rather different: We want to discover which equivalence relations (transformations thereof) affect the label.
C ONCLUSIONS
This work considered an out-of-distribution (OOD) classiﬁcation task where the shift between train and test environments is through different symmetry transformations of the input, where symmetry transformations are deﬁned via equivalence relations over the input space.
We described the task of ﬁnding symmetries that affect the label as a causal structure discovery task and show that, under certain conditions, we can use the predictive performance of invariant representations on the observational data to predict whether an edge exists in the causal DAG (Theorem 1).
We then proposed an MDLbased scoring for this causal structure discovery.
Finally, we test our approach in two simulated physics tasks and six image classiﬁcation tasks.
A CKNOWLEDGMENTS
This work was funded in part by the National Science Foundation (NSF) Awards CAREER IIS1943364 and CCF-1918483, the Purdue Integrative Data Science Initiative, and the Wabash Heartland
Innovation Network.
Any opinions, ﬁndings and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reﬂect the views of the sponsors.
A A PPENDIX A. J USTIFICATION FOR A SSUMPTION 1.
The above assumption is inspired by the deep relationship between symmetries and intelligence.
Young children, unlike monkeys and baboons, assume that a conditional stimulus F given another stimulus D extrapolates to a symmetric relation D given F without ever seeing any such examples (Sidman et al., 1982).
That is, if given D, action F produces a treat, the child assumes that given F, action
D also produces a treat.
Young children differ from primates in their ability to use symmetries to build conceptual relations beyond visual patterns (Sidman & Tailby, 1982; Westphal-Fitch et al., 2012), allowing extrapolations from intelligent reasoning.
However, forcing symmetries against data evidence is undesirable, since symmetries can provide valuable information when they are broken.
Unsurprising, humans are generally able to quickly ﬁnd and pay attention to some types of asymmetries.
A. P ROOF OF T HEOREM
Theorem 1 (Counterfactual invariance & causal DAG identiﬁcation) .
Let X { „ i be lumpable given every T p j q , j ‰ i as in Deﬁnition 2.
Then, the structural causal DAG implied by Equations (2) to (4) (depicted in Figure 2a(i)) does not contain the edge U
i Ñ Y
iff
[Equation]
P p X : q , @ P p U 1 q , . . .
P p U m q , where Γ i is a most-expressive representation that is invariant with respect to „ i .
Proof.
Notation (following
Equation (3)):
The observed input
X is
X : “ t p U 1 , . . .
, U
i ´ 1 , u i , U i ` 1 , . .
, U m q ˝
X : where t p U 1 , . .
, U m q is obtained by interleaving the transformation sequences from each individual U 1 , . . .
, U m
and we have set U i “ u i .
Necessity: We wish to show that if the SCM does not contain edge U i Ñ Y , then Equation (5) holds for all P p X : q , P p U 1 q , . .
, P p U m q .
By this assumption, Y outputs the same label for any value of U i .
Consider the collection of equivalence classes
X { „ i .
By the lumpability condition of Deﬁnition 2, all transformations t p j q P T p j q , j ‰ i , map all points in one equivalence class of „ i to points in a different one.
On the other hand, all transformations t p i q P T p i q map points to other points within the same equivalence class under „ i .
Now, consider the equivalence class of X after all the transformations have been applied to X : .
The equivalence class of X “ t p U 1 , . .
, U
i ´ 1 , u i , U i ` 1 , . .
, U m q ˝
X : is the same as that of X ˚ “ t p U 1 , . .
, U
i ´ 1 , u id i , U i ` 1 , . .
, U m q ˝
X : where U i “ u id i always selects identity transformations.
This is because changing u i to u id i only impacts the transformations chosen from T p i q , and these transformations do not change the equivalence class under „ i .
Thus, we have shown that we reach the same equivalence class under „ i for both X and X ˚ .
Now let Γ i be a most-expressive representation that is invariant with respect to „ i .
By deﬁnition, Γ i outputs the same value within an equivalence class, thus, Γ i p X q “ Γ i p X ˚ q .
But since by assumption
U i Ñ Y does not exist, X and X ˚ have the same label always.
Thus, there is no loss of information incurred by Γ i in predicting Y with the additional restraint Γ i p X q “ Γ i p X ˚ q .
Since Γ i is mostexpressive, we have P p Y “ y | Γ i p X q , U Y q “ P p Y “ y | X, U Y q for all y P Y .
This holds for all values of u i , and hence we get the desired result for any distribution
P p
U
i q .
Sufﬁciency: We wish to show that if Equation (5) holds for all P p X : q and P p U 1 q , . .
P p U m q , then there is no edge U
i Ñ Y in the causal graph.
We will prove by contrapositive: Assume there is an edge U i Ñ Y , then we will show there exists distributions P
p
X : q and P p U 1 q , . .
P p U m q such that
Equation (5) does not hold.
Deﬁne P p X : q “ δ x : for some x : P X where δ denotes a Dirac-delta function.
Deﬁne P p
U i “ u id i q “ 0 . 5 and P p
U i “ u i q “ 0 .
for u id i , u i P supp p U
i q .
As usual, u id i always selects the identity transformation, and u i selects a single transformation t
u i P T p i q .
Similarly, for all j ‰ i , deﬁne
P p
U j q “ δ u id j for u id j P supp p
U j q that only select identity transformations.
Now, there are two possible observed inputs: x “ t p u id 1 , . .
, u id m q ˝ x : “ x : and x 1 “ t p u id 1 , . . .
, u i , . .
, u id m q ˝ x : “ t u i ˝ x : .
Finally, deﬁne Y : “ 1 p U i “ u id i q , thus x and x 1 have different labels.
But, any invariant representation Γ i by deﬁnition has Γ i p x q “ Γ i p x 1 q since they belong to the same equivalence class.
Thus, even if Γ i is most-expressive, we have | P p Y | Γ i p X q , U Y q ´ P p Y | X, U Y q| TV “ 0 . 5 .
A. A DDITIONAL R ELATED W ORK
Counterfactual invariances.
Wang & Jordan (2021) use counterfactual language to formally deﬁne and learn non-spurious, disentangled representations from a single environment.
Our work is different in the following ways.
In the structural causal model (SCM) of their work, the authors assume that there are no confounders between the observed X and the label Y .
However, in our
SCM (Figure 2a(i)), we allow unobserved confounders X : and U i , i P D .
The hidden transformation variables U
i , i P D are confounders because they affect both the observed input X and the labels Y .
We leverage the fact that the confounders are related to symmetries (and do not affect X arbitrarily) to resolve the issue with unobserved confounding.
Wang & Jordan (2021) also require pinpointability of the cause of the observed X .
In our setting, this is typically not possible since there are multiple paths of transformations from X : to the same observed X .
Thus, all the parents of X may not be pinpointable, speciﬁcally the transformation variables U 1 , . . .
, U m .
Kaushik et al.
(2020; 2021) propose counterfactual data augmentation for text datasets where human annotators are asked to make minimal modiﬁcations to the input document so as to change its label (for example, by changing a few positive words to negative words) while keeping style, etc.
ﬁxed.
This type of augmentation essentially asks the labelers to identify all the causal features in the document and make modiﬁcations to those features alone.
This can be seen as obtaining new counterfactual examples by simulating the causal model and requires knowing the true function that describes how the features affect the labels.
We consider the more realistic setting where we do not have access to such a collection of counterfactual examples.
In this work, we consider the traditional automated data augmentations under a mostly unknown data generation process, as opposed to the counterfactual data augmentation (Kaushik et al., 2020) that either considers a fully-speciﬁed toy
SCM or relies on humans-in-the-loop to generate counterfactual data.
In Figure 1(c) we show that the standard data augmentation is not sufﬁcient for the OOD task.
However, if one had access to the fully-speciﬁed causal model, one could generate the counterfactual data shown in Figure 1(d) and learn an OOD classiﬁer with the counterfactually augmented data (as done by Kaushik et al.
But our work does not assume access to these counterfactual examples.
Additionally, we prove that a counterfactual invariant classiﬁer can be constructed from traditional data augmentation alone if the lumpability condition (Deﬁnition 2) is satisﬁed.
This is not the case in Figure 1(d).
Veitch et al.
(2021) deﬁne counterfactual invariant predictors f p X q when X has a single parent Z and provide conditions such predictors must satisfy over the observed distribution (given an SCM).
Note also that Veitch et al.
(2021) assume that part of the observed input X
( X K Z ) is not causally inﬂuenced by the confounder Z .
In our scenarios this is not generally true.
For example, under a color change, the entire observed image X changes.
Still, we show that the notion of a counterfactual invariant predictor exists.
Hence, the deﬁnition of Veitch et al.
(2021, Lemma 3.1) of a counterfactually invariant predictor that requires a segment of X to not causally depend on Z , a fundamental result of their work, unfortunately does not apply to our setting (since X may have no such segment).
A. MNIST- t 3 , 4 u EXPERIMENTS WITH FINITE TRANSFORMATION GROUPS
We test our proposed method on out-of-distribution tasks on images where the equivalence relations (symmetries) are provided as transformation groups (e.g., 90 ˝ rotations).
We use the MNIST- t 3 , 4 u (colored) dataset (Mouli & Ribeiro, 2021) that only contains digits 3 and 4, and follow their experimental setup.
MNIST- t 3 , 4 u is used to avoid any confounding factors while testing if the proposed method can learn the correct invariances, not for any practical considerations (e.g., rotated 6 is aand would interfere with some experiments, etc.).
We consider equivalence relations obtained from 3 different transformation groups: rotations by 90 ˝ (denoted G rot ), vertically ﬂipping the image (denoted G v-ﬂip ), and permuting the RGB color channels of the image (denoted G col ).
The 3 corresponding equivalence relations are lumpable (Deﬁnition 2) with respect to the transformations in the other two groups in almost all the cases.
Only exception [Table] [Table] is the equivalence relation „ v-ﬂip , which is not lumpable with respect to the transformations in G rot .
Consequently, we do not consider a task with invariance to vertical ﬂip alone.
We test our method on the same 4 classiﬁcation tasks proposed by Mouli & Ribeiro (2021) where each task represents the case where the target Y has different invariances, i.e., invariant to all three groups, to two, to one, invariant to none (the task is sensitive to the remaining groups).
We use the VGG architecture (Simonyan & Zisserman, 2014) for image classiﬁcation and construct a collection of function classes F : “ t F S : S Ď t rot , col , v-ﬂip uu corresponding to various invariant representations.
For example, F t rot , col u is a space of functions (CNNs) that are G-invariant to the rotation and color-permutation groups ( G rot and G col ), and F H is the space of functions with no invariance (standard CNN).
Results.
Our results are shown in Tables 2 to 5 for the four tasks respectively where the label is (i) invariant to all three groups, (ii) invariant to only rotation and vertical ﬂips, (iii) invariant to colorpermutation, and (iv) invariant to none.
We show the values for R p F q , { COMP p F , D q and S p F , D q as as discussed in Section 4.2.
Bold values in the tables indicate the function class chosen by GES method with the proposed scoring criterion (minimizing S p F , D q ).
Test accuracy is computed on the extrapolated dataset after shifting the distribution of P pt U
i u
i P
s D q (i.e., by applying the transformations that the label is invariant to).
In Tables 2 and 3, we see that the proposed method selects the correct model class in training and achieves the best OOD test accuracy.
In Tables 4 and 5, the method is excessively invariant (to vertical ﬂip) but still achieves within 1% of the best OOD test accuracy. The OOD test accuracy of a [Table] [Table]
[Table] [Table] [Table] standard CNN with no invariance ( F H ) is typically very low except in Table 5 where sensitivity to all groups is required.
We can also see the importance of R p F q for tie-breaking in these experiments.
As discussed in Section 4.2, { COMP p F , D q is unable to distinguish between the different function classes because the training data contains a single example per equivalence class (see Figure 2c).
A. CIFAR10 EXPERIMENTS WITH INFINITE / NONGROUP TRANSFORMATION SETS
In this section, we test our proposed method on out-of-distribution tasks on CIFAR10 images (Krizhevsky et al., 2009) where the equivalence relations are provided as inﬁnite sets of transformations that may not form a group.
We used (a) arbitrary rotation transformations over an image (denoted T rot ), and (b) shifting the hue of an image (denoted T col ).
Note that for a bounded image, arbitrary rotation is not a group due to cropping.
Further, transformations from the respective sets commute with each other, and hence, the lumpability condition is satisﬁed (Deﬁnition 2) for the corresponding equivalence relations.
We tested our method on 2 classiﬁcation tasks: (i) invariant to both sets of transformations (arbitrary rotations and hue shifts), and (ii) invariant to arbitrary rotations, but sensitive to hue shifts.
As before, we use the VGG architecture (Simonyan & Zisserman, 2014) for image classiﬁcation and construct a collection of function classes F : “ t F S : S Ď t rot , col uu corresponding to the various invariant representations.
We use data augmentation to construct these invariant representations (this is possible since the lumpability condition holds).
For example, F t rot , col u refers to CNNs that were trained by [Table] [Table] augmenting both arbitrarily rotated images and hue-shifted images.
Once again, F H is the space of functions with no invariance (standard CNN with no data augmentations).
Results.
We show in Tables 6 and 7 that our method is able to ﬁnd the correct invariance and achieves the best OOD test accuracy whereas the standard CNN with no invariance has poor OOD performance.
A. M ORE ON LUMPABILITY (D EFINITION 2)
We show that the lumpability condition of Deﬁnition 2 is equivalent to the normal subgroup condition of Mouli & Ribeiro (2021, Theorem 2) when the given equivalence relations are obtained from transformation groups.
However, unlike the normal subgroup condition, the lumpability condition applies in the general case when the equivalence relations are not necessarily obtained via transformation groups.
Proposition 1.
Let „ G 1 and „ G 2 be two equivalence relations on the input space X obtained as orbits under transformation groups G 1 and G 2 respectively, i.e., for i “ 1 , 2 , x „ G i x 1 iff there exists t p i
q P G i with x 1 “ t p i q
˝ x .
Then, „ G 1 is lumpable with respect to the transformations G(Deﬁnition 2) if and only if G 1 is a normal subgroup of G 1 _ G 2 , where _ is the join operator.
Proof.
First, given „ G 1 is lumpable with respect to G 2 , we wish to prove that G 1 is a normal subgroup of G 1 _ G 2 .
By deﬁnition of the join operator on transformation groups, G 1 is a subgroup of G 1 _ G 2 .
Next, consider an equivalence class r x s G 1 P X { „ G 1 .
Then, by the lumpability of „ G 1 with respect to G 2 , we have that for all t p 2 q P G 2 , there exists r x 1 s G 1 with x ˚ P r x s G 1 ùñ t p 2 q ˝ x ˚ P r x 1 s G 1 .
In other words, each t p 2 q maps all points in one equivalence class r x s G 1 to another equivalence class r x 1 s G 1 .
Speciﬁcally, t p 2 q maps x P r x s G 1 to t p 2 q ˝ x
P r x 1 s G 1 .
Thus, we can set x 1 “ t p 2 q ˝ x without loss of generality.
Then, for all t p 2 q P G 2 , we have from the lumpability condition that
[Equation]
Recall from the deﬁnition of the equivalence class derived from a transformation group (i.e., the orbit) that x ˚ P r x s G 1 means that there exists a transformation t p 1 q P G 1 that maps x to x ˚ , i.e., x ˚ “ t p 1 q ˝ x .
Similarly, t p 2 q ˝ x ˚ P r t p 2 q ˝ x s G 1 means that there exists another transformation ˜ t p 1 q such that t p 2 q ˝ x ˚ “ ˜ t p 1 q ˝ t p 2 q ˝ x .
Equation (9) then becomes
[Equation]
for all t p 2 q P G 2 .
Since Equation (10) holds for all x ˚ P r x s G 1 and for all x P X , we have @
t p 2 q P G 2 , @
t p 1 q P
G 1 , D ˜ t p 1 q P G 1 such that,
[Equation]
which implies that G 1 is a normal subgroup of G 1 _ G 2 .
The converse can be proved trivially by reversing the steps of the above proof.