Reinforcement learning improves dense retrieval models and question answering

Learning to Retrieve via Reinforcement Learning in Embedding Space

Information RetrievalArtificial IntelligenceComputation and LanguageMachine Learning

Summary

Dense retrieval models usually learn to find relevant documents by comparing how close queries and documents are in a mathematical space, but they don't learn directly to improve final retrieval results or answer quality. The authors created RELER, a method that uses reinforcement learning to directly optimize these retrieval results by sampling from special probability distributions and updating the model based on how well it performs. They also introduced a trick to reduce noise during the learning process by focusing on a smaller space of possibilities. Testing on a difficult benchmark and several question answering tasks showed RELER improved retrieval accuracy and the quality of generated answers.

What this means in practice

  • For search engine developers: Improve document retrieval accuracy by optimizing query embeddings directly for retrieval and downstream tasks using reinforcement learning.
  • For question answering system teams: Enhance question answering quality by fine-tuning query encoders with rewards that jointly consider retrieval and answer generation performance.

Authors

Qi Liu, Fengming Liang, Yiqun Chen, Erhan Zhang, Jiaxin Mao

Abstract

Dense retrieval models are typically trained with contrastive objectives that learn effective representations but do not directly optimize retrieval metrics or downstream task performance. To address this problem, we introduce RELER (REinforcement LEarning for Retrieval), a reinforcement learning framework that enables existing embedding models to learn to retrieve directly in embedding space and align to task-specific rewards. We train RELER by sampling unit-length query and document embedding actions from von Mises-Fisher (vMF) distributions centered on normalized encoder outputs, scoring the resulting retrieval or downstream outcomes as rewards, and updating the encoder with REINFORCE using a leave-one-out baseline (RLOO). As exploration in the high-dimensional embedding space is prone to sampling noise, we further propose conditional-mean projection (CMP), which projects each sampled embedding onto the low-dimensional subspace spanned by its encoder output and the candidate embeddings it is compared against, reducing noise in the policy gradient while preserving its expectation. We evaluate RELER on BRIGHT, a benchmark with reasoning-intensive queries that remain challenging for existing embedding models. RELER consistently outperforms InfoNCE and LambdaLoss in average nDCG@10 when post-training BGE-M3 and Qwen3-Embedding backbones. We further evaluate downstream utility through retrieval-augmented generation (RAG), where we adapt only the query encoder while keeping the document index and generator fixed. Across seven QA datasets, jointly optimizing retrieval and answer rewards improves both average retrieval performance and answer quality in RAG.