NuTrea: Neural Tree Search
for Context-guided Multi-hop KGQA

Here, we first define the problem settings for Neural Tree Search (NuTrea). The multi-hop KGQA task is primarily a natural language processing problem that receives a human language question $x_{q}$ as input, and requires retrieving the set of nodes from $\mathcal{G}=(\mathcal{V},\mathcal{E})$ that answer the question. Following the standard protocol in KGQA [11], the subject entities in $x_{q}$ are given and assumed always to be mapped to a node in $\mathcal{V}$ via entity-linking algorithms [31]. These nodes are called seed nodes, denoted $v_{s}\in\mathcal{V}_{s}$, which are used to extract a subgraph $\mathcal{G}_{q}=(\mathcal{V}_{q},\mathcal{E}_{q})$ from $\mathcal{G}$ so that $\mathcal{V}_{q}$ is likely to contain answer nodes $v_{a}\in\mathcal{V}_{a}$. Then, the task reduces to a binary node classification problem of whether each node $v\in\mathcal{V}_{q}$ satisfies $v\in\mathcal{V}_{a}$.

Following prior works [14, 1], we first compute different question representation vectors with a language model based module, called Instruction Generators (IG). We have two IG modules, each for the Expansion (section 3.2.1) and Backup (section 3.2.2) step, to compute

$$\{\mathbf{q}_{\text{exp}}^{(i)}\}_{i=1}^{N}=\text{IG}_{\text{exp}}(x_{q}),\;\;\;\{\mathbf{q}_{\text{bak}}^{(j)}\}_{j=1}^{M}=\text{IG}_{\text{bak}}(x_{q}),$$

(1)

where we name $\mathbf{q}_{\text{exp}}^{(i)},\mathbf{q}_{\text{bak}}^{(j)}\in\mathbb{R}^{D}$ as the expansion instruction and backup instruction, respectively. Detailed description on the IG module is in the supplement. Then, the learnable edge (relation) type embeddings $\boldsymbol{R}\in\mathbb{R}^{|\mathcal{R}|\times D}$ are randomly initialized or computed with a pretrained language model. On the other hand, node embeddings $\boldsymbol{H}\in\mathbb{R}^{|\mathcal{V}_{q}|\times D}$ of $\mathcal{V}_{q}$ are initialized using the edges in $\mathcal{E}_{q}$ and their relation types by function $\mathcal{H}$ as $\boldsymbol{H}=\mathcal{H}(\mathcal{E}_{q},\boldsymbol{R})$ (e.g., $\mathcal{H}\equiv$ arithmetic mean of incident edge relations). This is because the node entities are often uninformative proper nouns or encrypted codes. Then, our NuTrea model $\mathcal{F}$ function is defined as

$$\hat{\mathbf{y}}=\mathcal{F}(\{\mathbf{q}_{\text{exp}}^{(i)}\}_{i=1}^{N},\{\mathbf{q}_{\text{bak}}^{(j)}\}_{j=1}^{M},\mathcal{V}_{s}\,;\,\mathcal{G}_{q},\boldsymbol{H},\boldsymbol{R}),$$

(2)

where $\hat{\mathbf{y}}\in\mathbb{R}^{|\mathcal{V}_{q}|}$ is the predicted node score vector normalized across nodes in $\mathcal{V}_{q}$, whose ground-truth labels are $\mathbf{y}=[\mathbbm{1}(v\in\mathcal{V}_{a})]_{v\in\mathcal{V}_{q}}\in\mathbb{B}^{|\mathcal{V}_{q}|}$. Then, the model is optimized with the KL divergence loss between $\hat{\mathbf{y}}$ and $\mathbf{y}$. In this work, we claim our contributions in $\mathcal{F}$ (section 3.2) and $\mathcal{H}$ (section 3.3).

Our Neural Tree Search (NuTrea) model consists of multiple layers that each performs message passing in three consecutive steps: (1) Expansion, (2) Backup, and (3) Node ranking. The Expansion step propagates information outwards to expand the search tree, which is followed by the Backup step, where the depth-$K$ subtree content is aggregated to each of their root nodes, to enhance the past-oriented messages from the Expansion step. Then, the nodes are scored based on how likely they answer the given question.

