Source-linked AI summary

Graph Neural Networks for Link Prediction with Subgraph Sketching

Benjamin Paul Chamberlain, Sergey Shirobokov, Emanuele Rossi, Fabrizio Frasca, Thomas Markovich, Nils Hammerla, Michael M. Bronstein, Max Hansmire

arXiv:2209.15486v3cs.LGcs.IR

TL;DR

GNNs often underperform heuristics on link prediction, in part because they cannot count triangles or distinguish automorphic nodes. The paper proposes ELPH, which passes subgraph sketches as messages, and BUDDY, which uses precomputation. ELPH or BUDDY achieves the best performance on five of seven datasets, while ELPH runs out of memory on the two largest and BUDDY performs well on both.

  • Problem

    GNNs often fail to outperform heuristics on common link-prediction benchmarks, and MPNNs cannot count triangles or distinguish nodes in the same automorphism orbit.

    Source: §1 · p. 2
  • Method

    ELPH passes subgraph sketches as messages, while BUDDY uses precomputed sketches and node features to scale link prediction.

    Source: §1 · pp. 2–3, §5 · p. 8
  • Results

    ELPH or BUDDY achieves the best performance on five of seven datasets; ELPH runs out of memory on the two largest, while BUDDY performs extremely well on both.

    Source: §7 · p. 9
  • Takeaways & Limitations

    ELPH is strictly more expressive than MPNNs, and BUDDY empirically outperforms common MPNNs while remaining extremely scalable.

    Source: §4.2 · p. 8, §5 · p. 9

    The work is limited to undirected graphs or directed graphs preprocessed as undirected, leaving extensions to directed and temporal or dynamically evolving graphs for future work.

    Source: §8 · p. 10

Abstract

from arXiv · show

Many Graph Neural Networks (GNNs) perform poorly compared to simple heuristics on Link Prediction (LP) tasks. This is due to limitations in expressive power such as the inability to count triangles (the backbone of most LP heuristics) and because they can not distinguish automorphic nodes (those having identical structural roles). Both expressiveness issues can be alleviated by learning link (rather than node) representations and incorporating structural features such as triangle counts. Since explicit link representations are often prohibitively expensive, recent works resorted to subgraph-based methods, which have achieved state-of-the-art performance for LP, but suffer from poor efficiency due to high levels of redundancy between subgraphs. We analyze the components of subgraph GNN (SGNN) methods for link prediction. Based on our analysis, we propose a novel full-graph GNN called ELPH (Efficient Link Prediction with Hashing) that passes subgraph sketches as messages to approximate the key components of SGNNs without explicit subgraph construction. ELPH is provably more expressive than Message Passing GNNs (MPNNs). It outperforms existing SGNN models on many standard LP benchmarks while being orders of magnitude faster. However, it shares the common GNN limitation that it is only efficient when the dataset fits in GPU memory. Accordingly, we develop a highly scalable model, called BUDDY, which uses feature precomputation to circumvent this limitation without sacrificing predictive performance. Our experiments show that BUDDY also outperforms SGNNs on standard LP benchmarks while being highly scalable and faster than ELPH.

1 Introduction

Link prediction matters in applications such as recommendation, drug discovery, and knowledge-graph construction, yet GNNs often fail to outperform heuristics on common benchmarks. The paper analyzes subgraph GNNs and proposes ELPH and BUDDY to address expressiveness and scalability challenges.

1 Introduction

  • MPNNs cannot count triangles or common neighbors, limiting their ability to represent several one- and two-hop link-prediction heuristics.The paper attributes this to message passing's equivalence to the Weisfeiler-Leman test.Source: §1 · p. 2
  • Automorphic nodes receive equal representations, so a GNN can assign the same link probability even when the candidate pairs have different graph distances.Figure 1 illustrates this for links (1,2) and (1,4).Source: §1 · p. 2, Figure 1 · p. 2
  • Structural labeling features can address the automorphic-node problem, but conditioning node representations on an edge prevents efficient parallel computation.Prior subgraph methods make link prediction into binary subgraph classification.Source: §1 · p. 2
  • SGNNs face expensive subgraph construction, inefficient GPU batching of irregular subgraphs, and inference costs nearly as high as training.These drawbacks constrain applications requiring scalability or efficient inference.Source: §1 · p. 2
  • The paper analyzes SGNN components and proposes ELPH, which passes subgraph sketches as messages, and BUDDY, which precomputes sketches and node features for scalability.The authors also report favorable accuracy and speed against state-of-the-art methods and provide an open-source PyTorch library for graph sketching.Source: §1 · pp. 2–3

