Struct-X: Enhancing Large Language Models Reasoning with Structured Data

We first implement “read-model-fill” process and we start by processing input KGs using a graph attention encoder (GAE) that consists of $L$ layers Xu et al. (2021). The initial node features, denoted as $h_{v}^{(0)}$, are set up using information obtained from the KG completion module Fei et al. (2021). After processing through $L$ layers, we obtain the final node embeddings, $h_{v}^{(L)}$, which effectively represent both the semantic and structural information of the KGs. These encoded graph embeddings, $h_{v}^{(L)}$, are then partially masked at a specific rate, $p_{mask}$, to assist in learning about missing knowledge. This masking process can be mathematically represented as $\tilde{h}_{v}=M(h_{v}^{(L)})$, where $M(\cdot)$ symbolizes the masking operation. The masked nodes, denoted as $\tilde{h}_{v}$, are then fed into the knowledge retrieval module, $R(\tilde{h}_{v})$, which is explained in the following section. This module plays a crucial role in supplementing the missing information, thereby facilitating the generation of complete graph embeddings $\bar{h}_{v}$ Reda et al. (2022).

To address the gaps in entity information within the structured knowledge graph, we have developed a knowledge learning module, denoted as $F$. This module is designed to retrieve pertinent facts from the knowledge base to enhance the masked node embeddings, $\tilde{h}_{v}$ Yasunaga et al. (2022). More specifically, for each masked node, we calculate a similarity score between its embedding $\tilde{h}_{v}$ and all tail entities $t$ that are part of the set $E$. This is achieved using the scoring function $f_{score}$, which can be represented as:

This process enables us to efficiently fill in the missing information in the knowledge graph.

The scoring function evaluates both feature and topological similarities within the graph in Figure 2. It selects the top $K$ entities $t$ based on the highest scores and retrieves the related facts $(h,r,t)$ from the knowledge base. To incorporate these facts into the node embeddings, a relation-aware aggregation function $f_{agg}$ is used. This function accumulates relevant knowledge for each node, using a score threshold $\tau$ to filter out irrelevant facts Yu et al. (2022). The aggregation function adeptly manages various relations in structured knowledge by considering information from retrieved triples in a relation-aware manner. Additionally, before being input into the Transformer encoder, one linear layer $o$ concatenates and processes embeddings from all GAT layers.

The $o$ merges inputs and reduces dimensionality. This process retains rich multi-scale structural and semantic features at various depths. The output $e_{v}$ is flattened and prepared into sequences to replace token embeddings for the Transformer encoder input, as suggested by Wang et al. (2021). The refined node embeddings $\bar{h}_{v}$, enriched with retrieved entity information, supply additional knowledge for reasoning in the downstream LLMs.

Here we perform the “reflect” process. The module $R(\tilde{h}_{v})$ retrieves relevant knowledge absent in masked graph node inputs $\tilde{h}_{v}$. Related entities can be dynamically discovered by matching tail entities $t$ to each masked node using a similarity scoring function $f_{score}(\tilde{h}_{v},t)$ considering both feature and topological similarity Lewis (2020). Related facts $(h,r,t)$ are recalled to fill gaps. After concatenating retrieved knowledge sequences for all nodes ordered by similarity scores, we employ a pruning algorithm leveraging multi-head self-attention distillation and thresholds to filter out lower-weighted tokens. The remaining dense sequence provides supplemental external knowledge to complete masked graph node inputs.

To filter and verify the relevance of retrieved knowledge, we design a self-retrieved generation module $\text{$SelfReg$}_{\psi}(k)$ parameterized by $\psi$ that takes as input the retrieved knowledge sequences $k$ and outputs a filtered subset $\hat{k}$ containing only the most valuable tokens Shuster (2021). Specifically, $\text{$SelfReg$}_{\psi}$ first encodes the knowledge sequence $k=(x_{1},x_{2},\ldots,x_{N})$ using a Transformer encoder to obtain representations $h_{i}=f_{\text{$enc$}}(x_{i})$. Next, we compute an importance score for each token $s_{i}=\sigma(f_{\text{score}}(h_{i}))$, where $f_{\text{$score$}}$ is a scoring network and $\sigma$ is a sigmoid activation function. To train the scoring network in a self-supervised manner, we create corrupted knowledge sequences $\tilde{k}$ by randomly masking or shuffling some tokens. A contrastive loss is implemented to assign higher scores $s_{i}$ to tokens from the original $k$ versus corrupted $\tilde{k}$:

