PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning to Retrieve via Reinforcement Learning in Embedding Space

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

arXiv:2610.07731v1cs.IRcs.AIcs.CLcs.LG

TL;DR

Contrastive training does not directly optimize retrieval metrics or downstream task performance. RELER samples embedding actions and updates encoders with task rewards, using conditional-mean projection to reduce gradient noise. It improves retrieval on BRIGHT and, with a frozen index and generator, improves retrieval and answer quality across seven QA datasets.

  • Problem

    Contrastive objectives learn representations but do not directly optimize retrieval metrics or downstream task performance such as RAG answer accuracy.

    Source: §1 · p. 1

  • Method

    RELER samples query and document embeddings from vMF policies, scores retrieval or downstream outcomes as rewards, and updates the encoder with REINFORCE and a leave-one-out baseline.

    Source: §1 · p. 2

  • Results

    RELER generally improves average retrieval over InfoNCE across BRIGHT backbones, and query-only adaptation improves retrieval and answer quality with a fixed index and generator.

    Source: §1 · p. 2, §4.1.2 · p. 7, §4.2.2 · p. 9

  • Takeaways & Limitations

    These results show the potential of embedding-space RL for retrieval post-training, including query-only adaptation with a fixed document index and generator.

    Source: §5 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

RELER addresses the mismatch between contrastive training and retrieval or downstream task objectives by treating embeddings as retrieval actions that can be optimized with task feedback. On BRIGHT and across seven QA datasets, post-training improves retrieval and RAG answer quality.

1 INTRODUCTION

  • Contrastive objectives train relevant query–document pairs to score above irrelevant ones but are only indirectly related to retrieval metrics and downstream answer accuracy.Lower contrastive loss therefore does not necessarily translate into better retrieval quality or downstream utility.Source: §1 · p. 1
  • RL can optimize retrieval and downstream task metrics by sampling retrieval decisions and estimating policy gradients from their outcomes without differentiating through retrieval or the evaluator.Earlier retrieval policies explored generated text or discrete document selections rather than sampling the query and document vectors used for similarity scoring.Source: §1 · pp. 1–2
  • RELER samples query and document embedding actions, evaluates retrieval or downstream outcomes as rewards, and updates the encoder to favor higher-reward actions.It supports query-only adaptation against a fixed document index and retains deterministic single-vector retrieval at inference.Source: §1 · p. 2
  • RELER uses vMF policies, product rollouts, and REINFORCE with a leave-one-out baseline; conditional-mean projection reduces reward-irrelevant gradient noise while preserving the expected policy gradient.The policy samples around normalized encoder outputs on the unit sphere used for cosine retrieval.Source: §1 · p. 2
  • Across BGE-M3 and Qwen3-Embedding, RELER outperforms same-data supervised baselines on BRIGHT and improves average retrieval and answer quality across seven QA datasets.The RAG evaluation adapts only the query encoder while keeping the document index and generator fixed.Source: §1 · p. 2

2 RELATED WORK

Prior work chiefly trains dense retrievers with contrastive learning, while retrieval-oriented RL has explored actions such as document selections, ranked lists, and generated reasoning. RELER instead samples query and document embeddings as actions and uses retrieval rewards with conditional-mean projection.

2 RELATED WORK

  • Contrastive learning dominates dense retrieval, with later work extending its training recipe through hard negatives, weak supervision, multi-stage training, and specialized LLM adaptation.The cited approaches provide pretrained dual encoders that RELER subsequently adapts with task-level rewards.Source: §2 · p. 2
  • RL retrieval methods have sampled personal documents for generation rewards (Salemi et al., 2024), dynamically retrieved products (Liu et al., 2025c), and ordered document lists (Zhang et al., 2026).Other approaches use generated reasoning, including rationale-conditioned embeddings (Sun et al., 2026) and feedback from a separate embedder (Jiang et al., 2026).Source: §2 · pp. 2–3
  • RELER independently samples query and document embeddings and uses retrieval rewards with CMP, whereas prior approaches reshape retrieval through document selections, scores, or generated reasoning.At deployment, RELER uses deterministic single-vector encoder outputs.Source: §2 · p. 3

3 METHOD

The method uses embedding-space policies and task rewards to train retrieval models, with product-rollout RLOO and conditional-mean projection for learning. It also supports query-only adaptation against a fixed index and rewards based on retrieval, answer quality, or both.

