Tree-of-Traversals: A Zero-Shot Reasoning Algorithm for Augmenting Black-box Language Models with Knowledge Graphs (2024)

Abstract: Knowledge graphs (KGs) complement Large Language Models (LLMs) by providing reliable, structured, domain-specific, and up-to-date external knowledge. However, KGs and LLMs are often developed separately and must be integrated after training. We introduce Tree-of-Traversals, a novel zero-shot reasoning algorithm that enables augmentation of black-box LLMs with one or more KGs. The algorithm equips a LLM with actions for interfacing a KG and enables the LLM to perform tree search over possible thoughts and actions to find high confidence reasoning paths. Tree-of-Traversals significantly improves performance on question answering and KG question answering tasks. Code is available at https://github.com/amazon-science/tree-of-traversals

## introduction large language models (llms) are used for a range of knowledge-intensive tasks such as information retrieval @xcite , summarization @xcite , and question answering @xcite . trained on large amounts of textual data, these models learn a wide breadth of information. however, llms suffer from several limitations -they produce hallucinated information @xcite , lack deep domain-specific knowledge @xcite , and have a static knowledge cutoff when training ends. knowledge graphs (kgs) naturally complement llm weaknesses. kgs contain up-to-date information covering general @xcite and/or domainspecific topics (abu-salih, 2020; @xcite @xcite in highly structured and interpretable format. augmenting an llm's ability to reason and respond in natural language with an external kg's up-to-date knowledge presents a path toward a reliable and factual llm. the rise of powerful llms with new capabilities has renewed interest in combining llms with kgs. numerous survey and position papers have recently emphasized their combined potential @xcite @xcite . existing works augmented llms with kgs in multiple ways, such as integration into pretraining @xcite , fine-tuning @xcite , or later adaptation with subsequently trained components @xcite . all of these carry some limitations. in particular, training or fine-tuning a large scale llms is computationally expensive. in some cases, model weights are unavailable publicly. finally, the largest kgs require their own servers and cannot be integrated in-memory with llms. additionally, previous works do not consider the case of augmenting with multiple kgs. an algorithm that allows the augmentation of a powerful black-box llm with any number of internal or external kgs without training from scratch or fine-tuning the model is valuable. such a zeroshot algorithm would enable several innovative use cases such as (1) customers using a black-box llm api in conjunction with an internal domain-specific kg, (2) integration of personalized kgs into an llm without the risks associated with training a model with such personal data, (3) integration of deep domain knowledge via an array of api accessible kgs (e.g., imdb foot_0 , musicbrainz foot_1 ). we introduce tree-of-traversals, a novel algorithm that addresses the above issues by allowing the augmentation of any powerful llm with arbitrary number of kgs in a zero-shot fashion. it rethe "expand kg" option had higher value than "answer: no one". after choosing to expand kg, the model next selects which entities from the kg to expand. we conduct detailed experiments on three models of varied sizes hosted on amazon bedrock and present detailed ablation studies. ## tree-of-traversals the tree-of-traversals algorithm maintains a local kg subgraph that is repeatedly expanded until it contains all the information required by an llm to answer the given query. at the start, a local kg subgraph is initialized to contain the entities present in the original query. it is then expanded using a tree search algorithm to choose actions and thoughts generated by an llm in order to obtain relevant knowledge from the kg using a kg interface. the algorithm halts when the llm is able to answer the original query using the local kg subgraph as context. tree-of-traversals consists of three major components. (1) a knowledge graph interface implemented to interact with one or more required kgs. (2) an action state machine (asm) which is a finite state machine that defines the feasible space of actions, states, and prompt templates when the llm interacts with a kg to expand the local kg subgraph. (3) a tree search algorithm which defines the overall llm search trajectory such as best first search, backtrack upon making a mistake, and termination condition when an answer is found. ## experiments we evaluate tree-of-traversals using three different models available on amazon bedrock 3 : claude-instant (claude-instant-v1), llama2 70b (llama2-70b-chat-v1), and llama2 13b (llama2-13b-chat-v1). aws bedrock provides on-demand access through a single api to various foundation models, including both open source and black box ones. this is precisely a use case tree-of-traversals is designed for. ## results and discussion on musicbrainz-x-wikidata ( ## related works knowledge base question answering there is a large history of work that has looked into answering questions using a knowledge graph @xcite . the approaches can broadly be broken into those that attempt to parse the question into a logical form @xcite @xcite , and information retrieval approaches @xcite knowledge enhancement. recently, researchers have looked into enhancing pretrained language models and llms using knowledge graphs. many works have focused on incorporating knowledge graph data into llms through training, often with architectural changes to the model @xcite @xcite @xcite . some of these have found success using specialized graph encoder layers @xcite . some have sought to mix this process with pretraining of the language model @xcite . the limitations of including kgs during llm training are: the models are incapable of incorporating kg updates without retraining, are not able to change kg source (e.g., a domain-specific one), and may have low explainability resulting from learning knowledge in model weights rather than explicit retrieval. in addition, these methods add complexity to the training process, increase cost, and do not work with llm without access to model weights. as a result, scaling these methods has often been seen as too risky. other methods treat a kg as a more complete source of truth and use the llm to generate structured queries for the kg. @xcite and @xcite teach the llm to generate a logical kg query which then returns matching entities from the kg. these methods share similarities with the semantic parsing based approaches mentioned earlier. however, by returning a logical query for the kg, these methods lose the reasoning and commonsense abilities of the llm. some approaches have looked at injecting kg or knowledge base data into llm prompts. many methods do a single round of retrieval based on the query @xcite using a retrieval mechanism, such as dense passage retrieval @xcite . these methods cannot answer more complex multi-hop questions as the initial retrieval is unlikely to contain the secondary or tertiary information that will ultimately be required. later methods have made the retrieval process iterative. e.g. flare @xcite uses prediction of the upcoming sentence to retrieve relevant documents and regenerates until the sentence contains high-confidence tokens. in @xcite multi-round qa format is used for multi-round retrieval. react @xcite does multiple rounds of thoughts and actions to query a text-based knowledge base. @xcite repeatedly interfaces with a kg. in other parallel works, @xcite explore methods for building kg context for llms using path and neighborhood based search methods. the above methods do not incorporate tree search or forms of advanced reasoning beyond the llm's innate ability. they cannot explore multiple reasoning paths or solutions and cannot backtrack if making a mistake. this results in lesser capabilities. additionally, none of the above methods explore reasoning over multiple kgs. multiple knowledge base qa. some works studied qa on questions derived from differing domains. these works learn to pick between different domain-specific qa models, each trained on a single knowledge base @xcite @xcite . besides requiring training, such systems cannot answer questions requiring the synthesis of information from different domains (musicbrainz-x-wikidata). ## conclusion tree-of-traversals is a powerful algorithm that enables llms to utilize kg reasoning with zero examples, no schema dependency, and no training. we demonstrate its efficacy experimentally on multiple llms and datasets. we hope tree-of-traversals continues to be developed and see value in studying integration with personalized user kgs as well as with other domain-specific datasets.