PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Chi-Min Chan, Chunpu Xu, Ruibin Yuan, Hongyin Luo, Wei Xue, Yike Guo, Jie Fu

arXiv:2404.00610v1cs.CL

TL;DR

RQ-RAG targets retrieval-augmented generation systems that do not adequately handle ambiguous or complex queries and may rely on unsuitable retrieval context. It trains a 7B Llama2 model to refine queries through rewriting, decomposition, and disambiguation, using contextually grounded answer generation and controllable search trajectories. The method surpasses prior state of the art across three single-hop QA tasks and performs strongly on complex multi-hop QA scenarios.

  • Problem

    Existing retrieval-augmented generation approaches primarily retrieve context from the initial query, limiting their handling of ambiguous or complex questions requiring clarification or decomposition.

  • Method

    RQ-RAG trains a 7B Llama2 model end-to-end to dynamically rewrite, decompose, or disambiguate queries, using retrieved contexts and control tokens during answer generation.

  • Results

    RQ-RAG surpasses the previous state-of-the-art method across three single-hop QA tasks and demonstrates superior performance across three multi-hop QA tasks.

  • Takeaways & Limitations

    Query refinement and contextually grounded answer regeneration provide the framework with strong performance across single-hop and complex multi-hop QA settings.

  • Takeaways & Limitations

    Trajectory selection and retrieved-context quality remain improvement opportunities, including LLM-based scoring, context reranking, and explicit noise reduction.

Abstract

from arXiv · show

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.

1 Introduction

RQ-RAG addresses limitations of retrieval-augmented generation by teaching a 7B Llama2 model to search on demand and refine queries through rewriting, decomposition, and disambiguation. The method uses contextually grounded answer regeneration and improves performance over prior methods on single-hop and multi-hop question answering tasks.

  • LLMs’ static, parametric knowledge can leave them without up-to-date information and prone to hallucinations.
  • Retrieval can help access external information, but indiscriminate retrieval may introduce irrelevant context and reduce answer quality.
  • RQ-RAG trains a 7B Llama2 model end-to-end to search on demand and dynamically rewrite, decompose, or disambiguate queries.
  • RQ-RAG outperforms the previous state-of-the-art method on three single-hop QA tasks and demonstrates superior performance on three multi-hop QA tasks.
  • The method regenerates answers from retrieved contexts and uses control tokens to select query-refinement actions or terminate search before answering.
  • The framework also shows a high upper bound and resilience to different data sources compared with previous methods.

2 RQ-RAG: Learning to Refine Query for Retrieval Augmented Generation

RQ-RAG constructs training data that mirrors inference-time query refinement and retrieval, then trains a 7B Llama2 model to generate and select among refined-query trajectories. The approach supports rewriting, decomposition, and disambiguation, with trajectories selected using model-internal scoring methods.

  • 2.1 Dataset Construction: The dataset transforms original input-output pairs into iterative refinement actions, refined queries, retrieved documents, and a final context-conditioned answer.Each iteration specifies a refinement type, retrieves top-k documents, and ultimately generates a new answer.
  • 2.1 Dataset Construction: Training data covers multi-turn dialogue, decomposition-required queries, and ambiguity-resolution queries.The task pool is transformed through an automated annotation workflow using ChatGPT to generate refined queries and renewed responses from retrieved contexts.
  • 2.2 Generator Training: The generator is trained autoregressively to maximize the likelihood of responses conditioned on the input, refined queries, and retrieved documents.The objective averages response-generation likelihood over the constructed dataset.
  • 2.3 Sampling Strategies: At inference time, tree decoding explores trajectories produced by different query-refinement actions before selecting a final response.A trajectory contains the input, intermediate query-context steps, and the final answer.
  • 2.3 Sampling Strategies: RQ-RAG selects trajectories using perplexity, final-answer confidence, or cumulative-confidence ensembling, without relying on an external evaluator.The upper bound counts an instance as correct when any generated trajectory contains the correct answer.
  • 2.3 Sampling Strategies: Unlike prior sampling methods, RQ-RAG does not require a larger model to evaluate trajectories and is not limited to fixed-option final answers.These differences distinguish its sampling strategy from the cited approaches.