3.1 RETRIEVAL AS AN EMBEDDING-SPACE POLICY

  • A vMF policy samples query and document embeddings around normalized encoder outputs, while a task-defined scalar reward scores retrieved lists or downstream outputs against supervision.The reward need not be differentiable, allowing the objective to cover retrieval and downstream task metrics.Source: §3.1 · p. 3

3.2 PRODUCT ROLLOUTS WITH REINFORCE-LOO

  • Product rollouts evaluate every pairing of G query actions and G document bundles, reusing samples for G^2 rewards without extra encoder passes or action samples.REINFORCE-LOO compares each action’s mean reward with the other actions in its role; its baseline is independent of that action and the estimator is unbiased.Source: §3.2 · pp. 3–4

3.3 VARIANCE REDUCTION BY CONDITIONAL-MEAN PROJECTION

  • CMP projects sampled actions onto subspaces retaining their policy means and reward-relevant candidate directions, averaging out residual randomness while preserving candidate scores.For bounded rewards determined by retained scores, CMP preserves the expected gradient and does not increase any individual summand’s variance along any parameter direction.Source: §3.3 · pp. 4–5

3.4 QUERY ADAPTATION WITH A FIXED INDEX

  • For fixed-index retrieval, RELER freezes document embeddings and adapts a separate query encoder; query-only RLOO samples query vectors and rewards their full-index top-k lists.The method anchors the query encoder to its pretrained initialization to limit drift; it omits CMP because preserving retrieved-document scores need not preserve the full-index top-k list.Source: §3.4 · p. 5

3.5 TRAINING REWARDS

  • Training rewards use relevance labels for retrieval and reference answers for RAG, allowing query-only updates without differentiating through retrieval or generation.RAG reward variants include nDCG@10, answer token F1, and their equal mixture; the retrieval reward also combines graded ordering near the cutoff with a pairwise term over the candidate set.Source: §3.5 · pp. 5–6

4 EXPERIMENTS

The experiments test RELER against supervised baselines on reasoning-intensive BRIGHT, then assess retrieval and answer quality when only the query encoder adapts in a fixed-index RAG pipeline. Across these settings, RELER improves retrieval and answer outcomes, while ablations examine its rewards, rollouts, and gradient estimator.

4.1.1 EXPERIMENTAL SETUP

  • Training adapts Qwen3-Embedding and BGE-M3 on 4,963 ReasonRank queries; BRIGHT scores average equally over 12 subsets and three seeds.RELER uses product rollouts, RLOO+CMP, and graded nDCG@10 augmented by pairwise feedback.Source: §4.1.1 · p. 6

4.1.2 MAIN RESULTS

  • RELER achieves the highest average nDCG@10 across all three backbones and both BRIGHT query settings, ahead of InfoNCE, LambdaLoss, and pretrained encoders.The advantage over LambdaLoss, which uses the same teacher grades, persists with GPT-4 reasoning queries.Source: §4.1.2 · p. 7, Table 1 · p. 7

4.1.3 ABLATIONS AND TRAINING ANALYSIS

  • With the combined reward, CMP raises BRIGHT from 21.91 ± 0.70 to 23.38 ± 0.17 and yields 291–590-fold lower gradient variance in matched probes.CMP also adds 1.87 points with graded nDCG alone; the authors report improvement for every paired seed in both comparisons.Source: §4.1.3 · pp. 7–8, Table 2 · p. 8
  • Product rollouts beat one-to-one pairing by 2.14 points; teacher grades raise nDCG-based results from 17.45 to 22.72, and pairwise feedback adds 0.66 points.The best tested pairwise weight is λ = 0.5; BRIGHT is stable across ρ = 0.60–0.90, with a 0.29-point full sweep. Figure 2 also reports similar training nDCG curves with and without pairwise feedback.Source: §4.1.3 · p. 8, Figure 2 · p. 9

4.2.1 EXPERIMENTAL SETUP

  • For fixed-index QA, the Q3E-0.6B document index and Qwen2.5-7B-Instruct generator remain frozen while query-only RLOO is trained on two datasets.Five out-of-domain datasets provide no training queries; evaluation averages answer EM/F1 and answer-containing-passage Hit@10 across seven datasets.Source: §4.2.1 · p. 9, Table 3 · p. 9

4.2.2 RESULTS AND ANALYSIS

  • All three RLOO rewards improve Hit@10 and answer EM/F1 over InfoNCE, while the mixed reward exceeds the pretrained encoder on all seven-dataset averages.Answer-F1 reward slightly improves average EM/F1 over graded nDCG, and the mixture gives the best overall answer scores; Hit@10 does not capture passage order, distracting context, or generator input limits.Source: §4.2.2 · pp. 9–10, Table 3 · p. 9

