PaperlayerView exact revision on arXiv

Source-linked AI summary

Neural Bellman-Ford Networks: A General Graph Neural Network Framework for Link Prediction

Zhaocheng Zhu, Zuobai Zhang, Louis-Pascal Xhonneux, Jian Tang

arXiv:2106.06935v4cs.LG

TL;DR

Link prediction needs methods that combine the inductive generalization and interpretability of path-based approaches with the capacity of GNNs. The paper introduces NBFNet, which learns path-aggregation operators through generalized Bellman-Ford computation, and reports strong results across homogeneous and knowledge-graph settings. Its main limitations are the lack of a theoretical guarantee under neural operators and evaluation restricted to simple edge prediction.

  • Problem

    Existing path-based methods are interpretable, inductive, and scalable but handcrafted, whereas GNN approaches can lack inductivity and interpretability or compromise scalability.

  • Method

    NBFNet represents node pairs through generalized path aggregation and learns the generalized Bellman-Ford operators with INDICATOR, MESSAGE, and AGGREGATE functions.

  • Results

    NBFNet outperforms existing methods across homogeneous-graph and knowledge-graph link prediction in both transductive and inductive settings.

  • Takeaways & Limitations

    The framework covers traditional path-based methods and supports interpretable path visualization while retaining learned neural operators.

  • Takeaways & Limitations

    NBFNet lacks a theoretical guarantee for the relaxation of the semiring assumption and has only been verified on simple edge prediction.

Abstract

from arXiv · show

Link prediction is a very fundamental task on graphs. Inspired by traditional path-based methods, in this paper we propose a general and flexible representation learning framework based on paths for link prediction. Specifically, we define the representation of a pair of nodes as the generalized sum of all path representations, with each path representation as the generalized product of the edge representations in the path. Motivated by the Bellman-Ford algorithm for solving the shortest path problem, we show that the proposed path formulation can be efficiently solved by the generalized Bellman-Ford algorithm. To further improve the capacity of the path formulation, we propose the Neural Bellman-Ford Network (NBFNet), a general graph neural network framework that solves the path formulation with learned operators in the generalized Bellman-Ford algorithm. The NBFNet parameterizes the generalized Bellman-Ford algorithm with 3 neural components, namely INDICATOR, MESSAGE and AGGREGATE functions, which corresponds to the boundary condition, multiplication operator, and summation operator respectively. The NBFNet is very general, covers many traditional path-based methods, and can be applied to both homogeneous graphs and multi-relational graphs (e.g., knowledge graphs) in both transductive and inductive settings. Experiments on both homogeneous graphs and knowledge graphs show that the proposed NBFNet outperforms existing methods by a large margin in both transductive and inductive settings, achieving new state-of-the-art results.

1 Introduction

Link prediction is important but existing approaches trade off handcrafted simplicity, neural capacity, inductive generalization, interpretability, and scalability. The paper proposes NBFNet to learn path-based operators while retaining broad applicability.

  • Link prediction supports applications including recommender systems, knowledge graph completion, and drug repurposing.
  • Traditional path-based methods are interpretable, inductive, and scalable but rely on handcrafted metrics that may not be optimal.
  • GNN-based methods provide high expressiveness but are often transductive and lack interpretability, while inductive subgraph methods compromise scalability.
  • The proposed path formulation sums generalized products of edge representations and can be efficiently solved by a generalized Bellman-Ford algorithm.
  • NBFNet learns the generalized Bellman-Ford operators through INDICATOR, MESSAGE, and AGGREGATE functions.
  • 18% average relative gain in knowledge graph completion HITS@1 and 22% in inductive relation prediction HITS@10 were reported against existing state-of-the-art methods.

2 Related Work

