OOD Link Prediction Generalization Capabilities of Message-Passing GNNs in Larger Test Graphs
Anonymous Author(s)
AfﬁliationAddress email
Abstract
This work provides the ﬁrst theoretical study on the ability of graph Message Passing Neural Networks (gMPNNs) —such as Graph Neural Networks (GNNs)— toachieve counterfactually-invariant representations for inductive out-of-distribution(OOD) link prediction tasks, where deployment (test) graph sizes are larger thantraining graphs.
We ﬁrst prove non-asymptotic bounds showing that link predictors based on permutation-equivariant (structural) node embeddings obtainedby gMPNNs can converge to a random guess as test graphs get larger.
We thenpropose a theoretically-sound gMPNN that outputs structural pairwise (2-node)embeddings and prove non-asymptotic bounds showing that, as test graphs grow,these embeddings converge to embeddings of a continuous function that retains itsability to predict links OOD.
Empirical results on random graphs show agreementwith our theoretical results.
Introduction
Link prediction is the task of predicting whether two nodes likely have a missing link [1, 9, 20, 25, 44].
Link prediction tasks arise in many settings, ranging from predicting edges on bipartite graphs betweenusers and products or content in recommender systems [4, 8, 21, 22, 26, 40], to knowledge graphreconstruction [3, 11, 14, 34, 44, 45], to predicting protein-protein interactions [35].
In recent years, there has been growing interest in applying neural network models to inductive linkprediction tasks.
Inductive link prediction considers methods trained on a graph G tr and deployedat test time on another graph G te .
It also encompasses the task of training the method on a smallerinduced subgraph G tr of a larger graph G te , then deploying it on the entire graph.
In particular, ourwork focuses on graph message-passing Neural Networks (gMPNNs) [15, 38] or, more precisely, thewidely used Graph Neural Network (GNN) framework [6, 7, 10, 16, 17, 19, 39, 42, 47].
Our work asks the following questions: Are link prediction methods able to produce approximatecounterfactually-invariant representations that can cope with the task of inductive out-of-distribution(OOD) link prediction, where (unseen) test graphs are signiﬁcantly larger than training graphs?
Counterfactually-invariant representations are representations that are invariant to counterfactualchanges in spurious variables [33, 46].
How can these OOD link prediction tasks be theoreticallydeﬁned?
Can we obtain non-asymptotic bounds on the generalization capabilities of these methods?
The majority of today’s link prediction methods are based on a similar principle.
Consider anattributed graph G =
( V, E ) , with node set V = { 1 , ..., N } , edge set E ✓ V ⇥ V , and node features
F 2 R N ⇥ F 0 , F 0  1 .
Then, given a pair of nodes i, j 2 V , after T  1 iterations over G , thesemethods produce associated node embeddings (representation vectors) ⇥
• i , ⇥ • j 2 R F T , F T  1 ,which are then used in a link function ⌘ • : R F T ⇥ R F T !
[0 , 1] such that ⌘ • ( ⇥ • i , ⇥ • j ) predicts theprobability that i and j have a missing link in G .
In our notation we will denote all node embeddings and associated functions with the superscript “ • ”.
Henceforth we denote gMPNNs that outputstructural node embeddings as gMPNNs • .
Node embeddings.
The ﬁrst part of our work considers a subset of these methods, where the outputnode embeddings are permutation equivariant (a.k.a.
structural node embeddings [43]).
Informally,a sequence of node embeddings ⇥ • 2 R N ⇥ F T given by an embedding method is permutationequivariant if for any arbitrary graph G and any permutation ⇡ 2 S N of the node indices, where S Nis the symmetric group, the resulting isomorphic graph G 0 = ( ⇡  V, ⇡  E, ⇡  F ) gets permutednode embeddings ⇥ •
= ⇡  ⇥ • , where ⇡  M deﬁnes the action of ⇡ on M (we will provide a formaldeﬁnition in Section 2).
We leave the study of OOD link prediction with positional node embeddings(a.k.a.
permutation-sensitive node embeddings [43]) to future work.
The application of GNNs to link prediction tasks is made difﬁcult by the fact that, by construction,permutation-equivariant GNNs give the same embeddings ⇥ • i , ⇥ • j to any isomorphic nodes i, jin G , as noted by You et al.
[52] and Srinivasan and Ribeiro [43].
Isomorphic nodes are nodesthat are structurally indistinguishable in G (even when considering node features) except by their(assumed arbitrary) node indices i, j 2 V .
That is, if a graph has isomorphic pairs, permutationequivariant GNN link prediction can fail. The recent link prediction literature has signiﬁcantly reliedon isomorphic nodes for theoretical results (e.g., Zhang et al.
[54, Theorem 2] uses isomorphic nodesto prove that, uniformly, graphs are likely to have many isomorphic nodes and hence are not amenableto accurate link prediction). However, isomorphic nodes are rare in both real-world graphs (see Figure 3 in the Appendix) and in large random graphs (Proposition 1). An important open question is whether equivariant GNN would be able to predict links in asymmetricgraphs.
That is, the concerns of [43, 52] may not be of practical importance. Our work also answersthis question: We see that in-distribution (where graph test sizes are the same as training sizes),permutation-equivariant GNNs are able to predict links by tapping into the graph asymmetries. However, we show theoretically and empirically that tapping into asymmetries can fail OOD evenwhen it works in-distribution.
Pairwise embeddings.
Taking a different route, Srinivasan and Ribeiro [43] provides an existenceproof that the link prediction task between i and j can always be performed by a pairwise embedding⇥ •• ij ( G ) , i.e., for any pair of nodes i, j in a graph G , there exists a pairwise embedding ⇥ •• ij ( G )and a link function ⌘ •• : R F T !
[0 , 1] such that ⌘ •• ( ⇥ •• ij ) approximates the probability that i andj have a hidden link.
In our notation we will denote all pairwise (joint 2-node) embeddings andassociated functions with the superscript “ •• ”.
Unfortunately, as the test graph grows, we wereunable to prove existing pairwise embedding methods [31, 49, 53–55] are able to perform OOD linkprediction tasks.
Hence, we propose a novel family of gMPNNs for pairwise embeddings, denotedgMPNNs •• henceforth.
The second part of our work considers the OOD generalization capability ofthese gMPNNs •• .
Contributions.
In this work we study inductive OOD link prediction tasks for larger test graphsusing permutation-equivariant node and pairwise embeddings, ⇥ • and ⇥ •• , respectively.
Our workmakes the following contributions:
We provide a theoretical framework deﬁning OOD inductive link prediction tasks, where testgraphs are signiﬁcantly larger than training graphs.
We show that structural node embeddings from message-passing GNNs can fail in OOD linkprediction tasks if the test graph (from the same graph family) is signiﬁcantly larger than thetraining graph.
Our work ﬁlls an important gap in the literature , where Bevilacqua et al.
[5]studied the OOD capabilities of GNNs for graph classiﬁcation using random graph models.
Ourwork studies the OOD capabilities of GNNs for inductive link prediction in a similar setting.
We propose a new family of structural pairwise embeddings, denoted gMPNNs •• , that canprovably perform the above OOD task
We provide non asymptotic
bounds on the convergence of node and pairwise gMPNNs embeddings to an easy-to-analyze continuous model.
Extensive empirical experiments using stochastic blockmodels (SBMs [41]) validate our theoretical results.
Our work focuses on providing a theoreticalunderstanding of the challenges of OOD link prediction tasks rather than propose real-world linkprediction tasks and compare baselines.
However, we believe that our work lays the theoreticalfoundation (and challenges) for future application-focused works.
Preliminaries
Given an attributed graph G =
( V, E ) , with node set V = { 1 , ..., N } , edge set E ✓ V ⇥ V , adjacencymatrix A 2 { 0 , 1 } N ⇥ N , where A ij = 1 { ( i,j ) 2 E } , and node features F 2 R N ⇥ F 0 , F 0 > 0 .
Let 92 P ⇡ 2 B N be a permutation matrix associated with permutation ⇡ 2 S N (where S N is the symmetric 93 group), where
B N denotes the Birkhoff polytope of N ⇥ N doubly-stochastic matrices.
Doublysctochastic matrices are non-negative square matrices whose rows and columns sum to one.
Thematrix P ⇡ deﬁnes the action of permutation ⇡ on these matrices, e.g., ⇡  A = P ⇡ AP T ⇡ .
We denotea pair of nodes i, j 2 V as isomorphic in G if exists ⇡ 2 S N such that ⇡ i = j , A = P ⇡ AP T ⇡ , and F = P ⇡ F .
Node features can be deﬁned by the graph signal f : V !
R F 0 as a function that maps anode to an F 0 -dimensional feature in R F 0 .
Then the signal of the graph F can be represented by amatrix F =
[ f 1 , ..., f N ] T 2 R N ⇥ F 0 , where f i 2 R F 0 are the features of node i 2 V .
Random graph model for G .
Let ( X , d, µ ) be a metric-measure space, where X is a set, d is a metric,and µ is a probability Borel measure.
A graphon is deﬁned as a mapping W : X ⇥ X !
In what follows we deﬁne how the graph G is sampled from the graphon models.
The signal deﬁnitionfollows Maskey et al.
[29] and the edge samples follow Airoldi et al.
[2], Lawrence and Hyvärinen [23].
[Figure]
- ables are shaded white and gray, respectively.
Deﬁnition 1 (Random graph model) .
A random graph model for
G on ( X , d, µ ) is deﬁned as a pair ( W, f ) of a graphon W : X ⇥
X !
[0 , 1] and a metric-space signal f : X !
R F 0 , where f 2 108 L 1 ( X ) is an essentially bounded measurable function with theessential supremum norm.
We obtain ( G, F ) by ﬁrst sampling Ni.i.d.
random points
X 1 , ..., X N from X with probability densityµ , as the nodes of G .
Then the edge ( i, j ) between nodes i and j issampled with probability W
( X i , X j ) , i.e., the adjacency matrix
A =
( A i,j ) i,j of G is deﬁned as A i,j = ( Z i,j < W
( X i , W j ))for i, j = 1 , ..., N , where { Z i,j } Ni,j =1 are sampled i.i.d.
from
Uniform (0 , 1) .
The graph signal F =
[ f 1 , ..., f N ] T 2 R N ⇥ F 0is deﬁned as f i = f
( X i ) .
We say ( G, F ) ⇠ ( W, f ) .
Further,we restrict our attention to graphons W such that there exists aconstant d min satisfying the graphon degree
d W ( x ) :=
R X
W ( x, y ) dµ ( y )  d min > 0 , 8 x 2 X .
In an abuse of notation we identify node i 2 V with the sampled value X i ⇠ µ, 8 i 2 { 1 , ..., N } , 120 since generally µ is such that P ( X i = X j ) =
almost everywhere for i 6 =
j (e.g., µ is uniform).
The causal DAG of the data generation process of G is given in Figure 1.
Our goal is to produceclassiﬁers that survive the distribution shift implied by a change in the distribution of graph sizes
N during test time.
Furthermore, we note that all proofs are relegated to the Appendix due to spaceconstraints.
Henceforth use the terms node embeddings and node representations interchangeably.
Graph messagepassing Neural Network (gMPNN • ) is deﬁned by realizing a message-passing Neural Network(MPNN) on a graph.
The degree-normalized aggregator deﬁnitions follow Maskey et al.
[29] whilethe sum aggregator deﬁnition is new.
Deﬁnition 2 (MPNN) .
Let T 2 N denote the number of layers.
For t = 1 , ..., T , let  ( t ) :
R 2 F t  1 !
R H t  1 and ( t ) :
R F t  1 + H t  1 !
R F t be functions, where F t 2 N is called the feature dimension oflayer t .
The corresponding MPNN ⇥ is deﬁne by the sequence of message functions (  ( t ) )
Tt =1 andupdate functions ( ( t ) )
Tt =1 , i.e. ⇥ =
((  ( t ) )
Tt =1 , ( ( t ) )
Tt =1 ) .
The message and update functions in Deﬁnition 2 are usually given by MLPs.
We now introducethe gMPNN • with T message-passing layers.
For each node i 2 V , f • ( t ) i at layer t 2 { 1 , . . .
, T } isdeﬁned recursively using (a) its own representation at layer t  1 ( f • ( t  1) i ) and (b) an aggregatedrepresentation of its neighbors m ( t ) i (commonly) deﬁned via two distinct permutation-invariantaggregation functions as follows:
Deﬁnition 3 (gMPNN • ) .
Let ( G, F ) be a graph with graph signals as in Deﬁnition 1 and ⇥ be a
MPNN as in Deﬁnition 2.
For layer t = 1 , ..., T , deﬁne ⇥ • ( t ) A and ⇥ • ( t ) A as maps from the input graph
G and graph signals F (0) =
F 2 R N ⇥ F 0 to the features in the t -th neural layer by
[Equation]
where F ( t ) is deﬁned by either the neighbor-average aggregation procedure (a),
8 i 2 V , for ⇥ • ( t ) Am ( t ) i :
= 1 N
P
{+Nj =
 A i,j d i  ( t ) ( f • ( t  1) i , f • ( t  1) j ) , f • ( t ) i := ( t ) ( f • ( t  1) i , m ( t ) i ) ,+} where the N -normalizeddegree of node i is d i = 1 N