5 CONCLUSION

The conclusion presents RELER as embedding-space reinforcement learning for retrieval and downstream rewards, with conditional-mean projection reducing reward-invisible gradient noise. Results support its potential for retrieval post-training, while adaptive exploration and multi-step retrieval remain future directions.

5 CONCLUSION

  • RELER uses vMF exploration during training to align embedding models with retrieval and downstream rewards, while retaining deterministic vector retrieval at inference.Conditional-mean projection removes reward-invisible gradient noise without changing the expected fixed-candidate gradient.Source: §5 · p. 10
  • Across BRIGHT backbones, RELER generally improves average retrieval over InfoNCE, though gains vary by query form and domain.With a fixed document index and generator, query-encoder adaptation improves both retrieval performance and end-to-end answer quality.Source: §5 · p. 10
  • The authors identify adaptive exploration policies and multi-step retrieval for search agents as directions for future work.Source: §5 · p. 10

A THEORETICAL ANALYSIS

The theoretical analysis derives RELER’s product-rollout gradient and establishes conditional unbiasedness and per-term variance guarantees for conditional-mean projection (CMP). The implementation specifies vMF exploration, retrieval rewards, and practical CMP projections.

A.1 GRADIENT ESTIMATION FOR RELER

  • Product rollouts pair query actions with document bundles, and leave-one-out feedback yields an unbiased policy gradient under fixed-record, candidate-mask, and concentration assumptions.The reward-table row and column means estimate inner expectations; 2G component actions produce G² correlated evaluations.Source: §A.1 · pp. 14–15

A.2 CONDITIONAL-MEAN PROJECTION: UNBIASEDNESS AND VARIANCE

  • CMP is unbiased and does not increase any individual summand’s variance along any parameter direction when bounded rewards depend on retained candidate scores.The guarantee requires retaining all reward-relevant directions; it does not order the covariance of the summed gradient.Source: §A.2 · p. 15

A.3 PAIRWISE CONDITIONAL-MEAN PROJECTION

  • The pairwise CMP estimator is unbiased for the pairwise reward gradient, but its termwise variance guarantee does not order the covariance of the combined reward gradient.Documents outside a positive–negative pair have zero expected score-weighted contribution and can be omitted.Source: §A.3 · p. 16

B.1 EXPLORATION AND EMBEDDING SAMPLING

  • RELER samples unit-sphere actions from vMF policies using rejection sampling, decomposing each action into alignment with its encoder mean and an orthogonal tangent direction.The query and document means are sampled independently, and sampled actions are held fixed during the policy-gradient update.Source: §B.1 · p. 16

B.2 REWARD DEFINITIONS

  • The main retrieval reward combines graded nDCG with binary pairwise feedback at λ = 0.5, evaluated exactly without a differentiable rank relaxation.Teacher ranks map to grades 3 for rank 1, 2 for ranks 2–5, 1 for ranks 6–10, and 0 otherwise; pairwise feedback scores correctly ordered positives and negatives 1, ties 1/2, and reversals 0.Source: §B.2 · p. 17
  • If every pairing in a reward table has the same value, RLOO and CMP produce zero updates for that reward component; broader exploration can change orderings but move behavior farther from deterministic deployment.Projection can reduce gradient noise but cannot create a missing reward difference.Source: §B.2 · p. 17

B.3 CONDITIONAL-MEAN PROJECTION IMPLEMENTATION

  • The listwise query projection spans the policy mean and sampled documents, while each pairwise query projection spans the mean and positive–negative difference; CMP requires no extra encoder forwards or reward evaluations.The listwise query subspace has dimension at most 21 for up to 20 candidates, versus embedding dimension 1024.Source: §B.3 · pp. 17–18

B.4 QUERY–DOCUMENT POST-TRAINING ALGORITHM

The post-training algorithm samples query and document embeddings, computes leave-one-out feedback, and combines listwise and pairwise CMP gradients. The accompanying BRIGHT material specifies preprocessing, evaluation settings, and supervised baseline objectives.

B.4 QUERY–DOCUMENT POST-TRAINING ALGORITHM

  • Algorithm 2 encodes each minibatch once, samples query and document actions, computes cellwise leave-one-out weights and listwise and pairwise CMP gradients, then updates the shared encoder.The combined batch gradient uses one backward pass; an empty pair set leaves only the listwise contribution.Source: §B.4 · p. 18

