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

Graph neural networks (GNNs) [1, 2, 3, 4, 5, 6, 7, 8, 9, 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.

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.

We first give a simple example to show when GAE fails. In Figure 1, $v_{2}$ and $v_{3}$ have symmetric positions in the graph—from their respective views, they have exactly the same $h$-hop neighborhood for any $h$. Thus, without node features, GAE will learn the same representation for $v_{2}$ and $v_{3}$. However, when we want to predict which one of $v_{2}$ and $v_{3}$ is more likely to form a link with $v_{1}$, GAE will aggregate the representations of $v_{1}$ and $v_{2}$ as the link representation of $(v_{1},v_{2})$, and aggregate the representations of $v_{1}$ and $v_{3}$ to represent $(v_{1},v_{3})$, thus giving $(v_{1},v_{2})$ and $(v_{1},v_{3})$ the same representation and prediction. The failure to distinguish links $(v_{1},v_{2})$ and $(v_{1},v_{3})$ that have apparently different structural roles in the graph reflects one key limitation of GAE-type methods: by computing $v_{1}$ and $v_{2}$’s representations independently of each other, GAE cannot capture the dependence between two end nodes of a link. For example, $(v_{1},v_{2})$ has a much shorter path between them than that of $(v_{1},v_{3})$; and $(v_{1},v_{2})$ has both nodes in the same hexagon, while $(v_{1},v_{3})$ 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 cannot capture such neighborhood-overlap-based features. This can be seen from Figure 1 too. There is 1 common neighbor between $(v_{1},v_{2})$ and 0 between $(v_{1},v_{3})$, but GAE always gives $(v_{1},v_{2})$ and $(v_{1},v_{3})$ 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, $v_{2}$ and $v_{3}$ will have different node representations, thus $(v_{1},v_{2})$ and $(v_{1},v_{3})$ may also have different link representations after aggregation, enabling GAE to discriminate $(v_{1},v_{2})$ and $(v_{1},v_{3})$. However, using node-discriminating features loses GNN’s inductive learning ability to map nodes and links with identical neighborhoods (such as nodes $v_{2}$ and $v_{3}$, and links $(v_{1},v_{2})$ and $(v_{4},v_{3})$) 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 $(v_{1},v_{2})$ and $(v_{1},v_{3})$ can be distinguished), while maintaining their inductive learning ability (so that links $(v_{1},v_{2})$ and $(v_{4},v_{3})$ 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 $(v_{1},v_{2})$ and link $(v_{4},v_{3})$ 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 $(v_{1},v_{2})$ and $(v_{1},v_{3})$). According to our discussion above, GAE-type methods that directly aggregate node representations cannot 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.

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 ${\mathcal{G}}=(V,E,{\bm{\mathsfit{A}}})$, where $V=\{1,2,\ldots,n\}$ is the set of $n$ vertices, $E\subseteq V\times V$ is the set of edges, and ${\bm{\mathsfit{A}}}\in\mathbb{R}^{n\times n\times k}$ is a 3-dimensional tensor containing node and edge features. The diagonal components ${\bm{\mathsfit{A}}}_{i,i,:}$ denote features of node $i$, and the off-diagonal components ${\bm{\mathsfit{A}}}_{i,j,:}$ denote features of edge $(i,j)$. For heterogeneous graphs, the node/edge types can also be expressed in ${\bm{\mathsfit{A}}}$ using integers or one-hot encoding vectors. We further use ${\bm{A}}\in\{0,1\}^{n\times n}$ to denote the adjacency matrix of ${\mathcal{G}}$ with ${\bm{A}}_{i,j}=1$ iff $(i,j)\in E$. We let ${\bm{A}}$ be the first slice of ${\bm{\mathsfit{A}}}$, i.e., ${\bm{A}}={\bm{\mathsfit{A}}}_{:,:,1}$. Since ${\bm{\mathsfit{A}}}$ contains the complete information of a graph, we sometimes directly use ${\bm{\mathsfit{A}}}$ to denote the graph.