P
Nj =
A i,j , or the ( N -normalized) sum aggregation procedure (b),8 i 2 V , for ⇥ • ( t ) A , m ( t ) i :
= 1 N
P
Nj {+=
+}
A i,
[Equation]
( [Equation] ) ( f • (
[Equation]
[Equation] ) i , f • ( t  1) j ) , and f • ( t ) i := ( t ) ( f • ( t  1) i , m ( t ) i ) .
Given a gMPNN • , ⇥ • ( T ) A or ⇥ • ( T ) A , with T  1 layers as in Deﬁnition 3, their outputs are the graphsignals ⇥ • ( T ) A ( F ) 2 R N ⇥ F T for the neighbor-average aggregation and ⇥ • ( T ) A ( F ) 2 R N ⇥ F T for the( N -normalized) sum aggregation, and are henceforth denoted as node embedding outputs of thegMPNN • .
We denote ⇥ • ( T )
A ( F ) i as the node embedding for node
i 2 V .
2.2 Node embeddings with continuous message passing neural networks
Here we deﬁne continuous message passing neural networks for structural node embeddings.
Theintegral pooling deﬁnition is new while the degree-averaged one follows Maskey et al.
Deﬁnition 4 (Continuous message-passing) .
Given a MPNN ⇥ as in Deﬁnition 2, the node continuousmessage passing neural network (cMPNN • ) on graphons and metric-space signals f : X !
R F 0can be deﬁned by replacing the graph node features and the aggregation scheme in Deﬁnition 3 bythe following continuous counterparts.
Using a message signal U : X ⇥ X !
R H , the continuousmean aggregation of U is deﬁned by M • W ( U )( x ) =
R X
W ( x,y )
d W ( x ) U ( x, y ) dµ ( y ) , while the integralaggregation is deﬁned as M • W ( U )( x ) =
R X
W ( x, y ) U ( x, y ) dµ ( y ) , where W is a graphon andd W ( x ) =
R X
W ( x, y ) dµ ( y ) is the graphon degree as in Deﬁnition 1.
By replacing the discrete aggregations in Deﬁnition 3 by the above continuous aggregations, the same
MPNN ⇥ can process metric-space signals instead of graph signals in the following deﬁnition.
Deﬁnition 5 (cMPNN • ) .
Let ( W, f ) be a random graph model as in Deﬁnition 1 and ⇥ be a MPNN as in Deﬁnition 2.
For t = 1 , ..., T , deﬁne ⇥ • ( t ) W and ⇥ • ( t ) W as maps from input metric-space signal f • (0) =
f
: X !
R F 0 to the features in the t -th layer by
[Equation]
where f • ( t ) are deﬁned sequentially either through the degree-average aggregation (a) for ⇥ • ( t ) W :g • ( t ) ( x ) :
= M • W (  ( t ) ( f • ( t  1) , f • ( t  1) ))
( x ) = R X
W ( x,y )
d W ( x )  ( t ) ( f • ( t  1) ( x ) , f • ( t  1) ( y )) dµ ( y ) ,f • ( t ) ( x ) :
= ( t ) ( f • ( t  1) ( x ) , g • ( t ) ( x )) , or through the integral aggregation (b) for ⇥ • ( t ) W :g • ( t ) ( x ) :
= M • W (  ( t ) ( f • ( t  1) , f • ( t  1) ))
( x ) =
R X
W ( x, y )  ( t ) ( f • ( t  1) ( x ) , f • ( t  1) ( y )) dµ ( y ) ,with f • ( t ) ( x ) :
= ( t ) ( f • ( t  1) ( x ) , g • ( t ) ( x )) .
Size-stability of node representation and its drawbacks
We now present our main theorems about convergence of gMPNN • to cMPNN • for test graphs G tesampled from the graphon random graph model (see Deﬁnition 1), and how it leads to size-stability ofgMPNN • for nodes that have the same representation under cMPNNs • .
In what follows we will focuson the neighbor-average aggregation procedure (a) of Deﬁnition 3, since this is the more difﬁcultcase to prove.
Similar results for the ( N -normalized) sum aggregation procedure (Deﬁnition 3(b))are shown in the Appendix due to space constraints.
Moreover, common deﬁnitions (e.g., Lipschitzcontinuous functions) are also relegated to the Appendix to save space.
3.1 Convergence of gMPNNs towards cMPNNs as test graph size increase
We now prove that, with high probability, the maximum inﬁnity difference between the gMPNN • andcMPNN • node representations
Theorem 1.
Let ⇥ =
((  ( l ) )
Tl =1 , ( ( l ) )
Tl =1 ) be a MPNN as in Deﬁnition 2 with T layers such that ( l ) : R 2 F l  1 !
R
H
l  1 and ( l ) :
R F l  1 + H l  1 !
R F l are Lipschitz continuous with Lipschitzconstants L ( l )  and L ( l ) .
Let gMPNN • ⇥ • ( T ) A and cMPNN • ⇥ • ( T ) W be as in Deﬁnitions 3 and 5.
For arandom graph model ( W, f ) satisfying Deﬁnition 1, consider training ( G tr , F tr ) ⇠ ( W, f ) and test( G te , F te )
⇠ ( W, f ) graphs, where N tr and N te are their respective number of nodes, N te > N tr .
Let X te 1 , ..., X te N te and A te be as in Deﬁnition 1.
Let p 2 (0 ,
P
[Equation]
[Equation]
l
[Equation]
Then, if
[Equation]
we have with probability at least 1
[Equation]
where the constants C 1 and C 2 are deﬁned in the Appendix and depend on { L ( l )  , L ( l ) } Tl =1 and( G tr , F tr ) , N tr .
Theorem 1 above shows that as the test graph size N te grows, the node representations from thediscrete gMPNNs • converge to the continuous cMPNNs • .
Our OOD statement 1 in Theorem 1 hasprofound consequences when it comes to predicting links using the node representations obtainedby a gMPNN • .
Next, Corollary 1 shows that for any two nodes i, j 2 V te that are indistinguishablein the degree-aggregated cMPNN • (deﬁned as ⇥ • ( T ) W ( f )( X te i ) = ⇥
• ( T ) W ( f )( X te j ) ), they will getincreasingly similar representations in the discrete gMPNN • as N te grows.
Corollary 1.
Let ⇥ =
((  (
[Equation]
Tl =1 , ( ( l ) )
Tl =1 ) , ⇥ • ( T ) A , ⇥ • ( T ) W , p, ( W, f ) , ( G tr , F tr ) , ( G te , F te ) , N tr ,
N te , A te , and X te 1 , ..., X te N te be as in Theorem 1.
If there exists i, j 2 V te , i 6 = j, s.t. ⇥ • ( T ) W ( X i ) =⇥
• ( T ) W ( X j ) and Equation (1) is satisﬁed, then, with C 1 and C 2 as in Theorem 1, we have that withprobability at least 1
[Equation]
[Equation]
Implications of Corollary 1 on Stochastic Block Models (SBMs).
In what follows, we willdiscuss circumstances where two nodes i, j 2 V get the same cMPNN • representations (i.e., both⇥ • ( T ) W ( f )( X i ) = ⇥
• ( T ) W ( f )( X j ) and ⇥ • ( T ) W ( f )
( X i ) = ⇥
• ( T ) W ( f )( X j ) ).
In what follows we restrictour results to an important family of graphon models: Stochastic Block Models (SBMs) [41], wherewe also model node attributes.
SBMs were chosen because they can consistently model large graphsgenerated by any piecewise Lipschitz graphon model [2].
SBMs are also intuitive models, whichmakes them useful to illustrate our results.
Deﬁnition 6 (Stochastic Block Model (SBM)) .
An SBM ( W, f ) is a random graph model (Definition 1) with cluster structures in W and f .
Partition the node set into r  2 disjoint subsets
S 1 , S 2 , ..., S r ✓ V (known as blocks or communities) with an associated r ⇥ r symmetric matrix S , where the probability of an edge ( i, j ) , i 2 S a and j 2 S b is S ab , for a, b 2 { 1 , . . .
, r } .
Let X =
[0 , 1] , and µ be the uniform distribution on [0 , 1] .
By dividing X =
[0 , 1] into disjoint convex sets [ t 0 , t 1 ) , [ t 1 , t 2 ) , . . .
, [ t r  1 , t r ] , where t 0 = 0 and t r = 1 , node i belongs toblock S a if X i ⇠ Uniform (0 , 1) satisﬁes
X i 2 [ t a  1 , t a ) .
The graphon function W is deﬁned 210 as W ( X i , X j ) =
P a,b 2 { 1 ,...,r } S ab 1 ( X i 2 [ t a  1 , t a )) 1 ( X j 2 [ t b  1 , t b )) .
We take the liberty toalso deﬁne node signals in our SBM model, where for B =
[ B 1 , ..., B r ] T 2 R r ⇥ F
the metric-spacesignal f : X !
R F 0 is deﬁned as f ( x ) =
P a 2 { 1 ,...,r } 1 ( x 2 [ t a  1 , t a ))
B a . We deﬁne the action of permutation ⇡ on B of Deﬁnition 6 as ⇡  B , where ( ⇡  B ) ⇡ a = B a
Deﬁnition 7 (Isomorphic SBM blocks) .
For the SBM model ( W, f ) in Deﬁnition 6, we say two blocksa, b 2 { 1 , . . .
, r } are isomorphic if the SBM satisﬁes the following two conditions: (a) t a  t a  1 = 216 t b  t b  1 , and (b) for ⇡ 2 S r , such that ⇡ a = b , ⇡ b = a and ⇡ c
= c, 8 c 2 { 1 , ..., r } , c 6 = a, b , 217 S =
⇡  S , and B = ⇡  B .
A similar deﬁnition can be obtained for the general graphons in Deﬁnition 1 using the isomorphicgraphon deﬁnition of Lovász and Szegedy [27].
Now that we have the deﬁnition for isomorphic blocks in SBM models, we can prove that all nodesin these isomorphic blocks will obtain the same representations under degree-average aggregationand integral aggregation cMPNNs • .
Lemma 1.
Let ⇥ =
((  ( l ) )
Tl =1 , ( ( l ) )
Tl =1 ) be a MPNN as in Deﬁnition 2, and ⇥ • ( T ) W , ⇥ • ( T ) W as in
Deﬁnition 5.
For the SBM model ( W, f ) in Deﬁnition 6 with N te nodes X 1 , . . .
, X N te .
If there existsi, j 2 V te such that X
te
i , X te j are nodes that belong to isomorphic SBM blocks (Deﬁnition 7), then⇥ • ( T ) W ( f )
( X te i ) = ⇥
• ( T ) W ( f )( X te j ) and ⇥ • ( T ) W ( f )
( X te i ) = ⇥
• ( T ) W ( f )( X te j ) .
Note that even though any two nodes in isomorphic SBM blocks get the same cMPNN • representations per Lemma 1, these nodes are likely not isomorphic in G te (as shown in Proposition 1 in
Appendix) and, hence, they get different gMPNN • representations.
However, Corollary 1 shows thatthese representations become increasingly similar as the test graph size grows ( N te  1) .
We usethis observation to understand the ability of gMPNNs • to perform link prediction tasks next.
3.2 The hardness of OOD inductive link prediction using structural node embeddings
The convergence of gMPNNs • to cMPNNs • as the test graph size N te grows (Theorem 1) impliesthrough Corollary 1 and Lemma 1 that node representations of distinct SBM blocks can becomeincreasingly similar as the test graph size grows ( N te  1) , even though these nodes are notisomorphic in G te with high probability (see Proposition 1 in the Appendix).
Deﬁnition 8 (Link prediction function from structural node embeddings) .
An inductive link predictionfunction ⌘ • : R F T ⇥ R F T !
[0 , 1] takes the gMPNN • node representations of two nodes i, j 2 V teand predicts the edge probability P ( A te ij = 1) .
We assume ⌘ • is Lipschitz continuous with Lipschitzconstant L ⌘ .
In the context of graphon random graph models (Deﬁnition 1), we aim to learn⌘ • ( ⇥ • ( T )
A te ( F te ) i , ⇥ • ( T )
A te ( F te ) j ) ⇡ W
( i, j ) .
We further assume we predict a link if ⌘ • ( · , · ) > ⌧ , whileno link if ⌘ • ( · , · ) < ⌧ , for some (arbitrary) threshold ⌧ 2 [0 , 1] chosen by the user of such system.
The next corollary showcases the difﬁculty in OOD predicting links using structural node representations as N te grows.
Corollary 2.
Let ⇥ =
((  ( l ) )
Tl =1 , ( ( l ) )
Tl =1 ) be the MPNN with T layers and ⇥ • ( T ) A , ⇥ • ( T ) W asin Theorem 1.
Let ⌘ • : R F T ⇥ R F T !
[0 , 1] be as in Deﬁnition 8.
Consider the SBM ( W, f )in Deﬁnition 6 with isomorphic blocks (Deﬁnition 7).
Let ( G tr , F tr ) ⇠ ( W, f ) and ( G te , F te )
⇠( W, f ) be the training and test graphs with N tr and N te nodes, respectively.
Consider any two testnodes i, j 2 { 1 , ..., N te } , i 6 = j , for which we can make a link prediction decision with ⌘ • (i.e.,⌘ • ( ⇥ • ( T )
A te ( F te ) i , ⇥ • ( T ) A ( F te ) j ) 6 = ⌧ ).
Let G te be large enough to satisfy both Equation (1) and
[Equation]
where p , C 1 , and C 2 are as given in Corollary 1.
Then, if i and j belong to isomorphic blocks (i.e.,⇥ • ( T ) W ( f )( X te i )
= ⇥
• ( T ) W ( f )( X te j ) ), with probability at least 1
P
Tl = [Equation] 2( H l +1) p the link predictionmethod in Deﬁnition 8 will make the same link prediction regardless of the SBM probability matrix
S (Deﬁnition 6) and whether i and j are in the same block or distinct isomorphic blocks .
Corollary 2 proves that link prediction with structural node embeddings form gMPNNs • is unreliable.
That is, for any link prediction method satisfying Deﬁnition 8, as the test graph grows N te  1 , themethod will increasingly struggle to give different predictions within and across isomorphic SBMblocks, even when these probabilities are arbitrarily different in the underlying graph model.
In whatfollows we show that pairwise embeddings can address this challenge.
Size-stability of structural pairwise embeddings and its advantages
We have discussed the limitation of gMPNNs • on node representation for link prediction.
Now weclaim that a joint continuous message passing graph neural network is capable of link predictionin graphon random graph models (Deﬁnition 1).
We deﬁne the joint continuous message passinggraph neural network inspired by the cMPNNs • for node representations (Deﬁnition 5).
First, weneed to deﬁne the graphon fraction of common neighbors for graphon nodes x and y , c W ( x, y ) :=
R X
W ( x, z ) W ( y, z ) dµ ( z ) .
We only consider graphons W such that there exists d cmin satisfyingc W ( x, y )  d cmin > 0 , 8 x, y 2 X in this section.
Since we do not have edge feature as in Deﬁnition 1,we deﬁne the metric-space pair-wise signal as f •• ( x, y )
= 1 , 8 x, y 2 X .
Deﬁnition 9 (cMPNN •• ) .
Let ( W, f ) be a random graph model as in Deﬁnition 1 and ⇥ be a MPNNas in Deﬁnition 2.
For t = 1 , ..., T , deﬁne the continuous (pairwise) cMPNN •• ⇥ •• ( t ) W as the mappingthat maps input pairwise metric-space signals
f •• (0) =
f •• to the features in the t -th layer by
[Equation]
where f •• ( t ) are deﬁned recursively by
[Equation]
1) ( x, z )) + W ( x, z ) c W ( x, y )  ( t ) ( f •• ( t  1) ( x, y ) , f •• ( t  1) ( y, z ))) dµ ( z ) , f •• ( t ) ( x, y ) := ( t ) ( f •• ( t  1) ( x, y ) , g •• ( t ) ( x, y )) .
The intuition of the aggregation function is that two edges with one same node is considered neighborsin a higher-order graph [32], and to go from ( x, y ) to ( x, z ) , we need to transition from y to z , whichhas probability W ( y, z ) .
The same holds for going from ( x, y ) to ( y, z ) .
Lemma 2.
If  ( x, y ) =
y and ( x, y ) =
x/y , then f •• ( t ) ( x, y )
=
W ( x, y ) , 8 x, y 2 X is a stationary point in the cMPNN •• , i.e. if f •• ( t  1) ( x, y )
=
W ( x, y ) , then f •• ( t ) ( x, y ) =
W ( x, y ) , 8 x, y
X .
We deﬁne the corresponding gMPNN •• as follows.
First we deﬁne the fraction of common neighborsbetween nodes i and j as c Ai,j = 1 N
P
Nz = 1 A i,z · A j,z .
If two nodes do not have common neighbors,then we set c Ai,j = 1 N to avoid computation error.
Further, we deﬁne f •• i,j = 1 8 i, j 2 V for anygraph G , and F •• = ( f •• i,j ) i,j 2 V as the pair-wise graph signals.
Deﬁnition 10 (gMPNN •• ) .
Let ( G, F ) be a graph with graph signals as in Deﬁnition 1 and ⇥ be a
MPNN as in Deﬁnition 2.
For t = 1 , ..., T layers we deﬁne the gMPNN •• ⇥ •• ( t ) A as the mapping that maps input pairwise graph signals F •• (0) =
F •• to the features in the t  th layer by
[Equation]
where f •• ( t ) are deﬁned recursively by the following function,
[Equation]
Next, Theorem 2 that under the causal DAG of Figure 1, cMPNN •• is approximately counterfactuallyinvariant [33] to larger graph sizes.
Theorem 2.
The gMPNN •• representation ⇥ •• ( T ) A in Deﬁnition 10 is approximately counterfactuallyinvariant to larger graph sizes in the random graph model ( W, f ) as in Deﬁnition 1 with d cmin > 0 .
More precisely, let ( G, F ) ⇠ ( W, f ) be a graph of size n from the training distribution and deﬁne thetest distribution ( G te , F te ) through the causal graph in Figure 1 as an interventional change to obtainlarger test graph sizes N te > N (which means any test graph is larger than the largest possible training graph).
Let ⇥ =
((  ( l ) )
Tl =1 , ( ( l ) )
Tl =1 ) be a MPNN as in Deﬁnition 2 with T layers such that ( l ) and ( l ) are Lipschitz continuous with Lipschitz constants L ( l )  and L ( l ) .
Let cMPNN •• ⇥ ••
( T ) Wbe as in Deﬁnition 9.
Let X te 1 , ..., X te N te
and A te be as in Deﬁnition 1.
Let p 2 (0 ,
P
Tl =1 2( H l +1) ) .
Then, if [Equation] , we have with probability at least 1
[Equation]
[Equation]
where the constants C 3 and C 4 are deﬁned in the Appendix and depend on { L ( l )  , L ( l ) } Tl =1 and( G, F ) , N .
Hence, as the test graph size N te gets counterfactually larger (that is, as we intervene on the causal DAG of Figure 1 to change the support of the distribution of N in order to obtain larger testgraphs), the link predictor learned in the training data using gMPNN •• will converge to a continuousmethod (cMPNN •• ) that can predict links in OOD tasks (i.e., W ( X te i , X te j ) is a stationary solution ofcMPNN •• per Lemma 2).
This convergence is observed in our empirical results.
Further Related Work
In what follows we describe works related to learning transferability in GNNs.
The concept oftransferability of GNN is introduced by Levie et al.
[24], Ruiz et al.
[36], which state that if twographs represent same phenomena (e.g., are sampled from the same distribution), then a transferable
GNN has approximately the same predictive performance on both graphs.
This is closely relatedto in-distribution generalization capabilities of GNNs to unseen test data, i.e., generalization errorwhen train and test data come from the same distribution.
Existing works [18, 28, 36, 37] provethe transferability for spectral-based GCNs under graphon models, and Maskey et al.
[29] extendsthese results to more general message passing GNNs.
Our results are also based on general messagepassing GNNs.
Our goal (OOD link prediction) is, however, signiﬁcantly different than these priorworks, which focus on in-distribution graph and node classiﬁcation.
Related works relating to therepresentation power, higher order structural and positional link prediction methods (not alreadycovered in our introduction) can be found in Appendix A due to space constraints.
Empirical Evaluation
In what follows we empirically validate our theoretical results in two parts.
We implement all ourmodels in Pytorch Geometric [13] and make it available 2 .
Due to space constraints we relegate adetailed description of our experiments to the Appendix.
Convergence and stability.
First we will empirically validate Theorems 1 and 2 and Corollary 1.
Details of this experiment can be found in Appendix B.2.
Figures 2(a-c) (in Appendix) showlog-log plots of the convergence of gMPNNs to their continuous cMPNN counterparts as the testgraph size N te increases.
The empirical approximation errors  • A-W (Theorem 1) (Figure 2(a))and  •• A-W (Theorem 2) are shown as a function of the test graph size N te = 2 n , n = 5 , ..., 13 .
The empirical results show agreement with the theory since  • A-W and
•• A-W are bounded aboveby O ( p log N te / p N te ) , which is approximated by the slope  1 / 2 in a log-log plot.
Figures 2(de) (in Appendix) show histograms of the difference between gMPNN • embeddings of differentnodes in G te .
Let  • i,j := ⇥ • ( T )
A ( F ) i  ⇥ • ( T ) A ( F ) j for i, j 2 V te ,  • i,j 2 R F T and furtherdeﬁne  • iso (resp.
non-iso) i,j :=
(  • i,j ) argmax k | (  • i,j ) k | , where k 2 { 1 , . . .
, F T } is the dimension ofthe embedding.
We use subscript iso (resp.
non-iso ) when i, j 2 V te are in isomorphic (resp.non-isomorphic) SBM blocks (Deﬁnition 7).
As N te increases, Figure 2(d) shows that embeddingsbetween isomorphic blocks converge, validating Corollary 1, while Figure 2(e) shows that nonisomorphic blocks do not.
Link prediction performance evaluation with SBMs (in-distribution and OOD).
In what followswe introduce empirical results using the SBM described earlier.
We start by sampling the traininggraph ( G tr , F tr ) with N tr = 10 3 nodes.
We randomly split E tr into positive train (80%) and validation(10%) edges (we reserve 10% of E tr for the transductive test scenario), and uniformly sample the same number of across-block non-edges as negative train and validation edges.
The embedding methodgMPNN • (resp. gMPNN •• ) along with link predictor ⌘ • (resp. ⌘ •• ) are trained in an end-to-endmanner for predicting positive and negative edges in training using cross-entropy loss.
The modelachieving the best validation accuracy is then chosen to be applied in test.
Our experiments considerthree scenarios (in all scenarios we use the same number of negative test edges as positive test edges,sampled from non-edges in G te with endpoints in different isomorphic blocks): (i) (In-distribution)transductive scenario where G te =
G tr , where positive test edges are the 10% reserved in E tr not usedin training or validation; (ii) In-distribution inductive scenario where G te is sampled from the same
SBM with N te = N tr , where we sample 0 .
1 | E tr | positive test edges from E te (for fair comparisonacross all scenarios); (iii) OOD inductive scenario where G te is sampled from the same SBM with
N te = 10 ⇥ N tr , where we sample 0 .
1 | E tr | positive test edges from E te .
For structural node embeddings we consider GraphSAGE [17], GCN [19] (without positionalfeatures), GAT [48] and GIN [51] as the representatives of gMPNN • models.
The link predictor⌘ • is as feedfoward network (with 3 hidden layers and 10 neurons each) that receives the two nodeembeddings as input, and has link prediction threshold ⌧ = 0 .
5 (see Deﬁnition 8 for details).
For structural pairwise embeddings we choose our proposed gMPNN •• method of Deﬁnition 10,since
we can prove that our approach is theoretically sound in Lemma 2.
We test gMPNN •• in twoversions: The  and functions in Lemma 2 (denoted ﬁxed ) and a feedforward neural networkfor with 2 hidden layers and 5 neurons each (denoted learn ).
The link predictor ⌘ •• is the sameas ⌘ • except it just takes one pairwise embedding as input, rather than two node embeddings.
Table 1 (in Appendix) presents our empirical results.
The oracle predictor knows the graphonvalues W ( X te i , X te j ) .
Our evaluation metrics include the Matthews correlation coefﬁcient (mcc) [30],balanced accuracy, and Hits@
K for K = 10 , 50 , 100 that counts the ratio of positive edges ranked atthe k -th place or above against all negative edges.
Note that gMPNN • structural node representationscan very accurately predict links in the transductive tasks, and still performs reasonably well ininductive in-distribution tasks.
However, as expected from Corollary 2, this performance sufferssigniﬁcantly as N te becomes 10 ⇥ larger than N tr .
Now all gMPNN • methods produce predictorsthat are no better than a random guess over all metrics (e.g., see OOD mcc and accuracy (in red)).
Incontrast, the gMPNN •• is able to consistently offer good performance on both in-distribution and OOD tasks.
Link prediction performance evaluation with ogbl-ddi (in-distribution and OOD).
In whatfollows we introduce empirical results using the ogbl-ddi dataset, which represents a drug-druginteraction network.
For the purpose of performing OOD tasks, we start by sampling 10% of thenodes ( 427 nodes) and its induced subgraph to be the training graph.
Further experimental details canbe found in Appendix B.4.
The in-distribution inductive scenario has G te constructed as an inducedsubgraph with N te =
N tr nodes from the remaining ogbl-ddi graph.
Our OOD inductive scenariohas G te as the induced subgraph without the training nodes ( N te = 3840 nodes).
The test edges areobtained by applying the original edge split on the newly induced test subgraph, where we furtherdown-sample to the same amount of test edges as in our in-distribution scenarios for fair comparisonacross all scenarios.
Table 3 in the Appendix presents our empirical results on the ogbl-ddi linkprediction task.
All gMPNN • methods performs worse in inductive settings than transductive settings,and suffer much worse performance in OOD transductive setting except GCNs.
In contrast, thegMPNN •• is able to consistently offer good performance on both in-distribution and OOD tasks,showing that the theoretical results are not limited to SBM models.
Conclusions
This work studied and provided the ﬁrst theoretical framework for the task of out-of-distribution(OOD) link prediction, where test graphs are larger than training graphs.
Using non-asymptoticbounds, this work showed that OOD link prediction methods using structural node embeddingsgiven by message-passing GNNs converge to link predictors that may perform no better than randomguesses.
The work also proposed a theoretically-sound structural pairwise embedding with a messagepassing algorithm which is able to perform our OOD link prediction task by being approximatelyinvariant to interventions on test graph sizes.
Extensive empirical evaluation showed agreement withthese theoretical results.
We do not foresee adverse social impacts for this theoretical work.