Another notable challenge with KGs is embedding nodes. Many KG entities (nodes) are proper nouns that are not informative, and several KGQA datasets [16, 17] consist of encrypted entity names. Thus, given no proper node features, the burden is on the model layers to learn meaningful node embeddings. To alleviate this, we propose a novel node embedding method, Relation Frequency-Inverse Entity Frequency (RF-IEF), which grounds on the local and global topological information of the KG nodes.

Term frequency-inverse document frequency (TF-IDF) is one effective feature that characterizes a textual document [32, 33, 34]. The bag-of-words model computes the frequency of terms in a document and penalizes frequent but noninformative words, e.g., ‘a’, ‘the’, ‘is’, ‘are’, by the frequency of the term across documents. Motivated by the idea, we represent a node on a KG as a bag of relations. An entity node is characterized by the frequencies of rare (or informative) relations. Similar to TF-IDF, we define two functions: Relation Frequency (RF) and Inverse Entity Frequency (IEF). The RF function is defined for node $v\in\mathcal{V}_{q}$ and relation type $r\in\mathcal{R}$ as

$$\text{RF}(v,r)=\underset{e\in\mathcal{I}(v)}{\sum}\mathbb{1}\{f(e)=r\},$$

(11)

where $\mathcal{I}(v)$ is the set of incident edges of node $v$, $\mathbb{1}$ is an indicator function, and $f:\mathcal{E}\rightarrow\mathcal{R}$ is a function that retrieves the relation type of an edge. Then, the output of the RF function is a matrix $RF\in\mathbb{R}^{|\mathcal{V}_{q}|\times|\mathcal{R}|}$ that counts the occurrence of each relation type incident to each node in the KG subgraph. We used raw counts for relation frequency (RF) to reflect the local degree information of a node. On the other hand, the IEF function is defined as

$$\text{IEF}(r)=\log\frac{|\mathcal{V}_{q}|}{1+\text{EF}(r)},$$

(12)

where

$$\text{EF}(r)=\underset{v\in\mathcal{V}_{q}}{\sum}\mathbb{1}\{\exists\;e\in\mathcal{I}(v)\text{ s.t. }f(e)=r\}.$$

(13)

EF counts the global frequency of nodes across KG subgraphs that have relation $r$ within its incident edge set. With $IEF\in\mathbb{R}^{|\mathcal{R}|}$, the RF-IEF matrix $\boldsymbol{F}\in\mathbb{R}^{|\mathcal{V}_{q}|\times|\mathcal{R}|}$ is computed as

$$\boldsymbol{F}=RF\,\text{diag}(IEF),$$

(14)

where $\text{diag}(IEF)\in\mathbb{R}^{|\mathcal{V}_{q}|\times|\mathcal{R}|}$ denotes a diagonal matrix constructed with the elements of $IEF$.

The RF-IEF matrix $\boldsymbol{F}$ captures both the local and global KG structure and it can be further enhanced with the rich semantic information on the edges of KGs. Unlike entity nodes with uninformative text, e.g., proper nouns and encrypted entity names, edges generally are accompanied by linguistic descriptions (relation types). Hence, the relations are commonly embedded by a pre-trained language model in KGQA. Combining with the relation embeddings $\boldsymbol{R}\in\mathbb{R}^{|\mathcal{R}|\times D}$, our final RF-IEF node embeddings $\boldsymbol{H}$ is computed as

$$\boldsymbol{H}=\boldsymbol{F}\,\boldsymbol{R}\,\boldsymbol{W}_{h},$$

(15)

where $\boldsymbol{H}\in\mathbb{R}^{|\mathcal{V}_{q}|\times D}$, and $\boldsymbol{W}_{h}\in\mathbb{R}^{D\times D}$. The RF-IEF node embeddings can be viewed as the aggregated semantics of relations, which are represented by a language model, based on graph topology as well as the informativeness (or rareness) of relations. A row in $\boldsymbol{H}$ is a node embedding vector $\mathbf{h}_{v}$, which is used in Eq. (LABEL:eq:prop_emb) at the first NuTrea layer.