Next, we define set isomorphism, which generalizes graph isomorphism to arbitrary node sets.

When $(V,{\bm{\mathsfit{A}}})\simeq(V^{\prime},{\bm{\mathsfit{A}}}^{\prime})$, we say two graphs ${\mathcal{G}}$ and ${\mathcal{G}}^{\prime}$ are isomorphic (abbreviated as ${\bm{\mathsfit{A}}}\simeq{\bm{\mathsfit{A}}}^{\prime}$ because $V=\pi(V^{\prime})$ for any $\pi$). 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^{\prime}$. In practice, when $S\neq V$, we are often more concerned with the case of ${\bm{\mathsfit{A}}}={\bm{\mathsfit{A}}}^{\prime}$, where isomorphic node sets are defined in the same graph (automorphism). For example, when $S=\{i\},S^{\prime}=\{j\}$ and $(i,{\bm{\mathsfit{A}}})\simeq(j,{\bm{\mathsfit{A}}})$, we say nodes $i$ and $j$ are isomorphic in graph ${\bm{\mathsfit{A}}}$ (or they have symmetric positions/same structural role in graph ${\bm{\mathsfit{A}}}$). An example is $v_{2}$ and $v_{3}$ in Figure 1.

We say a function $f$ defined over the space of $(S,{\bm{\mathsfit{A}}})$ is permutation invariant (or invariant for abbreviation) if $\forall\pi\in\Pi_{n}$, $f(S,{\bm{\mathsfit{A}}})=f(\pi(S),\pi({\bm{\mathsfit{A}}}))$. Similarly, $f$ is permutation equivariant if $\forall\pi\in\Pi_{n}$, $\pi(f(S,{\bm{\mathsfit{A}}}))=f(\pi(S),\pi({\bm{\mathsfit{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.

For simplicity, we will briefly use structural representation to denote most expressive structural representation in the rest of the paper. We will omit ${\bm{\mathsfit{A}}}$ if it is clear from context. We call $\Gamma(i,{\bm{\mathsfit{A}}})$ a structural node representation for $i$, and call $\Gamma(\{i,j\},{\bm{\mathsfit{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 $(v_{1},v_{2})$ and $(v_{4},v_{3})$ in Figure 1, are always predicted the same, while non-isomorphic links, such as $(v_{1},v_{2})$ and $(v_{1},v_{3})$, 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.

In this section, using GAE for link prediction as an example, we show the key limitation of directly aggregating node representations as a node set representation.

In this section, we first define the general form of labeling trick, and use a specific implementation, zero-one labeling trick, to intuitively explain why labeling trick helps GNNs learn better link representations. Next, we present our main theorem showing that labeling trick enables a node-most-expressive GNN to learn structural representations of node sets, which formally characterizes GNN’s maximum multi-node representation learning ability. Then, we review SEAL and show it exactly uses one labeling trick. Finally, we discuss other labeling trick implementations in previous works.

The concept of most expressive structural representation is based on assigning node sets the same representation if and only if they are isomorphic to each other in the graph. However, exact isomorphism is not very common. For example, Babai and Kucera [36] prove that at least $(n-\log n)$ nodes in almost all $n$-node graphs are non-isomorphic to each other. In practice, 1-WL-GNN also takes up to $\mathcal{O}(n)$ message passing layers to reach its maximum power for discriminating non-isomorphic nodes, making it very hard to really target on finding exactly isomorphic nodes/links.

In this regard, we propose a more practical concept, called local isomorphism.

Local $h$-isomorphism only requires the $h$-hop enclosing subgraphs around $S$ and $S^{\prime}$ are isomorphic, instead of the entire graphs. We argue that this is a more useful definition than isomorphism, because: 1) Exact isomorphism is rare in real-world graphs. 2) Algorithms targeting on exact isomorphism are more likely to overfit. Only assigning the same representations to exactly isomorphic nodes/links may fail to identify a large amount of nodes/links that are not isomorphic but have very similar neighborhoods. Instead, nodes/links locally isomorphic to each other may better indicate that they should have the same representation. With local $h$-isomorphism, all our previous conclusions based on standard isomorphism still apply. For example, GAE (without node-discriminating features) still cannot discriminate locally $h$-non-isomorphic links. And a node-most-expressive GNN with labeling trick can learn the most expressive structural representations of node sets w.r.t. local $h$-isomorphism, i.e., learn the same representation for two node sets if and only if they are locally $h$-isomorphic:

Corollary 1 demonstrates labeling trick’s power in the context of local isomorphism. To switch to local $h$-isomorphism, all we need to do is to extract the $h$-hop enclosing subgraph around a node set, and apply labeling trick and GNN only to the extracted subgraph. This is exactly what SEAL does.

There is emerging interest in studying the representation power of graph neural networks recently. Xu et al. [35] and Morris et al. [40] first show that the discriminating power of GNNs performing neighbor aggregation is bounded by the 1-WL test. Many works have since been proposed to increase the power of GNNs by simulating higher-order WL tests [40, 41, 42]. However, most previous works focus on improving GNN’s whole-graph representation power. Little work has been done to analyze GNN’s node/link representation power. Srinivasan and Ribeiro [26] first formally studied the difference between structural representations of nodes and links. Although showing that structural node representations of GNNs cannot perform link prediction, their way to learn structural link representations is to give up GNNs and instead use Monte Carlo samples of node embeddings learned by network embedding methods. In this paper, we show that GNNs combined with labeling trick can as well learn structural link representations, which reassures using GNNs for link prediction.

Many works have implicitly assumed that if a model can learn node representations well, then combining the pairwise node representations can also lead to good link representations [43, 19, 44]. However, we argue in this paper that simply aggregating node representations fails to discriminate a large number of non-isomorphic links, and with labeling trick the aggregation of structural node representations leads to structural link representations. Li et al. [21] proposed distance encoding (DE), whose implementations based on $S$-discriminating distances can be shown to be specific labeling tricks. They proved that DE can improve 1-WL-GNNs’ discriminating power, enabling them to differentiate almost all $(S,{\bm{\mathsfit{A}}})$ tuples sampled from $r$-regular graphs. Our paper contributes to an important aspect that Li et al. [21] overlooked: 1) Our theory focuses on the gap between a GNN’s single-node and multi-node representation power. We show even a GNN has maximum node representation power, it still fails to learn structural representations of node sets unless combined with a labeling trick. However, the theory of DE cannot explain this. 2) Our theory is not restricted to $r$-regular graphs, but applies to any graphs. 3) Our theory points out that a valid labeling trick is not necessarily distance based—it need only be permutation equivariant and $S$-discriminating. More discussion on the difference between DE’s theory and the theory in this paper is given in Appendix E.

You et al. [45] also noticed that structural node representations of GNNs cannot capture the dependence (in particular distance) between nodes. To learn position-aware node embeddings, they propose P-GNN, which randomly chooses some anchor nodes and aggregates messages only from the anchor nodes. In P-GNN, nodes with similar distances to the anchor nodes, instead of nodes with similar neighborhoods, have similar embeddings. Thus, P-GNN cannot learn structural node/link representations. P-GNN also cannot scale to large datasets. You et al. [38] later proposed ID-GNN. As discussed in Section 4.3, ID-GNN’s node coloring scheme can be seen as a partial labeling trick.

Finally, although labeling trick is formally defined in this paper, various forms of specific labeling tricks have already been used in previous works. To our best knowledge, SEAL [20] proposes the first labeling trick, which is designed to improve GNN’s link prediction power. It is later adopted in inductive knowledge graph completion [28] and matrix completion [29], and is generalized into DE [21] which works for $|S|>2$ cases. Wan et al. [46] use labeling trick for hyperedge prediction.

In this section, we use a two-node task, link prediction, to empirically validate the effectiveness of labeling trick for multi-node representation learning. We use four link existence prediction datasets in Open Graph Benchmark (OGB) [27]: ogbl-ppa, ogbl-collab, ogbl-ddi, and ogbl-citation2. These datasets are open-sourced, large-scale (up to 2.9M nodes and 30.6M edges), adopt realistic train/validation/test splits, and have standard evaluation procedures, thus providing an ideal place to benchmark an algorithm’s realistic link prediction power. The evaluation metrics include Hits@$K$ and MRR. Hits@$K$ counts the ratio of positive edges ranked at the K-th place or above against all the negative edges. MRR (Mean Reciprocal Rank) computes the reciprocal rank of the true target node against 1,000 negative candidates, averaged over all the true source nodes. Both metrics are higher the better. We include more details and statistics of these datasets in Appendix F. Our code is available at https://github.com/facebookresearch/SEAL_OGB.

Baselines. We use the following baselines for comparison. We use 5 non-GNN methods: CN (common neighbor), AA (Adamic-Adar), MLP, MF (matrix factorization) and Node2vec. Among them, CN and AA are two simple link prediction heuristics based on counting common neighbors, which are used for sanity checking. We use 3 plain GAE baselines: GraphSAGE [44], GCN [19], and GCN+LRGA [47]. These methods use the Hadamard product of pairwise node representations output by a GNN as link representations, without using a labeling trick. Finally, we compare 3 GNN methods using labeling tricks: GCN+DE [21], GCN+DRNL, and SEAL [20]. GCN+DE/GCN+DRNL enhance GCN with the DE/DRNL labeling trick. SEAL uses a GCN and the DRNL labeling trick, with an additional subgraph-level readout SortPooling [9]. More details are in Appendix G. Moreover, we test the zero-one labeling trick in our ablation experiments. Results can be found in Appendix H.

Results and discussion. We present the main results in Table 1. Firstly, we can see that GAE methods without labeling trick do not always outperform non-GNN methods. For example, on ogbl-ppa and ogbl-collab, simple heuristics CN and AA outperform plain GAE methods by large margins. This suggests that GAE methods cannot even learn simple neighborhood-overlap-based heuristics, verifying our argument in Introduction. In contrast, when GNNs are enhanced by labeling trick, they are able to beat heuristics. With labeling trick, GNN methods achieve new state-of-the-art performance on 3 out of 4 datasets. In particular, we observe that SEAL outperforms GAE and positional embedding methods, sometimes by surprisingly large margins. For example, in the challenging ogbl-ppa graph, SEAL achieves an Hits@100 of 48.80, which is 87%-195% higher than GAE methods without using labeling trick. On ogbl-ppa, ogbl-collab and ogbl-citation2, labeling trick methods also achieve state-of-the-art results.

Despite obtaining the best results on three datasets, we observe that labeling trick methods do not perform well on ogbl-ddi. ogbl-ddi is considerably denser than the other graphs. It has 4,267 nodes and 1,334,889 edges, resulting in an average node degree of 500.5. In ogbl-ddi, labeling trick methods fall behind GAE methods using trainable node embeddings. One possible explanation is that ogbl-ddi is so dense that a practical GNN with limited expressive power is hard to inductively learn any meaningful structural patterns. In comparison, the transductive way of learning free-parameter node embeddings makes GAEs no longer focus on learning inductive structural patterns, but focus on learning node embeddings. The added parameters also greatly increase GAEs’ model capacity. An interesting future topic is to study how to improve labeling tricks’ performance on dense graphs. Appendix H presents more ablation experiments to study the power of different labeling tricks, the effect of subgraph pooling, and the number of hops/layers.

In this paper, we proposed a theory of using GNNs for multi-node representation learning. We first pointed out the key limitation of a common practice in previous works that directly aggregates node representations as a node set representation. To address the problem, we proposed labeling trick which gives target nodes distinct labels in a permutation equivariant way. We proved that labeling trick enables GNNs to learn most expressive structural representations of node sets, which formally characterizes GNNs’ maximum multi-node representation learning ability. Experiments on four OGB datasets verified labeling trick’s effectiveness.

The authors greatly thank the actionable suggestions from the reviewers. Li is partly supported by the 2021 JP Morgan Faculty Award and the National Science Foundation (NSF) award HDR-2117997.