RQ-RAG: Learning to Refine Queries for Retrieval Augmented Generation

TL;DR

Proposes RQ-RAG, which refines queries via rewriting, decomposition, and disambiguation, boosting 7B Llama2 performance on single/multi-hop QA, surpassing SOTA.

cs.CL 🔴 Advanced 2024-03-31 35 views
Chi-Min Chan Chunpu Xu Ruibin Yuan Hongyin Luo Wei Xue Yike Guo Jie Fu
NLP retrieval-augmented generation question answering query rewriting multi-task learning

Key Findings

Methodology

RQ-RAG employs control-token-guided query rewriting, decomposition, and disambiguation within an end-to-end framework. It leverages ChatGPT for automated multi-scenario data generation, training the model to select optimal query transformations via multi-strategy (perplexity, confidence, ensemble) filtering. The core algorithm involves autoregressive training (formula 2) and multi-path sampling, enabling dynamic query refinement. Experiments on datasets like Arc-Challenge, PopQA, HotpotQA, and others demonstrate significant improvements over baselines, especially in complex multi-hop tasks.

Key Results

  • On three single-hop QA datasets, RQ-RAG with 7B Llama2 outperforms previous SOTA by 1.9% on average, with an increase of 33.5% over zero-shot models. In multi-hop QA, it improves performance by 22.6%, with highest success rate reaching 84.0%. The multi-strategy sampling yields the best trajectory selection, significantly boosting answer accuracy.
  • Automated data construction using ChatGPT enables the model to learn query rewriting, decomposition, and disambiguation across diverse scenarios, enriching training data and improving generalization. Multi-path sampling allows the model to iteratively retrieve and refine information, leading to more accurate answers.
  • Robustness tests across data sources (Bing, Wikipedia, DuckDuckGo) show minimal performance variance, indicating high system stability. Upper bound analysis reveals a success rate of up to 76.8%, highlighting the potential for further improvements through better trajectory evaluation and context reranking.

Significance

This work advances retrieval-augmented generation by integrating learned query refinement, addressing limitations in handling complex, ambiguous, or multi-hop questions. It enhances the ability of large models to dynamically adapt queries, leading to more accurate, contextually grounded answers. The approach bridges the gap between static knowledge and real-time information, promising significant impact on practical AI applications like customer service, knowledge retrieval, and intelligent assistants. Its scalable framework offers a new paradigm for multi-task, multi-modal AI systems, fostering future innovations.

Technical Contribution

Key innovations include: 1) a control-token-based end-to-end training framework for query rewriting, decomposition, and disambiguation; 2) automated multi-scenario data generation via ChatGPT; 3) multi-strategy trajectory filtering (perplexity, confidence, ensemble) to select optimal responses. Compared to SOTA methods like Self-RAG and SAIL, this work introduces explicit query optimization, improving retrieval relevance and answer accuracy, especially in complex multi-hop scenarios.

Novelty

This is the first comprehensive system integrating learned query rewriting with multi-strategy trajectory filtering in retrieval-augmented generation. Unlike prior works that focus solely on retrieval or generation, this approach actively refines queries during inference, enabling the model to handle complex, multi-step reasoning tasks more effectively. It bridges the gap between static retrieval and dynamic query adaptation, setting a new standard for QA systems.

Limitations

  • The reliance on automatically generated data introduces noise, which may limit generalization. Training costs are high due to multi-path sampling and large-scale data creation.
  • Trajectory selection strategies, while effective, can still misjudge the best path, requiring more sophisticated evaluation methods.
  • Data generation depends on ChatGPT, which has knowledge cutoff and may not reflect real-time information, affecting temporal relevance.

Future Work

Future directions include developing more intelligent trajectory evaluation mechanisms, integrating real-time data sources, and extending the framework to multi-modal inputs such as images and videos. Additionally, optimizing computational efficiency and exploring reinforcement learning for dynamic query refinement will further enhance system performance and applicability in real-world scenarios.

AI Executive Summary

