SeaKR: Self-aware Knowledge Retrieval for Adaptive Retrieval Augmented Generation

Zijun Yao,Weijian Qi,Liangming Pan,Shulin Cao,Linmei Hu,Weichuan Liu,Lei Hou,Juanzi Li

2024-06-27

Abstract:This paper introduces Self-aware Knowledge Retrieval (SeaKR), a novel adaptive RAG model that extracts self-aware uncertainty of LLMs from their internal states. SeaKR activates retrieval when the LLMs present high self-aware uncertainty for generation. To effectively integrate retrieved knowledge snippets, SeaKR re-ranks them based on LLM's self-aware uncertainty to preserve the snippet that reduces their uncertainty to the utmost. To facilitate solving complex tasks that require multiple retrievals, SeaKR utilizes their self-aware uncertainty to choose among different reasoning strategies. Our experiments on both complex and simple Question Answering datasets show that SeaKR outperforms existing adaptive RAG methods. We release our code at <a class="link-external link-https" href="https://github.com/THU-KEG/SeaKR" rel="external noopener nofollow">this https URL</a>.

Computation and Language

What problem does this paper attempt to address?

This paper proposes a new approach called Self-aware Knowledge Retrieval (SEAKR) for adaptive retrieval enhanced generation (RAG) model. RAG is an effective strategy to alleviate the incorrect information (hallucination) generated by large language models (LLMs), which dynamically determines when external knowledge is needed for generation and performs retrieval. SEAKR activates retrieval by extracting the self-aware uncertainty of the internal state of LLMs, and triggers retrieval when the model exhibits high self-aware uncertainty during the generation process. In order to effectively integrate the retrieved knowledge fragments, SEAKR re-ranks these fragments based on the self-aware uncertainty of LLMs and selects them for retention based on the degree of uncertainty reduction. In addition, for complex tasks that require multiple retrievals, SEAKR utilizes self-aware uncertainty to select different reasoning strategies.

Compared with existing adaptive RAG methods, SEAKR performs better on complex and simple question-answering tasks. The paper mentions that most existing methods default to knowledge retrieval for every input query, which may lead to misleading or conflicting information. SEAKR accurately determines knowledge needs by utilizing the self-aware uncertainty of LLMs, thereby improving efficiency and accuracy. Experimental results show that dynamically integrating retrieved knowledge brings greater performance improvement than self-aware retrieval alone, emphasizing the necessity of dynamic integration in adaptive RAG.