Existing link-prediction research spans path-based methods, embedding methods, and graph neural networks, each using a distinct representation strategy.

  • Link-prediction research is generally organized into path-based, embedding, and graph-neural-network paradigms.
  • Path-based Methods: Path-based methods score node pairs using paths, random walks, shortest paths, relational paths, or learned symbolic rules.
  • Embedding Methods: Embedding methods learn distributed node and edge representations while preserving graph structure, with separate examples for homogeneous and knowledge graphs.
  • Graph Neural Networks: Common GNN link-prediction frameworks encode nodes and decode node pairs, while subgraph-based approaches explicitly encode neighborhoods around node pairs.

3 Methodology

The methodology represents node pairs through generalized path aggregation and computes these representations with Bellman-Ford-style message passing. NBFNet replaces handcrafted operators with neural functions for source-conditioned pair representations.

  • 3.1 Path Formulation for Link Prediction: The path formulation represents a node pair as a generalized sum of path representations, each formed by a generalized product of ordered edge representations.
  • 3.1 Path Formulation for Link Prediction: The formulation can model traditional link-prediction methods and graph algorithms, including Katz index, personalized PageRank, graph distance, widest path, and most reliable path.
  • 3.1 Path Formulation for Link Prediction: The generalized Bellman-Ford algorithm avoids enumerating exponentially many paths by exploiting multiplication's distributivity over summation.
  • 3.2 Neural Bellman-Ford Networks: NBFNet learns the operators with INDICATOR, MESSAGE, and permutation-invariant AGGREGATE functions, relaxing the semiring assumption.
  • 3.2 Neural Bellman-Ford Networks: NBFNet initializes the source node and propagates source-conditioned representations, then reads out pair representations at target nodes.
  • 3.2 Neural Bellman-Ford Networks: With a small constant number of iterations, NBFNet has time complexity O(|E|d + |V|d2).

4 Experiment

NBFNet is evaluated across transductive and inductive link-prediction settings, where it generally outperforms strong baselines. Ablations examine learned operators, depth, relation categories, and path-based interpretations.

  • Experiment Setup: NBFNet is evaluated on knowledge graph completion, homogeneous graph link prediction, and inductive relation prediction using established datasets and splits.The evaluation includes FB15k-237, WN18RR, Cora, Citeseer, and PubMed.
  • Main Results: NBFNet significantly outperforms existing methods on both knowledge graph completion datasets and improves HITS@1 by 21% over DRUM and 18% over LowFER on average.NBFNet uses 3M parameters on FB15k-237, compared with 30M for TransE.
  • Main Results: NBFNet performs best on Cora and PubMed and competitively on sparse Citeseer, despite not using the datasets’ node features.Combining node features with structural representations is left for future work.
  • Main Results: NBFNet achieves the best inductive relation-prediction results on all splits, with a 22% average relative gain in HITS@10 over GraIL.Unlike GraIL, which usually encodes at most a 2-hop subgraph, NBFNet can efficiently explore longer paths.
  • Ablation Study: PNA aggregation improves over sum, mean, and max, while performance increases with depth and saturates after six layers.The results support learning more powerful aggregation functions and suggest paths up to length six may suffice for link prediction.
  • Path Interpretations of Predictions: Path interpretations identify influential routes through edge importance, including soft logical entailment, analogical reasoning, and longer paths.Table 7 visualizes the top-two path interpretations and their weights for each query triplet.

5 Discussion and Conclusion

NBFNet is evaluated across graph settings and achieves strong performance, while the paper identifies theoretical, task-scope, and societal limitations.

  • The generalized Bellman-Ford semiring assumption does not hold for NBFNet because neural-network activations are nonlinear, leaving the relaxation without a theoretical loss guarantee.
  • NBFNet is verified only on simple edge prediction, not complex logical queries involving conjunctions or disjunctions.
  • NBFNet may encode training-data bias and produce stereotyped predictions in social or e-commerce applications.
  • Powerful link prediction models could augment harmful activities such as spamming, phishing, and social engineering.
  • NBFNet outperforms existing state-of-the-art methods by a large margin across homogeneous and knowledge graphs in transductive and inductive settings.