The rapid development of large language models (LLMs) like GPT-3 and GPT-4 has revolutionized natural language understanding and generation. However, these models face inherent limitations: static knowledge bases that cannot incorporate real-time information, and difficulties in handling complex or ambiguous queries. Retrieval-augmented generation (RAG) emerged as a promising solution, combining external knowledge retrieval with generative models to improve accuracy and updateability. Yet, existing RAG systems often rely on static queries, which can be insufficient for complex questions requiring multi-step reasoning or clarification.

This paper introduces RQ-RAG, a novel framework that endows models with the ability to learn query refinement strategies. By training a 7B Llama2 model with a carefully curated dataset, the authors enable the system to dynamically rewrite, decompose, and disambiguate user queries based on context. This process is guided by control tokens, allowing the model to choose among different query transformation paths during inference. The training data is generated automatically using ChatGPT, covering diverse scenarios such as multi-turn dialogues, complex decompositions, and ambiguous questions. Multiple trajectory sampling strategies—perplexity, confidence, and ensemble—are employed to select the most relevant retrieval path, significantly improving answer accuracy.

Experimental results across multiple datasets demonstrate that RQ-RAG surpasses previous state-of-the-art methods, including Self-RAG and SAIL, with an average improvement of 1.9% on single-hop QA and 22.6% on multi-hop QA. The model exhibits high robustness across different data sources, maintaining performance stability. Furthermore, the upper bound analysis indicates a potential success rate of up to 84%, highlighting the system’s strong capacity for query optimization and knowledge retrieval.

Overall, RQ-RAG represents a significant step forward in retrieval-augmented language modeling, offering a flexible, scalable approach to complex question answering. Its ability to actively refine queries paves the way for more intelligent, context-aware AI systems capable of real-time knowledge integration, with broad implications for industry and research. Future work will focus on enhancing trajectory evaluation, incorporating multi-modal data, and reducing computational costs, aiming to realize fully autonomous, real-time intelligent assistants.

Deep Analysis

Background

Recent advances in LLMs like GPT-3, BERT, and their后续版本极大推动了自然语言理解与生成技术的发展。这些模型在多任务、多领域表现出色,但其知识静态、更新缓慢的问题逐渐显现。检索增强生成(RAG)技术应运而生,通过引入外部知识库,解决模型知识过时和信息不足的问题。早期工作如Karpukhin等提出的Dense Retriever,显著优于传统的BM25检索方法。随后,Self-RAG、SAIL等模型结合检索与生成,提升了多任务适应性。然而,面对复杂、多义或多跳问题,现有方法多依赖静态检索策略,难以动态应对查询歧义与结构复杂性,限制了其实际应用效果。

Core Problem

当前检索增强模型在处理复杂、多义或多跳问答时,存在检索不准确、查询歧义未能有效消解的问题。传统方法多采用原始查询直接检索,忽略了查询的潜在歧义与结构复杂性,导致检索结果偏离用户意图,影响答案的准确性。此外,模型缺乏主动优化查询的能力,难以在多轮交互中逐步精炼信息,限制了推理深度。如何让模型自主学习重写、拆解与消歧查询,成为提升问答性能的关键。

Innovation

本文的核心创新包括:1)引入基于控制标记的端到端训练机制,使模型能自主选择重写、拆解或消歧查询,提升多场景适应性;2)利用ChatGPT自动生成多场景、多轮对话数据,丰富训练样本,增强模型泛化能力;3)提出多策略轨迹筛选方法(困惑度、置信度、集成),有效提升答案准确率。这些创新突破了以往只优化单一检索或生成环节的局限,实现了查询的主动优化与动态调整,为复杂问答提供了新思路。

Methodology

  • �� 数据集构建:采集多场景(多轮对话、拆解、歧义)任务,利用ChatGPT自动生成重写、拆解、消歧查询及对应答案。
  • �� 训练目标:采用自回归模型(公式2),端到端学习查询重写、拆解与答案生成。
  • �� 查询策略:利用控制标记引导模型选择不同路径(重写、拆解、消歧、终止),实现多轨迹生成。
  • �� 轨迹筛选:基于困惑度、置信度及集成策略筛选最优路径,提升答案质量。
  • �� 训练流程:结合自动生成的数据,进行多轮训练,优化模型对复杂查询的处理能力。

