SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xingyue Huang
Oxford, UK.
xingyue.huang@cs.ox.ac.uk
Ismail Ilkan Ceylan
ismail.ceylan@cs.ox.ac.uk
Miguel Romero
Universidad Catolica de Chile
& CENIA Chile
mgromero@uc.cl
Pablo Barceló
Inst. for Math. and Comp. Eng.
&IMFD Chile & CENIA Chile
pbarcelo@uc.cl
Graph neural networks are prominent models for representation learning over graph-structured data. While the capabilities and limitations of these models are well-understood for simple graphs, our understanding remains incomplete in the context of knowledge graphs. Our goal is to provide a systematic understanding of the landscape of graph neural networks for knowledge graphs pertaining to the prominent task of link prediction. Our analysis entails a unifying perspective on seemingly unrelated models and unlocks a series of other models. The expressive power of various models is characterized via a corresponding relational Weisfeiler-Leman algorithm. This analysis is extended to provide a precise logical characterization of the class of functions captured by a class of graph neural networks. The theoretical findings presented in this paper explain the benefits of some widely employed practical design choices, which are validated empirically.
Graph neural networks (GNNs) [27, 12] are prominent models for representation learning over graph-structured data, where the idea is to iteratively compute vector representations of nodes of an input graph through a series of invariant (resp., equivariant) transformations. While the landscape of GNNs is overwhelmingly rich, the vast majority of such models are instances of message passing neural networks [11] which are well-studied, leading to a theoretical understanding of their capabilities and limitations [34, 21]. In turn, our understanding is rather limited for GNN models dedicated to learning over knowledge graphs, which are applied in a wide range of domains.
To make our context precise, we first consider an extension of message passing neural networks with relation-specific message functions, which we call relational message passing neural networks. Two prominent examples of this framework are RGCN [28] and CompGCN [32], and their expressive power has recently been characterized through a dedicated relational Weisfeiler-Leman test [4].
While offering principled means for learning over knowledge graphs, the standard relational message passing framework is tailored for computing unary node representations and therefore models of this class are better suited for node-level tasks (e.g., node/entity classification). Actually, it is well-known that even a good node-level representation might not necessarily induce a good edge representation,
hinding the applicability of such models for the crucial task of link prediction [39]. This has led to the design of GNN architectures specifically tailored for link prediction over knowledge graphs [30, 40, 33, 18], for which our understanding remains limited.
The goal of this paper is to offer a theory of the capabilities and limitations of a class of relational GNN architectures which compute pairwise node representations to be utilized for link prediction. Although most such architectures can be seen to be subsumed by higher-order message passing neural networks that compute pairwise representations on nodes, the inherently quadratic behavior of the latter justifies local approximations which align better with models used in practice. Of particular interest to us is Neural Bellman-Ford Networks (NBFNets) [40], which define a message passing approach inspired by the Bellman-Ford algorithm. We argue in salient detail that the crucial insight of this approach is in leveraging the idea of computing conditional pairwise-node representations, which leads to more expressive models at a more reasonable computational cost, given its local nature.
Building on this fundamental aspect, we define conditional message passing neural networks, which extend traditional ones by a conditional message passing paradigm: every node representation is conditional on a source node and a query relation, which allows for computing pairwise node representations. This framework strictly contains NBFNets and allows for a systematic treatment of various other models. Through a careful study of this framework, we can explain the conceptual differences between different models along with their respective expressive power.
- We introduce conditional message passing neural networks which encode representations of nodes  $v$  conditioned on a (source) node  $u$  and a query relation  $q$ , yielding pairwise node representations. We discuss the model design space, including a discussion on different initialization regimes and the presence of global readout functions in each layer of the network.
- We define a relational Weisfeiler-Leman algorithm (building on similar works such as Barceló et al. [4]), and prove that conditional message passing neural networks can match the expressive power of this algorithm. This study reveals interesting insights about NBFNets, suggesting that their strong empirical performance is precisely due to the expressive power, which can be matched by other instances of this framework.
- Viewing conditional message passing neural networks (with or without global readouts) as classifiers over pairs of nodes, we give logical characterizations of their expressive power based on formulas definable in some binary variants of graded modal logic [8, 19]. This provides us with a declarative and well-studied formal counterpart of C-MPNNs.
- We conduct an experimental analysis to verify the impact of various model choices, particularly pertaining to initialization, history, message computation, and global readout functions. We also conduct both inductive and transductive experiments on various real-world datasets, empirically validating our theoretical findings.
2 Related work and motivation
Early GNNs for knowledge graphs are relational variations of message passing neural networks. A prototypical example is the RGCN architecture [28], which extends graph convolutional networks (GCNs) [17] with relation-specific message functions. CompGCN [32] and several other architectures [37] follow this line of work with differences in their aggregate, update, and message functions. These architectures encode unary node representations and typically rely on a pairwise decoder function to predict the likelihood of a link which is known to be suboptimal for link prediction [39]. There is a good understanding of the expressive power of these architectures [4]: we generalize these results in Section 3, since they form the basis for the rest of our results.
A different approach is given for single-relational graphs by SEAL [38], where the idea is to encode (labeled) subgraphs (instead of nodes). GraIL [30] extends this idea to knowledge graphs, and one important virtue of these models is that they are inductive even if there are no node features in the input graph. The idea is to use a form of labeling trick [39] based on pairwise shortest path distances in sampled subgraphs, but these architectures suffer from scalability issues. More recent inductive architectures integrate ideas from earlier path-based link prediction approaches [24, 14] into modern GNN architectures, resulting in proposals such as PathCon [33], Geodesic GNNs [18], and NBFNets [40]. Our study is very closely related to NBFNets which is inspired by the generalized
version of the Bellman-Ford algorithm for finding shortest paths. These architectures aggregate over relational paths by keeping track of conditional pairwise-node representations. While NBFNets can be intuitively seen as the neural counterpart of the generalized Bellman-Ford algorithm, they do not provably align with this algorithm since the "semiring assumption" is invalidated through the use of non-linearities (which is explicit in Zhu et al. [40]). This leaves open many questions regarding the capabilities and limitations of these architectures.
In this paper, we argue that the key insight behind architectures such as NBFNets is in locally computing pairwise representations through conditioning on a source node, and this has roots in earlier works, such as ID-GNNs [36]. To formally study this, we introduce conditional message passing neural networks as a strict generalization of NBFNets [40] and related models such as NeuralLP [35], or DRUM [25]. Through this abstraction, we theoretically study the properties of a large class of models in relation to local variants of relational Weisfeiler-Leman algorithms. Broadly, our study can be seen as the relational counterpart of the expressiveness studies conducted for GNNs [34, 21, 3], particularly related to higher-order GNNs [21], which align with higher-order dimensional variants of the WL test. Our characterization relies on local versions of higher-order WL tests [22], albeit not in a relational context. This can be seen as a continuation and generalization of the results given for relational message passing neural networks [4] to a broader class of models.
3.1 Knowledge graphs and invariants
Knowledge graphs. A knowledge graph is a tuple  $G = (V, E, R, c)$ , where  $V$  is a set of nodes,  $E \subseteq R \times V \times V$  is a set of labeled edges, or facts,  $R$  is the set of relation types and  $c: V \to D$  is a node coloring. When  $D = \mathbb{R}^d$ , we also say that  $c$  is a  $d$ -dimensional feature map, and typically use  $\pmb{x}$  instead of  $c$ . We write  $r(u, v)$  to denote a labeled edge, or a fact, where  $r \in R$  and  $u, v \in V$ . The neighborhood of a node  $v \in V$  relative to a relation  $r \in R$  is defined as  $\mathcal{N}_r(v) := \{u \mid r(u, v) \in E\}$ .
Graph invariants. We define  $k$ -ary graph invariants following the terminology of Grohe [13], for which we first define isomorphism over knowledge graphs. An isomorphism from a knowledge graph  $G = (V, E, R, c)$  to a knowledge graph  $G' = (V', E', R, c')$  is a bijection  $f: V \to V'$  such that  $c(v) = c'(f(v))$  for all  $v \in V$ , and  $r(u, v) \in E$  if and only if  $r(f(u), f(v)) \in E'$ , for all  $r \in R$  and  $u, v \in V$ . A 0-ary graph invariant is a function  $\xi$  defined on knowledge graphs such that  $\xi(G) = \xi(G')$  for all isomorphic knowledge graphs  $G$  and  $G'$ . For  $k \geq 1$ , a  $k$ -ary graph invariant is a function  $\xi$  that associates with each knowledge graph  $G = (V, E, R, c)$  a function  $\xi(G)$  defined on  $V^k$  such that for all knowledge graphs  $G$  and  $G'$ , all isomorphisms  $f$  from  $G$  to  $G'$ , and all  $k$ -tuples of nodes  $v \in V^k$ , it holds that  $\xi(G)(v) = \xi(G')(f(v))$ . If  $k = 1$ , this defines a node invariant, or unary invariant, and if  $k = 2$ , this defines a binary invariant, which is central to our study.
Refinements. A function  $\xi(G): V^k \to D$  refines a function  $\xi'(G): V^k \to D$ , denoted as  $\xi(G) \preceq \xi'(G)$ , if for all  $\pmb{v}$ ,  $\pmb{v}' \in V^k$ ,  $\xi(G)(\pmb{v}) = \xi(G)(\pmb{v}')$  implies  $\xi'(G)(\pmb{v}) = \xi'(G)(\pmb{v}')$ . We call such functions equivalent, denoted as  $\xi(G) \equiv \xi'(G)$ , if  $\xi(G) \preceq \xi'(G)$  and  $\xi'(G) \preceq \xi(G)$ . A  $k$ -ary graph invariant  $\xi$  refines a  $k$ -ary graph invariant  $\xi'$ , if  $\xi(G)$  refines  $\xi'(G)$  for all knowledge graphs  $G$ .
3.2 Relational message passing neural networks
We introduce relational message passing neural networks (R-MPNNs), which encompass several known models such as RGCN [28] and CompGCN [32]. The idea is to iteratively update the feature of a node  $v$  based on the different relation types  $r \in R$  and the features of the corresponding neighbors in  $\mathcal{N}_r(v)$ . In our most general model we also allow readout functions, that allow further updates to the feature of  $v$  by aggregating over the features of all nodes in the graph.
Let  $G = (V, E, R, \boldsymbol{x})$  be a knowledge graph, where  $\boldsymbol{x}$  is a feature map. An  $R$ -MPNN computes a sequence of feature maps  $\boldsymbol{h}^{(t)}: V \to \mathbb{R}^{d(t)}$ , for  $t \geq 0$ . For simplicity, we write  $\boldsymbol{h}_v^{(t)}$  instead of  $\boldsymbol{h}^{(t)}(v)$ . For each node  $v \in V$ , the representations  $\boldsymbol{h}_v^{(t)}$  are iteratively computed as:
$$
\begin{array}{l} \boldsymbol {h} _ {v} ^ {(0)} = \boldsymbol {x} _ {v} \\ \boldsymbol {h} _ {v} ^ {(t + 1)} = \operatorname {U P D} \left(\boldsymbol {h} _ {v} ^ {(f (t))}, \operatorname {A G G} \left(\left\{\left\{\mathbf {M S G} _ {r} \left(\boldsymbol {h} _ {w} ^ {(t)}\right) \right| w \in \mathcal {N} _ {r} (v), r \in R \right\}\right, \operatorname {R E A D} \left(\left\{\left\{\boldsymbol {h} _ {w} ^ {(t)} \mid w \in V \right\}\right)\right)\right), \\ \end{array}
$$
where  $\mathrm{UPD}$ ,  $\mathrm{AGG}$ ,  $\mathrm{READ}$ , and  $\mathrm{MSG}_r$  are differentiable update, aggregation, global readout, and relation-specific message functions, respectively, and  $f: \mathbb{N} \to \mathbb{N}$  is a history function $^1$ , which is always non-decreasing and satisfies  $f(t) \leq t$ . An R-MPNN has a fixed number of layers  $T \geq 0$ , and then, the final node representations are given by the map  $h^{(T)}: V \to \mathbb{R}^{d(T)}$ .
The use of a readout component in message passing is well-known [5] but its effect is not well-explored in a relational context. It is of interest to us since it has been shown that standard GNNs can capture a larger class of functions with global readout [3].
An R-MPNN can be viewed as an encoder function  $\mathsf{enc}$  that associates with each knowledge graph  $G$  a function  $\mathsf{enc}(G):V\to \mathbb{R}^{d(T)}$ , which defines a node invariant corresponding to  $\pmb{h}^{(T)}$ . The final representations can be used for node-level predictions. For link-level tasks, we use a binary decoder  $\mathsf{dec}_q:\mathbb{R}^{d(T)}\times \mathbb{R}^{d(T)}\to \mathbb{R}$ , which produces a score for the likelihood of the fact  $q(u,v)$ , for  $q\in R$ .
In Appendix A.1 we provide a useful characterization of the expressive power of R-MPNNs in terms of a relational variant of the Weisfeiler-Leman test [4]. This characterization is essential for the rest of the results that we present in the paper.
4 Conditional message passing neural networks
R-MPNNs have serious limitations for the task of link prediction [39], which has led to several proposals that compute pairwise representations directly. In contrast to the case of R-MPNNs, our understanding of these architectures is limited. In this section, we introduce the framework of conditional MPNNs that offers a natural framework for the systematic study of these architectures.
Let  $G = (V, E, R, \mathbf{x})$  be a knowledge graph, where  $\mathbf{x}$  is a feature map. A conditional message passing neural network (C-MPNN) iteratively computes pairwise representations, relative to a fixed query  $q \in R$  and a fixed node  $u \in V$ , as follows:
$$
\boldsymbol {h} _ {v | u, q} ^ {(0)} = \operatorname {I N I T} (u, v, q)
$$
$$
\boldsymbol {h} _ {v | u, q} ^ {(t + 1)} = \mathrm {U P D} (\boldsymbol {h} _ {v | u, q} ^ {f (t)}, \mathrm {A G G} (\{\{\mathrm {M S G} _ {r} (\boldsymbol {h} _ {w | u, q} ^ {(t)}, \boldsymbol {z} _ {q}) | w \in \mathcal {N} _ {r} (v), r \in R \}), \mathrm {R E A D} (\{\{\boldsymbol {h} _ {w | u, q} ^ {(t)} | w \in V \})),
$$
where  $\mathrm{INIT}$ ,  $\mathrm{UPD}$ ,  $\mathrm{AGG}$ ,  $\mathrm{READ}$ , and  $\mathrm{MSG}_r$  are differentiable initialization, update, aggregation, global readout, and relation-specific message functions, respectively, and  $f$  is the history function. We denote by  $\pmb{h}_q^{(t)}:V\times V\to \mathbb{R}^{d(t)}$  the function  $\pmb{h}_q^{(t)}(u,v)\coloneqq \pmb{h}_{v|u,q}^{(t)}$ , and denote  $\pmb{z}_q$  to be a learnable vector representing the query  $q\in R$ . A C-MPNN has a fixed number of layers  $T\geq 0$ , and the final pair representations are given by  $\pmb{h}_q^{(T)}$ . We sometimes write C-MPNNs without global readout to refer to the class of models which do not use a readout component.
Intuitively, C-MPNNs condition on a source node  $u$  in order to compute representations of  $(u,v)$  for all target nodes  $v$ . To further explain, we show a visualization of C-MPNNs and R-MPNNs in Figure 1 to demonstrate the differences in the forward pass. Contrary to the R-MPNN model where we carry out relational message passing first and rely on the binary decoder to compute a query fact  $q(u,v)$ , the C-MPNN model first initializes all node representations with the zero vector except the representation of the node  $u$  which is assigned a vector with non-zero entry. Following the initialization, we carry out relational message passing and decode the hidden state of the target node  $v$  to obtain the output, which yields the representation of  $v$  conditioned on  $u$ .
Observe that C-MPNNs compute binary invariants, provided that the initialization  $\mathrm{INIT}$  is a binary invariant. To ensure that the resulting model computes pairwise representations, we require  $\mathrm{INIT}(u,v,q)$  to be a nontrivial function in the sense that it needs to satisfy target node distinguishability: for all  $q\in R$  and  $v\neq u\in V$ , it holds that  $\mathrm{INIT}(u,u,q)\neq \mathrm{INIT}(u,v,q)$ .
This is very closely related to the Labeling Trick proposed by Zhang et al. [39], which is an initialization method aiming to differentiate a set  $\{u, v\}$  of target nodes from the remaining nodes in a graph. However, the Labeling Trick only applies when both the source  $u$  and target nodes  $v$  are labeled. Recent state-of-the-art models such as ID-GNN [36] and NBFNet [40] utilize a similar method, but only with the source node  $u$  labeled differently in initialization. Our definition captures precisely this, and we offer a theoretical analysis of the capabilities of the aforementioned architectures accordingly.
The typical choice is  $f(t) = t$ , but other options are considered in the literature, as we discuss later.
Figure 1: Visualization of R-MPNN and C-MPNN. The dashed arrow is the target query  $q(u, v)$ . Arrow colors indicate distinct relation types, while node colors indicate varying hidden states. R-MPNN considers a unary encoder and relies on a binary decoder, while C-MPNN first initializes binary representation based on the target query  $q(u, v)$ , and then uses a unary decoder.
One alternative is to directly learn pairwise representations following similar ideas to those of higher-order GNNs, but these algorithms are not scalable. Architectures such as NBFNets represent a trade-off between computational complexity and expressivity. The advantage of learning conditional representations  $h_{v|u,q}$  is to be able to learn such representations in parallel for all  $v \in V$ , amortizing the computational overhead; see Zhu et al. [40] for a discussion. We have also carried out a runtime analysis comparison among different classes of models in Appendix B.
4.1 Design space and basic model architectures
To specify a C-MPNN architecture, we need to specify the functions INIT, AGG,  $\mathrm{MSG}_r$ ,  $f$ , and READ. In the following, we consider three initialization functions  $\{\mathrm{INIT}^1, \mathrm{INIT}^2, \mathrm{INIT}^3\}$ , two aggregation functions  $\{\mathrm{sum}, \mathrm{PNA}\}$ , three message functions  $\{\mathrm{Msg}_r^1, \mathrm{Msg}_r^2, \mathrm{Msg}_r^3\}$ , two history functions  $\{f(t) = t, f(t) = 0\}$ , and either sum global readout or no readout term.
Initialization. We consider the following natural variants for initialization:
$$
\mathrm {I N I T} ^ {1} (u, v, q) = \mathbb {1} _ {u = v} * \mathbf {1}, \quad \mathrm {I N I T} ^ {2} (u, v, q) = \mathbb {1} _ {u = v} * \boldsymbol {z} _ {q}, \quad \mathrm {I N I T} ^ {3} (u, v, q) = \mathbb {1} _ {u = v} * (\boldsymbol {z} _ {q} + \epsilon_ {u}),
$$
where  $*$  represents element-wise multiplication, the function  $\mathbb{1}_{u = v}(v)$  is the indicator function which returns the all-ones vector  $\mathbf{1}$  if  $u = v$  and the all-zeros vector  $\mathbf{0}$  otherwise with corresponding size.
Clearly, both  $\mathrm{INIT}^1$  and  $\mathrm{INIT}^2$  satisfy target node distinguishability assumption if we assume  $z_{q}$  has no zero entry, where  $\mathrm{INIT}^2$ , in addition, allows query-specific initialization to be considered. Suppose we further relax the condition to be target node distinguishability in expectation. Then,  $\mathrm{INIT}^3$  can also distinguish between each conditioned node  $u$  given the same query vector  $z_{q}$  by adding an error vector  $\epsilon_{u}$  sampled from  $\mathcal{N}(0,1)$  to the conditioned node's initialization.
Aggregation. We consider sum aggregation and Principal Neighborhood Aggregation (PNA) [7].
Message. We consider the following variations of message functions:
$$
\operatorname {M S G} _ {r} ^ {1} \left(\boldsymbol {h} _ {w | u, q} ^ {(t)}, \boldsymbol {z} _ {q}\right) = \boldsymbol {h} _ {w | u, q} ^ {(t)} * \boldsymbol {W} _ {r} ^ {(t)} \boldsymbol {z} _ {q},
$$
$$
\operatorname {M S G} _ {r} ^ {2} \left(\boldsymbol {h} _ {w | u, q} ^ {(t)}, \boldsymbol {z} _ {q}\right) = \boldsymbol {h} _ {w | u, q} ^ {(t)} * \boldsymbol {b} _ {r} ^ {(t)},
$$
$$
\mathbf {M} \mathrm {S G} _ {r} ^ {3} \left(\boldsymbol {h} _ {w | u, q} ^ {(t)}, \boldsymbol {z} _ {q}\right) = \boldsymbol {W} _ {r} ^ {(t)} \boldsymbol {h} _ {w | u, q} ^ {(t)},
$$
where  $\mathbf{W}_r^{(t)}$  are relation-specific transformations.  $\mathrm{Msg}_r^1$  computes a query-dependent message, whereas  $\mathrm{Msg}_r^2$  and  $\mathrm{Msg}_r^3$  are analogous to message functions of CompGCN and RGCN, respectively.
History. In addition, we can set  $f$ , which intuitively is the function that determines the history of node embeddings to be considered. By setting  $f(t) = t$ , we obtain a standard message-passing algorithm where the update function considers the representation of the node in the previous iteration. We can alternatively set  $f(t) = 0$ , in which case we obtain (a generalization of) NBFNets.
Readout. We consider a standard readout which sums the representations and applies a linear transformation on the resulting representations. Alternatively, we consider the special case, where we omit this component (we discuss a dedicated readout operation in our empirical analysis later).
A simple architecture. Consider a basic C-MPNN architecture with global readout, which, for a query relation  $q \in R$  and a fixed node  $u$ , updates the representations as:
$$
\begin{array}{l} \boldsymbol {h} _ {v | u, q} ^ {(0)} = \mathbb {1} _ {u = v} * \boldsymbol {z} _ {q} \\ \boldsymbol {h} _ {v | u, q} ^ {(t + 1)} = \sigma \Bigg (\boldsymbol {W} _ {0} ^ {(t)} \big (\boldsymbol {h} _ {v | u, q} ^ {(t)} + \sum_ {r \in R} \sum_ {w \in \mathcal {N} _ {r} (v)} \operatorname {M S G} _ {r} ^ {1} (\boldsymbol {h} _ {w | u, q} ^ {(t)}, \boldsymbol {z} _ {q}) \big) + \boldsymbol {W} _ {1} ^ {(t)} \sum_ {w \in V} \boldsymbol {h} _ {w | u, q} ^ {(t)} \Bigg), \\ \end{array}
$$
where  $\pmb{W}_0^{(t)},\pmb{W}_1^{(t)}$  are linear transformations followed by a non-linearity  $\sigma$ .
5 Characterizing the expressive power
5.1 A relational Weisfeiler-Leman characterization
To analyze the expressive power of C-MPNNs, we introduce the relational asymmetric local 2-WL, denoted by  $\mathsf{rawl}_2$ . In this case, we work with knowledge graphs of the form  $G = (V,E,R,c,\eta)$ , where  $\eta : V \times V \to D$  is a pairwise coloring. We say that  $\eta$  satisfies target node distinguishability if  $\eta(u,u) \neq \eta(u,v)$  for all  $u \neq v \in V$ . The notions of isomorphism and invariants extend to this context in a natural way. For each  $t \geq 0$ , we update the coloring as:
$$
\begin{array}{l} \operatorname {r a w l} _ {2} ^ {(0)} (u, v) = \eta (u, v), \\ \left. \mathbf {r a w l} _ {2} ^ {(t + 1)} (u, v) = \tau \left(\mathbf {r a w l} _ {2} ^ {(t)} (u, v), \{\{\left(\mathbf {r a w l} _ {2} ^ {(t)} (u, w), r\right) \mid w \in \mathcal {N} _ {r} (v), r \in R) \} \}\right), \right. \\ \end{array}
$$
where  $\tau$  injectively maps the above pair to a unique color, which has not been used in previous iterations. Observe that  $\mathsf{rawl}_2^{(t)}$  defines a binary invariant, for all  $t\geq 0$ .