SAGE: Strategy-Adaptive Generation Engine for Query Rewriting

TL;DR

This paper introduces SAGE, a strategy-guided reinforcement learning framework for query rewriting that improves retrieval performance, reduces exploration costs, and enhances interpretability by incorporating expert strategies and novel reward mechanisms.

Contribution

The paper presents SAGE, a novel RL-based query rewriting system that integrates expert strategies and innovative reward shaping to outperform existing methods and learn optimal strategies.

Findings

Achieves state-of-the-art NDCG@10 results.

Learns to select optimal strategies and reduces unnecessary exploration.

Generates concise rewrites, lowering inference costs.

Abstract

Query rewriting is pivotal for enhancing dense retrieval, yet current methods demand large-scale supervised data or suffer from inefficient reinforcement learning (RL) exploration. In this work, we first establish that guiding Large Language Models (LLMs) with a concise set of expert-crafted strategies, such as semantic expansion and entity disambiguation, substantially improves retrieval effectiveness on challenging benchmarks, including HotpotQA, FEVER, NFCorpus, and SciFact. Building on this insight, we introduce the Strategy-Adaptive Generation Engine (SAGE), which operationalizes these strategies in an RL framework. SAGE introduces two novel reward shaping mechanisms-Strategic Credit Shaping (SCS) and Contrastive Reward Shaping (CRS)-to deliver more informative learning signals. This strategy-guided approach not only achieves new state-of-the-art NDCG@10 results, but also uncovers a compelling emergent behavior: the agent learns to select optimal strategies, reduces unnecessary exploration, and generates concise rewrites, lowering inference cost without sacrificing performance. Our findings demonstrate that strategy-guided RL, enhanced with nuanced reward shaping, offers a scalable, efficient, and more interpretable paradigm for developing the next generation of robust information retrieval systems.