[Paper Review] Neural Bellman-Ford Networks: A General Graph Neural Network Framework for Link Prediction

Zhaocheng Zhu, Zuobai Zhang+2 more|arXiv (Cornell University)|Jun 13, 2021

Advanced Graph Neural NetworksComputer Science75 references112 citations

TL;DR

NBFNet is a neural framework that learns a generalized Bellman-Ford path formulation for link prediction, enabling inductive, interpretable, and scalable predictions on both knowledge graphs and homogeneous graphs by using learnable Indicator, Message, and Aggregate components.

ABSTRACT

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.

Motivation & Objective

• Motivate a general path-based representation learning framework for link prediction that combines interpretability of traditional path metrics with neural network capacity.
• Define a generalized path formulation where a pair representation sums over all paths and each path representation multiplies edge representations.
• Introduce NBFNet, a neural parameterization of the generalized Bellman-Ford solver using Indicator, Message, and Aggregate functions.
• Demonstrate that NBFNet achieves strong performance in both transductive and inductive settings on homogeneous graphs and knowledge graphs, with competitive efficiency and interpretability.

Proposed method

• Represent a pair of nodes h_q(u,v) as the generalized sum of all path representations between u and v, with each path represented as the generalized product of edge representations along the path.
• Solve the path formulation efficiently via a generalized Bellman-Ford algorithm under a semiring framework with summation and multiplication operators.
• Parameterize the generalized Bellman-Ford algorithm with three neural components: Indicator (boundary condition), Message (multiplication operator), and Aggregate (summation operator).
• Edge representations w_q(x,r,v) are learned and can be related to relational operators in knowledge graph embeddings; the model supports inductive generalization.
• Train using negative sampling (PCA) with a binary prediction objective to score triplets in knowledge graphs and edges in homogeneous graphs.
• Achieve amortized inference time of O(|E| d / |V| + d^2) per group of triplets, enabling scalable predictions.

Experimental results

Research questions

• RQ1 Can a path-based representation framework capture and unify traditional path metrics and modern graph neural networks for link prediction?
• RQ2 Is it possible to perform accurate link prediction in both inductive and transductive settings on knowledge graphs and homogeneous graphs using a learnable Bellman-Ford-based approach?
• RQ3 Do neuralized operators (Indicator, Message, Aggregate) improve performance and interpretability compared to handcrafted operators in path-based formulations?

Key findings

• NBFNet significantly outperforms state-of-the-art methods on knowledge graph completion across datasets, achieving a notable average relative gain in HITS@1 compared with the best path-based method.
• Compared with embedding-based methods, NBFNet attains substantial gains (e.g., 18% average relative improvement in HITS@1) while using far fewer parameters (around 3M vs. 30M for TransE on FB15k-237).
• NBFNet delivers strong results in homogeneous graph link prediction, outperforming several baselines on Cora and PubMed, and remaining competitive on CiteSeer despite its sparsity.
• In inductive relation prediction, NBFNet achieves the best results across all inductive splits, with meaningful improvements over prior methods (e.g., averaging 22% relative gain in HITS@10 over GraIL).
• Ablation studies show that advanced edge representations (e.g., RotatE, DistMult) and the learnable Aggregate (PNA) yield noticeable gains, and deeper architectures improve performance up to a saturation around 6 layers.
• The framework provides path-level interpretability by allowing extraction of top contributing paths for predictions.