From Complex to Atomic: Enhancing Augmented Generation via Knowledge-Aware Dual Rewriting and Reasoning

Jinyu Wang, Jingjing Fu, Rui Wang, Lei Song, Jiang Bian

TL;DR: We present an advanced RAG system with knowledge-aware dual rewriting and reasoning capabilities, designed to improve knowledge extraction and rationale formulation within specialized datasets.

Abstract: Recent advancements in Retrieval-Augmented Generation (RAG) systems have significantly enhanced the capabilities of large language models (LLMs) by incorporating external knowledge retrieval. However, the sole reliance on retrieval is often inadequate for mining deep, domain-specific knowledge and for performing logical reasoning from specialized datasets. To tackle these challenges, we present an approach, which is designed to extract, comprehend, and utilize domain knowledge while constructing a coherent rationale. At the heart of our approach lie four pivotal components: a knowledge atomizer that extracts atomic questions from raw data, a query proposer that generates subsequent questions to facilitate the original inquiry, an atomic retriever that locates knowledge based on atomic knowledge alignments, and an atomic selector that determines which follow-up questions to pose guided by the retrieved information. Through this approach, we implement a knowledge-aware task decomposition strategy that adeptly extracts multifaceted knowledge from segmented data and iteratively builds the rationale in alignment with the initial query and the acquired knowledge. We conduct comprehensive experiments to demonstrate the efficacy of our approach across various benchmarks, particularly those requiring multihop reasoning steps. The results indicate a significant enhancement in performance, up to 12.6\% over the second-best method, underscoring the potential of the approach in complex, knowledge-intensive applications.

Lay Summary:

Current Retrieval-Augmented Generation (RAG) methods encounter significant challenges in domain-specific reasoning. A primary issue is that the corpus available may not always possess the necessary information to support each step of the complex reasoning processes that sophisticated language models might generate.

Our approach involves tagging each original corpus document with atomic knowledge units and iteratively rewriting the user query based on the information already retrieved. This dual-rewriting mechanism bridges the gap between the user query and the available corpus. Moreover, the iterative process evolves step-by-step, culminating in a comprehensive response to the query. This mimics the meticulous approach researchers take in refining questions and gathering evidence, ensuring a thoroughly reasoned and substantiated conclusion.

In public benchmarks featuring challenging multi-hop questions, our method demonstrated a 20% increase in accuracy over the second-best performing method. Moving forward, we plan to integrate in-context learning capabilities to dynamically select suitable examples and allow the components of our approach to learn from user feedback.

Application-Driven Machine Learning: This submission is on Application-Driven Machine Learning.

Link To Code: https://github.com/microsoft/PIKE-RAG

Primary Area: Applications->Language, Speech and Dialog

Keywords: RAG, Knowledge atomizing, Knowledge-aware task decomposition, Multihop QA

Submission Number: 4599