where $\Delta$ is a margin hyperparameter. This drives the model to identify the most valuable knowledge. Finally, we filter the sequence by discarding tokens scoring below a threshold of $\tau$ to retain only the most relevant phrases, significantly reducing the learning burden when provided as supplements to the LLMs.

The filtered relevant knowledge $\hat{k}$ provides targeted assistance to improve reasoning without overwhelming the LLMs with extraneous and irrelevant information.

Here we perform the “reason” phase. To capture semantic and structural interactions between entities within the KGs, we use a specialized graph encoder in Figure 3, denoted as $E_{\theta}(G)$, which is parameterized by $\theta$. This KG is represented as $G=(V,E)$, where $V$ is the set of node entities and $E$ is the set of relation edges Li (2022). For each entity node $v_{i}$ in $V$, we first derive its initial feature representation $h_{v_{i}}^{(0)}$, which is a vector in a high-dimensional space.

The graph encoder works through a series of $L$ layers, each layer enhancing the node representations through message passing. This process can be described as:

where $\mathcal{N}(v_{i})$ refers to the neighboring nodes of $v_{i}$, and $f_{\theta}(\cdot)$ is a function that aggregates information from these neighbors to update the node’s embedding.

To focus on the most relevant semantic connections, we use a graph self-attention layer within $f_{\theta}(\cdot)$. This layer calculates attention weights as follows $a_{ij}=\frac{\exp(\langle q_{i},k_{j}\rangle)}{\sum_{v_{k}\in\mathcal{N}(v_{i})}\exp(\langle q_{i},k_{k}\rangle)}$, where $q_{i}$ and $k_{j}$ are derived from the embeddings of the nodes. This method allows the model to selectively emphasize the most informative signals from neighboring nodes Cui et al. (2020).

After processing through $L$ layers, we obtain refined node embeddings $z_{v_{i}}=h_{v_{i}}^{(L)}$, which encapsulate both semantic and structural information of the graph. To make these embeddings more manageable for downstream tasks, we compress them through a trainable down-projection layer:

This step reduces the dimensionality of the embeddings to $d_{e}$, which is smaller than the original $d_{z}$. The resulting condensed embeddings $e_{v_{i}}$ still retain crucial token-level interactions but are more concise, making them better suited for training models for specific tasks. This approach ensures that while the size of the input sequence is significantly reduced, the essential semantic and structural features of the knowledge graph are preserved for subsequent reasoning.

To further guide the LLMs in effectively reasoning over the structured input with a knowledge graph, we have developed an Auxiliary Module. This module is designed to create dynamic prompts that enhance the coherence of answers generated by LLMs. It functions by analyzing the LLM’s predicted answer, denoted as $\hat{y}$, along with the current loss, $L$. Based on these inputs, it generates a refined prompt, $p^{\prime}$, which is then used for a new round of inference. We use the pre-trained Bert model (i.e., bert-base-NER) Devlin et al. (2018) to construct this Auxiliary Module, symbolized as $G$ and parameterized by $\theta_{g}$. This generator crafts the prompt text, taking into account the input values $p^{\prime}=G(L,\hat{y};\theta_{g})$. The generator is trained jointly with the overall system using policy gradient methods to maximize the expected reward $R$ of producing coherent answers: $J(\theta_{g})=\mathbb{E}{p^{\prime}\sim G}[R(p^{\prime})]$, $\nabla{\theta_{g}}J(\theta_{g})=\mathbb{E}{p^{\prime}\sim G}[\nabla{\theta_{g}}\log G(p^{\prime}|L,\hat{y};\theta_{g})R(p^{\prime})]$.

This reward function is designed to encourage the LLMs to generate responses that are not only fluent but also logically consistent, particularly when using the updated prompt. This feature enables the module to dynamically adjust prompts based on the current performance, thereby offering new approaches to improve the quality of answers. Throughout the training process, the module progressively learns to produce more effective prompts, leading to enhanced accuracy and coherence in the LLM’s reasoning. For more detailed experimental testing and analysis, please refer to Appendix G.

We assess the performance of our proposed Struct-X framework on four open-source benchmark datasets designed for knowledge graph reasoning and multi-hop reasoning abilities on graphs.

Task1:WebQSP contains 4,737 QA pairs where the questions require logical reasoning over a knowledge graph derived from Wikipedia to infer the correct answer. The knowledge graph consists of 5,719 entities and 2,150 relations.

Task2:MetaQA comprises a set of more complex compositional questions constructed from an underlying knowledge graph with a vocabulary of 300 entities and 100 relations. It has a total of 1,200 unique questions that test the multi-hop, logical, and comparative reasoning abilities of models.