A Path Formulations for Traditional Methods

The path formulation represents node pairs by generalized sums of generalized products over paths, unifying several traditional methods and enabling generalized Bellman-Ford computation.

  • A Path Formulations for Traditional Methods: The path formulation models a pair representation as a generalized sum of path representations, with each path formed by an ordered generalized product of edge representations.
  • A.1 Katz Index: Katz index is a path formulation using ⊕= +, ⊗= ×, and damped edge representations βwe.
  • A.2 Personalized PageRank: Personalized PageRank is a path formulation using ⊕= +, ⊗= ×, and edge representations given by random-walk probabilities scaled by α.
  • A.3 Graph Distance: Graph distance is a path formulation using ⊕= min, ⊗= +, and edge representations equal to edge lengths.
  • A.4 Widest Path: Widest path is a path formulation using ⊕= max and ⊗= min, selecting the path with the largest minimum edge weight.
  • A.5 Most Reliable Path: Most reliable path is a path formulation using ⊕= max and ⊗= ×, with edge representations equal to edge probabilities.
  • B Generalized Bellman-Ford Algorithm: When the operators form a semiring, generalized Bellman-Ford efficiently solves the path formulation and aggregates path representations through iterative updates.
  • B Generalized Bellman-Ford Algorithm: After t iterations, intermediate representations aggregate paths up to t edges; simple paths require at most |V| iterations, while practice often converges by T = 6.

B.3 Traditional Methods

Traditional path-based methods satisfy semiring structures, so Katz index, personalized PageRank, graph distance, widest path, and most reliable path can use generalized Bellman-Ford.

  • Katz index, personalized PageRank, graph distance, widest path, and most reliable path can all be solved via generalized Bellman-Ford.
  • The proof strategy reduces solvability to showing that each method’s operator pair satisfies the semiring requirements.
  • Katz index and personalized PageRank use natural summation + and multiplication ×, which satisfy a semiring.
  • Graph distance uses min and + with identities +∞ and 0, satisfying the required semiring properties.
  • Widest path uses max and min with identities −∞ and +∞, and satisfies the semiring properties through distributivity.
  • Most reliable path uses max and × over non-negative probabilities, with identities 0 and 1 and the required distributive properties.

C.1 NBFNet

NBFNet has relatively low inference complexity and parameter growth that depends on relations rather than entities. With a small iteration count, its complexity becomes linear in graph size for fixed representation dimensions.

  • Time complexity: O(T(|E|d + |V|d2)) is the time complexity of one NBFNet run.The bound accounts for MESSAGE and AGGREGATE computation across T iterations.
  • Time complexity: A small constant T works well in practice, reducing NBFNet complexity to O(|E|d + |V|d2).The paper reports that T = 6 can suffice for convergence in link prediction.
  • Inference: NBFNet amortizes inference over triplets sharing the same condition by executing the algorithm once for each group.This contrasts with running message passing separately for every candidate triplet.
  • Comparison: RGCN-based inference has O(T|E|d2) complexity, while GraIL and SEAL have amortized complexity O(|E|d2).The comparison attributes the difference to linear transformations in RGCN message functions and local-subgraph processing in GraIL and SEAL.
  • Comparison: Restricting GraIL to very small local subgraphs can severely harm link-prediction performance, while moderate-radius neighborhoods may still be expensive on small-world graphs.The limitation arises because moderate-radius neighborhoods can contain many nodes and edges.
  • Parameters: NBFNet requires 3M parameters on FB15k-237 versus 30M for TransE, with parameter count scaling linearly with the number of relations.Its parameter count does not depend on the number of entities or edges.

E Statistics of Datasets

