TWO-DIMENSIONAL WEISFEILER-LEHMAN GRAPH NEURAL NETWORKS FOR LINK PREDICTION

Abstract

Link prediction is one important application of graph neural networks (GNNs). Most existing GNNs for link prediction are based on one-dimensional Weisfeiler-Lehman (1-WL) test. As pointed out by previous works, 1-WL-GNNs by nature learn node-level representations thereby have poor expressive power on links. Some node labeling methods relieve this weakness but introduce low efficiency. In this paper, we study a completely different approach which directly obtain node pair (link) representations based on two-dimensional Weisfeiler-Lehman (2-WL) tests. 2-WL tests directly use links (2-tuples) as message passing units instead of nodes, and thus can directly obtain link representations. We theoretically analyze the expressive power of 2-WL tests to discriminate non-isomorphic links, and prove their superior link discriminating power than 1-WL. Based on different 2-WL variants, we propose a series of novel 2-WL-GNN models for link prediction. Experiments on a wide range of real-world datasets demonstrate their competitive performance to state-of-the-art baselines.

1. INTRODUCTION

Link prediction is a key problem of graph-structured data (Al Hasan et al., 2006; Liben-Nowell & Kleinberg, 2007; Menon & Elkan, 2011; Trouillon et al., 2016) . It refers to utilizing node characteristics and graph topology to measure how likely a link exists between a pair of nodes. Due to the importance of predicting pairwise relations, it has wide applications in various domains, such as recommendation in social networks (Adamic & Adar, 2003) , knowledge graph completion (Nickel et al., 2015) , and metabolic network reconstruction (Oyetunde et al., 2017) . One class of traditional link prediction methods are heuristic methods, which use manually designed graph structural features of a target node pair such as number of common neighbors (CN) (Liben-Nowell & Kleinberg, 2007) , preferential attachment (PA) (Barabási & Albert, 1999) , and resource allocation (RA) (Zhou et al., 2009) to estimate the likelihood of link existence. Another class of methods, embedding methods, including Matrix Factorization (MF) (Menon & Elkan, 2011) and node2vec (Grover & Leskovec, 2016) , learn node embeddings from the graph structure in a transductive manner, which cannot generalize to unseen nodes or new graphs. Recently, with the popularity of GNNs, their application to link prediction brings a number of cutting-edge models (Kipf & Welling, 2016; Zhang & Chen, 2018; Zhang et al., 2021; Zhu et al., 2021) . Most existing GNN models for link prediction are based on one-dimensional Weisfeiler-Lehman (1-WL) test (Weisfeiler & Leman, 1968; Shervashidze et al., 2011) . 1-WL test is a popular heuristic for detecting non-isomorphic graphs. In each update, it obtains all nodes' new colors by hashing their own colors and multisets of their neighbors' colors. Vanilla GNNs simulate 1-WL test by iteratively aggregating neighboring node features to the center node to update node representations, which we call 1-WL-GNNs. With the node representations, 1-WL-GNNs compute link prediction scores by aggregating pairwise node representations. Graph Auto-encoder (GAE, and its variant VGAE) (Kipf & Welling, 2016 ) is such a model. However, 1-WL-GNNs can only discriminate links on the "node" level. This is illustrated by Figure 1 left: v 2 and v 3 are symmetric nodes in the graph thus having the same representation by 1-WL-GNN, but links (v 1 , v 2 ) and (v 1 , v 3 ) are not symmetric. However, 1-WL-GNNs are unable to discriminate links (v 1 , v 2 ) and (v 1 , v 3 ), though (v 1 , v 2 ) has a shorter path between them than (v 1 , v 3 ). Although positional node embeddings or random features can alleviate this problem, they fail to guarantee symmetrical links (such as (v 1 , v 2 ) and (v 4 , v 3 )) to have the same representation. In order to surpass 1-WL, plenty of link prediction models apply node labeling inherently, including SEAL (Zhang & Chen, 2018), Distance Encoding (Li et al., 2020) , ID-GNN (You et al., 2021) , and some models for matrix completion (Zhang & Chen, 2020) and knowledge graph completion (Teru et al., 2020) . It raises expressive power from "node" to "link" level by breaking the symmetry between the target node pair and other nodes during the message passing. Figure 1 middle and right illustrated this effect. However, labeling also introduces a challenge. It requires repeatedly applying GNN to a labeled subgraph for every link to predict thereby being inefficient. Therefore, we aim to develop novel GNN models with both full-batch link prediction ability and higher expressive power than 1-WL. 𝑣 ! 𝑣 " 𝑣 # 𝑣 ! 𝑣 " 𝑣 # 𝑣 $ 𝑣 # 𝑣 ! 𝑣 " Figure 1: 1-WL-GNNs cannot distinguish links (v 1 , v 2 ) and (v 1 , v 3 ) in the left graph. With labeling trick, 1-WL-GNNs can distinguish them in their respective labeled graphs (middle and right). We propose a completely different paradigm for link prediction. We construct GNNs based on two-dimensional Weisfeiler-Lehman (2-WL) tests, which we call 2-WL-GNNs. In 2-WL-GNNs, node pairs are used as the elemental message passing units so that link representations are directly obtained. Figure 2 gives an illustration for a particular 2-WL algorithm. We first theoretically study the link discriminating power of different 2-WL test variants, including the plain 2-WL, 2-FWL (Folklore WL), and their newly proposed local version. We show that 2-WL, 2-FWL and local 2-FWL are strictly more expressive than 1-WL for link prediction, while local 2-WL has equivalent power to 1-WL. Based on these 2-WL tests, we construct a series of 2-WL-GNN models. Despite all using node pairs to propagate messages, these models have different aggregation schemes, link discriminating power, time/space complexity, as well as drastically different implementations, which we discuss in Section 4. Extensive experiments on multiple benchmark datasets verify 2-WL-GNNs' power for link prediction. 2-WL-GNNs achieve highly competitive link prediction performance to state-of-the-art models including SEAL (Zhang & Chen, 2018) and NBFNet (Zhu et al., 2021) , while using significantly less time. In this paper, we aim to develop novel GNN models with both full-batch link prediction ability and higher expressive power than 1-WL.

