Towards Distribution Shift of Node-Level Prediction on Graphs: An Invariance Perspective

1 INTRODUCTION . As the demand for handling in-the-wild unseen instances draws increasing concerns , out-ofdistribution ( OOD ) generalization ( Mansour et al. , 2009 ; Blanchard et al. , 2011 ; Muandet et al. , 2013 ; Gong et al. , 2016 ) occupies a central role in the ML community . Yet , recent evidence suggests that deep neural networks can be sensitive to distribution shifts , exhibiting unsatisfactory performance within new environments , e.g. , Beery et al . ( 2018 ) ; Su et al . ( 2019 ) ; Recht et al . ( 2019 ) ; Mancini et al . ( 2020 ) . A more concerning example is that a model for COVID-19 detection exploits undesired ‘ shortcuts ’ from data sources ( e.g. , hospitals ) to boost training accuracy ( DeGrave et al. , 2020 ) . Recent studies of the OOD generalization problem like Rojas-Carulla et al . ( 2018 ) ; Bühlmann ( 2018 ) ; Gong et al . ( 2016 ) ; Arjovsky et al . ( 2019 ) treat the cause of distribution shifts between training and testing data as a potential unknown environmental variable e. Assuming that the goal is to predict target label y given associated input x , the environmental variable would impact the underlying data generating distribution p ( x , y|e ) = p ( x|e ) p ( y|x , e ) . With E as the support of environments , f ( · ) as ∗This work was done during authors ’ internship at AWS Shanghai AI Lab . a prediction model and l ( · , · ) as a loss function , the OOD problem could be formally represented as min f max e∈E E ( x , y ) ∼p ( x , y|e=e ) [ l ( f ( x ) , y ) |e ] . ( 1 ) Such a problem is hard to solve since the observations in training data can not cover all the environments in practice . Namely , the actual demand is to generalize a model trained with data from p ( x , y|e = e1 ) to new data from p ( x , y|e = e2 ) . Recent research opens a new possibility via learning domain-invariant models ( Arjovsky et al. , 2019 ) under a cornerstone data-generating assumption : there exists a portion of information in x that is invariant for prediction on y across different environments . Based on this , the key idea is to learn a equipredictive representation model h that gives rise to equal conditional distribution p ( y|h ( x ) , e = e ) for ∀e ∈ E . The implication is that such a representation h ( x ) will bring up equally ( optimal ) performance for a downstream classifier under arbitrary environments . The model p̂ ( y|x ) with such a property is called as invariant model/predictor . Several up-to-date studies develop new objective designs and algorithms for learning invariant models , showing promising power for tackling OOD generalization ( Chang et al. , 2020 ; Ahuja et al. , 2020 ; Krueger et al. , 2021 ; Liu et al. , 2021 ; Creager et al. , 2021 ; Koyama & Yamaguchi , 2021 ) . While the OOD problem is well-established in certain settings ( where the dataset can be obviously modeled as a set of i.i.d . generated pairs ( x , y ) from p ( x , y|e ) ) , its formulation for graph-structured data , especially for node-level tasks on graphs ( where we note that each node in the graph corresponds to an instance ) , is not clear and remains as an open problem . Compared with classic data format ( e.g . vision or texts ) , graph-structured data have two fundamental differences : 1 ) the non-independent and non-identically distributed nature exists in data generation even within the same environment , embodied with the inter-connection among data points in one graph ; 2 ) the structural information also plays a role for prediction beside the node features . These differences bring up unique technical challenges for handling distribution shifts of node-level tasks on graphs . Distribution shifts indeed widely exist in real-world graphs . For instance , in citation networks , the distributions for paper citations ( the input ) and subject areas/topics ( the label ) would go through significant change as time goes by ( Hu et al. , 2020b ) . In social networks , the distributions for users ’ friendships ( the input ) and their activity ( the label ) would highly depend on when/where the networks are collected ( Fakhraei et al. , 2015 ) . In financial networks ( Pareja et al. , 2020 ) , the payment flows between transactions ( the input ) and the appearance of illicit transactions ( the label ) would have strong correlation with some external contextual factors ( like time and market ) . In these cases , neural models built on graph-structured data , particularly , Graph Neural Networks ( GNNs ) which are the common choice , need to effectively deal with OOD data during test time . Moreover , as GNNs have become popular and easy-to-implement tools for modeling relational structures in broad AI areas ( vision , texts , audio , etc . ) , enhancing its robustness to distribution shifts is a pain point for building general AI systems , especially applied to high-stake applications like autonomous driving ( Dai & Gool , 2018 ) , medical diagnosis ( AlBadawy et al. , 2018 ) , criminal justice ( Berk et al. , 2018 ) , etc . In this paper , we endeavor to 1 ) formulate the OOD problem for node-level tasks on graphs , 2 ) develop a new learning approach based on an invariance principle , 3 ) provide theoretical results to dissect its rationale , and 4 ) design comprehensive experiments to show its practical efficacy . Concretely : 1 . To accommodate the non-IID nature of nodes in a graph , we fragment a graph into a set of egographs for centered nodes and decompose the data-generating process into : 1 ) sampling a whole input graph and 2 ) sampling each node ’ s label conditioned the ego-graph . Based on this , we can inherit the spirit of Eq . 1 to formulate the OOD problem for node-level tasks over graphs ( see Section 2.1 ) . 2 . To take into account structural information that is useful for prediction , we first re-formulate the invariant assumption used in prior arts with recursive computation on the induced BFS trees of ego-graphs , inspired by the Weisfeiler-Lehman test ( Weisfeiler & Lehman , 1968 ) . Then , to endow GNNs with enough ability for handling distribution shifts , we devise a new learning approach , entitled Explore-to-Extrapolate Risk Minimization , that aims GNNs at minimizing the mean and variance of risks from multiple environments that are simulated by adversarial context generators ( instantiated as graph editers ) , as shown in Fig . 1 ( a ) ( see Section 2.2 and 3 ) . 3 . To shed more insights on the rationales of the proposed approach and its relationship with the formulated OOD problem , we prove that our objective can guarantee a valid solution for the formulated OOD problem given some mild conditions and furthermore , an upper bound on the OOD error can be effectively controlled when minimizing the training error ( see Section 4 ) . 4 . To evaluate the approach , we design a comprehensive set of experiments on diverse real-world nodelevel prediction datasets that entail distribution shifts from artificial spurious features , cross-domain transfers and dynamic graph evolution . We also apply our approach to distinct GNN backbones ( GCN , GAT , GraphSAGE , GCNII and GPRGNN ) , and the results show that it consistently outperforms standard empirical risk minimization with promising improvements on OOD data ( see Section 5 ) . 2 PROBLEM FORMULATION . In this section , we present our problem formulation for the OOD node-level prediction problem on graphs and then introduce a cornerstone invariance assumption for data generation . All the random variables are denoted as bold letters while the corresponding realizations are denoted as thin letters . 2.1 OUT-OF-DISTRIBUTION PROBLEM FOR GRAPH-STRUCTURED DATA . An input graph G = ( A , X ) contains two-fold information1 : an adjacency matrix A = { avu|v , u ∈ V } and node features X = { xv|v ∈ V } where V denotes node set . Apart from these , each node in the graph has a label , which can be represented as a vector Y = { yv|v ∈ V } . We define G as a random variable of input graphs and Y as a random variable of node label vectors . Such a definition takes a global view and treat the input graph as a whole . Based on this , one can adapt the definition of general OOD problem Eq . 1 via instantiating the input as G and the target as Y , and then the data generation can be characterized as p ( G , Y|e ) = p ( G|e ) p ( Y|G , e ) where e is a random variable of environments that is a latent variable and impacts data distribution . However , the above definition makes little sense in node-level problems where in most cases there is a single input graph that contains a massive number of nodes . To make the problem-solving reasonable , we instead take a local view and investigate each node ’ s ego-graph that has influence on the centered node . Assume v as a random variable of nodes . We define node v ’ s L-hop neighbors as Nv ( where L is an arbitrary integer ) and the nodes in Nv form an ego-graph Gv which consists of a ( local ) node feature matrix Xv = { xu|u ∈ Nv } and a ( local ) adjacency matrix Av = { auw|u , w ∈ Nv } . Use Gv as a random variable of ego-graphs2 whose realization is Gv = ( Av , Xv ) . Besides , we define y as a random variable of node labels . In this way , we can fragment a whole graph as a set of instances { ( Gv , yv ) } v∈V where Gv denotes an input and yv is a target . Notice that the ego-graph can be seen as a Markov blanket for the centered node , so the conditional distribution p ( Y|G , e ) can be decomposed as a product of |V | independent and identical marginal distributions p ( y|Gv , e ) . Therefore , the data generation of { ( Gv , yv ) } v∈V from a distribution p ( G , Y|e ) can be considered as a two-step procedure : 1 ) the entire input graph is generated via G ∼ p ( G|e ) which can then be fragmented into a set of ego-graphs { Gv } v∈V ; 2 ) each node ’ s label is generated via y ∼ p ( y|Gv = Gv , e ) . Then the OOD node-level prediction problem can be formulated as : given training data 1Our formulation and method can be trivially extended to cover edge features which we omit here for brevity . 2We use a subscript v here to remind that it is an ego-graph from the view of a target node . { Gv , yv } v∈V from p ( G , Y|e = e ) , the model needs to handle testing data { Gv , yv } v∈V ′ from a new distribution p ( G , Y|e = e′ ) . Denote E as the support of environments , f as a predictor model with ŷ = f ( Gv ) and l ( · , · ) as a loss function . More formally , the OOD problem can be written as : min f max e∈E EG∼p ( G|e=e ) [ 1 |V | ∑ v∈V Ey∼p ( y|Gv=Gv , e=e ) [ l ( f ( Gv ) , y ) ] ] . ( 2 ) We remark that the first-step sampling G ∼ p ( G|e = e ) can be ignored since in most cases one only has a single input graph in the context of node-level prediction tasks .

Summary: This paper studies the problem of distribution shifts as out-of-distribution generalization. Specifically, it formulates the OOD problem as invariant risk minimization under different environments. The relation between these two has been extensively discussed in the paper. Multiple environment is done by graph editing using policy gradient. Experiments on three different kind of distribution shifts are presented and the effectiveness of EERM is validated.