PaperlayerView exact revision on arXiv

Source-linked AI summary

Query Rewriting for Retrieval-Augmented Large Language Models

Xinbei Ma, Yeyun Gong, Pengcheng He, Hai Zhao, Nan Duan

arXiv:2305.14283v3cs.CL

TL;DR

Retrieval-augmented LLMs face a mismatch between input text and the knowledge needed for retrieval, especially when the reader is a black-box model. The paper introduces Rewrite-Retrieve-Read, which rewrites queries before web retrieval and optionally trains a small rewriter with reader feedback. Evaluations on open-domain and multiple-choice QA show consistent performance improvement, while the framework adapts frozen retriever-reader systems with a smaller trainable component.

  • Problem

    Fixed retrieval queries can fail to express the knowledge needed from an input, while black-box LLM readers limit adaptation of existing retrieval-augmentation methods.

  • Method

    Rewrite-Retrieve-Read adds query rewriting before web search, with a small rewriter trained by reinforcement learning from the frozen LLM reader's performance.

  • Results

    Query rewriting consistently improves retrieval-augmented LLM performance across open-domain and multiple-choice QA, and smaller language models can perform the rewriting.

  • Takeaways & Limitations

    The framework enhances retrieval augmentation by adapting the search query while retaining a frozen retriever and LLM reader.

  • Takeaways & Limitations

    Training introduces a trade-off between downstream-task specialization and generalization, compromising scalability to direct transfer compared with few-shot in-context learning.

Abstract

from arXiv · show

Large Language Models (LLMs) play powerful, black-box readers in the retrieve-then-read pipeline, making remarkable progress in knowledge-intensive tasks. This work introduces a new framework, Rewrite-Retrieve-Read instead of the previous retrieve-then-read for the retrieval-augmented LLMs from the perspective of the query rewriting. Unlike prior studies focusing on adapting either the retriever or the reader, our approach pays attention to the adaptation of the search query itself, for there is inevitably a gap between the input text and the needed knowledge in retrieval. We first prompt an LLM to generate the query, then use a web search engine to retrieve contexts. Furthermore, to better align the query to the frozen modules, we propose a trainable scheme for our pipeline. A small language model is adopted as a trainable rewriter to cater to the black-box LLM reader. The rewriter is trained using the feedback of the LLM reader by reinforcement learning. Evaluation is conducted on downstream tasks, open-domain QA and multiple-choice QA. Experiments results show consistent performance improvement, indicating that our framework is proven effective and scalable, and brings a new framework for retrieval-augmented LLM.

1 Introduction

Retrieval augmentation can improve factuality, but black-box LLM readers and fixed queries leave a gap between input text and retrieved knowledge. Rewrite-Retrieve-Read addresses this gap by rewriting queries and can further train a small rewriter using reader feedback.

  • Retrieval augmentation supplies external knowledge to mitigate hallucination and temporal misalignment in knowledge-intensive tasks.
  • Black-box LLM access makes methods requiring complete retriever-reader access infeasible, motivating adaptation around frozen modules.
  • Existing approaches leave retrieval queries fixed, despite a gap between the input text and the knowledge needed for retrieval.
  • Rewrite-Retrieve-Read inserts query rewriting before retrieval and uses an internet search engine to access up-to-date knowledge without maintaining a search index.
  • A small trainable rewriter is optimized with reinforcement learning from LLM-reader performance to adapt queries to frozen retriever and reader modules.
  • Experiments on open-domain and multiple-choice QA show consistent performance improvement, while the smaller language model can perform query rewriting.
  • The framework adapts input text for frozen retriever and reader modules and achieves performance gains with less resource consumption.

2 Related Work

Prior retrieval-augmented systems use retrievers, prompting, or trainable components to improve LLM responses, but this work focuses on adapting the search query before retrieval. Its framework adds query rewriting ahead of retrieve-then-read and introduces a small rewriting model.

  • 2.1 Retrieval Augmentation: Retrieval augmentation provides factual context such as common sense and real-time news for contextualized language-model prediction.
  • 2.1 Retrieval Augmentation: Earlier retrieve-then-read systems commonly use sparse or dense retrievers with trainable pretrained language-model readers in end-to-end frameworks.
  • 2.2 Cooperation with Black-box LLMs: LLMs are often accessed as black boxes because some are private and their scale requires substantial computational resources.
  • 2.2 Cooperation with Black-box LLMs: Prompt-based methods and interactive pipelines leverage LLMs through generated context, chain-of-thought, web APIs, multi-hop decomposition, and retrieval.
  • 2.2 Cooperation with Black-box LLMs: RePlug fine-tunes a dense retriever under supervision from a frozen LLM, exemplifying adaptation of retrieval to the reader.
  • 2.2 Cooperation with Black-box LLMs: Unlike prior approaches, the proposed pipeline adds query rewriting before retrieve-then-read and uses a small rewriting model.