C BRIGHT EXPERIMENTS

  • The BRIGHT experiments section collects training and evaluation protocols, reference baselines, ablations, and gradient measurements.The supplied passages describe these components but do not report their results.Source: C BRIGHT EXPERIMENTS · p. 19

C.1.1 DATA AND PREPROCESSING

  • Preprocessing reduces 6,721 records to 4,963 queries, with 4,896 retained under source-wise batching; there is no internal development split, and overlap screening does not establish complete semantic decontamination.The remaining queries contain 15,830 positive occurrences.Source: §C.1.1 · p. 19

C.1.2 REPRESENTATION AND EVALUATION SETTINGS

  • Qwen3-Embedding uses instruction-prefixed queries and terminal-token pooling, whereas BGE-M3 uses CLS pooling without query or document prefixes; both allow evaluation up to 8,192 tokens.The evaluation settings specify deterministic retrieval using unscaled cosine similarity.Source: §C.1.2 · p. 19

C.1.3 BASELINE OBJECTIVES

  • InfoNCE treats each positive separately and uses the full valid cross-device pool, while LambdaLoss uses teacher grades and detached absolute nDCG@10 swap weights.LambdaLoss normalizes its weighted logistic sum by each query’s total swap weight.Source: §C.1.3 · p. 19

C.1.4 TRAINING HYPERPARAMETERS

This appendix reports BRIGHT training configurations, per-seed results, and reference-baseline context. With GPT-4 reasoning queries, RELER beats InfoNCE at every paired seed across all three backbones; literature results use a different protocol.

C.1.4 TRAINING HYPERPARAMETERS

  • BRIGHT adaptation disables gradient clipping, reference-policy KL, and auxiliary contrastive losses; its main reward uses λ = 0.5, while a shared step budget fixes query exposure across methods.All 4B BRIGHT methods use LoRA (rank 16) with a learning rate of 2 × 10−4.Source: §C.1.4 · p. 19, Table 4 · p. 20

C.1.5 PER-SEED MAIN RESULTS

  • RELER exceeds InfoNCE at every paired seed on all three backbones with GPT-4 reasoning queries.Table 5 reports per-seed BRIGHT nDCG@10 averaged over 12 subsets; pretrained encoders were evaluated once.Source: §C.1.5 · p. 20, Table 5 · p. 20

C.2 BRIGHT REFERENCE BASELINES

  • Published BRIGHT reference scores provide context, not controlled comparisons, because their training and evaluation protocols differ.Table 6 omits MS MARCO fine-tuning controls and pretrained Qwen3-Embedding rows; its literature results come from Chen et al. (2026).Source: §C.2 · p. 20, Table 6 · p. 21

C.3 COMPLETE BRIGHT ABLATION RESULTS

The appendix details BRIGHT ablation and sensitivity evaluations. Among the tested pairwise weights, λ = 0.5 performs best, while the average changes little across the tested alignment values.

C.3 COMPLETE BRIGHT ABLATION RESULTS

  • The ablation tables use Q3E-0.6B, original queries, G = 64, 113 steps, and three seeds; Table 7 groups comparisons by reward and rollout design.Its blocks compare full RELER with an unprojected RLOO control, standalone graded nDCG variants, and binary-reward variants.Source: §C.3 · p. 21, Table 7 · p. 22

C.4 PAIRWISE WEIGHT AND ALIGNMENT

  • 0.08 points separate ρ = 0.80 from the main setting, but that setting does not improve every paired seed.The alignment sweep has a smaller effect on average than the pairwise-weight sweep.Source: §C.4 · p. 22, Table 8 · p. 22
  • In own-query candidate product-rollout sweeps with RLOO+CMP, λ=0.5 at ρ=0.70 gains come from StackExchange, while Coding and Theorem-based retrieval score lower than at λ=0; alignment changes have a smaller average effect.Source: §C.4 · p. 22

C.5 GRADIENT VARIANCE: PROTOCOL AND RESULTS

The appendix measures CMP’s gradient-noise reduction and documents a fixed-index QA setup. It also specifies the reward, anchoring, and generation procedures used for those experiments.

C.5 GRADIENT VARIANCE: PROTOCOL AND RESULTS

  • The probes compare paired estimators on shared actions, candidates, and rewards across three fixed microbatches, with dropout off and no optimizer updates.The estimated squared mean-gradient difference is nonpositive in five probes and at most 2.81% of CMP’s estimated squared signal in the other four.Source: §C.5 · p. 23
  • 291–590-fold lower full-gradient variance appears in all nine fixed-state probes, while CMP adds 11–14% forward/backward time per draw.The BRIGHT seed standard deviation also reflects the subsequent optimization trajectory, unlike these fixed-state rollout-noise measurements.Source: §C.5 · p. 23, Table 9 · p. 23