2 Preliminaries

The preliminaries define the graph and enclosing-subgraph notation, describe sketch-based set-intersection estimation, and explain how message passing and subgraph GNNs produce link predictions.

2 Preliminaries

  • A k-hop subgraph enclosing link (u, v) contains the union of the k-hop neighborhoods of both endpoints; node and derived structural features are denoted X_uv and Z_uv.The graph is defined as undirected in the stated notation.Source: §2 · p. 3
  • HyperLogLog estimates union cardinality and MinHash estimates Jaccard similarity, and their sketches support permutation-invariant set operations with costs controlled by parameter p.The sketch operations' running times are constant with respect to set size.Source: §2 · p. 3
  • MPNNs use local permutation-invariant aggregation, while their expressiveness is bounded by the Weisfeiler-Leman test and automorphic nodes receive identical representations.A node-representation readout can therefore assign equal probabilities to links whose endpoints are automorphic, even at different distances.Source: §2 · p. 3
  • SGNNs classify links from enclosing subgraphs, augmenting node features with structural features before GNN propagation and graph-level readout.Their subgraph workflow is extraction, feature augmentation, propagation, and link prediction.Source: §2 · pp. 3–4
  • Figure 2 distinguishes structure features from LP readouts that pool all nodes in the enclosing subgraph or only the two endpoints.The text reports better predictive performance for endpoint edge pooling than for sum or mean pooling over all subgraph nodes.Source: Figure 2 · p. 4

3 Analyzing Subgraph Methods for Link Prediction

The analysis finds that SGNN link prediction depends strongly on structural information and link-level readout, while subgraph feature propagation can often be simplified. It also identifies per-link subgraph and feature computation as an efficiency burden.

3 Analyzing Subgraph Methods for Link Prediction

  • An MLP on structure-feature counts outperforms SGNN propagation on two of three datasets, while preprocessing node features can nearly recover SGNN propagation performance.Figure 4 contrasts count-only MLP performance with raw versus pre-propagated node features.Source: Figure 4 · p. 5, §3 · p. 5
  • DRNL features must be computed for each link at training and inference, requiring two subgraph traversals and up to O(|E|) time on complex networks.For regular graphs, the stated complexity is O(deg^k).Source: §3 · p. 5
  • Most DRNL structure-feature importance is concentrated at low distances.Importance is assessed from logistic-regression weights, grouped by maximum distance and normalized.Source: Figure 3 · p. 4, §3 · p. 5
  • Edge pooling outperforms mean or sum pooling over subgraph nodes across all datasets.The authors note that edge pooling can also be formulated without explicit subgraph construction.Source: §3 · p. 5
  • Structure features produce very large performance improvements across all datasets.The analysis identifies these features as central to the predictive value of SGNNs.Source: §3 · pp. 5–6

4 Link Prediction with Subgraph Sketching

The paper replaces explicit link subgraphs with node-wise sketches that estimate structural features for link prediction. ELPH uses these sketches in a full-graph MPNN, with a proof of greater expressive power than MPNNs.

4 Link Prediction with Subgraph Sketching

  • ELPH uses node-wise subgraph sketches to approximate structure features such as DE and DRNL label counts, avoiding explicit subgraph construction.The paper motivates these features as capturing salient link patterns.Source: §4 · p. 6, §4.2 · p. 7

4.1 Structure Features Counts

  • For receptive field k, the method uses k^2 A counts and 2k B counts, totaling k(k + 2) structure-count features.The counts can be computed without constructing the whole subgraph for each edge.Source: §4.1 · p. 6
  • HyperLogLog and MinHash sketches approximate neighborhood cardinalities and intersections, with speed-accuracy tradeoffs controlled by sketch parameters independently of graph size.The resulting estimates are used for the structure-feature counts.Source: §4.1 · pp. 6–7

4 Link Prediction with Subgraph Sketching

  • Precomputing SGNN subgraphs does not remove the memory burden: their combined size can exceed available memory on moderately sized datasets.This is the stated constraint on that alternative, rather than a limitation specific to ELPH.Source: §4.1 · p. 6