Experiments

在单跳(Arc-Challenge、PopQA、OpenbookQA)和多跳(HotpotQA、2WikiMultiHop、Musique)问答任务上,采用公开数据集,比较无检索、检索增强、以及本文提出的RQ-RAG模型。指标包括准确率、F1值,超参数设置遵循原始论文,进行多轮验证。模型训练采用Adam优化器,批次大小为32,训练时间约为48小时。多策略筛选显著提升多场景问答性能。

Results

在三项单跳问答中,RQ-RAG平均超越SOTA 1.9%,多跳任务提升22.6%。在检索增强设置中,超越无检索模型33.5%,优于自监督的Self-RAG 20.3%。多策略筛选达成最高84.0%的成功轨迹比例。模型在不同数据源(Bing、Wikipedia)表现稳定,显示出强大的鲁棒性与潜力。

Applications

该技术适用于企业客服、智能问答、知识库检索等场景,提升系统对复杂、多义问题的理解与回答能力。未来结合多模态信息(图像、视频)将拓展其应用范围,推动智能助理、自动问答的产业升级。

Limitations & Outlook

模型依赖大量自动生成数据,存在噪声与偏差,训练成本较高。轨迹筛选策略仍有误判风险,未来需引入更智能的评估机制。多场景数据由ChatGPT生成,可能受限于其知识截止时间,影响时效性。未来应结合实时信息源,优化模型效率与准确性。

Plain Language Accessible to non-experts

想象你在厨房做菜,面对一道复杂菜谱,单纯按照原始食谱操作,可能会出现味道不佳或步骤繁琐。现在,如果你能提前准备好不同的调料和步骤,甚至根据味道调整配料,做菜就会变得更简单、更好吃。这就像这篇论文里的模型,它学会了根据不同的情况,主动调整“菜谱”——也就是查询内容——以获得更准确的答案。模型会像厨师一样,先理解问题的复杂性,然后拆解成更简单的部分,或者用不同的方式表达,确保每一步都精准无误。这样,不管问题多复杂,模型都能灵活应对,做出满意的“菜肴”。

ELI14 Explained like you're 14

想象你在学校里问老师一个很难的问题,比如“为什么天是蓝色的?”如果老师只用一句话回答,可能会不够详细。可是如果老师能先问你一些问题,弄清楚你具体想知道的内容,再用更合适的方式回答,那答案就会更清楚、更贴合你的需要。这篇论文的模型就像那个聪明的老师,它学会了在回答问题前,先“重写”问题,把复杂的问题拆成更简单的部分,或者用不同的表达方式,确保理解清楚。它还会根据不同情况,选择不同的“回答策略”,让答案更准确、更完整。这样,无论问题多复杂,模型都能像老师一样,给出最合适的答案,帮助我们更好地理解世界。

Abstract

Large Language Models (LLMs) exhibit remarkable capabilities but are prone to generating inaccurate or hallucinatory responses. This limitation stems from their reliance on vast pretraining datasets, making them susceptible to errors in unseen scenarios. To tackle these challenges, Retrieval-Augmented Generation (RAG) addresses this by incorporating external, relevant documents into the response generation process, thus leveraging non-parametric knowledge alongside LLMs' in-context learning abilities. However, existing RAG implementations primarily focus on initial input for context retrieval, overlooking the nuances of ambiguous or complex queries that necessitate further clarification or decomposition for accurate responses. To this end, we propose learning to Refine Query for Retrieval Augmented Generation (RQ-RAG) in this paper, endeavoring to enhance the model by equipping it with capabilities for explicit rewriting, decomposition, and disambiguation. Our experimental results indicate that our method, when applied to a 7B Llama2 model, surpasses the previous state-of-the-art (SOTA) by an average of 1.9\% across three single-hop QA datasets, and also demonstrates enhanced performance in handling complex, multi-hop QA datasets. Our code is available at https://github.com/chanchimin/RQ-RAG.

cs.CL