D.1 CORPUS AND TRAINING DATA

  • The QA experiments freeze a 21,015,324-passage Wikipedia index and Qwen2.5-7B-Instruct generator while adapting the query encoder.Training uses NQ and HotpotQA for one epoch; recipes differ in candidate access and objective.Source: §D.1 · p. 24, Table 10 · p. 24

D.2 TRAINING OBJECTIVES AND HYPERPARAMETERS

  • Graded rewards assign qrel positives grade 3, other supporting passages grade 2, and remaining answer-containing passages grade 1; mixed reward averages answer F1 and graded nDCG@10.InfoNCE uses a once-mined pool, whereas query-only RLOO searches the frozen full-corpus index.Source: §D.2 · p. 24

D.3 FIXED-INDEX QUERY ANCHOR

  • The fixed-index training anchor penalizes cosine distance from pretrained query embeddings, but does not alter rewards, the frozen index, or the generator.Initial embeddings are stored as constants; inference uses the current query embedding without an anchor lookup.Source: §D.3 · p. 25

D.4 GENERATION AND EVALUATION

  • RAG generation uses greedy decoding at temperature zero; evaluation normalizes answers for EM/F1 and separately measures reference-answer containment in Hit@10.Prompts drop the lowest-ranked passages first when they exceed the input budget.Source: §D.4 · p. 25

D.5 PER-DATASET RESULTS

Per-dataset QA results vary behind the aggregate averages. The mixed reward’s F1 advantage over graded nDCG narrows without Bamboogle, while its margin over answer-F1 training reverses.

D.5 PER-DATASET RESULTS

  • −0.25 points is the mixed reward’s F1 margin over answer-F1 training without Bamboogle, versus +0.24 across all seven datasets.Without Bamboogle, the mixed reward’s F1 margin over graded nDCG is +0.13, down from +0.30 across seven datasets.Source: §D.5 · p. 25
  • InfoNCE is best on NQ at 47.98 F1, while answer-F1 rewards beat graded nDCG in F1 on five of seven datasets.NQ and Bamboogle are the exceptions to that five-dataset pattern.Source: §D.5 · p. 25
  • Table 11 reports per-dataset results and query counts, with each dataset weighted equally in group means despite differing query counts.The table uses the same index, generator, and query counts across recipes.Source: §D.5 · p. 25, Table 11 · p. 26

E EMBEDDING TRAINING FROM A GENERAL-PURPOSE LLM

Starting from Qwen3-0.6B, RELER shows no clear advantage over InfoNCE, while the extended related work distinguishes its embedding-action approach and conditional-mean estimator from prior retrieval and ranking methods.

E EMBEDDING TRAINING FROM A GENERAL-PURPOSE LLM

  • Both methods train Qwen3-0.6B embeddings on E2Rank for 1,200 steps and are evaluated across 41 MTEB English v2 tasks in seven types.RELER uses graded nDCG@10 plus pairwise feedback, own-query candidates, product rollouts, RLOO+CMP, G = 64, and ρ = 0.70.Source: E EMBEDDING TRAINING FROM A… · p. 26
  • InfoNCE leads RELER by 1.48 points in retrieval and has the higher task-type average, while their all-task averages are nearly tied.The comparison uses one seed and complete recipes with different supervision and candidate handling.Source: E EMBEDDING TRAINING FROM A… · p. 26, Table 12 · p. 26

F EXTENDED RELATED WORK

  • Unlike prior MDP-based ranking policies and stochastic rankings, RELER studies a one-step embedding action with offline relevance feedback and assesses its value against matched InfoNCE and LambdaLoss controls.The related work situates contrastive training and pairwise or listwise learning-to-rank as complementary approaches.Source: F EXTENDED RELATED WORK · p. 27
  • Unlike approaches that select profile documents or generate reasoning for retrieval, RELER samples query and document embeddings during training and retains deterministic single-vector inference.In fixed-index mode, it samples a query vector and searches the full index, using retrieval or answer feedback in QA experiments.Source: F EXTENDED RELATED WORK · p. 27
  • CMP conditions its Rao–Blackwell averaging on vMF means, reward cells, and LOO baselines; its variance guarantee requires retaining every reward-relevant direction.The guarantee does not extend to using only returned full-corpus top-k hits.Source: F EXTENDED RELATED WORK · p. 27
Loading 2610.07731v1…