RAG query rewriting method based on multi-stage retrieval feedback
By introducing a query rewriter and multi-stage feedback training, the query rewriting method of the RAG framework is optimized, which solves the problems of user query brevity and semantic gap, improves the accuracy of retrieval and generation, and is applicable to a variety of large-scale language models.
Patent Information
- Application Number
- CN202510964540.X
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2025-07-14
- Publication Date
- 2025-09-12
- Estimated Expiration
- 2045-07-14
AI Technical Summary
In the existing RAG framework, the user's original query is brief, vague, or does not match the standard retrieval format, causing the retriever to return a large number of irrelevant or low-quality documents, affecting the accuracy of generated answers. Existing query rewriting methods fail to effectively address the semantic gap between query expression and retrieval requirements.
A query rewriter is introduced, an improved RAG framework is constructed, the query rewriter is trained through two stages of static feedback and dynamic feedback, an information alignment strategy is designed, the query rewriting process is optimized, and the matching degree of retrieval requirements is improved.
It improves the coverage and accuracy of retrieval, enhances the accuracy of the generator's answers, and is suitable for large language models of different scales.
Smart Images

Figure CN120492613B_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to the field of retrieval enhancement generation technology, and in particular to a RAG query rewriting method based on multi-stage retrieval feedback. Background Art
[0002] In recent years, large language models (LLMs) have made significant progress in natural language processing (NLP). These models possess powerful language understanding and knowledge storage capabilities, and in many cases are capable of answering questions directly based on parameterized knowledge. However, due to their parameterized knowledge storage, LLMs still face challenges in open-domain question answering tasks, such as knowledge staleness, hallucinations, and lack of interpretability. To address these challenges with LLMs, retrieval-augmented generation (RAG) is an effective solution. It consists of two key components: a retriever and a generator. The retriever retrieves relevant documents from a large corpus to answer the question, while the generator uses this contextual information to generate accurate answers. This approach not only relies on the LLM's own extensive training data, but also requires retrieving and extracting relevant information from large-scale external knowledge sources (such as internet documents, encyclopedias, and domain-specific knowledge bases) to generate answers.
[0003] In the existing technology, the research on RAG mainly focuses on the optimization of the retriever or generator, ignoring the semantic gap between the query expression and the retrieval requirements. The user's original query is often brief, vague or does not match the standard retrieval format. There is an information gap between the query expression and the retrieval requirements, which causes the retriever to return a large number of irrelevant or low-quality documents, thereby affecting the accuracy of the final generated answer. Therefore, optimizing only the retriever or generator is not enough to improve the overall performance of the RAG framework.
[0004] Furthermore, existing research on query rewriting methods in RAG is mainly divided into two categories: human-friendly query rewriting methods and generator-friendly query rewriting methods. Human-friendly query optimization methods use carefully designed prompt scripts to guide large language models (LLMs) to generate queries that conform to human language habits based on their own powerful language capabilities. Although easy for humans to understand, they may not be completely suitable for retrieval systems. Generator-friendly query rewriting methods use improving the generation quality of the generator as a training goal, which improves the generation effect to a certain extent. However, in a loosely coupled framework like RAG, methods that use long-distance indirect signals as optimization targets have unstable effects on retrieval tasks, which in turn affects the stability of generation. Summary of the Invention
[0005] In view of this, the present invention provides a RAG query rewriting method based on multi-stage retrieval feedback, which is used to optimize the query and improve the matching degree of the retrieval requirements, so as to enhance the coverage and accuracy of the retrieval.
[0006] In a first aspect, the present invention provides a RAG query rewriting method based on multi-stage retrieval feedback, the method comprising:
[0007] Step 1: Build an improved RAG framework by introducing a query rewriter;
[0008] Step 2: Based on step 1, the query rewriter is trained in two stages based on static feedback and dynamic feedback of the retrieval task; and an information alignment strategy is designed to calibrate the rewritten query to the information scope and semantics of the retrieval requirement.
[0009] Optionally, step 1 includes:
[0010] Build a RAG framework consisting of a query rewriter, retriever, and generator, and introduce the T5-large model as a query rewriter ; First, the query rewriter rewrites the original query to obtain the rewritten query; second, the retriever searches based on the rewritten query to obtain the retrieved document; finally, the generator combines the rewritten query and the retrieved document to predict the answer.
[0011] Optionally, step 2 includes:
[0012] The static feedback phase is based on supervised learning, using labeled queries and annotated documents to guide the query rewriting process. The query rewriter uses static feedback to perform query rewriting and information alignment. The dynamic feedback phase is based on reinforcement learning, using the accuracy and coverage of the retrieved documents as reward signals and dynamically feeding them back to the query rewriter. Through the interaction between query rewriting and retrieval tasks, the query rewriter generates queries that meet retrieval requirements.
[0013] Reformulate the original query q as a set of annotated documents that is quantitatively and semantically close to the set The rewritten query is:
[0014] ;
[0015] in, , Represents annotated documents; Represents the rewritten query generated by the query rewriter, Indicates rewriting the query set, that is Generated by the query rewriter, , Represents a query rewriter;
[0016] The information alignment strategy includes two sub-strategies: quantity alignment and semantic alignment. Quantity alignment: Based on the annotated document The number of queries determines the query rewriter generates rewritten queries The number of , and Represents the number of rewritten queries and annotated documents respectively; Semantic alignment: During the supervised learning training process, the annotated documents The semantic representation of and The generated query information contains the information of the annotated document.
[0017] Optionally include:
[0018] The query rewriter is trained in two stages: supervised learning and reinforcement learning. The process is as follows:
[0019] I. Supervised learning: First, we use the Large Language Model (LLM) to generate labeled queries based on annotated documents. Then, we combine the labeled query supervision signal and the annotated document supervision signal to jointly optimize the query rewriting task.
[0020] Dataset construction, through the prompt instruction to guide the Tongyi Qianwen model Qwen-max, for the original query q combined with the annotated document collection Generate K semantically distinct annotated documents Matching rewritten queries , forming a rewritten query set; then select The set of retrieval indicators that are better than the original query in terms of both R-Precision and R-Precision is used as the label query , ; Construct a bigram As training data for supervised learning;
[0021] The query rewriter is initialized based on the query supervision signal of the label query. The query rewriter learns to generate the label query based on the input original query q. The query generation loss is calculated using the cross entropy loss function, which is expressed as:
[0022] ;
[0023] in, Indicates tag query At the sequence position The words on T express the label query The length of the sequence;
[0024] To label document supervision signals, we first deploy a context encoder of a pre-trained dense paragraph retriever DPR to generate each labeled document The hidden state , and then perform average pooling on the hidden state to obtain the semantic representation of the annotated document ; When the encoder of the query rewriter generates the hidden state of the original query q , and similarly for the hidden state Perform average pooling to obtain the query representation ; Then, by minimizing the query representation and annotated document semantic representation The mean square error MSE between the query and the document is used to align the semantics of the query and the document. The expression of the information alignment loss is:
[0025] ;
[0026] Joint optimization objective, the overall training objective loss is generated by the query loss and information alignment loss Composition, its expression is:
[0027] ;
[0028] in, represents the weight factor used to balance the impact of query generation and information alignment strategies;
[0029] II. Reinforcement Learning: We use a policy gradient-based reinforcement learning framework to model the rewriting process as a Markov decision process (MDP). This process is defined by the state space, action space, reward function, and optimization strategy. We use retrieval effect feedback to drive the iterative improvement of the query rewriter.
[0030] Optionally, the reinforcement learning process includes:
[0031] Task modeling, based on the policy gradient reinforcement learning framework, the rewriting process is modeled as a Markov decision process MDP, which consists of five tuples Definition: The state space S is defined by the vocabulary and the maximum sequence length, where each state Represents the original query q and the generated partially rewritten query The action space A is equivalent to the vocabulary space, and each action Select the next token to be generated ; The state transition probability P is determined by the policy network The policy network outputs a probability distribution over the vocabulary to guide token generation; in generating rewrite queries After obtaining the retrieved document D, calculate the reward function R; discount factor Used to balance immediate rewards and long-term benefits; initial strategy Initialized by a supervised pre-trained rewriter; at each step t, the policy observes the current state and perform actions , generate tokens , until the terminator is generated to end the current round; the final reward is evaluated based on the retrieval effect of the complete rewritten query;
[0032] Design reward function, given the original query q, policy network Generate rewritten queries , the reward function designed is expressed as:
[0033] ;
[0034] in, 、 represents the weight coefficient, represents the regularization strength, 、 、 These are all hyperparameters, which are tuned through experiments; Represents the retrieval accuracy reward signal, reflecting the proportion of annotated documents in the retrieval results; Represents the retrieval coverage reward signal, which is used to evaluate the retrieval system's recall ability for all annotated documents. The expressions of the above two reward signals are:
[0035] ;
[0036] ;
[0037] Introduce the KL divergence term as a policy stability constraint to limit the current policy and supervised learning initial strategy The distribution difference is used to prevent semantic deviation, and its expression is:
[0038] ;
[0039] Strategy optimization, using the proximal policy optimization PPO algorithm to maximize the expected reward, the expression of its objective function is:
[0040] ;
[0041] in, represents the advantage value calculated by the generalized advantage estimate GAE, represents the shear threshold; value network Accumulate discounted rewards through regression For training, the optimization goal is to minimize the prediction error, which is expressed as:
[0042] ;
[0043] Training process, policy network and value network Update alternately through the following steps: Based on the current strategy Generate rewritten queries , obtain the document set D through the retriever R, and calculate the composite reward including retrieval accuracy, coverage and semantic similarity ;use Predict state value and calculate the advantage value at each moment through generalized advantage estimation GAE ; Through the objective function renew , while optimizing Minimize the prediction error ; Regularly Parameter synchronization to the old policy .
[0044] In a second aspect, an embodiment of the present invention provides a computer-readable storage medium, which includes a stored program, wherein when the program is running, the device where the computer-readable storage medium is located is controlled to execute the RAG query rewriting method based on multi-stage retrieval feedback in the first aspect or any possible implementation of the first aspect.
[0045] In a third aspect, an embodiment of the present invention provides a method comprising: one or more processors; a memory; and one or more computer programs, wherein the one or more computer programs are stored in the memory, and the one or more computer programs include instructions that, when executed, enable the device to execute the RAG query rewriting method based on multi-stage retrieval feedback in the first aspect or any possible implementation of the first aspect.
[0046] In the technical solution provided by the present invention, the method includes constructing an improved RAG framework by introducing a query rewriter; training the query rewriter in two stages based on static feedback and dynamic feedback of the retrieval task; and designing an information alignment strategy to calibrate the rewritten query to the information scope and semantics of the retrieval requirements. This method optimizes the query, improves the matching degree of the retrieval requirements, and enhances the coverage and accuracy of the retrieval. BRIEF DESCRIPTION OF THE DRAWINGS
[0047] In order to more clearly illustrate the technical solutions of the embodiments of the present invention, the following briefly introduces the drawings required for use in the embodiments. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying any creative work.
[0048] Figure 1 A flowchart of a RAG query rewriting method based on multi-stage retrieval feedback provided by an embodiment of the present invention;
[0049] Figure 2 A schematic diagram of an improved RAG framework provided in an embodiment of the present invention;
[0050] Figure 3 Schematic diagram of query rewriter training provided by an embodiment of the present invention, wherein (a) is a schematic diagram of generating a label query using LLM before supervised learning; (b) is a schematic diagram of the supervised learning process; (c) is a schematic diagram of the reinforcement learning process;
[0051] Figure 4 A schematic diagram of an electronic device provided by an embodiment of the present invention. DETAILED DESCRIPTION
[0052] To make the objectives, technical solutions, and advantages of the embodiments of the present invention more clear, the technical solutions in the embodiments of the present invention will be clearly and completely described below in conjunction with the accompanying drawings in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making creative efforts shall fall within the scope of protection of the present invention.
[0053] It should be understood that the embodiments described are only a portion of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by persons of ordinary skill in the art without creative work are within the scope of protection of the present invention.
[0054] The terms used in the embodiments of the present invention are only for the purpose of describing specific embodiments and are not intended to limit the present invention. The singular forms "a", "the" and "the" used in the embodiments of the present invention are also intended to include plural forms, unless the context clearly indicates other meanings.
[0055] It should be understood that the term "and / or" as used herein simply describes a relationship between associated objects, indicating that three possible relationships exist. For example, "A and / or B" can represent: A alone, A and B together, or B alone. Furthermore, the character " / " in this document generally indicates an "or" relationship between the associated objects.
[0056] The word "if," as used herein, may be interpreted as "at the time of" or "when" or "in response to determining" or "in response to detecting," depending on the context. Similarly, the phrases "if it is determined" or "if (stated condition or event) is detected" may be interpreted as "when it is determined" or "in response to the determination" or "when detecting (stated condition or event)" or "in response to detecting (stated condition or event)," depending on the context.
[0057] Figure 1 A flowchart of a RAG query rewriting method based on multi-stage retrieval feedback provided by an embodiment of the present invention is shown in FIG. Figure 1 As shown, the method includes:
[0058] Step 1: Build an improved RAG framework by introducing a query rewriter.
[0059] Retrieval-enhanced generation is a technical approach that combines retrieval and generative models. First, the retriever searches for content related to the current task from a large amount of stored information (such as knowledge bases, historical task records, etc.), and then uses the generative model to generate answers or instructions for specific tasks based on the retrieved information.
[0060] The traditional RAG framework includes a retriever and a generator. When an original query q is input, the system will return a set of retrieved documents D through the retriever. ; The generator uses retrieved documents to help the model complete the open domain question answering task.
[0061] In the embodiment of the present invention, Figure 2 As shown, step 1 includes:
[0062] Construct a RAG framework consisting of a query rewriter, a retriever, and a generator, namely the Query Rewriting with Multi-Stage Retrieval Feedback in Retrieval-Augmentation Generation (QR-MSRF), and introduce the T5-large model as the query rewriter. ; First, the query rewriter rewrites the original query to obtain the rewritten query; second, the retriever searches based on the rewritten query to obtain the retrieved document; finally, the generator combines the rewritten query and the retrieved document to predict the answer.
[0063] In an embodiment of the present invention, the query rewriter includes an encoder and a decoder, which rewrites the original query to, firstly, compensate for problems such as the user's original query being brief, vague, or not matching the standard retrieval format, and secondly, rewrite the original query into a retrieval-friendly query that better meets the retrieval requirements, thereby improving the coverage and accuracy of the retrieved documents.
[0064] In the embodiment of the present invention, Figure 2 The red dotted arrows in the middle represent the two training stages of static feedback and dynamic feedback respectively; the solid arrows represent the inference process.
[0065] Step 2: Based on step 1, the query rewriter is trained in two stages based on static feedback and dynamic feedback of the retrieval task; and an information alignment strategy is designed to calibrate the rewritten query to the information scope and semantics of the retrieval requirement.
[0066] In this embodiment of the present invention, step 2 includes:
[0067] The static feedback phase is based on supervised learning, using labeled queries and annotated documents to guide the query rewriting process. The query rewriter uses static feedback to perform query rewriting and information alignment. However, a rewriter trained solely on the static signals provided by labeled queries and annotated documents produces queries that are not sufficiently generalizable and do not significantly improve RAG performance. Therefore, the dynamic feedback phase is based on reinforcement learning, using the accuracy and coverage of the retrieved documents as reward signals, which are dynamically fed back to the query rewriter. Through the interaction between query rewriting and retrieval tasks, the query rewriter generates queries that meet retrieval requirements.
[0068] Reformulate the original query q as a set of annotated documents that is quantitatively and semantically close to the set The rewritten query is:
[0069] ;
[0070] in, , Represents annotated documents; Represents the rewritten query generated by the query rewriter, Indicates rewriting the query set, that is Generated by the query rewriter, , Represents a query rewriter;
[0071] There is an information gap between the original query and the retrieval target in terms of both quantity and semantics, and rewriting queries is an important bridge to bridge this gap. Therefore, in the static feedback stage, not only should the rewriter learn high-quality label queries, but the annotated documents should also be used to guide the query rewriter to have the ability to understand and bridge this gap. To this end, the present invention designs an information alignment strategy, which includes two sub-strategies: quantity alignment and semantic alignment. Quantity alignment: too few rewritten queries may result in missing valuable retrieval documents, while too many rewritten queries may result in the retrieval of irrelevant documents, interfering with the generation quality of the generator. The present invention balances the characteristics of single-query and multi-query rewriting, and according to the annotated documents The number of queries determines the query rewriter generates rewritten queries The number of , and Represents the number of rewritten queries and annotated documents respectively; Semantic alignment: The core of the retrieval task is to calculate the semantic similarity between the query and the candidate documents. This paper uses reverse thinking to The semantic representation of and The generated query information contains the information of the annotated document.
[0072] In the embodiment of the present invention, Figure 3 Shown, including:
[0073] Based on two stages of supervised learning and reinforcement learning, such as Figure 3 (b) and Figure 3 As shown in (c) in the figure, the query rewriter is trained as follows:
[0074] I. Supervised learning, such as Figure 3 As shown in (a), we first use the large language model (LLM) to generate label queries based on the annotated documents. Then, we combine the label query supervision signal and the annotated document supervision signal to jointly optimize the query rewriting task.
[0075] Dataset construction, through the prompt instruction to guide the Tongyi Qianwen model Qwen-max, for the original query q combined with the annotated document collection Generate K semantically distinct annotated documents Matching rewritten queries , forming a rewritten query set; then select The set of retrieval indicators that are better than the original query in terms of (Recall at 5) and R-Precision is used as the tag query , ; Construct a bigram As training data for supervised learning;
[0076] The T5-large model is selected as the initial query rewriter for the label query supervision signal; the supervised fine-tuning of the rewriter is to enable the query rewriter to learn the high-quality rewritten queries generated by the Large Language Model (LLM) Specifically: the query rewriter generates and labels queries based on the input original query q learning The query generation loss is calculated using the cross entropy loss function, which is expressed as:
[0077] ;
[0078] in, Indicates tag query At the sequence position The words on T express the label query The length of the sequence;
[0079] To label document supervision signals, we first deploy a pre-trained dense passage retrieval (DPR) context encoder to generate each labeled document. The hidden state , and then perform average pooling on the hidden state to obtain the semantic representation of the annotated document ; When the encoder of the query rewriter generates the hidden state of the original query q , and similarly for the hidden state Perform average pooling to obtain the query representation ; Then, by minimizing the query representation and annotated document semantic representation The mean square error MSE between the query and the document is used to align the semantics of the query and the document. The expression of the information alignment loss is:
[0080] ;
[0081] Joint optimization objective, the overall training objective loss is generated by the query loss and information alignment loss Composition, its expression is:
[0082] ;
[0083] in, represents the weight factor used to balance the impact of query generation and information alignment strategies;
[0084] II. Reinforcement Learning: We use a policy gradient-based reinforcement learning framework to model the rewriting process as a Markov decision process (MDP). This process is defined by the state space, action space, reward function, and optimization strategy. We use retrieval effect feedback to drive iterative improvements in the query rewriter.
[0085] In the embodiment of the present invention, the process of reinforcement learning includes:
[0086] Task modeling, based on the policy gradient reinforcement learning framework, the rewriting process is modeled as a Markov decision process MDP, which consists of five tuples Definition: The state space S is defined by the vocabulary and the maximum sequence length, where each state Represents the original query q and the generated partially rewritten query The action space A is equivalent to the vocabulary space, and each action Select the next token to be generated ; The state transition probability P is determined by the policy network (i.e. query rewriter ) determines that the policy network outputs a probability distribution over the vocabulary to guide token generation; in generating rewriting queries After obtaining the retrieved document D, calculate the reward function R; discount factor Used to balance immediate rewards and long-term benefits; initial strategy Initialized by a supervised pre-trained rewriter; at each step t, the policy observes the current state and perform actions , generate tokens , until the terminator is generated to end the current round; the final reward is evaluated based on the retrieval effect of the complete rewritten query;
[0087] Design reward function, the core of reinforcement learning task is to maximize the cumulative reward function, and its design focuses on retrieval performance. Given the original query q, the policy network Generate rewritten queries , the reward function designed is expressed as:
[0088] ;
[0089] in, 、 represents the weight coefficient, represents the regularization strength, 、 、 These are all hyperparameters, which are tuned through experiments; Represents the retrieval accuracy reward signal, reflecting the proportion of annotated documents in the retrieval results; Represents the retrieval coverage reward signal, which is used to evaluate the retrieval system's recall ability for all annotated documents. The expressions of the above two reward signals are:
[0090] ;
[0091] ;
[0092] Introduce the KL divergence term as a policy stability constraint to limit the current policy and supervised learning initial strategy The distribution difference is used to prevent semantic deviation, and its expression is:
[0093] ;
[0094] Policy optimization uses the Proximal Policy Optimization (PPO) algorithm to maximize the expected reward. The expression of its objective function is:
[0095] ;
[0096] in, represents the advantage value calculated by generalized advantage estimation (GAE), represents the shear threshold; value network Accumulate discounted rewards through regression For training, the optimization goal is to minimize the prediction error, which is expressed as:
[0097] ;
[0098] Training process, policy network and value network Update alternately through the following steps: Based on the current strategy Generate rewritten queries , obtain the document set D through the retriever R, and calculate the composite reward including retrieval accuracy, coverage and semantic similarity ;use Predict state value and calculate the advantage value at each moment through generalized advantage estimation GAE ; Through the objective function renew , while optimizing Minimize the prediction error ; Regularly Parameter synchronization to the old policy , maintaining sampling stability.
[0099] We conducted experiments using three representative open-domain question answering datasets, representing open-domain question answering tasks of varying difficulty: (I) Natural Questions (NQ), a question answering dataset designed for end-to-end problem solving; (II) HotPotQA, a dataset containing complex questions requiring multi-hop reasoning; and (III) TriviaQA, a challenging reading comprehension dataset containing complex questions, diverse evidential documents, and cross-sentence reasoning requirements. We randomly collected 60,000 instances from the training sets of these three datasets to train the query rewriter. During the evaluation phase, we used the test sets of NQ and TriviaQA, as well as the development set of HotpotQA, as internal evaluation datasets.
[0100] This paper evaluates both retrieval and generation quality metrics. Specifically, the top-5 recall (R@5) and R-Precision (R-Pre) of the retrieved documents are used to assess the effectiveness of the rewriting for retrieval. The generator's answer quality is evaluated by calculating the F1 score of the answer.
[0101] We use standard retrieval enhancement as the comparison baseline, which directly uses the original query to retrieve and then generate answers. In addition, we select the human-friendly query rewriting method and the generator-friendly query rewriting method as the comparison baseline methods of the present invention.
[0102] For human-friendly query rewriting methods: (1) Large Language Model Rewriting (LLM-Rewrite) uses hints to leverage the inherent capabilities of LLM for general query rewriting; (2) RAG-Fusion leverages LLM to rewrite the user's original query into multiple different queries, supplementing and expanding the user's question from different perspectives, and reordering documents through RRF to generate a fused document list.
[0103] For generation-friendly query rewriting methods: (1) Rewrite-Retrieve-Read (RRR) uses reinforcement learning technology to fine-tune the T5 model using the accuracy of LLM answers as a reward signal; (2) RQ-RAG optimizes queries by constructing a search query dataset containing various scenarios and training LLM to rewrite, decompose, and clarify the original queries.
[0104] The retrieval component in the RAG system is built using FAISS and a pre-trained embedding model, DPR. During the training phase, DPR is frozen. KILT provides a unified knowledge source and evaluation benchmark, and manually removes duplicate related documents from the original annotations in the training data before using them as annotation documents for training. A series of hyperparameters in the training phase were tested and verified and ultimately determined as: weight factor Set to 0.5; and Both are 1, is 0.2; the clipping threshold To facilitate comparison with the baseline method, both the present invention and the baseline method use Qwen-max as the generator and utilize the first five documents recalled by the retriever as additional context to generate answers.
[0105] The specific performance of various query rewriting methods in three open domain question answering datasets is shown in Table 1.
[0106] Table 1 Results of different query rewriting methods
[0107] ;
[0108] Table 1 shows the results obtained using different query rewriting methods on three representative datasets. QR-MSRF significantly improves both retrieval and generation performance on the NQ, HotpotQA, and TriviaQA datasets, and outperforms other baseline methods on most metrics. The improvement is particularly pronounced for the HotpotQA and TriviaQA datasets. On the HotpotQA dataset, QR-MSRF improves the R-Precision metric by 11.3% over the second-best result, and on the TriviaQA dataset, it improves the R-Precision metric by 5.8% over the second-best result. This indicates that QR-MSRF can more effectively rank relevant paragraphs at the top, thereby improving retrieval performance. The experimental results verify the superiority of the method of using retrieval tasks to guide query rewriting.
[0109] The best result is in bold and the second best result is underlined in Table 1. Since RQ-RAG uses an iterative query rewriting and retrieval approach, its retrieval quality cannot be directly evaluated. Therefore, only the evaluation of the generated results is provided here.
[0110] In order to prove that QR-MSRF (i.e., the improved RAG framework) can be applied to other LLMs, the present invention also tests the Llama3-8B model and the Qwen2-7B model, and selects RAG-Fusion and RRR as comparison methods from the two query rewriting baseline methods. The experimental results are shown in Table 2.
[0111] Table 2 Framework universality test results
[0112] ;
[0113] As shown in Table 2, the performance of different rewriting methods was tested using three different-scale LLMs: Llama3-8B, Qwen2-7B, and Qwen-max. The experimental results demonstrate that the proposed method is adaptable to LLMs of varying scales and achieves significantly better query rewriting performance than other baseline methods, effectively demonstrating the versatility of QR-MSRF.
[0114] Compared with the prior art, the present invention has the following beneficial effects:
[0115] 1. The existing technology only focuses on optimizing and improving the retriever and generator in the original RAG framework. In addition, the user's original query is often brief, vague, or does not match the standard retrieval format, causing the retriever to return a large number of irrelevant or low-quality documents, which in turn affects the accuracy of the final generated answer. The present invention proposes to construct an improved RAG framework (QR-MSRF): inserting a lightweight and specialized query rewriter, namely the T5-large model, into the traditional RAG framework, improving the quality of queries, thereby improving the quality of retrieved documents, and ultimately enhancing the accuracy of the generator's answers.
[0116] 2. To address the problems of generating human-friendly and generator-friendly queries, this paper proposes query rewriting with multi-stage feedback based on the retrieval task, including static and dynamic feedback. This maximizes the guiding role of the retrieval task in the query rewriting process, rewriting the original query into a retrieval-friendly query that better meets retrieval requirements.
[0117] 3. To address the information gap between query expression and retrieval requirements, the present invention proposes an information alignment strategy in the static feedback stage. This strategy includes two sub-strategies: quantity alignment and semantic alignment. It achieves adaptive alignment of the rewritten query and the retrieval target in both quantity and semantic dimensions, calibrates the rewritten query to the required information range and semantics, and thus improves the coverage and accuracy of the retrieval.
[0118] In the technical solution provided by the present invention, the method includes constructing an improved RAG framework by introducing a query rewriter; training the query rewriter in two stages based on static feedback and dynamic feedback of the retrieval task; and designing an information alignment strategy to calibrate the rewritten query to the information scope and semantics of the retrieval requirements. This method optimizes the query, improves the matching degree of the retrieval requirements, and enhances the coverage and accuracy of the retrieval.
[0119] Each step of the embodiment of the present invention may be performed by an electronic device, including but not limited to a tablet computer, a portable PC, a desktop computer, etc.
[0120] An embodiment of the present invention provides a computer-readable storage medium, which includes a stored program, wherein when the program runs, the electronic device where the computer-readable storage medium is located is controlled to execute the above-mentioned embodiment of the RAG query rewriting method based on multi-stage retrieval feedback.
[0121] Figure 4 A schematic diagram of an electronic device provided for an embodiment of the present invention, as shown in the figure, the electronic device 21 includes: a processor 211, a memory 212, and a computer program 213 stored in the memory 212 and executable on the processor 211. When the computer program 213 is executed by the processor 211, the RAG query rewriting method based on multi-stage retrieval feedback in the embodiment is implemented. To avoid repetition, they are not described here one by one.
[0122] The electronic device 21 includes, but is not limited to, a processor 211 and a memory 212. Those skilled in the art will appreciate that the figure is merely an example of the electronic device 21 and does not limit the electronic device 21. The electronic device 21 may include more or fewer components than shown, or may combine certain components or different components. For example, the electronic device may also include input and output devices, network access devices, buses, and the like.
[0123] The processor 211 may be a central processing unit (CPU), or other general-purpose processors, digital signal processors (DSP), application-specific integrated circuits (ASIC), field-programmable gate arrays (FPGA), or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, etc. A general-purpose processor may be a microprocessor or any conventional processor.
[0124] The memory 212 can be an internal storage unit of the electronic device 21, such as the hard drive or memory of the electronic device 21. The memory 212 can also be an external storage device of the electronic device 21, such as a plug-in hard drive, a Smart Media Card (SMC), a Secure Digital (SD) card, a Flash Card, etc. equipped on the electronic device 21. Furthermore, the memory 212 can include both the internal storage unit of the electronic device 21 and an external storage device. The memory 212 is used to store computer programs and other programs and data required by the network device. The memory 212 can also be used to temporarily store data that has been output or is about to be output.
[0125] Those skilled in the art will clearly understand that, for the convenience and brevity of description, the specific working processes of the systems, devices and units described above can refer to the corresponding processes in the aforementioned method embodiments and will not be repeated here.
[0126] The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, improvements, etc. made within the spirit and principles of the present invention should be included in the scope of protection of the present invention.
Claims
1. A RAG query rewriting method based on multi-stage retrieval feedback, characterized in that: The method comprises: Step 1: Build an improved RAG framework by introducing a query rewriter; Step 2: Based on step 1, the query rewriter is trained in two stages based on static and dynamic feedback of the retrieval task. An information alignment strategy is designed to calibrate the rewritten query to the information scope and semantics of the retrieval requirement. The step 1 comprises: Build a RAG framework consisting of a query rewriter, retriever, and generator, and introduce the T5-large model as a query rewriter First, the query rewriter rewrites the original query to obtain the rewritten query; second, the retriever searches based on the rewritten query to obtain the retrieved document; finally, the generator combines the rewritten query and the retrieved document to predict the answer; The step 2 includes: The static feedback phase is based on supervised learning, using labeled queries and annotated documents to guide the query rewriting process. The query rewriter uses static feedback to perform query rewriting and information alignment. The dynamic feedback phase is based on reinforcement learning, using the accuracy and coverage of the retrieved documents as reward signals and dynamically feeding them back to the query rewriter. Through the interaction between query rewriting and retrieval tasks, the query rewriter generates queries that meet retrieval requirements. Reformulate the original query q as a set of annotated documents that is quantitatively and semantically close to the set The rewritten query is: ; in, , Represents annotated documents; Represents the rewritten query generated by the query rewriter, Indicates rewriting the query set, that is Generated by the query rewriter, , Represents a query rewriter; The information alignment strategy includes two sub-strategies: quantity alignment and semantic alignment. Quantity alignment: Based on the annotated document The number of queries determines the query rewriter generates rewritten queries The number of , and Represents the number of rewritten queries and annotated documents respectively; Semantic alignment: During the supervised learning training process, the annotated documents The semantic representation of and The semantic similarity of the generated query information contains the information of the annotated document; The query rewriter is trained in two stages: supervised learning and reinforcement learning. The process is as follows: I. Supervised learning: First, we use the Large Language Model (LLM) to generate labeled queries based on annotated documents. Then, we combine the labeled query supervision signal and the annotated document supervision signal to jointly optimize the query rewriting task. Dataset construction, through the prompt instruction to guide the Tongyi Qianwen model Qwen-max, for the original query q combined with the annotated document collection Generate K semantically distinct annotated documents Matching rewritten queries , forming a rewritten query set; then select The set of retrieval indicators that are better than the original query in terms of both R-Precision and R-Precision is used as the label query , ; Construct a bigram As training data for supervised learning; The query rewriter is initialized based on the query supervision signal of the label query. The query rewriter learns to generate the label query based on the input original query q. The query generation loss is calculated using the cross entropy loss function, which is expressed as: ; in, Indicates tag query At the sequence position The words on T express the label query The sequence length of To label document supervision signals, we first deploy a context encoder of a pre-trained dense paragraph retriever DPR to generate each labeled document The hidden state , and then perform average pooling on the hidden state to obtain the semantic representation of the annotated document ; When the encoder of the query rewriter generates the hidden state of the original query q , and similarly for the hidden state Perform average pooling to obtain the query representation ; Then, by minimizing the query representation and annotated document semantic representation The mean square error MSE between the query and the document is used to align the semantics of the query and the document. The expression of the information alignment loss is: ; Joint optimization objective, the overall training objective loss is generated by the query loss and information alignment loss Composition, its expression is: ; in, represents the weight factor used to balance the impact of query generation and information alignment strategies; II. Reinforcement Learning: We use a policy gradient-based reinforcement learning framework to model the rewriting process as a Markov decision process (MDP). This process is defined by the state space, action space, reward function, and optimization strategy. We use retrieval effect feedback to drive the iterative improvement of the query rewriter.
2. The method according to claim 1, characterized in that The reinforcement learning process includes: Task modeling, based on the policy gradient reinforcement learning framework, the rewriting process is modeled as a Markov decision process MDP, which consists of five tuples Definition: The state space S is defined by the vocabulary and the maximum sequence length, where each state Represents the original query q and the generated partially rewritten query The action space A is equivalent to the vocabulary space, and each action Select the next token to be generated ; The state transition probability P is determined by the policy network The policy network outputs a probability distribution over the vocabulary to guide token generation; in generating rewrite queries After obtaining the retrieved document D, calculate the reward function R; discount factor Used to balance immediate rewards and long-term benefits; initial strategy Initialized by a supervised pre-trained rewriter; at each step t, the policy observes the current state and perform actions , generate tokens , until the terminator is generated to end the current round; the final reward is evaluated based on the retrieval effect of the complete rewritten query; Design reward function, given the original query q, policy network Generate rewritten queries , the reward function designed is expressed as: ; in, 、 represents the weight coefficient, represents the regularization strength, 、 、 These are all hyperparameters, which are tuned through experiments; Represents the retrieval accuracy reward signal, reflecting the proportion of annotated documents in the retrieval results; Represents the retrieval coverage reward signal, which is used to evaluate the retrieval system's recall ability for all annotated documents. The expressions of the above two reward signals are: ; ; Introduce the KL divergence term as a policy stability constraint to limit the current policy and supervised learning initial strategy The distribution difference is used to prevent semantic deviation, and its expression is: ; Strategy optimization, using the proximal policy optimization PPO algorithm to maximize the expected reward, the expression of its objective function is: ; in, represents the advantage value calculated by the generalized advantage estimate GAE, represents the shear threshold; value network Accumulate discounted rewards through regression For training, the optimization goal is to minimize the prediction error, which is expressed as: ; Training process, policy network and value network Update alternately through the following steps: Based on the current strategy Generate rewritten queries , obtain the document set D through the retriever R, and calculate the composite reward including retrieval accuracy, coverage and semantic similarity ;use Predict state value and calculate the advantage value at each moment through generalized advantage estimation GAE ; Through the objective function renew , while optimizing Minimize the prediction error ; Regularly Parameter synchronization to the old policy .
3. A computer-readable storage medium, characterized in that The computer-readable storage medium includes a stored program, wherein when the program is executed, the device where the computer-readable storage medium is located is controlled to execute the RAG query rewriting method based on multi-stage retrieval feedback according to any one of claims 1 to 2.
4. An electronic device, characterized in that: include: one or more processors; Memory; and one or more computer programs, wherein the one or more computer programs are stored in the memory, and the one or more computer programs include instructions that, when executed by the device, enable the device to perform the RAG query rewriting method based on multi-stage retrieval feedback as described in any one of claims 1 to 2.
Citation Information
Patent Citations
Generative large language model training method and model-based search method
CN116127020A
Information search method and device, electronic equipment, storage medium and program product
CN117009621A