4.2 Efficient Link Prediction with Hashes (ELPH)

  • ELPH augments MPNN messages with subgraph sketches and uses a link-level readout, achieving greater expressiveness at the same asymptotic complexity.Its nodewise sketch computation and edge features remove the need to construct subgraphs explicitly; the paper describes the resulting model as as efficient as GCN.Source: §4.2 · p. 7
  • With exact structure-count estimates, ELPH does not suffer from the automorphic node problem and is strictly more powerful than MPNNs.The proof establishes that ELPH can distinguish some links that MPNNs cannot.Source: §4.2 · p. 8

5 Scaling ELPH with Preprocessing (BUDDY)

BUDDY addresses ELPH’s GPU-memory boundary by precomputing node features and sketches, then scoring links with an MLP over edge features. Its preprocessing avoids keeping the full dataset in GPU memory, but its non-parametric propagation does not subsume MPNNs.

5 Scaling ELPH with Preprocessing (BUDDY)

  • BUDDY avoids requiring the full dataset in GPU memory by precomputing node features and sketches, then scoring links with an MLP over edge features.Cached edge features can further reduce repeated computation when the same edge sets are used.Source: §5 · p. 8
  • Fixed node-feature propagation and sketch generation can be performed once during preprocessing, using efficient sparse scatter operations.The fixed propagation almost recovers the performance of learnable SGNN propagation.Source: §5 · p. 8
  • BUDDY avoids the automorphic node problem, but its non-parametric feature propagation does not subsume MPNNs.The authors cannot exclude links that MPNNs distinguish but BUDDY does not.Source: §5 · p. 9

6 Related Work

Related work spans subgraph-based link prediction, distance-based neural methods, heuristic learning, and scalable embeddings. Table 1 frames time complexity in terms of graph size, feature dimension, propagation hops, and sketch size.

6 Related Work

  • WLNM introduced subgraph link prediction; SEAL [48] replaced its MLP with a GNN, and subsequent work developed labeling, subgraph-generation, and application variants.DRNL originated in SEAL; other labeling schemes and equivariant positional encodings followed [21] [45] [38].Source: §6 · p. 9
  • NBFNet [53] models single-source distances with Bellman-Ford iterations and shares graph augmentation across destinations, making it faster amortized than SGNNs but memory-intensive.The passage characterizes its memory consumption as impractically high.Source: §6 · p. 9
  • Neo-GNN [46] learns heuristic functions with edgewise and nodewise networks, while #GNN [41] hashes node features and requires set-membership interpretations.Source: §6 · p. 9
  • TransE, DistMult [43], and complEx [37] are embedding-based link predictors; decoupled feature propagation appeared in SGC [40] and was scaled in SIGN [32].Source: §6 · p. 9
  • Table 1 summarizes time complexity using N and E for node and edge counts, d for feature dimension, k for propagation hops, and h for sketch size.Source: Table 1 · p. 9

7 Experiments

Across the link-prediction benchmarks, ELPH or BUDDY leads on five of seven datasets, with SEAL the closest competitor. ELPH runs out of memory on the two largest datasets, where BUDDY performs extremely well; their results show no clear overall winner when both are evaluated.

7 Experiments

  • Evaluation covers Cora, Citeseer, Pubmed, and OGB link-prediction datasets, comparing heuristics, GCN, GraphSAGE, Neo-GNN, SEAL, and NBFNet.Source: §7 · p. 9
  • ELPH or BUDDY achieves the best performance on five of seven datasets; SEAL is closest, and BUDDY performs extremely well on the two largest datasets where ELPH runs out of memory.Source: §7 · p. 9, Table 2 · p. 10
  • ELPH outperforms BUDDY on three of the five datasets where both have results, but there is no clear overall performance winner.Source: §7 · p. 10
  • Wall-time reporting defines training time as one epoch and inference time as the full test set.Source: Table 3 · p. 10

8 Conclusion

The conclusion presents a model derived from analyzing existing state-of-the-art methods, with improved time and space complexity and predictive performance on standard benchmarks. The preliminaries explain how graph symmetries create an automorphic-node problem for link prediction.

