Less is More: Making Smaller Language Models Competent Subgraph Retrievers for Multi-hop KGQA (2024)

Abstract: Retrieval-Augmented Generation (RAG) is widely used to inject external non-parametric knowledge into large language models (LLMs). Recent works suggest that Knowledge Graphs (KGs) contain valuable external knowledge for LLMs. Retrieving information from KGs differs from extracting it from document sets. Most existing approaches seek to directly retrieve relevant subgraphs, thereby eliminating the need for extensive SPARQL annotations, traditionally required by semantic parsing methods. In this paper, we model the subgraph retrieval task as a conditional generation task handled by small language models. Specifically, we define a subgraph identifier as a sequence of relations, each represented as a special token stored in the language models. Our base generative subgraph retrieval model, consisting of only 220M parameters, achieves competitive retrieval performance compared to state-of-the-art models relying on 7B parameters, demonstrating that small language models are capable of performing the subgraph retrieval task. Furthermore, our largest 3B model, when plugged with an LLM reader, sets new SOTA end-to-end performance on both the WebQSP and CWQ benchmarks. Our model and data will be made available online: https://github.com/hwy9855/GSR.

## introduction large language models (llms) have demonstrated tremendous capabilities in various natural language processing tasks @xcite . despite their success, their hallucination tendencies still limit their performance across the involved tasks, and in question answering, in particular @xcite . retrieval-augmented generation methods are widely used to enhance llms @xcite figure 1 : our proposed gsr architecture facilitates training a smaller lm with less parameters by reconsidering the subgraph retrieval task as a subgraph id generation task, leading to shorter sequences of less target tokens. in addition, we propose two ways of obtaining less training samples from weakly supervised retrieval data to reduce the noise (e.g., path gendergender is not informative and brings noise in training). 2023) and provide better safeguards against hallucination issues @xcite . knowledge graphs (kgs) @xcite have been recognised as valuable knowledge sources due to their compact triple representation, clear and noise-free knowledge format, and rich domain-specific information @xcite . consequently, they have attracted significant attention by researchers who seek to propose efficient solutions that leverage information enclosed within a kg for question answering (kgqa) @xcite @xcite . traditional semantic parsing (sp) based kgqa @xcite @xcite initiate the process by generating a sparql query derived from the natural language question. this query is subsequently executed against the kg of interest. recently, subgraph retrieval (sr) based kgqa becomes popular as they eliminate the need of extensive sparql annotations (i.e. demonstra-tions consisting of (question, sparql) pairs). typical sr-based kgqa methods @xcite @xcite operate within a retrieval-augmented generation framework. they first retrieve a subgraph from the knowledge graph, which is then sent to a reader to generate answers. these methods achieve competitive performance compared with complex sp-based methods. since state-of-the-art (sota) methods for srbased kgqa, such as those by @xcite and @xcite , typically employ llms to generate relevant relation chains, they can be expensive in both training and reasoning due to the complexity of relation names (e.g., government.-government_position_held.office_position_or_title is tokenized into 18 tokens, with the llama2 tokenizer). in addition, the complexity of relation names induces an unnecessary task in mapping relations to relation names, besides mapping questions to relation chains. these drawbacks inspire us to consider a generative subgraph retriever architecture specified with smaller language models that have less parameters. this forms the primary research question of this work: how can small language models be utilized to accomplish the subgraph retrieval task with better efficiency and comparable or superior effectiveness compared to large language models? to simplify the task of generating relation chain given the complex relation names and to improve the representation of relations in language models, we treat each relation as a special token (relation id) in the small language model. this transforms the subgraph retrieval task into predicting a sequence of relation ids, which we define as the identifier of a subgraph (subgraph id). as demonstrated in figure 1 , our redefined task requires much less tokens (2 tokens vs. 35 tokens) as target, makes the training easier and increases the inference efficiency. our subgraph retriever, named generative subgraph retriever (gsr), is trained jointly using two types of data: indexing data and subgraph retrieval data. in addition, we observe a large amount of noise inside the subgraph retrieval data used by most previous works @xcite . to mitigate this issue, we proposed two data pruning methods to obtain denoised retreival data with less training samples. our comprehensive experiment showcase that employing large language model for (generative) sub-graph retrieval is an unnecessary expense. our best setting, even when a 220m parameters model is used (30× less parameters), can achieve +9.2% and +5.3% f1 score improvement (more effective) on the webqsp and cwq benchmarks respectively over the previous sota subgraph retrieval work @xcite consisting of 7b parameters. by integrating with our llm reader, which is finetuned from the same base model used in the sota approaches, we achieve f1 score improvements of +6.3% on webqsp and +4.9% on cwq in the endto-end evaluation. our 3b model with llama3 reader reaches new sota performance among srbased kgqa on both webqsp and cwq dataset with f1 80.1% and 64.4% respectively. our contributions can be summarised as follows: 1) we introduce gsr, a method utilizing small language models to accomplish the subgraph retrieval task. 2) we propose a training framework comprising an indexing step and a retrieval step for training gsr, including: a) an automatic method for collecting indexing data; b) two distinct methods to enhance the quality of the retrieval data. 3) comprehensive experimental results demonstrate the effectiveness of our work. our best model achieves an average improvement of +5.6% in f1 score on two kgqa datasets compared to the previous sota sr-based kgqa models, while being 7.7 times more efficient during the subgraph retrieval step. ## related works our method draws inspiration from works on knowledge graph question answering (kgqa) and generative retrieval. ## problem statement a knowledge graph @xmath0, r, o)∥s, @xmath1, @xmath2 james k. polk office_holder office_position_or_title m.04j60kh m.04j60kc m.04469y8 united states representative president of the united states what did james k polk do before he was president? generative subgraph retriever office_holder --office_position_or_title government_positions_held --office_position_or_title profession question united states representative prediction subgraph id (relation chain) subgraph llm reader fetch subgraph textualize generate beam search figure 2: overall pipeline of proposed rag framework for kgqa task with the generative subgraph retriever. eral (s, r, o) triples, where e is the entity set and r is the relation set, and s, r, o are instances of a subject entity, a relation, and an object entity respectively, for a @xmath3. knowledge graph question answering (kgqa) aims to figure out the answer set â given the question q and topic entity e t . kgqa methods based on subgraph retrieval model the question answering task in a retrievalaugment generation framework by first retrieving relevant subgraph sg ⊆ kg given the question q, and, subsequently, using a reader to generate the predicted answer set a given the question q and the retrieved subgraph sg. 4 methodology 4.1 generative subgraph retriever subgraph definition a kg subgraph sg ⊆ kg is a subset of all triples in the original kg. specially, in this work, we made a further constraint on the definition of subgraph, which we called path constrained subgraph 1 . in short, path constrained subgraph is a set of special subgraph that can be identified with a simple identifier (e, c), where e is the entity and @xmath4, r 2 , ..., r n is an ordered relation chain that indicates n-hop reasoning. by moving from the entity e through the relation chain c, we can identify the whole subgraph with the simple identifier. the middle part of figure 2 shows an example of identify subgraph with topic entity james k. polk and relation chain (office_holder, office_position_or_title), which can be identified with identifier [james k. polk, office_holder, office_position_or_title]. in addition, for better flexibility, we decouple the topic entity e and relation chain c in the subgraph identifier, which means that the two part can be retrieved separately. we assume the topic entity e is known in this work for aligning other sr-based kgqa works @xcite . subgraph retrieval with the subgraph definition above, we can model the subgraph retrieval task as a subgraph id generation task. given a natural language query q, subgraph retrieval aims to find the relevant subgraph, as a relation chain @xmath5, r 2 , ..., r n }, that forms an n-hop reasoning chain from the topic entity e t to the answer entity e a . the relation chain is ordered, and is predicted by the probability of the above equation can be decomposed into the conventional auto-regressive language modelling objective, assuming relations are mapped into the involved model's token space. subgraph id since the major task of our gsr model is to map questions to a subgraph id (relation chain), it is essential to build an efficient and effective way of representing the relations in the model. @xcite simply ask llms to generate the whole name of the relation (e.g., tv.regular_tv_appearance.actor), but it is not efficient and hard for lms to learn the mapping. instead, we adapt the atomic document representation methods from dsi-based works @xcite . on relation level, we assign each relation an atomic id, i.e., each relation is mapped to a special token in the generative language model. on subgraph level, we build a hierarchical indexing, i.e., each subgraph id is mapped to the relation tokens which forms the reasoning chain. ## experimental setup datasets we use webqsp @xcite and cwq @xcite datasets that are both constructed from freebase for all the experiments. following previous works that seek to increase the reasoning efficiency on freebase, we choose to use the same simplified freebase (only contains multi-hop triples of all topic entities appeared in webqsp and cwq) @xcite @xcite . specially, webqsp is a relatively easier dataset as its questions require at most 2-hop reasoning from their topic entity. cwq is more challenging as the questions involve up to 4-hop reasoning. detailed statistics of both datasets are provided in appendix b. baselines we compare our proposed rag with gsr pipeline with several baselines. for evaluating subgraph retrieval performance, we compare our methods with rog @xcite . for end-to-end kgqa comparison, we choose both sp-based approaches including unikgqa @xcite , decaf @xcite , chatkbqa @xcite and sr-based approaches including sr @xcite , kaping @xcite , rog @xcite and tog @xcite . metrics for evaluating subgraph retrieval, we choose precision, recall and f1 to measure the answer coverage of the retrieved subgraphs. for the end-to-end evaluation, we use f1, hits@1 @xcite . rog, in particular, measures hits@1 in an uncommon way @xcite . in order to align their results with ours, we further added a hits metric, which treats all answers generated by llm as the top-1 answer and measures the hits@1 score accordingly. for the original hits@1, we use the first answer generated by llm. ## results we evaluate our approach across both sugraph retrieval and the end-to-end qa. ## analysis we analyse in detail the effectiveness of several variants of our proposed method, to answer the following research questions. rq1: how to efficiently and effectively prompt llms with a retrieved subgraph? (sec 7.1) rq2: how does the indexing data contribute to the performance on subgraph retrieval? rq3: what is the performance upper-bound of the proposed methods? (sec 7.3) ## conclusion we propose a retrieval-augmented generation framework for kgqa comprising a subgraph retriever and an llm reader. to efficiently and effectively retrieve a subgraph from kg, we propose a novel generative subgraph retrieval method that transform the subgraph retrieval task into a sequential subgraph id generation task. our base gsr model with 220m parameters is capable of outperforming the previous 7b sota baseline in subgraph retrieval, with around 30× less parameters. combined with our llm reader, we achieve new sota performance in subgraph retrievalbased kgqa on both the webqsp and cwq datasets. in summary, our proposed gsr model with less model parameters, less training samples, and learned to generate less tokens, achieves more efficiency and more effectiveness compared sota sr-based kgqa works.