3 Methodology

Rewrite-Retrieve-Read inserts query rewriting before retrieval so the generated query better matches the knowledge needed for the task. Its trainable rewriter is optimized with reader feedback through reinforcement learning while the retriever and reader remain frozen.

  • Rewrite-Retrieve-Read: Rewrite-Retrieve-Read adds a query-rewriting step that generates a search query before retrieving context and reading the input with that context.The pipeline maps input x to rewritten query x̃, retrieved document doc, and prediction ŷ.
  • Rewrite-Retrieve-Read: An LLM can use few-shot prompting to generate one or more queries for information potentially needed by the task.The prompted LLM performs the rewriting step in the frozen pipeline.
  • Trainable Rewriter: A small trainable language model, initialized from T5-large, takes over rewriting and is first trained on pseudo data before reinforcement learning.Pseudo labels are generated by prompting an LLM, and samples receiving correct reader predictions form the warm-up dataset.
  • Trainable Rewriter: Warm-up performance is modest because it depends on pseudo-data quality, prompt reliance, and the relatively small rewriter size, motivating subsequent reinforcement learning.The reinforcement-learning stage is intended to align the rewriter with the retriever and LLM reader.
  • Reinforcement Learning: The rewriter is optimized as a policy-gradient agent whose generated query is evaluated after retrieval and reading through the LLM reader's final prediction.The formulation uses a Markov Decision Process, PPO, and a value network with Generalized Advantage Estimation.
  • Reinforcement Learning: The reward combines task performance measures such as exact match and F1 with KL-divergence regularization against the initialization policy.The regularization is intended to prevent the rewriter from deviating too far from its initialization.

4 Implementation

The implementation uses few-shot prompting for query rewriting and fine-tunes T5 as the trainable rewriter. Prompt wording and demonstration selection are recognized as factors affecting in-context performance but are not further optimized.

  • Rewriter: The prompted rewriter uses an instruction, constant few-shot demonstrations, and an input, while fine-tuning uses T5 as the trainable rewriter.Demonstrations illustrate task-specific output formats for HotpotQA and MMLU.
  • Prompting: Prompt-line phrasing and demonstration selection affect in-context learning performance, but prompt editing is outside this work's focus.

5 Experiments

The experiments evaluate query rewriting for open-domain and multiple-choice QA using frozen or trainable rewriting settings with LLM readers. Query rewriting consistently improves performance across open-domain QA, while gains on multiple-choice QA vary by reader and category.

  • Evaluation setup: Evaluation covers HotpotQA, AmbigNQ, PopQA, and MMLU, using open-domain QA and four multiple-choice QA categories.HotpotQA tests multi-hop reasoning, AmbigNQ disambiguates questions, PopQA emphasizes long-tail knowledge, and MMLU covers Humanities, STEM, Social Sciences, and Other.
  • Evaluation setup: Open-domain QA uses Exact Match, F1, and Hit, with the reinforcement-learning reward defined as a weighted sum of these measures.Hit rewards retrieved content containing the answer and penalizes misses.
  • Evaluation setup: The study compares direct answering, retrieve-then-read, an LLM as frozen rewriter, and a trainable rewriter.In the rewriting settings, generated queries are passed to the retriever and reader; Table 1 provides the prompt lines used for the LLMs.
  • Open-domain QA results: Query rewriting consistently improves performance across the three open-domain QA datasets with both frozen and trainable rewriters.Standard retrieval helps on AmbigNQ and PopQA but hurts on HotpotQA; rewriting increases scores, although the trainable rewriter remains below the LLM rewriter on PopQA.
  • Multiple-choice QA results: On multiple-choice QA, query rewriting improves most ChatGPT-reader settings except Social Sciences, while Vicuna shows gains across all four categories.The reported MMLU metrics are Exact Match, which is also used as the reward.

6 Analysis

