Labeling Trick: A Theory of Using Graph Neural Networks for Multi-Node Representation Learning

1 Introduction . Graph neural networks ( GNNs ) [ 1–10 ] have achieved great successes in recent years . While GNNs have been well studied for single-node tasks ( such as node classification ) and whole-graph tasks ( such as graph classification ) , using GNNs to predict a set of multiple nodes is less studied and less understood . Among such multi-node representation learning problems , link prediction ( predicting the link existence/class/value between a set of two nodes ) is perhaps the most important one due to its wide applications in practice , including friend recommendation in social networks [ 11 ] , movie recommendation in Netflix [ 12 ] , protein interaction prediction [ 13 ] , drug response prediction [ 14 ] , knowledge graph completion [ 15 ] , etc . In this paper , we use link prediction as a medium to study GNN ’ s multi-node representation learning ability . Note that although our examples and experiments are all around link prediction , our theory applies generally to all multi-node representation learning problems such as triplet [ 16 ] , motif [ 17 ] and subgraph [ 18 ] prediction tasks . ∗Corresponding author : Muhan Zhang ( muhan @ pku.edu.cn ) . Work done as a research scientist at Facebook . †Pan Li inspires the study of labeling tricks , proves Theorem 2 , and helps check the theoretical framework . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . There are two main classes of GNN-based link prediction methods : Graph AutoEncoder ( GAE ) [ 19 ] and SEAL [ 20 , 21 ] . GAE ( and its variational version VGAE [ 19 ] ) first applies a GNN to the entire network to compute a representation for each node . The representations of the two end nodes of the link are then aggregated to predict the target link . GAE represents a common practice of using GNNs to learn multi-node representations . That is , first obtaining individual node representations through a GNN as usual , and then aggregating the representations of those nodes of interest as the multi-node representation . On the contrary , SEAL applies a GNN to an enclosing subgraph around each link , where nodes in the subgraph are labeled differently according to their distances to the two end nodes before applying the GNN . Despite both using GNNs for link prediction , SEAL often shows much better practical performance than GAE . As we will see , the key lies in SEAL ’ s node labeling step . 𝑣2 𝑣1 𝑣3 𝑣4 apparently different structural roles in the graph reflects one key limitation of GAE-type methods : by computing v1 and v2 ’ s representations independently of each other , GAE can not capture the dependence between two end nodes of a link . For example , ( v1 , v2 ) has a much shorter path between them than that of ( v1 , v3 ) ; and ( v1 , v2 ) has both nodes in the same hexagon , while ( v1 , v3 ) does not . Take common neighbor ( CN ) [ 22 ] , one elementary heuristic feature for link prediction , as another example . CN counts the number of common neighbors between two nodes to measure their likelihood of forming a link , which is widely used in social network friend recommendation . CN is the foundation of many other successful heuristics such as Adamic-Adar [ 11 ] and Resource Allocation [ 23 ] , which are also based on neighborhood overlap . However , GAE can not capture such neighborhood-overlapbased features . This can be seen from Figure 1 too . There is 1 common neighbor between ( v1 , v2 ) and 0 between ( v1 , v3 ) , but GAE always gives ( v1 , v2 ) and ( v1 , v3 ) the same representation . The failure to learn common neighbor demonstrates GAE ’ s severe limitation for link prediction . The root cause still lies in that GAE computes node representations independently of each other—when computing the representation of one end node , it is not aware of the other end node . One way to alleviate the above failure is to use one-hot encoding of node indices or random features as input node features [ 24 , 25 ] . With such node-discriminating features , v2 and v3 will have different node representations , thus ( v1 , v2 ) and ( v1 , v3 ) may also have different link representations after aggregation , enabling GAE to discriminate ( v1 , v2 ) and ( v1 , v3 ) . However , using node-discriminating features loses GNN ’ s inductive learning ability to map nodes and links with identical neighborhoods ( such as nodes v2 and v3 , and links ( v1 , v2 ) and ( v4 , v3 ) ) to the same representation , which results in a great loss of generalization ability . The resulting model is no longer permutation invariant/equivariant , violating the fundamental design principle of GNNs . Is there a way to improve GNNs ’ link discriminating power ( so that links like ( v1 , v2 ) and ( v1 , v3 ) can be distinguished ) , while maintaining their inductive learning ability ( so that links ( v1 , v2 ) and ( v4 , v3 ) have the same representation ) ? In this paper , we analyze the above problem from a structural representation learning point of view . Srinivasan and Ribeiro [ 26 ] prove that the multi-node prediction problem on graphs ultimately only requires finding a most expressive structural representation of node sets , which gives two node sets the same representation if and only if they are isomorphic ( a.k.a . symmetric , on the same orbit ) in the graph . For example , link ( v1 , v2 ) and link ( v4 , v3 ) in Figure 1 are isomorphic . A most expressive structural representation for links should give any two isomorphic links the same representation while discriminating all non-isomorphic links ( such as ( v1 , v2 ) and ( v1 , v3 ) ) . According to our discussion above , GAE-type methods that directly aggregate node representations can not learn a most expressive structural representation . Then , how to learn a most expressive structural representation of node sets ? To answer this question , we revisit the other GNN-based link prediction framework , SEAL , and analyze how node labeling helps a GNN learn better node set representations . We find out that two properties of a node labeling are crucial for its effectiveness : 1 ) target-nodes-distinguishing and 2 ) permutation equivariance . With these two properties , we define labeling trick ( Section 4.1 ) , which unifies previous node labeling methods into a single and most general form . Theoretically , we prove that with labeling trick a sufficiently expressive GNN can learn most expressive structural representations of node sets ( Theorem 1 ) , which reassures GNN ’ s node set prediction ability . It also closes the gap between GNN ’ s node representation learning nature and node set tasks ’ multi-node representation learning requirement . We further extend our theory to local isomorphism ( Section 5 ) . And finally , experiments on four OGB link existence prediction datasets [ 27 ] verified our theory . Note that the labeling trick theory allows the presence of node/edge features/types , thus is not restricted to non-attributed and homogeneous graphs . Previous works on heterogeneous graphs , such as knowledge graphs [ 28 ] and recommender systems [ 29 ] have already seen successful applications of labeling trick . Labeling trick is also not restricted to two-node link representation learning tasks , but generally applies to any multi-node representation learning tasks . 2 Preliminaries . In this section , we introduce some important concepts that will be used in the analysis of the paper , including permutation , set isomorphism and most expressive structural representation . We consider a graph G = ( V , E , A ) , where V = { 1 , 2 , . . . , n } is the set of n vertices , E ⊆ V × V is the set of edges , and A ∈ Rn×n×k is a 3-dimensional tensor containing node and edge features . The diagonal components Ai , i , : denote features of node i , and the off-diagonal components Ai , j , : denote features of edge ( i , j ) . For heterogeneous graphs , the node/edge types can also be expressed in A using integers or one-hot encoding vectors . We further use A ∈ { 0 , 1 } n×n to denote the adjacency matrix of G with Ai , j = 1 iff ( i , j ) ∈ E. We let A be the first slice of A , i.e. , A = A : , : ,1 . Since A contains the complete information of a graph , we sometimes directly use A to denote the graph . Definition 1 . A permutation π is a bijective mapping from { 1 , 2 , . . . , n } to { 1 , 2 , . . . , n } . Depending on the context , π ( i ) can mean assigning a new index to node i ∈ V , or mapping node i to node π ( i ) of another graph . All n ! possible π ’ s constitute the permutation group Πn . For joint prediction tasks over a set of nodes , we use S to denote the target node set . For example , S = { i , j } if we want to predict the link between i , j . We define π ( S ) = { π ( i ) |i ∈ S } . We further define the permutation of A as π ( A ) , where π ( A ) π ( i ) , π ( j ) , : = Ai , j , : . Next , we define set isomorphism , which generalizes graph isomorphism to arbitrary node sets . Definition 2 . ( Set isomorphism ) Given two n-node graphs G = ( V , E , A ) , G′ = ( V ′ , E′ , A′ ) , and two node sets S ⊆ V , S′ ⊆ V ′ , we say ( S , A ) and ( S′ , A′ ) are isomorphic ( denoted by ( S , A ) ' ( S′ , A′ ) ) if ∃π ∈ Πn such that S = π ( S′ ) and A = π ( A′ ) . When ( V , A ) ' ( V ′ , A′ ) , we say two graphs G and G′ are isomorphic ( abbreviated as A ' A′ because V = π ( V ′ ) for any π ) . Note that set isomorphism is more strict than graph isomorphism , because it not only requires graph isomorphism , but also requires the permutation maps a specific node set S to another node set S′ . In practice , when S 6= V , we are often more concerned with the case of A = A′ , where isomorphic node sets are defined in the same graph ( automorphism ) . For example , when S = { i } , S′ = { j } and ( i , A ) ' ( j , A ) , we say nodes i and j are isomorphic in graph A ( or they have symmetric positions/same structural role in graph A ) . An example is v2 and v3 in Figure 1 . We say a function f defined over the space of ( S , A ) is permutation invariant ( or invariant for abbreviation ) if ∀π ∈ Πn , f ( S , A ) = f ( π ( S ) , π ( A ) ) . Similarly , f is permutation equivariant if ∀π ∈ Πn , π ( f ( S , A ) ) = f ( π ( S ) , π ( A ) ) . Permutation invariance/equivariance ensures representations learned by a GNN is invariant to node indexing , which is a fundamental design principle of GNNs . Now we define most expressive structural representation of a node set , following [ 26 , 21 ] . Basically , it assigns a unique representation to each equivalence class of isomorphic node sets . Definition 3 . Given an invariant function Γ ( · ) , Γ ( S , A ) is a most expressive structural representation for ( S , A ) if ∀S , A , S′ , A′ , Γ ( S , A ) = Γ ( S′ , A′ ) ⇔ ( S , A ) ' ( S′ , A′ ) . For simplicity , we will briefly use structural representation to denote most expressive structural representation in the rest of the paper . We will omit A if it is clear from context . We call Γ ( i , A ) a structural node representation for i , and call Γ ( { i , j } , A ) a structural link representation for ( i , j ) . Definition 3 requires that the structural representations of two node sets are the same if and only if the two node sets are isomorphic . That is , isomorphic node sets always have the same structural representation , while non-isomorphic node sets always have different structural representations . This is in contrast to positional node embeddings such as DeepWalk [ 30 ] and matrix factorization [ 31 ] , where two isomorphic nodes can have different node embeddings [ 32 ] . GAE using node-discriminating features also learns positional node embeddings . Why do we study structural representations ? Formally speaking , Srinivasan and Ribeiro [ 26 ] prove that any joint prediction task over node sets only requires a structural representation of node sets . They show that positional node embeddings carry no more information beyond that of structural representations . Intuitively speaking , it is because two isomorphic nodes in a network are perfectly symmetric and interchangeable with each other , and should be indistinguishable from any perspective . Learning a structural node representation guarantees that isomorphic nodes are always classified into the same class . Similarly , learning a structural link representation guarantees isomorphic links , such as ( v1 , v2 ) and ( v4 , v3 ) in Figure 1 , are always predicted the same , while non-isomorphic links , such as ( v1 , v2 ) and ( v1 , v3 ) , are always distinguishable , which is not guaranteed by positional node embeddings . Structural representation characterizes the maximum representation power a model can reach on graphs . We use it to study GNNs ’ multi-node representation learning ability .

Summary: This paper studies the problem of using GNN for link prediction. The contributions are mainly two folds. First, it studies the limitation of using GAE for link prediction, which directly aggregate a pair of nodes embedding. Secondly, it proposes the labeling trick which is sufficiently expressive and unifies some of the existing frameworks. The proposed method is general for high-order node set prediction where link prediction is a special case.