The paper uses standard transductive and inductive datasets for knowledge-graph and homogeneous-graph link prediction. Homogeneous graphs are undirected and split into training, validation, and test edges at an 85:5:10 ratio.

  • Dataset coverage: The datasets cover knowledge graph completion, homogeneous graph link prediction, and inductive relation prediction.Dataset statistics are reported separately for the two transductive settings and inductive relation prediction.
  • Splits: Knowledge graphs use standard transductive and inductive splits, while homogeneous graphs use an 85:5:10 train/validation/test edge split.All homogeneous graphs used in the paper are undirected.
  • Splits: Inductive relation prediction uses a transductive validation set sharing fact triplets with training for hyperparameter tuning.This reproduces the validation protocol of the referenced prior work.

F Implementation Details

NBFNet uses dataset-specific hyperparameter selection within a generally robust configuration, standard graph augmentations, normalization and skip connections, and fused message passing to reduce memory use.

  • Hyperparameters: The model configuration is reported as robust across datasets, with learning hyperparameters tuned separately using validation performance.For FB15k-237 and WN18RR, the same hyperparameters are used in transductive and inductive settings.
  • Data augmentation: Knowledge-graph triplets are augmented with flipped relations, while homogeneous-graph nodes receive self-loops.These augmentations follow previous work in each setting.
  • Architecture: Layer normalization follows each AGGREGATE function, and shortcut connections, an MLP feed-forward network, and ReLU activations are used.For undirected graphs, pair representations are symmetrized by summing both directions.
  • Training: Training uses edge dropout on query-node connections, validation-based checkpoint selection, and criteria of MRR for knowledge graphs or AUROC for homogeneous graphs.Experiments run with standard data parallelism on 4 Tesla V100 GPUs.
  • Efficiency: Fused MESSAGE and AGGREGATE operators reduce NBFNet memory complexity from O(|E|d) to O(|V|d).The implementation avoids explicitly creating all messages.
  • Evaluation: The paper reports wall time for NBFNet across datasets and settings, including all four inductive split versions.The corresponding results are summarized in Table 13.
  • Large-scale evaluation: Large-scale evaluation includes ogbl-biokg and WikiKG90M using OGB link-property-prediction protocols and MRR against 1,000 negative entities.These datasets extend evaluation to substantially larger knowledge graphs.

G.1 Results on ogbl-biokg

On large-scale knowledge graphs, NBFNet is evaluated with neighborhood sampling and reports strong validation or leaderboard performance. The sampling procedure preserves paths up to a specified length while limiting expansion of high-degree entities.

  • Results on ogbl-biokg: NBFNet achieves the best reported result on ogbl-biokg with much fewer parameters than compared embedding methods.The dataset contains 93,773 entities, 51 relations, and 5,088,434 triplets.
  • Results on WikiKG90M: WikiKG90M contains 87,143,637 entities, 1,315 relations, and 504,220,369 triplets.It is an extremely large, general-purpose knowledge graph used in the OGB large-scale challenge.
  • Sampling: Bidirectional BFS samples k-hop neighborhoods around the head and candidate tails, then unions them into a sampled graph.The figure illustrates the neighborhoods as dashed circles and ignores edge direction during BFS sampling.
  • Sampling: The sampled graph guarantees every path of length at most 2k between the head and any tail candidate is present.This guarantee follows from combining the sampled neighborhoods around the query entities.
  • Sampling: Neighbor expansion is downsampled to m incoming and m outgoing entities per visited entity to handle high-degree nodes.This controls the size of the sampled graph during BFS.
  • Results on WikiKG90M: On WikiKG90M validation, the best single model uses k = 2 and m = 100; training improves when negatives decrease from 1,000 to 20 and batch size increases from 4 to 64.The latter choice creates a distribution shift between sampled graphs used for training and validation.

H Ablation Study

The ablation study reports results for different MESSAGE and AGGREGATE functions and for varying numbers of NBFNet layers.

  • Table 15 reports full results for different MESSAGE and AGGREGATE functions.
  • Table 17 reports full results for NBFNet configurations with different numbers of layers.
  • The ablations examine both learned function choices and network depth.
Loading 2106.06935v4…