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: The paper adapts a recently proposed invariant risk minimization approach for tackling distribution shift on node-level predictions on graphs. The main idea is to decompose the graph into a set of ego-graphs rooted at each node, thus incorporating structural information. Since the learning objective requires data from different environments they authors introduce auxiliary context generators trained to maximize the variance loss.