Asymmetry Learning for Counterfactually-invariant Classification in OOD Tasks

Generalizing from observed to new related environments ( out-of-distribution ) is central to the reliability of classifiers . However , most classifiers 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 finding a robust OOD classifier can be defined as finding the simplest causal model that defines 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 identifies which symmetries the classifier must break in order to correctly predict Y in both train and test . Asymmetry learning performs a causal model search that , under certain identifiability conditions , finds classifiers 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 classification tasks . 1 INTRODUCTION . A significant challenge in classification 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 , finding a classifier invariant to the distribution shift can still be a data-driven task ( e.g. , classical domain adaptation Ben-David et al . ( 2007 ) ; Muandet et al . ( 2013 ) ; Zhao et al . ( 2019 ) ) . 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 pY tr , X trq , where X tr : “ T tr ˝ X : , with X : as a hidden variable with distribution P pX : 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 pT trq ‰ P pT teq . We are interested in learning an invariant classifier that generalizes well in held out examples from the training and test distributions . The definition of transformation matters in this task . We first 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 classifier is not , in general , solvable via traditional ( interventional ) 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 . ( 2020 ) ) . 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 classifier that performs well OOD is itself a Layer 1 classifier , since it tries to predict P pY te|X teq . Layer 2 : Without observations from P pX teq and/or P pY te|X teq , learning an OOD classifier 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 pY |X , dopT qq . Following the recent literature , we will define the do operator as Y pT “ tq : “ Y |dopT “ tq . Layer 3 : Counterfactuals are the most challenging task . We start our description with an example . Consider a random continuous transformation T tr2 ( in training ) which changes to random transformation T te2 ( in test ) . Let X : describe a hidden variable such that X tr : “ T1 ˝ T tr2 ˝ T3 ˝ X : and X te : “ T1 ˝ T te2 ˝ T3 ˝X : , where T1 and T3 are independent continuous random transformations and P pT tr2 q ‰ P pT te2 q . Assume the target variable Y depends only onX : , T1 , and T3 . To counterfactually ask what would have happened to the observed input x if we had forced dopT tr2 “ t̃2q , we are inquiring about XpT tr2 “ t̃2q|X tr “ x . Note that dopT tr2 “ t̃2q does not change Y . Also note that the knowledge of X tr “ x is an indirect statement about T tr2 since P pT tr2 |X tr “ xq ‰ P pT tr2 q . That is , for x , x1 P X , P pXpT tr2 “ t̃2q “ x1|X tr “ xq “ ż t P pXpT tr2 “ t̃2q “ x1|T tr2 “ t , X tr “ xqdP pT tr2 “ t|X tr “ xq . ( 1 ) Equation ( 1 ) and the difference between the causal hierarchy layers will be relevant for our results . Contributions . Our contributions can be described as follows : 1 . We introduce a generalization of transformation groups via symmetry transformations tied to equivalence classes that removes the requirement of invertible transformations common in definitions using transformation groups . 2 . 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 . 3 . Finally , we introduce asymmetry learning , which describes a representation regularization that , under a set of assumptions , learns the correct counterfactual invariant OOD classifier . 2 SYMMETRIES AND TRANSFORMATIONS . Geometrically , an object is called symmetric if there is a transformation on the object that does not change its shape ( in some definition 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 define a mathematical object as symmetric if there is a transformation on the object that returns another object equivalent to the first ( Rosen , 2008 , Chapter 10 ) . It is clear from this definition of symmetry that we first need to define 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 definition . We define 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 „ cq ùñ a „ c. Equivalence relations allow us to define 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 , defined as ras : “ tx P X : x „ au . Note that one can define m ě 2 equivalence relations on the same input space . The equivalence class of x with respect to equivalence relation k is denoted rxspkq , 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 rasp1q and b R rasp2q . 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 Dc P X , c ‰ a , s.t . c P rasp1q X rasp2q . We denote the collection of equivalence classes of X under the equivalence relation „ k as the quotient space X { „ k : “ trxspkq | x P X u . Transformation group example . Consider the bijective transformations t : X Ñ X of a transformation group G , t P G. We now define an equivalence relation over G as t ˝ x „ G x for all t P G. The equivalence class rxspGq is x ’ s orbit defined as rxspGq : “ tx1 : Dt P G , x1 “ t ˝ xu . For example , if G is the group that permutes the elements of vectors in R3 , then p1 , 2 , 3q „ G p2 , 1 , 3q . Property functions example . Another way of deriving an equivalence relation is via functions of the input space z : X Ñ Rp , where the output zpxq is a particular property of the vector x P X . For example , given an observation of length T from a dynamical system , x P RdˆT , a possible property function could be zenergyp¨q that computes the energy of the dynamical system . Assuming there are m known properties z1 , . . . , zm with zi : X Ñ Rpi , we can construct corresponding equivalence relations „ 1 , . . . , „ m such that for any x , x1 P X , x „ i x1 if zjpxq “ zjpx1q , @ j ‰ i . In words , two inputs are equivalent under „ i if they have the same properties for all zj , j ‰ i. Symmetry transformations . As seen above , symmetries can be defined without defining how the input is transformed to create the equivalence classes , although defining a set of transformations is useful when describing the equivalence class . Given an equivalence relation „ , we can define 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 tid ˝ 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 defined as : for all x , x1 , xp „ 1 _ „ 2qx1 if and only if there exists a chain of equivalence relations x „ k1 x1 , . . . , xh´1 „ kh x1 with all kj P t1 , 2u . It is easy to check that „ 1 _ „ 2 is an equivalence relation . We are now ready to define a general causal model that defines the training and test distributions in our setting .

Summary: This paper considers a class of out of distribution (OOD) problems where at test time there may be new symmetry transformations of the input X (i.e., they don't change the label Y). The authors explain why standard invariances learned by data augmentation may not be OOD invariant. Next, the paper presents a method for learning OOD-invariant representations through causal structure discovery. This hinges on the concept of being counterfactually invariant to the symmetry transformations that could appear at test time but don't affect Y. Next is an algorithm for discovering the structure of a causal DAG which largely revolves around deciding whether or not there is an edge U_i -> Y; this existence question corresponds to whether or not Y is invariant to the transformation U_i (IIUC).  The paper test this approach on tasks in a simple simulated physics environment.