Query Rewriting in Retrieval-Augmented Large Language Models (2023)

Abstract: 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.

## introduction large language models (llms) have shown remarkable abilities for human language processing and extraordinary scalability and adaptability in few-or zero-shot settings. @xcite @xcite . however, the training process depends on large-scale high-quality corpora but without the perception of the real world. thus, llms still have to face the issue of hallucination @xcite and temporal misalignment @xcite @xcite . this affects the reliability of llms and hinders wider practical application, because the consistency between the llm responses with the real world needs further validation. existing work has proved that incorporating external knowledge (i.e., non-parametric knowledge) with internal knowledge (i.e., parametric knowledge) can effectively alleviate hallucination, especially for knowledge-intensive tasks. in fact, retrievalaugmented llms have been shown so effective that they have been regarded as a standard solution to alleviate the factuality drawbacks in naive llm generations. retrieval augmentation is applied to select relative passages as external contexts for the language model, which is retrieve-then-read framework @xcite @xcite . take the open-domain question-answering task (open-domain qa) as an example, a retriever first searches for related documents for a question. then the llm receives the question and the documents, then predicts an answer. as most llms are only accessible through inference apis, they play the part of black-box frozen readers in the pipeline. this makes previous retrieval augmentation methods that require complete access @xcite @xcite no longer feasible. recent studies on retrieval-augmented language models lean more on the llm-oriented adaptation. an idea is to train a dense retrieval model to cater to the frozen language model @xcite . by using feedback from the llm as a training objective, the retrieval model is tuned for better llm input contexts. another research line focuses on the design of interactions between the retriever and the reader @xcite , where both the retriever and the reader are usually frozen. the idea is to trigger the emergent ability through carefully crafted prompts or a sophisticated prompt pipeline. multiple interactions with external knowledge allow the llm to approach the correct answer step by step. however, there are still problems remaining to be solved. existing approaches overlook the adaptation of the query, i.e., the input of the retrievethen-read pipeline. the retrieval query is either original from datasets or directly determined by the black-box generation, thus is always fixed. however, there is inevitably a gap between the input text and the knowledge that is really needed to query. this limits performance and places a burden on retrieval capability enhancement and prompt engineering. in consideration of this issue, this paper proposes rewrite-retrieve-read, a new framework for retrieval augmentation, which can be further tuned for adapting to llms. in front of the retriever, a step of rewriting the input is added, filling the gap between the given input and retrieval need, as is shown in figure 1 . we adopt the off-the-shelf tool, an internet search engine, as the retriever, which avoids the maintenance of the search index and can access up-to-date knowledge @xcite . different from previous studies @xcite that require the memory of multiple interaction rounds between the retriever and the llm for each sample, the motivation of our rewriting step is to clarify the retrieval need from the input text. we also propose a trainable scheme for our rewrite-retrieve-read framework (figure 1 (c) ). the black-box retriever and the reader form a frozen system. to further smooth the steps of our pipeline, we apply a small, trainable language model to perform the rewriting step, denoted as the rewriter. the rewriter is trained by reinforcement learning using the llm performance as a reward, learning to adapt the retrieval query to improve the reader on downstream tasks. our proposed methods are evaluated on knowledge-intensive downstream tasks including open-domain qa (hotpoqa @xcite , ambignq @xcite , popqa @xcite ) and multiple choice qa (mmlu @xcite ). the experiments are implemented on t5-large @xcite as the rewriter, chatgpt @xcite and vicuna-13b @xcite as the llm reader. the results show that query rewriting consistently improves the retrieve-augmented llm performance. the results also indicate that the smaller language model can be competent for query rewriting. to sum up, our proposed novel retrievalaugmentation method, rewrite-retrieve-read is the first framework where the input text is adapted for the frozen retriever and llm reader. we introduce a tuneable scheme with a small, trainable model, achieving performance gains with less resource consumption. ## methodology we present rewrite-retrieve-read, a pipeline that improves the retrieval-augmented llm from the perspective of query rewriting. figure 1 shows an overview. this section first introduces the pipeline framework in section 3.1, then the trainable scheme in section 3.2. ## implementation rewriter for the frozen pipeline in §3.1, we prompt an llm to rewrite the query with few-shot in-context learning @xcite . our prompt follows the formulation of [instruction, demonstrations, input], where the input is x. the instruction is straightforward and demonstrations are 1-3 random examples from training sets and are kept constant across all runs, mainly for the task-specific output format illustration, i.e., a short phrase as an answer for hotpotqa, and an option as an answer for mmlu. for the training scheme in §3.2, we fine-tuning a t5 as the rewriter. retriever we use the bing search engine as the retriever. it requires no candidate index construction like a dense retriever, nor candidates like a textbook. but it allows for a wide knowledge scope and up-to-time factuality. with bing api, the retrieval is performed in two approaches. (i) for all retrieved web pages, we concatenate the snippets that are related sentences selected by bing. this method is similar to using a search engine in a browser, input a query and press enter, then collect the texts shown on the search result page. (ii) for retrieved web pages, we request the urls and parser to get all the texts. this is similar to clicking on items on the search result page. then we use bm25 to keep those with higher relevance scores with the query, reducing the document length. reader the reader is a frozen llm, where we adopt chatgpt (gpt-3.5-turbo) and vicuna-13b. it performs reading comprehension and prediction with few-shot in-context learning. in our prompt, following the brief instruction and the demonstrations, the input is x or [doc, x] with retrieval augmentation. it has been proved that both the phrasing of prompt lines @xcite and the selection of demonstrations show effects on the in-context learning performance @xcite . as it is not the focus of this work, we pay no more attention to prompt editing. ## analysis 6.1 training process the training process includes two stages, warm-up and reinforcement learning. this section shows the validation scores of the three open-domain qa datasets for further analysis. figure 2 presents the metric scores through training iterations in the process of reinforcement learning. as the rewriting models have been warmed up on the pseudo data before rl, scores at "0 iteration" denote the ability acquired from the warm-up training. it can be observed that the curves show upward trends with some fluctuations on all the datasets. (i) for multi-hop questions in hotpotqa, the standard retrieval is relatively weaker. complex questions can be not specific search queries and show a larger gap from rewritten queries, i.e., the green and red lines. (ii) on ambignq and popqa, our method surpasses the baselines after several iterations (3 or 4). this indicates that the rl training stage can compensate for the insufficiency of the distillation on the pseudo data during warm-up training. (iii) in particular, on popqa, the trainable rewriter remains inferior to the llm rewriter. this can be explained as the dataset is constructed for adaptive retrieval @xcite , which only uses retrieval where it helps to avoid harmful redundant retrieval. thus, "none" is a possible query that means no retrieval. this causes more complexity and uncertainty. llm rewriter knows better when the retrieval is needed for itself as a reader, although the rewriting step is not concatenated as the input context of the reader. we calculate the performance of query "none". the questions that can be correctly answered without retrieval (i.e., the "direct" method) are those samples that need no more context. comparing this retrieval-free set with those that are rewritten to be"none" query, the f 1 score of the llm rewriter is 71.9% and the t5 rewriter score is 67.1%. if we consider the questions that can be correctly answered without retrieval but go wrong with retrieval as the retrieval-free set, the f 1 scores are 78.7% for llm rewriter and 77.4% for t5. ## conclusion this paper introduces the rewrite-retrieve-read pipeline, where a query rewriting step is added for the retrieval-augmented llm. this approach is applicable for adopting a frozen large language model as the reader and a real-time web search engine as the retriever. further, we propose to apply a tuneable small language model the rewriter, which can be trained to cater to the frozen retriever and reader. the training implementation consists of two stages, warm-up and reinforcement learning. evaluation and analyses on open-domain qa and multiple-choice qa show the effectiveness of query rewriting. our work proposes a novel retrieval-augmented black-box llm framework, proves that the retrieval augmentation can be enhanced from the aspect of query rewriting, and provides a new method for integrating trainable modules into black-box llms.