Rethinking On-policy Optimization for Query Augmentation

AI summary

Overview

Research area: Information retrieval (IR) and natural language processing, specifically query augmentation for retrieval systems using large language models (LLMs), reinforcement learning (RL), and test-time prompting.

Technical level: Advanced. The paper assumes familiarity with BM25 and dense retrieval, pseudo-document expansion methods such as HyDE and Query2Doc, retrieval metrics such as NDCG@10 and Recall@10, and on-policy RL (PPO). The high-level conclusion — that simple prompting often matches expensive RL training — is accessible to a broader audience.

Scope (one sentence): The paper runs the first systematic, compute-aware comparison of zero-shot prompting-based query augmentation against on-policy RL for query rewriting across evidence-seeking, ad hoc, and tool retrieval, then proposes a hybrid method that optimizes pseudo-document generation instead of query rewriting.

What This Paper Is About

When you cannot change the retriever — for example, when using an API-based embedding service — the only way to improve search results is to change the query itself. Two families of methods do this: prompt-based approaches that ask an LLM to generate a hypothetical document to search with, and reinforcement-learning approaches that fine-tune an LLM to rewrite queries by directly rewarding retrieval metrics. These two families had never been compared under consistent experimental conditions, so it was unclear which one is worth its cost. The paper supplies that comparison and, motivated by the result, introduces a hybrid method called OPQE that applies RL to pseudo-document generation rather than to query rewriting.

Key Contributions

1. A controlled comparison of DeepRetrieval against a simple zero-shot prompting baseline. The authors replicate DeepRetrieval's PPO-based query-rewriting pipeline (denoted DR-3B and DR-7B for the 3B and 7B policy sizes) and compare it against SPQE (Simple Pseudo-document Query Expansion), a training-free method that concatenates the original query with an LLM-generated hypothetical document.

2. Evidence that training-free prompting often matches or surpasses RL-trained rewriting under a compute-aware setting. Prompting is allowed to use larger inference-time models (GPT-4o-mini, Qwen3-32B-Instruct, GPT-OSS-120B) while RL policies are small Qwen2.5-{3B,7B}-Instruct models trained with substantial multi-GPU compute, and prompting still achieves the best overall Borda rank.

3. A novel hybrid method, On-policy Pseudo-document Query Expansion (OPQE). OPQE applies on-policy optimization to the generation of a pseudo-document that maximizes retrieval performance, rather than to rewriting the query, merging prompting's generative structure with RL's targeted optimization. The abstract reports that OPQE outperforms both standalone prompting and RL-based rewriting.

4. An extension of the comparison to tool retrieval, an underexplored setting, using the ToolRet dataset (over 200k training instances) across three official splits (ToolRet-Web, ToolRet-Code, ToolRet-Customized).

Main Findings

• Both paradigms help over no augmentation, but differently. The "Base" columns (no query augmentation) are improved on by both SPQE and DeepRetrieval across most settings, but the gains differ by task and retriever. On evidence-seeking with BM25, the average moves from 70.2 (base) to 75.6 (DR-7B) and 78.2 (SPQE).

• RL can be effective, especially for sparse retrieval and dense ad hoc retrieval. On ad hoc retrieval with the dense Contriever-msmarco retriever, DR-7B achieves the highest average of 57.5 versus the base of 51.6, with FEVER rising from 75.8 to 85.5. With BM25 on evidence-seeking, DR-7B improves the average from 70.2 to 75.6.

• Simple prompting delivers the best sparse-retrieval averages. SPQE achieves the highest averages for evidence-seeking (78.2) and ad hoc (51.6) with BM25.

• Prompting wins on aggregate ranking. Averaged Borda scores across all task–retriever combinations: SPQE with GPT-4o-mini ranks 1 (29.5), GPT-OSS-120B ranks 2 (24.5), Qwen3-32B-Instruct ranks 3 (22.0), DR-7B ranks 4 (19.0), DR-3B ranks 5 (16.0), and the base setting ranks 6 (15.0). Sparse-only Borda scores are 17.0 for GPT-4o-mini SPQE, 16.0 for GPT-OSS-120B, 12.0 for Qwen3-32B, and 7.0 for both DR-3B and DR-7B.