8 Conclusion

  • The work is limited to undirected graphs or directed graphs preprocessed as undirected, leaving directed and temporal or dynamic extensions for future work.The authors also identify graph curvature as a direction for future investigation.Source: §8 · p. 10

A.1 Preliminaries

  • Graph automorphisms preserve adjacency, and their orbits group nodes that share a structural role.Nodes in the same orbit are described as symmetric or automorphic.Source: §A.1 · p. 14
  • Automorphic nodes can form non-automorphic pairs, so node-level symmetry can force identical representations for links with different characteristics.The text gives shortest-path distance and common-neighbor count as examples of differing pair characteristics.Source: §A.1 · p. 15
  • Link discrimination measures whether a model assigns distinct representations to two node pairs, and expressiveness compares model families by that ability.A family is strictly more expressive when it can distinguish a pair that the other family cannot.Source: §A.1 · p. 15

A.2 Deferred theoretical results and proofs

The proofs establish ELPH’s greater expressiveness than MPNNs and show that both ELPH and BUDDY can distinguish certain pairs that MPNNs cannot. They also state estimator assumptions, a boundary on BUDDY’s expressiveness, and dataset subgraph-size challenges.

A.2 Deferred theoretical results and proofs

  • With exact cardinality estimates, one ELPH layer distinguishes the C6 pairs (v0, v2) and (v0, v3): their A and B features are (1, 1) and (0, 2), respectively.The graph’s nodes all belong to one orbit, while the two pairs belong to distinct orbits.Source: §A.2 · p. 16, Figure 7 · p. 16, Table 5 · p. 17
  • BUDDY does not suffer from the automorphic node problem when its cardinality estimates are exact, because its readout can reproduce ELPH’s distinguishing structural features.The proof applies this to the same C6 node-pairs.Source: §A.2 · p. 16
  • The estimator-based results assume exact A estimates; the authors say sufficiently large samples could preserve them when MinHash and HyperLogLog estimators are unbiased.They cite [10] for examples of unbiased cardinality estimators.Source: §A.2 · p. 16
  • ELPH is more expressive than MPNNs because it can simulate an MPNN and also distinguish some pairs no MPNN can distinguish.The proof’s counterexample uses non-automorphic C6 node-pairs whose constituent nodes are automorphic.Source: §A.2 · pp. 17–18
  • BUDDY does not generally subsume MPNNs, so the proofs do not rule out pairs that MPNNs distinguish but BUDDY does not.Source: §A.2 · p. 18

B.1 Datasets and Their Properties

  • Across the experimental datasets, sampled subgraphs grow faster than regular-graph scaling suggests and have highly variable sizes; dense DDI is the stated exception for two-hop subgraphs.For DDI, most two-hop subgraphs include almost every node.Source: §B.1 · p. 18

B.2 Experimental Setup

The experimental setup specifies graph splits, hyperparameter selection, space-complexity comparisons, and implementation details. It uses validation-selected hyperparameters and reports test results from a test set used once.

B.2 Experimental Setup

  • Training message-passing links equal supervision links, while validation and test supervision links are held out from training.Planetoid and ogbl-collab test-time message-passing edges also include validation supervision edges; OGB uses fixed splits and Planetoid random 70-10-20 splits.Source: §B.2 · p. 18

B.3 Hyperparameters

  • Hyperparameters were selected by random search for highest validation accuracy, and the test set was used only once.Search ranges covered hidden dimension, learning rate, dropout, layers, and weight decay.Source: §B.3 · pp. 18–19

B.4 Space Complexity

  • ELPH’s space complexity adds node sketches to GCN’s representation-learning storage, while BUDDY has the same space complexity as ELPH and precomputes node representations.BUDDY caches cardinality estimates and does not use weight matrices in preprocessing.Source: §B.4 · p. 19

B.5 Implementation Details

  • The implementation uses PyTorch and PyTorch Geometric, with experiments run on AWS p2 or p3 machines.The machines used Tesla K80 or Tesla V100 GPUs.Source: §B.5 · p. 19

C.1 Heuristic Methods for Link Prediction

The appendices describe link-prediction heuristics, sketching methods, structural labels and features, and runtime and ablation analyses. They explain why subgraph growth and redundancy motivate sketch-based alternatives.

