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: The authors propose an approach for constructing classifiers that achieve out-of-distribution (OOD) generalization using a new learning paradigm they call _asymmetry learning_. They consider OOD tasks where the test input is obtained from the training input by applying a sequence of (random) input transformations. To obtain an invariant OOD classifier that generalizes well to both in-distribution and out-of-distribution samples, the authors introduce the concept of _counterfactual invariant representations over symmetric transformations_. They show how learning the invariant representations can be cast as a causal structure discovery task and propose a score-based (GES-based) algorithm for finding the causal directed acyclic graph (DAG) that best describes the invariances.