3 Experiments

The experiments evaluate RQ-RAG on single-hop and multi-hop question-answering tasks against no-retrieval and retrieval-based baselines. The benchmark suite includes three datasets in each QA category, with additional multi-hop comparisons using Chain-of-Thought and Chain-of-Note.

  • 3 Experiments: RQ-RAG is evaluated on two QA categories: single-hop and multi-hop question answering.The experiments explicitly assess performance across both task types.
  • 3 Experiments: The single-hop suite comprises Arc-Challenge, PopQA, and OpenbookQA.These are the three single-hop datasets used in evaluation.
  • 3 Experiments: The multi-hop suite comprises HotpotQA, 2WikiMultiHopQA, and Musique.These datasets require multi-hop question answering.
  • 3 Experiments: Comparisons include no-retrieval and retrieval-based baselines using zero-shot and task-specific fine-tuned Llama2-7B models.The baseline groups differ in whether they use contexts retrieved from external databases.
  • 3 Experiments: Multi-hop comparisons additionally include Chain-of-Thought and Chain-of-Note using ChatGPT and GPT-4 as underlying models.These methods extend the multi-hop baseline set beyond Llama2 configurations.

4 Results and Analysis

RQ-RAG outperforms baseline methods on both single-hop and multi-hop QA, while query refinement and search-augmented training contribute to its performance. Additional analyses examine sampling strategies, answer regeneration, and retrieval-source robustness.

  • 4.1 Single-hop QA: 33.5% average improvement over Llama2-7B Zero Shot occurs in retrieval settings, while search-augmented training adds value beyond supervised baselines.RQ-RAG also surpasses baselines in non-retrieval settings.
  • 4.1 Single-hop QA: 20.3% average improvement over SAIL-7B and 1.9% over Self-RAG are reported across three QA tasks, despite using about 40k versus 150k supervised training examples.Self-RAG is identified as the former state-of-the-art.
  • 4.2 Multi-hop QA: 22.6% average enhancement across three multi-hop QA datasets results from autonomous query refinement, which enables query decomposition beyond direct retrieval of the original query.RQ-RAG also outperforms Chain-of-Thought and Chain-of-Note baselines despite its smaller backbone than ChatGPT.
  • 4.3 Sampling strategies: Confidence-based sampling generally performs best on single-hop QA, whereas ensemble-based sampling performs best on multi-hop QA; PPL is moderate by comparison.The comparison covers six tasks.
  • 4.3 Upper bound: 76.8% in ARC_C, 65.6% in POPQA, 84.0% in OBQA, 80.5% in HOTPOTQA, 60.6% in 2WIKI, and 54.5% in MUSIQUE are the reported upper-bound success rates.Success means that any generated trajectory reaches the correct answer.
  • 4.6 Limitations: More effective trajectory selection, context reranking, and explicit noise reduction remain future improvement directions because the current upper bound is not the absolute limit.The paper specifically suggests leveraging LLMs to score generated trajectories.
  • 4.4 Answer regeneration: 0% retention of the original answer yields the best HotpotQA performance, and effectiveness declines as retained-answer proportion increases.The 0% setting regenerates the answer entirely from retrieved contexts.
  • 4.5 Retrieval sources: Inference retrieval-source changes have negligible impact, with RQ-RAG showing variance of 0.7 compared with Self-RAG’s 1.8.The evaluated sources are DuckDuckGo, Wikipedia, and Bing Search.

5 Related Works

RAG research has improved both retrieval and generation, while prior work has explored filtering, self-reflection, and query rewriting to address irrelevant contexts.

  • RAG systems use external data to supplement LLMs’ parametric knowledge during generation.
  • Retrieval quality matters because irrelevant contexts can harm LLM generation, while semantically relevant retrieval outperforms BM25-based retrieval.
  • SAIL trains LLMs to distinguish irrelevant contexts, whereas Self-RAG trains self-reflection over retrieved contexts.
  • Rewrite-Retrieve-Read trains a small model to rewrite queries for a black-box reader, while RQ-RAG trains the LLM to refine queries itself.

6 Conclusion