The analysis tracks reinforcement-learning progress, retrieval quality, and illustrative query effects across open-domain and multiple-choice settings. Rewriting improves retrieval and prediction, with performance varying by dataset and rewriter.

  • 6.1 Training Process: Validation curves show EM and F1 across HotpotQA, AmbigNQ, and PopQA, alongside EM baselines for standard retrieval and LLM rewriting.Scores are plotted through reinforcement-learning iterations; iteration 0 reflects warm-up training.
  • 6.1 Training Process: After three or four iterations, the proposed method surpasses baselines on AmbigNQ and PopQA, while standard retrieval is relatively weaker on HotpotQA.The HotpotQA gap is attributed to complex questions being less specific search queries.
  • 6.2 Retrieval Result: The pipeline evaluates retrieved-context quality with hit rate, measuring whether normalized retrieved context contains the correct answer.Because query rewriting affects context before reader output, QA metrics are indirect measurements.
  • 6.2 Retrieval Result: On AmbigNQ, BM25 content selection recalls better documents than snippets, and retrieval-augmented scores are analyzed against no-retrieval and upper-bound settings.The upper-bound scores use samples whose retrieved contexts contain the correct answers.

7 Conclusion

The conclusion presents Rewrite-Retrieve-Read as a query-rewriting framework for retrieval-augmented black-box LLMs. It combines frozen retrieval and reading components with a trainable rewriter and reports effectiveness across QA settings.

  • 7 Conclusion: Rewrite-Retrieve-Read adds query rewriting before retrieval, supporting a frozen LLM reader and a real-time web-search retriever.The framework adapts input text for frozen pipeline components.
  • 7 Conclusion: A small trainable language model performs rewriting through warm-up and reinforcement-learning stages to cater to the frozen retriever and reader.The reported evaluations cover open-domain QA and multiple-choice QA.
  • 7 Conclusion: Evaluations and analyses show the effectiveness of query rewriting and provide a way to integrate trainable modules into black-box LLM systems.The conclusion characterizes the framework as a retrieval-augmented black-box LLM approach.

Limitations

The paper identifies limits in transferability, interaction design, and retrieval control. These boundaries reflect trade-offs between task specialization, one-turn operation, and web-search retrieval.

  • Limitations: Training introduces a trade-off between downstream-task specialization and generalization, compromising direct-transfer scalability relative to few-shot in-context learning.This limitation concerns adding a training process to the framework.
  • Limitations: Unlike LLM-agent approaches that make multiple calls and follow-up questions, the proposed method enhances a one-turn retrieve-then-read framework with a trainable query rewriter.The paper explicitly distinguishes its motivation from multi-call agent systems.
  • Limitations: Using a web search engine may limit retrieval quality and control compared with neural dense retrievers built on professional, filtered knowledge bases.The paper states that such dense retrievers may potentially achieve better and more controllable retrieval.

A Warm-up Dataset

Warm-up training constructs a pseudo dataset for query rewriting from benchmark data, using official splits when available and random splits otherwise. PopQA receives a stratified train-test split.

  • A Warm-up Dataset: The pseudo rewriting dataset uses whole official training sets for HotpotQA and AmbigNQ, while PopQA and MMLU are randomly split because they lack official splits.The construction supports warm-up training of the tuneable rewriter.
  • A Warm-up Dataset: PopQA contains 16 question types and is split into 13k training examples and 714 testing examples using stratified sampling.The passage introduces the dataset statistics and warm-up-data split table.

B Setup Details

The pipeline uses specified warm-up and reinforcement-learning settings for the trainable rewriter, alongside BM25-based retrieval. Multiple-choice QA metrics are reported in Table 5.

  • Training: Warm-up training uses a 3e-5 learning rate, batch sizes of 16 or 20, and 6, 8, or 12 epochs.
  • Evaluation: Table 5 reports metrics for the multiple-choice QA evaluation.
  • Training: Reinforcement learning samples 5120 steps across 10 threads, with 512 steps per thread, then trains the policy network for 2, 3, or 4 epochs.
  • Retrieval: BM25 retrieval parses textbox paragraphs from searched URLs, scores them against the query, and retains high-scoring text until the context reaches its maximum length.

C Web Search: Tool Use

The framework uses an external web search engine to retrieve broad, current knowledge, while a trainable small model helps align language-model queries with external tools. This flexibility introduces practical and contextual risks from web retrieval.

  • Tool Use: The pipeline places an externally built web search engine in the retriever module.
  • Advantages: Web search provides access to a wide-ranging, up-to-date knowledge base and can alleviate temporal misalignment in fixed candidate databases.
  • Limitations: Web APIs require subscriptions, while retrieved Internet context can be difficult to control and occasionally inconsistent, redundant, or toxic.
  • Alignment: A trainable small model is proposed to align language models with tools that lack trainable parameters, including search engines.
Loading 2305.14283v3…