2. LINK-LEVEL TWO-DIMENSIONAL WEISFEILER-LEHMAN TESTS

In this section we introduce various 2-WL tests which directly use links as message passing unit, and define their link-level expressive power. We denote a set by {•}, an ordered set (tuple) by (•) and a multiset by { {•} } to have repeated elements. We use [n] to denote the set {1, 2, ..., n}.

2.1. k-DIMENSIONAL WEISFEILER-LEHMAN TESTS

k-dimensional WL test (k-WL) uses k-tuples of nodes as update unit. In each iteration, every k-tuple updates its color from its newly-defined neighboring k-tuples. There are two variants of k-WL algorithms: the plain k-dimensional WL (k-WL) and the k-dimensional Folklore WL (k-FWL) (Cai et al., 1992; Grohe, 2017; Maron et al., 2019) . Both k-WL and k-FWL update colors for k-tuples s := (s 1 , s 2 , ..., s k ) with s 1 , ..., s k being nodes. k-WL defines neighborhood of k-tuple s as N (s) = N 1 (s), N 2 (s), ..., N k (s) , where N j (s) = (s 1 , ..., s j-1 , s ′ , s j+1 , ..., s k )|s ′ ∈ [n] . (1) k-FWL has a different definition of neighborhood. k-FWL defines the jth neighborhood of s as N F j (s) = (j, s 2 , ..., s k ), (s 1 , j, ..., s k ), ..., (s 1 , ..., s k-1 , j) . (2) And the full neighborhood of s is given by N F (s) = { {N F j (s)|j ∈ [n]} }. Essentially, k-WL and k-FWL have the same nk neighbor tuples but differ in how these nk tuples are ordered and grouped.

