Q - PRM : Adaptive Query Rewriting for Retrieval-Augmented Generation via Step-level Process Supervision (2025)

Abstract: Query rewriting plays a pivotal role in Retrieval-Augmented Generation (RAG) by refining real-world queries of varying complexity. Existing approaches typically rely on outcome-supervised training or heuristic rules to guide the rewriting process. However, these paradigms often struggle to handle queries with varying levels of complexity, posing over- and under-refinement problems. We identify the root cause of these issues as the absence of supervision signals for intermediate steps. To fully construct and utilize such signals, we propose Q-PRM, a novel query rewriting framework. Q-PRM reformulates the rewriting process as a Markov Decision Process (MDP) composed of atomic rewriting steps. In this way, Q-PRM can apply process-level supervision to each atomic step according to the query type, offering more targeted and effective guidance. Q-PRM comprises three key stages: (1) applying Monte Carlo Tree Search to generate step-level process supervision signals; (2) performing reinforced self-training for progressive process refinement; and (3) employing PRM-guided decoding during inference. Experiments on several open-domain QA benchmarks demonstrate that Q-PRM consistently outperforms baselines across different levels of query complexity.

## introduction in real-world applications, retrieval-augmented generation (rag) systems need to handle user queries that vary widely in complexity (e.g., multihop reasoning, ambiguity) @xcite @xcite . this need arises from the users' diverse expression styles, intents, and background knowledge @xcite . to improve retrieval effectiveness and answer generation quality, it is essential to transform raw queries into syntactically clear and retrieval-friendly forms. this * corresponding author process, termed query rewriting, has emerged as a crucial pre-processing step in rag @xcite . existing query rewriting methods primarily rely on either heuristic rules @xcite or outcome-based training @xcite to perform single-step or multi-step rewriting @xcite . however, these approaches often fail to adapt to queries of varying complexity, posing over-and under-refinement challenges for realworld deployment. we conduct experiments to better illustrate these challenges. as shown in figure 1 (a), outcome-supervised approaches utilize downstream answer quality as the reward to train the rewriter. however, our experiment finds they tend to omit critical steps when handling complex queries (e.g., ambiguous, multihop), causing the under-refinement challenge. for example, the ambiguous query "which country won more medals in the 2008 olympics, china or the us?" requires disambiguation (e.g., distinguishing between the summer and winter olympics). however, the outcome-supervised rewriter merely decomposes the query without resolving the ambiguity, resulting in sub-queries that lack essential context @xcite . on the other hand, heuristic-based methods refine queries using fixed information expansion strategies @xcite . however, such approaches often cause step redundancy when rewriting simple queries, leading to overrefinement in the final rewritten query. as in figure 1 (b), the example simple query "who is elon musk?" is unnecessarily expanded with excessive background information, introducing redundant information into the rewritten version. the main reason for such over-and underrefinement issues stems from the lack of supervision signals under each intermediate rewriting step @xcite . instead of relying solely on the final answer, the process supervision paradigm provides step-level guidance using the process reward model (prm) @xcite . prm serves as a valuable reminder for the query rewriter to pay attention to essential steps when handling queries of varying complexity. in this way, the rewriter can better preserve essential steps when handling complex queries, while avoiding unnecessary ones for simpler queries. however, two primary challenges emerge when applying process supervision to the task: (1) the absence of process supervision signals to train the prm @xcite ; and (2) the lack of multistep reasoning ability tailored for rewriting tasks, which prevents rewriters from effective incorporation of prm signals. to address these challenges, we propose our method q-prm. specifically, q-prm consists of three key steps: (1) data collection: to address the high cost of manually annotating process supervision signals (challenge 1), we propose an uncertainty-aware monte carlo tree search (mcts) approach to automatically collect robust supervision signals for training the prm; (2) model training: we apply reinforced self-training @xcite to enhance the rewriter's reasoning ability to handle queries of varying complexity (challenge 2); and (3) inference: by combining the prm with the rewriter, we propose a decoding method that leverages the prm for fine-grained guidance, enhancing generalization to queries of varying complexity. our contribution can be summarized as follows: training-free methods leverage llms' world knowledge to heuristically expand queries for better retrieval. some works @xcite (e.g., query2doc) converted raw queries into pseudo-documents using the internal knowledge of llms. other works @xcite ) utilize heuristic rules to expand the query in a multi-step way. although efficient, these methods will inevitably introduce unnecessary and redundant information to the query due to hallucinations of llms @xcite . for training-based methods, existing works mainly use outcome supervision to fine-tune the rewriter model. some works @xcite @xcite utilize reinforcement learning (e.g., ppo @xcite ) to train the rewriter model with downstream task quality (i.e., generation, retrieval, re-ranking @xcite ) as the rewards. however, since downstream task performance depends not only on the quality of query rewriting but also on other factors (e.g., query difficulty @xcite , retrieval quality @xcite , and llm internal knowledge @xcite ), the reward is often unstable and sparse @xcite ), making it difficult for the rewrite model to effectively learn the optimal actions. process supervision for rag. rag is widely adopted for knowledge-intensive tasks @xcite @xcite , but it still faces limitations in complex multi-step reasoning. existing methods integrate rag with chainof-thought (cot) reasoning @xcite . since direct cot often falters on challenging tasks (e.g., competitive programming), cr-planner @xcite attempts to use prm to assist search and improve the reasoning ability of rag systems through test-time scaling. autoprm @xcite decomposed complex problems into subproblems while leveraging a self-supervisedly trained prm to provide feedback, thereby enhancing multi-step logical reasoning tasks. to improve the trustworthiness in the reasoning process, rearter @xcite introduces the process explanation model to provide explanations for the reasoning process. currently, most works use process supervision to enhance the in-retrieval and post-retrieval stages @xcite , but its potential in the pre-retrieval stage to refine user queries has not yet been explored. ## preliminary retrieval-augmented generation. consider a knowledge-intensive task (e.g., open-domain question answering), where the dataset comprises multiple query-answer instances. each instance consists of a raw query q paired with its corresponding ground-truth answer y. a rag system f (•) takes a query q as input and returns an answer ŷq = f (q). the rag process begins with a retriever that identifies a set of relevant documents d based on q. these documents are then passed to an llm reader, which generates the final response. the effectiveness of a rag system is typically assessed by evaluating the quality of the generated response ŷq using metrics such as exact match and f1 score @xcite , which compare the generated answer ŷq to the ground-truth answer y. in this paper, we denote this evaluation function as sim(ŷ q , y). ## our method: q-prm in this section, we present our proposed method, q-prm. the overall workflow of q-prm is illustrated in figure 2 . specifically, q-prm consists of three key steps. (1) data collection: leveraging mcts to gather process supervision signal for query rewriting; (2) model training: using the collected data to train the rewriting model m and the prm; (3) prm-guided inference: introducing a prm-guided search decoding mechanism for inference. ## conclusion in this paper, we identify a key challenge in query rewriting across varying levels of complexity: existing methods often overlook critical intermediate steps, leading to incorrect rewrites. we attribute this problem to the absence of process supervision signals. to bridge this gap, we emphasize the importance of process supervision signals and propose q-prm, a method that leverages a prm to guide the rewriting process. since such supervision signals are not readily available for this task, q-prm leverages uncertainty-aware mcts to alleviate reward instability in the rag process and to collect process signals for training the prm. offline experiments on four open-domain qa datasets and online evaluations show that q-prm effectively handles queries with diverse complexities.