RQ-RAG trains LLMs on a curated dataset to refine queries through rewriting, decomposition, and disambiguation. It outperforms prior SOTA methods on single-hop QA and performs strongly on complex multi-hop QA.

  • RQ-RAG trains LLMs on a curated dataset to rewrite, decompose, and disambiguate queries.
  • RQ-RAG surpasses prior SOTA methods across three single-hop QA tasks and performs better on complex multi-hop QA scenarios, including comparisons with ChatGPT.

A.1 Data Statistics

The curated data spans multiple QA categories and instruction-following tasks. Search augmentation substantially increases dataset length compared with the original data.

  • The data includes single-hop QA, multi-hop QA, ambiguous tasks, and instruction-following tasks.
  • The dataset uses Arc-Easy/Arc-Challenge and OpenbookQA for single-hop QA, HotpotQA and Musique for multi-hop QA, and ASQA for ambiguous tasks.
  • Search augmentation shifts most examples to 150–2000 tokens, whereas most original examples are shorter than 200 tokens.

A.2 Data Annotation

ChatGPT was used for reproducible data annotation, but examples with refusals or formatting failures were excluded. Prompt templates supported the annotation phases.

  • ChatGPT gpt-3.5-turbo-0125 was used for annotation with temperature set to 0 for reproducibility.
  • Annotation examples were excluded when ChatGPT refused to answer or failed to follow the required output format.
  • Tables 4–6 provide the prompt templates used across the annotation phases, with blue text marking input placeholders.

B.1 Training Hyperparameters

The method uses iterative query generation and retrieval, with queries tailored for rewriting, decomposition, or disambiguation. Data collection also covers multi-turn dialogue and complex query scenarios, while trajectory selection uses PPL, confidence, and ensemble strategies.

  • Training Hyperparameters: Models are trained for one epoch with a learning rate of 2e-5 and 3% warmup steps on 8 NVIDIA H800 GPUs.The maximum input length is 4096 because the dataset has an extended context length.
  • Iterative Query Refinement: The decoding process alternates between generating queries, retrieving contexts, and generating answers.Special tokens control expansion paths through the sequence generate → retrieve → generate → retrieve → answer.
  • Iterative Query Refinement: Queries are designed for rewriting, decomposition, or disambiguation to address different information needs.
  • Trajectory Selection: Final-answer trajectories are selected using perplexity, confidence, or an ensemble strategy without external LLM evaluation.The authors also estimate an upper bound by checking whether any generated trajectory contains the correct answer.

B.3 Evaluation Datasets and Metric

Evaluation covers single-hop and multi-hop question answering with task-specific metrics and retrieval settings. Single-hop tasks use DuckDuckGo contexts, while multi-hop tasks use candidate documents and select contexts at each step.

  • Single-Hop QA: Single-hop evaluation includes Arc-Challenge, PopQA, and OpenbookQA.Arc-Challenge and OpenbookQA contain four-choice questions, while PopQA uses a longtail subset.
  • Metrics: Accuracy measures Arc-Challenge and OpenbookQA, while PopQA uses a match score based on inclusion of ground truths.
  • Multi-Hop QA: Multi-hop evaluation samples 500 instances from HotpotQA-distractor, 2WikiMultihopQA, and MuSiQue-Ans.The experiments use a reading-comprehension setting with candidate documents from the original datasets.
  • Multi-Hop QA: HotpotQA questions link to 10 passages, with only 2 relevant passages; 2WikiMultihopQA questions have 2 or 4 relevant passages.
  • Retrieval: Inference retrieves three DuckDuckGo contexts for single-hop tasks and three embedding-selected contexts per step for multi-hop tasks.Multi-hop selection uses text-embedding-3-large over the candidate documents.

C Additional Results

Additional experiments examine how retention ratio affects two multi-hop QA tasks. The results indicate that regenerating answers grounded in provided contexts is important for downstream performance.

  • Interpretation: Regenerating answers grounded in provided contexts is crucial for downstream task performance.The same phenomenon is observed across the two multi-hop QA tasks shown in Figures 8 and 9.
  • Retention-Ratio Analysis: Retention ratio affects performance on 2WikiMultihopQA.
  • Retention-Ratio Analysis: Retention ratio affects performance on MuSiQue.
Loading 2404.00610v1…