• Open-weight prompting models are competitive with proprietary ones. SPQE averages with GPT-OSS-120B reach 80.0 on evidence-seeking with BM25 versus 78.2 for GPT-4o-mini, and 45.4 on tool retrieval with the dense retriever versus 44.8 for GPT-4o-mini. SPQE with an open-weight 32B model can be served on a single A100 40GB GPU.

• RL's gains are retriever-dependent and can be negative. On tool retrieval, DR-3B raises the sparse BM25 average Completeness@10 from 29.4 to 32.3 and NDCG@10 from 30.0 to 32.0, but on the dense retriever the same training does not surpass the base Completeness@10 of 29.6 (29.9 for DR-3B, 29.4 for DR-7B), and DR-7B lowers average NDCG@10 from 29.9 to 29.0.

• Prompting is consistently strong in tool retrieval. SPQE achieves the highest average NDCG@10 (33.4 sparse, 34.5 dense) and Completeness@10 (35.1 sparse, 34.8 dense). The paper's text attributes the Recall@10 advantage to "HyDE, another generative approach" even though the table reports SPQE values (44.4 sparse, 44.8 dense); HyDE does not appear as a row in that table.

• RL excels where domain adaptation matters. On the specialized FiQA dataset, the RL-trained model outperforms prompting, which the authors attribute to learning domain-specific financial terminology.

• Dense retrieval gains from RL are smaller or absent. Prompting is close to RL on dense evidence-seeking (SPQE 79.7 versus DR-7B 79.8 and DR-3B 79.6 with E5-base-v2) and on dense ad hoc retrieval (SPQE 56.6 versus DR-7B 57.5), while standing out on several individual datasets (for example, DL19 dense at 70.1 for SPQE versus 67.8 for DR-7B and 67.5 for the base).

• Compute profiles differ fundamentally. RL spends up-front training compute (8x A100 40GB for 3B policies, 16x A100 40GB for 7B policies), while SPQE incurs only per-query inference cost. The authors state that wallclock costs are not directly comparable, and that SPQE with the same Qwen2.5-{3B,7B}-Instruct backbones was not as competitive in preliminary experiments.

• OPQE's specific numbers are not reported in the available text. The abstract states that OPQE outperforms both standalone prompting and RL-based rewriting and yields the strongest overall performance, and the discussion refers to higher initial rewards when training OPQE (Fig. 1), but per-dataset results are not included in the supplied content.

Methodology in Plain English

The authors take an existing RL system, DeepRetrieval, and rebuild it faithfully. In that setup, a language model policy reads the original query, writes some intermediate reasoning, and then emits a final rewritten query. That rewritten query is sent to a retriever, the retriever returns its top passages, and a reward is computed from retrieval metrics against known relevance labels. The policy is then updated with PPO under a KL penalty to a reference policy, so the model is pushed toward queries that retrieve better without drifting into incoherent text. The retriever and corpus act as the "environment," meaning the model can be optimized even when the retrieval system is a black box.

Against this they place SPQE, which is much simpler: ask an instruction-following LLM to write a hypothetical document that answers the query, then search using the query and the generated document together. No training is involved. To make the comparison fair in terms of compute, they let the prompting baseline use larger models than the RL policies, since RL consumes heavy training compute.

They evaluate on three families of tasks. Evidence-seeking retrieval uses the Wikipedia-18 snapshot (English Wikipedia from 2018-12-20, chunked into non-overlapping 100-word passages) as the corpus, with NQ, TriviaQA, and SQuAD, and reports Hit@20 with BM25 and E5-base-v2. Ad hoc retrieval uses BEIR datasets (FEVER, HotpotQA, NFCorpus, SciFact) plus MS MARCO's dev set, DL19, and DL20, reports NDCG@10, and uses BM25 and Contriever-msmarco. Tool retrieval uses ToolRet with its three official splits, reports NDCG@10, Recall@10, and Completeness@10, and optimizes directly against Completeness@10. The tool-retrieval table caption lists BM25 and Contriever-msmarco as the retrievers, while the setup text names BM25 and E5-base-v2.

