On Single-environment Extrapolations in Graph Classification and Regression Tasks

1 INTRODUCTION . In some graph classification and regression applications , the graphs themselves are representations of a natural process rather than the true state of the process . Molecular graphs are built from a pairwise atom distance matrix by keeping edges whose distance is below a certain threshold and the choice impacts distinguishability between molecules ( Klicpera et al. , 2020 ) . Functional brain connectomes are derived from time series but researchers must choose a frequency range for the signals , which affects resulting graph structure ( De Domenico et al. , 2016 ) . Recent work ( e.g . Knyazev et al . ( 2019 ) ; Bouritsas et al . ( 2020 ) ; Xu et al . ( 2020 ) ) explore extrapolations in real-world tasks , showcasing a growing interest in the underexplored topic of graph extrapolation tasks . In this work , we refer to graph-processing environment ( or just environment ) as the collection of heuristics and other data curation processes that gave us the observed graph from the true state of the process under consideration . The true state alone defines the target variable . Our work is interested in what we refer as the graph extrapolation task : predict a target variable from a graph regardless of its environment . In this context , even graph sizes can be determined by the environment . Unsurprisingly , graph extrapolation tasks—a type of out-of-distribution prediction—are only feasible when we make assumptions about these environments . We define the graph extrapolation task as a counterfactual inference task that requires learning environment-invariant ( E-invariant ) representations . Unfortunately , graph datasets largely contain a single environment , while common E-invariant representation methods require training data from multiple environments , including Independence of causal mechanism ( ICM ) methods ( Bengio et al. , 2019 ; Besserve et al. , 2018 ; Johansson et al. , 2016 ; Louizos et al. , 2017 ; Raj et al. , 2020 ; Schölkopf , 2019 ; Arjovsky et al. , 2019 ) , Causal Discovery from Change ( CDC ) methods ( Tian & Pearl , 2001 ) , and representation disentanglement methods ( Bengio et al. , 2019 ; Goudet et al. , 2017 ; Locatello et al. , 2019 ) . Contributions . Our work contributes to a growing literature by providing , to the best of our knowledge , the first systematic counterfactual modeling framework for extrapolations in graph classification/regression tasks . Existing work , e.g. , the parallel work of Xu et al . ( 2020 ) , define extrapolations geometrically and , thus , have a different scope . Our work connects Lovász ’ s graph limit theory with graph-size extrapolation in a family of graph classification and regression tasks . Moreover , our experiments show that in these tasks , traditional graph classification/regression methods —including graph neural networks and graph kernels— are unable to extrapolate . 2 A FAMILY OF GRAPH EXTRAPOLATION TASKS . Geometrically , extrapolation can be thought as reasoning beyond a convex hull of a set of training points ( Hastie et al. , 2012 ; Haffner , 2002 ; King & Zeng , 2006 ; Xu et al. , 2020 ) . However , for neural networks—and their arbitrary representation mappings—this geometric interpretation is insufficient to describe a truly broad range of tasks . Rather , extrapolations are better described through counterfactual reasoning ( Neyman , 1923 ; Rubin , 1974 ; Pearl , 2009 ; Schölkopf , 2019 ) . Specifically we want to ask : After seeing training data from environment A , how to extrapolate and predict what would have been the model predictions of a test example from an unknown environment B , had the training data also been from B ? For instance , what would have been the model predictions for a large test example graph if our training data had also been large graphs rather than small ones ? A structural causal model ( SCM ) for graph classification and regression tasks . In many applications , graphs are simply representations of a natural process rather than the true state of the process . In what follows we assume all graphs are simple , meaning all pairs of vertices have at most one edge and no self-loops are present . Our work defines an n-vertex attributed graph as a sample of a random variable Gn : = ( X ( obs ) 1,1 , . . . , X ( obs ) n , n ) , where X ( obs ) i , j ∈ Ω ( e ) encodes edges and edge attributes and X ( obs ) i , i ∈ Ω ( v ) encodes vertex attributes ; we will assume Ω = Ω ( v ) = Ω ( e ) for simplicity . Consider a supervised task over a graph input Gn , n ≥ 2 , and its corresponding output Y . We describe the graph and target generation process with the help of a structural causal model ( SCM ) ( Pearl , 2009 , Definition 7.1.1 ) . We first consider a hidden random variableE with support in Z+ that describes the graph-processing environment ( see Introduction ) . We also consider an independent hidden random variableW ∈ DW that defines the true state of the data , which is independent of the environment variable E , with an appropriately defined space DW . In the SCM , these two variables are inputs to a deterministic graph-generation function g : Z+ ×DW ×DZ → Ωn×n , for some appropriately defined space DZ , that outputs G ( hid ) N ( obs ) : = ( X ( hid ) 1,1 , . . . , X ( hid ) N ( obs ) , N ( obs ) ) = g ( E , W , ZX ) , with N ( obs ) : = η ( E , W ) , ( 1 ) where ZX is another independent random variable that defines external noise ( like measurement noise of a device ) . Equation ( 1 ) gives edge and vertex attributes of the graph G ( hid ) N ( obs ) in some arbitrary canonical form ( Immerman & Lander , 1990 ) , where η is a function of both E and W that gives the number of vertices in the graph . To understand our definitions , consider the following simple example ( divided into two parts ) . Erdős-Rényi example ( part 1 ) : For a single environment e , let n = η ( e ) be the ( fixed ) number of vertices of the graphs in our training data , and p = W be the probability that any two vertices of the graph have an edge . Finally , the variable ZX can be thought as the seed of a random number generator that is drawn n ( n−1 ) 2 times to determine if two distinct vertices are connected by an edge . The above defines our training data as a set of Erdős-Rényi random graphs of size n with p = W . The data generation process in Equation ( 1 ) could leak information about W through the vertex ids ( the order of the vertices ) . Rather than restricting how W acts on ( X ( hid ) 1,1 , . . . , X ( hid ) N ( obs ) , N ( obs ) ) , we remedy this by modeling a random permutation to the vertex indices : G ( obs ) N ( obs ) : = ( X ( obs ) 1,1 , . . . , X ( obs ) N ( obs ) , N ( obs ) ) = ( X ( hid ) π ( 1 ) , π ( 1 ) , . . . , X ( hid ) π ( N ( obs ) ) , π ( N ( obs ) ) ) , ( 2 ) where π ∼ Uniform ( SN ( obs ) ) is an uniform permutation of the indices { 1 , . . . , N ( obs ) } and SN ( obs ) is the permutation group . The observed graph is the outcome of this joint permutation of vertex ids . SCM target variable . We now define our target variable Y . The true target of G ( obs ) N ( obs ) is Y = h ( W , ZY ) , ( 3 ) which is given by a deterministic function h that depends only on W and a random noise ZY independent of W and E. Our final structural causal model is summarized in the directed acyclic graph ( DAG ) of Figure 1 ( a ) . Erdős-Rényi example ( part 2 ) : The targets of the Erdős-Rényi graphs in our previous example can be , for instance , the value Y = W in Equation ( 3 ) , which is also the edge probability p. Graph extrapolation tasks over new environments . Equation ( 3 ) shows that our target variable Y is a function only of the true state W of the data , rather than the graph-processing environment E. Due to the reverse path between Y and E through G ( obs ) N ( obs ) in the DAG of Figure 1 ( a ) , Y is not independent of E given G ( obs ) N ( obs ) . These non-causal paths are called backdoor paths since they flow backwards from Y and G ( obs ) N ( obs ) . Hence , traditional ( interpolation ) methods can pick-up this corre- lation , which would prevent the learnt model from extrapolating over environments different than the ones provided in the training data ( or even over different P ( E ) distributions ) ( Arjovsky et al. , 2019 ; Schölkopf , 2019 ; de Haan et al. , 2019 ) . To address the challenge of predicting Y in spite of backdoor paths , we need a backdoor adjustment ( Pearl , 2009 , Theorem 3.3.2 ) . Instead of explicitly conditioning on the environment for the adjustment , we eliminate the need for conditioning with a graph representation that is invariant to the environment E. Before we proceed , we note that the existing counterfactual notation in the literature ( see Definition 7 of Bareinboim et al . ( 2020 ) ) could be cumbersome in our setting . Hence , we re-propose the powerful concept of random variable coupling from Markov chains ( Pitman , 1976 ; Propp & Wilson , 1996 ) to describe our counterfactual inference problem . The coupling of two independent variables D1 and D2 is a proof technique that creates a random vector ( D † 1 , D † 2 ) , such that Di , D † i have the same marginal distributions , i = 1 , 2 , butD†1 , D † 2 are structurally dependent . For instance , ifD1 , D2 are independent 6-sided die rolls , then D†1 = ( D + 2 ) mod 6 + 1 , D † 2 = ( D + 1 ) mod 6 + 1 are coupled variables corresponding to D1 and D2 , respectively , where D is a 6-sided die roll . Definition 1 ( Counterfactual coupling ( CFC ) ) . A counterfactual coupling of Equations ( 1 ) to ( 3 ) is P ( Y = y , G ( obs ) N ( obs ) = G ( obs ) n ( obs ) , G ( cf ) N ( cf ) = G ( cf ) n ( cf ) ) = EW , ZX , ZY , π , E , Ẽ [ 1 { y = h ( W , ZY ) } · 1 { G ( obs ) n ( obs ) =π ( g ( E , W , ZX ) ) } · 1 { G ( cf ) n ( cf ) =π ( g ( Ẽ , W , ZX ) ) } · 1 { n ( obs ) = η ( E , W ) } · 1 { n ( cf ) = η ( Ẽ , W ) } ] , ( 4 ) where G ( obs ) N ( obs ) : = ( X ( obs ) 1,1 , . . . , X ( obs ) N ( obs ) , N ( obs ) ) and G ( cf ) N ( cf ) : = ( X ( cf ) 1,1 , . . . , X ( cf ) N ( cf ) , N ( cf ) ) , π ( · ) is defined below , and E and Ẽ are independent random variables that sample environments , potentially with different distributions and supports , and 1 is the Dirac delta function . The counterfactual coupled variable G ( cf ) N ( cf ) asks what would have happened to G ( obs ) N ( obs ) if we had used the environment random variable Ẽ in place of E in Equation ( 1 ) . In an abuse of notation we have defined π ( G ( · ) N ) : = ( X ( · ) π ( 1 ) , π ( 1 ) , . . . , X ( · ) π ( N ) , π ( N ) ) above . Using Definition 1 we now prove that a graph representation function Γ ( · ) that is E-invariant is able to predict the targets of the counterfactual graphs . Proposition 1 . Let P ( Y |G ( obs ) N ( obs ) = G ( obs ) n ( obs ) ) and P ( Y |G ( cf ) N ( cf ) = G ( cf ) n ( cf ) ) be the conditional target distributions defined by the counterfactually-coupled random variables in Definition 1 . To simplify exposition , we consider the case where Y ∈ Y is discrete . The continuous case is similar but requires significantly more complex measure theory definitions . Consider a permutation-invariant graph representation Γ : ∪∞n=1Ωn×n → Rd , d ≥ 1 , and a function ρ ( · , · ) ∈ [ 0 , 1 ] ( e.g. , a feedforward network with softmax outputs ) such that , for some , δ > 0 , the interpolation error ( generalization error ) is defined as P ( |P ( Y = y|G ( obs ) N ( obs ) = G ( obs ) n ( obs ) ) − ρ ( y , Γ ( G ( obs ) n ( obs ) ) ) | ≤ ) ≥ 1− δ , ∀y ∈ Y. Γ is said environment-invariant ( E-invariant ) if Γ ( G ( obs ) N ( obs ) ) a.s. = Γ ( G ( cf ) N ( cf ) ) , where a.s. ( almost surely ) means Γ ( G ( obs ) n ( obs ) ) = Γ ( G ( cf ) n ( cf ) ) , except for a set of graphs { G ( obs ) n ( obs ) } and { G ( cf ) n ( cf ) } with zero probability ( measure ) . Then , the extrapolation error is the same as the interpolation error , i.e. , P ( |P ( Y = y|G ( cf ) N ( cf ) = G ( cf ) n ( cf ) ) − ρ ( y , Γ ( G ( cf ) n ( cf ) ) ) | ≤ ) ≥ 1− δ , ∀y ∈ Y . ( 5 ) Proposition 1 shows that an E-invariant representation will perform no worse on the counterfactual test data ( extrapolation samples from ( Y , G ( cf ) N ( cf ) ) ) than on a test dataset having the same environment distribution as the training data ( samples from ( Y , G ( obs ) N ( obs ) ) ) . Other notions of E-invariant representations are possible ( Arjovsky et al. , 2019 ; Schölkopf , 2019 ) , but ours —through coupling— provides a direct relationship with how we learn graph representations from a single training environment . Our task now becomes finding an E-invariant graph representation Γ that can describe well the training data distribution . Specifically , we are interested in single-environment extrapolations . Definition 2 ( Single-environment extrapolation ) . An extrapolation task is a single-environment extrapolation task if the observed training data is generated from a single-environment e ∈ Z+ , while the test data may come from a larger set of environments E ⊆ Z+ , E 6= { e } . In recent years , a crop of interesting research has analyzed the expressiveness of Γ . In what follows we explain why these are related to interpolations rather than extrapolations . A comment on most-expressive graph representations , interpolations , and extrapolations . The expressiveness of a graph classification/regression method is a measure of model family bias ( Morris et al. , 2019 ; Xu et al. , 2018a ; Gärtner et al. , 2003 ; Maron et al. , 2019a ; Murphy et al. , 2019 ) . That is , given enough training data , a neural network from a more expressive family can achieve smaller generalization error ( interpolation error ) than a neural network from a less expressive family , assuming appropriate optimization . However , this power is just a measure of interpolation capability , not extrapolation . Figure 1 ( b ) illustrates a space where each point is a set of neural network parameters from a most-expressive model family . The blue region ( ellipsoid i ) represents models that can perfectly interpolate over the training distribution ( i.e. , models with the smallest generalization error ) . The models in the blue region are mostly fitting spurious training environment E correlations with Y , that will cause poor extrapolations in new environments . The models illustrated in the red region of Figure 1 ( b ) ( ellipsoid ii ) are E-invariant and , thus , by Proposition 1 , can extrapolate across environments , since they can not fit these spurious environment correlations . The intersection between the blue and red regions contains models that are optimal both for test data from the same environment distribution as training ( interpolation test ) and test data from a different environment distribution ( extrapolation test ) . In our SCM in Equations ( 1 ) to ( 3 ) , the intersection between the blue and red ellipsoids is nonempty . We can denote the models in the red ellipsoid as the most-expressive family of E-invariant ( Proposition 1 ) . Our work focuses on a family of classifiers and regression models that reside inside the red ellipsoid . Summary . In this section we have defined a family of extrapolation tasks for graph classification and regression using counterfactual modelling , and connected it to the existing literature . Next , we show how these definitions can be applied to a family of random graph models ( graphons ) first introduced by Diaconis & Freedman ( 1981 ) .

Summary: This paper formulates a theoretical model for testing extrapolation abilities of graph learning tasks, and suggests some practical feature maps to achieve good extrapolation properties empirically. In more detail, this paper introduces a model, a so-called structural causal model, for graphs where the graph creation process is modeled as a random variable that depends on different independent factors:  environment $E$, which is used to model the graph size $n$; graph property $W$, which is used for example to model the probability of edge existence; and random seed $Z_X$. The created graph is also scrambled by a random permutation to yield the observed graph $G^{\text{obs}}$. The ground truth labels of the graph $Y$ are functions of two factors $W$, the graph property, and $Z_Y$ random seed.