C.1 Heuristic Methods for Link Prediction

  • Neighborhood-similarity heuristics include Common Neighbors, while path-based methods such as Katz and Personalized PageRank use multiple paths or random-walk landing probabilities.Adamic-Adar and Resource Allocation are described as second-order heuristics that penalize neighbors by degree.Source: §C.1 · pp. 19–20

C.2 Subgraph Generation Complexity

  • In complex networks, subgraph generation can scale with the number of graph edges, and sampled subgraphs in the benchmark datasets greatly exceed average degree with large variance.The text relates this to power-law degree distributions with 2 < γ < 3.Source: §C.2 · p. 20

C.3 Subgraph Sketches

  • HyperLogLog sketches estimate set-union cardinality, while MinHash estimates Jaccard similarity; together they can estimate set intersections.HyperLogLog merges sketches by elementwise maxima; MinHash estimates similarity from matching sketch entries.Source: §C.3 · p. 20, §C.3.1 · p. 20, §C.3.2 · p. 21

C.4 Labeling Schemes

  • DRNL labels nodes by distances to both link endpoints while masking the opposite endpoint and its edges, which improves expressiveness for positive edges.The masking can make distances large even in two-hop graphs.Source: §C.4 · p. 21

C.5 Substructure Counting and Curvature

  • The structural features count local substructures, including triangles and five-cycles, but cannot distinguish some different motifs such as a four-clique and two triangles.The authors note a planned investigation of links between link prediction and discrete Ricci curvature.Source: §C.5 · p. 23

C.6 Example Structure Features

  • The example computes A67[2, 1] as the intersection of node 6’s two-hop neighbors and node 7’s one-hop neighbors, yielding one shared node.The shared node is node 2; B67[2] counts one two-hop neighbor of node 6 absent from node 7’s k-hop neighbors.Source: §C.6 · pp. 23–24

D.1 Runtimes and Discussion

  • Table 7 compares method wall times against SEAL in dynamic and static modes, with one-epoch training and full-test-set inference times; OGB SEAL values are estimates.The experiments used a single Tesla K80 GPU for the reported runtime tables.Source: Table 7 · p. 24, §D.1 · p. 24

D.2 Ablation Studies

  • The ablations examine MinHash permutation and HyperLogLog precision settings, and remove node or structure features from BUDDY.Source: §D.2 · p. 24, Figure 9 · p. 25

D.2.1 Hashing Parameter Ablation

The appendix examines hashing-parameter and hop-count ablations and relates minhash permutation count to hash-generation runtime.

D.2.1 Hashing Parameter Ablation

  • The hashing-parameter ablation covers Planetoid datasets, including the hyperloglog p parameter.Source: Figure 10 · p. 25, §D.2.1 · p. 25
  • A separate ablation examines the number of hops.Source: Figure 11 · p. 26
  • Figure 12 plots hash-generation runtime against the number of minhash permutations.Source: Figure 12 · p. 26, §D.2.1 · p. 26
  • Hyperloglog sketch generation is negligible compared with minhash sketch generation, and parameter effects on runtime are appreciable only during hash preprocessing.At training and inference, generating structure features from hashes is negligible compared with an MLP forward pass.Source: §D.2.1 · p. 26

D.2.2 Feature Ablation

Feature ablations report dataset-dependent contributions from structure and node features; accompanying material gives BUDDY's procedure, learning curves, and a deployment-bias caution.

D.2.2 Feature Ablation

  • Structure features dominate BUDDY performance on Collab and PPA, contribute nearly as much as node features on Citation, and add smaller gains on Planetoid datasets.Combining node and structure features also reduces run-to-run variance on Planetoid and Collab; DDI has no node features, and PPA's one-hot species labels are not used.Source: §D.2.2 · p. 26, Table 9 · p. 27

E Full BUDDY Algorithm

  • BUDDY preprocessing initializes MinHash and hyperloglog sketches and propagates node features; prediction combines pairwise structure features and propagated endpoint features in an MLP.Source: E Full BUDDY Algorithm · p. 27

F Learning curves

G Societal Impact

  • The paper studies graph-structured datasets and methods rather than applications; recommendation-system deployment would require checks against amplifying existing biases.The authors report no evidence that their method enhances biases.Source: G Societal Impact · p. 29
Loading 2209.15486v3…