Training uses veRL 0.1 with PPO, the original DeepRetrieval codebase with engineering changes for throughput, and retrieval built on Pyserini and Faiss. Settings include learning rate 1e-6, global train batch size 128, actor and critic mini-batch size 32, and rollout temperature 0.6. Group-based alternatives such as GRPO were not evaluated.

Why This Matters

Impact on research. The paper reframes a common assumption — that tying training directly to retrieval metrics is the better path — by showing that under controlled, compute-aware conditions a training-free baseline is highly competitive and often better on aggregate. It also isolates a specific design choice by distinguishing on-policy optimization from off-policy preference tuning, deliberately scoping claims to PPO and leaving off-policy methods for future work. This gives the community a cleaner basis for deciding where RL effort is justified.

Real-world applications:

• API-based search over proprietary corpora. When the retriever is a hosted embedding service such as OpenAI's text-embedding-3 or a system like PubMed, the query is the only adjustable surface, and the paper shows a zero-shot pseudo-document prefix is a strong default.
• Tool and function retrieval for LLM agents. Agents must select the correct tool from large, functionally overlapping tool sets; the paper shows prompting substantially improves NDCG@10 and Recall@10 for both sparse and dense retrievers on ToolRet.
• Resource-constrained deployments. A zero-shot prompting setup with an open-weight 32B model can run on a single A100 40GB GPU, avoiding per-dataset RL training.
• Domain-specific search where adaptation pays off. Results on FiQA suggest that when terminology is specialized and training data exists, RL adaptation can outperform generic prompting.

Industry relevance. The findings give teams a cost-aware decision rule: start with prompt-based pseudo-document expansion, and invest in RL only when a stable domain, labeled relevance data, and the compute budget for per-dataset training are available. The observation that RL helps term-matching retrievers more than dense ones is directly actionable for hybrid retrieval stacks.

Future Directions

• Systematic comparison against off-policy methods. The paper explicitly restricts itself to on-policy PPO and leaves comparison with preference-tuning approaches such as Direct Preference Optimization to future work.
• More robust reward design. The authors optimize tool retrieval directly against Completeness@10 and leave fine-grained reward design to future work, especially given that sparse and dense retrievers respond differently to the same reward.
• Diagnosing why RL helps sparse but not dense retrievers. The authors suggest analyzing model outputs — length and verbosity, term distribution shifts, and query–tool coverage — and state that these statistics are not included in this paper.
• Comparison of on-policy algorithms. Group-based methods such as GRPO are not evaluated, and the authors leave a systematic comparison of on-policy algorithms to future work.

Target Audience

Researchers and practitioners in information retrieval, retrieval-augmented generation, and LLM agent design who need to decide between prompt-based query augmentation and RL-based query rewriting. It is most useful to those who already understand retrieval metrics and RL fine-tuning, and who want evidence about when the extra training cost of on-policy optimization is justified. Engineers building search over black-box or API-based retrievers will find the practical comparison and the tool-retrieval extension especially relevant.

Authors’ abstract

Recent advances in large language models (LLMs) have led to a surge of interest in query augmentation for information retrieval (IR). Two main approaches have emerged. The first prompts LLMs to generate answers or pseudo-documents that serve as new queries, relying purely on the model's parametric knowledge or contextual information. The second applies reinforcement learning (RL) to fine-tune LLMs for query rewriting, directly optimizing retrieval metrics. While having respective advantages and limitations, the two approaches have not been compared under consistent experimental conditions. In this work, we present the first systematic comparison of prompting-based and RL-based query augmentation across diverse benchmarks, including evidence-seeking, ad hoc, and tool retrieval. Our key finding is that under a compute-aware comparison setting, simple, training-free query augmentation often performs on par with, or even surpasses, more expensive RL-based counterparts, especially when using powerful LLMs. Motivated by this discovery, we introduce a novel hybrid method, On-policy Pseudo-document Query Expansion (OPQE), in which the LLM policy learns to generate a pseudo-document that maximizes retrieval performance, rather than rewriting the query, thus merging the flexibility and generative structure of prompting with the targeted optimization of RL. We show OPQE outperforms both standalone prompting and RL-based rewriting, demonstrating that a synergistic approach yields the best results. We open source our implementation to facilitate reproducibility.