Task3:Family Tree Age Consider a family tree $G=(V,E)$, where each individual $v_{i}$ in $V$ is associated with a description $d_{i}$ specifying their age. The objective of this task is to identify the triplet comprising an individual, one of their grandparents, and a grand-uncle/grand-aunt by marriage that collectively has the highest cumulative age.

Task4:Travel Route Optimization Let $G=(V,E)$ be a graph representing connected cities, where each city $v_{i}$ in $V$ has a description $d_{i}$ with the travel toll or tax. The LLM must plan the route from a source to a destination city that minimizes the total toll paid.

For all datasets, we incorporate the encoded graph representations into Llama2, which has been pre-trained on BookCorpus and English Wikipedia. Appendix C is a case analysis of the experimental results on the datasets.

Baseline

• •

  Embedding-based Model: We compare against representative embedding models for knowledge graph reasoning including TransE Bordes et al. (2013), DistMult Yang et al. (2015), EmbedKGQA Saxena et al. (2020), ComplEx Trouillon et al. (2016), and RotatE Sun et al. (2019).

• •

  Open-source LLM: We evaluate reasoning capabilities of widely-used pre-trained language models accessible through open APIs, including Llama2 [7B & 13B] Touvron et al. (2023) and Alpaca [7B & 13B] Yao et al. (2023a) which are openly available LLMs up to 13 billion parameters.

• •

  LLM-based Fine-tuning: To assess the performance of LM fine-tuning approaches, we include as baselines KG-LlaMA and KG-Alpaca Yao et al. (2023b), KG-BERT Yao et al. (2019), PKGC Lv et al. (2022) and vanilla IT Zhang (2023) which incorporate techniques to enhance LMs using annotated KG datasets or self-supervision.

Our implementation is in PyTorch and we run experiments on NVIDIA A100 GPUs. More details of training can be found in Appendix D.

The results presented in Table 2 indicate that Struct-X consistently outperforms existing baseline methods across various datasets. Specifically, in the WebQSP benchmark, Struct-X achieves an accuracy of 75.13%, which is 2.65% higher than the previously best-performing method, KoPA. Additionally, Struct-X shows modest improvements in precision and recall, with increases of 9.51% and 10.96%, respectively, compared to Vanilla IT. In the more challenging MetaQA dataset, Struct-X’s performance is notably better, surpassing the state-of-the-art accuracy scores by 1.84% and achieving a 1.68% higher precision. Furthermore, Struct-X demonstrates significant advancements in specialized tasks such as Family Tree and Travel Route, where it exceeds the top baseline results by 3.36% and 5.34% in accuracy, respectively.

Compared to embedding models such as TransE, DistMult, and EmbedKGQA, Struct-X also shows promising improvements in reasoning abilities by integrating both semantic and topological structures of knowledge graphs. For instance, against RotatE’s accuracy of 74.55% on the WebQSP dataset, Struct-X achieves higher performance with a 75.13% accuracy, an increase of 0.58%. The difference is slightly more pronounced on the MetaQA dataset, where Struct-X exceeds RotatE’s score of 78.19% by 1.44% in accuracy. In scenarios requiring complex reasoning inferences, Struct-X demonstrates enhanced capabilities, outperforming peak embedding model accuracy by a notable margin of 16.74% in Task4.

The results also show that Struct-X can enhance the capabilities of the Llama2, which itself achieves a 27.13% accuracy on the WebQSP benchmark. This enhancement is achieved through masking graph embeddings and using topology matching to retrieve relevant facts, thus addressing the gaps in factual knowledge that Llama2 requires. By overcoming these deficiencies in LLMs, Struct-X significantly improves performance, increasing accuracy by 47.4%. This indicates the effectiveness of structured augmentation, which is not present in the Llama2. Further, Struct-X filters out less important tokens using the Self-Reg module, ensuring focus on the most relevant information. In comparison to previous methods like KG-BERT fine-tuning, StructX offers essential enhancements, particularly in complex reasoning tasks, as evidenced by increases of up to 10.24% in accuracy and 5.61% in recall.

Based on the experimental results, the “reflect” process also plays a crucial role in enhancing reasoning capabilities. This process involves the IFReT $(x)$, which selectively gathers evidence as needed, and the IFReL $(x,p)$, which filters less relevant passages using relevance scores from Eq.3 to improve context for LLMs. Additionally, the IFSuP $(y,p)$ and IFUsE $(x,y)$ ensure passage-response consistency and assess overall utility, contributing to higher quality results. For further case studies and experiments on this topic, readers are directed to Appendix H.