Here, we present the experimental results of NuTrea. Following the common evaluation practice of previous works, we test the model that achieved the best performance on the validation set. In the WQP dataset experiments in Table 1, we achieved the best performance of 77.4 H@1 among strong KGQA baselines that take an information retrieval approach, as discussed in Section 2. Compared to the previous best, this is a large improvement of 0.6 points. In terms of the F1 score, which evaluates the answer set prediction, our method achieved a score of 72.7, exceeding the previously recorded value by a large margin of 1.8 points. In addition, we also improved the previous state-of-the-art performance on the CWQ dataset by achieving an 53.6 H@1, which is an improvement of 0.7 points.

We also experimented NuTrea on MetaQA to see if it performs reasonably well with easy questions as well. On three data splits, NuTrea achieved comparable performance with previous state-of-the-art methods for simple question answering. Evaluating with the average H@1 score of the three splits, NuTrea performs second best among all baseline models.

The KG is often human-made and the contents are prone to being incomplete. Hence, it is a norm to test a model’s robustness on incomplete KG settings where a certain portion of KG triplets, $\emph{i.e}.$, a tuple of $\langle$head, relation, tail$\rangle$, are dropped. This experiment evaluates the robustness of our model to missing relations in a KG. We follow the experiment settings in [1], and use the identical incomplete KG dataset which consists of WQP samples with [50%, 30%, 10%] of the original KG triplets remaining. In Table 3, NuTrea performs the best in most cases, among the GNN models designed to handle incomplete KGs. We believe that our model adaptively learns multiple alternative reasoning processes and can plan for future moves beforehand via our Backup step, so that it provides robust performance with noisy KGs.

Here, we evaluate the effectiveness of each component in NuTrea to answer Q1. An ablation study is performed on our two major contributions: the Backup step in our NuTrea layers and the RF-IEF node embeddings. By removing the RF-IEF node embeddings, we instead apply the common node initialization method used in [14, 1], which simply averages the relation embeddings incident to each node. Another option is to use zero embeddings, but we found it worse than the simple averaging method. For NuTrea layer ablation, we remove the Backup step which plays a key role in aggregating the future-oriented subtree information onto the KG nodes. Then, only the expansion-score ($s^{(e)}_{v}$) is computed, and the backup-score ($s^{(b)}_{v}$) is always 0. Also, the embedding $\mathbf{f}_{v}$ (Eq. (LABEL:eq:prop_emb)), is directly output from the NuTrea layer and no further updates are made via the Backup step.

In Table 3, we can see that the largest performance drop is observed when the Backup step is removed. Without it, the model has limited access to the broader context of the KG and cannot reflect the complex question constraints in node searching. Further discussion on this property of NuTrea’s message passing is provided in the next Section 5.2. Also, we observed a non-trivial 0.6 point drop in H@1 by removing the RF-IEF initialization method. By ablating both components, there was a significant degradation of 4.0 H@1 points.

To answer Q2, we highlight the key advantages of NuTrea over recent approaches. We analyze the efficiency of NuTrea, and provide qualitative results.

The RF-IEF node embedding is a simple method inspired by an effective text representation technique in natural language processing. Here, we disclose the specific effect of RF-IEF on the relation embedding aggregation weights, thereby answering Q3. In Figure 4 (left), the log-scaled $EF$ (Eq. 13) value of each relation type is sorted. The globally most frequent relation types, including “self_loop”s, are too general to provide much context in characterizing a KG node. Our RF-IEF suppresses such uninformative relation types for node embedding initialization, resulting in a weight distribution like Figure 4 (right). The pie charts display two examples on the difference between aggregation weights for a node entity before and after RF-IEF is applied. To illustrate, the weight after RF-IEF corresponds to a row of $\boldsymbol{F}$ in Eq. (14), while the weight before RF-IEF would be uniform across incident edges. As “Alexander Bustamante” is a politician, the relation types “orgainizations founded” and “founders” become more lucid via RF-IEF, while relations like “events_competed_in” and “competitors” are emphasized for an athlete like “Kemar Bailey-Cole”. Likewise, RF-IEF tends to scale up characteristic relation types in initializing the node features, thereby enhancing differentiability between entities. To further demonstrate RF-IEF’s general applicability, we also provide a plug-in experiment on another baseline model in the supplement.