State
Lexical query features (length, IDF mass, stop-word ratio) and BM25 retrieval-confidence statistics (top-1 score, score gap, score variance).
RAG · Contextual Bandit · Retrieval Routing
Selective and Cost-Aware Retrieval Routing for RAG
NYCU · Reinforcement Learning (535514) · Spring 2026
We learn a lightweight, per-query router over BM25, dense, hybrid, and
keyword-rewrite retrieval actions using offline contextual-bandit methods.
The policy improves cost-aware retrieval evidence on BEIR benchmarks without
touching any LLM weights.
A bandit router over BM25 / dense / hybrid / rewrite.
Better retrieval. Explicit cost. No LLM fine-tuning.
Abstract
Routing retrieval actions as an offline contextual bandit
Fixed retrieval strategies in Retrieval-Augmented Generation (RAG) apply the same action—BM25, dense, hybrid, or query rewrite—to every incoming query, regardless of difficulty or domain. We reframe per-query retrieval selection as an offline contextual-bandit problem: for each query, a lightweight policy observes lexical and BM25-confidence features and selects the retrieval action expected to maximize a cost-aware reward combining Recall@5 and MRR minus retrieval-call penalties.
We evaluate on two BEIR benchmarks—SciFact (scientific claim verification) and NFCorpus (biomedical/nutrition retrieval)—using full-corpus offline replay with 600 training and 300 test queries per dataset. The selective policy outperforms the best fixed action on both datasets, with bootstrap 95% CIs strictly above zero, and transfers across domains without retraining. The claim is intentionally narrow: cost-aware evidence retrieval, not generated-answer EM/F1, not LLM fine-tuning.
Problem
Fixed retrievers cannot adapt to query difficulty.
A single retrieval action is rarely optimal for every query. On SciFact, a scientific claim-verification corpus, hybrid retrieval and keyword rewrites improve coverage. On NFCorpus—a short-title biomedical corpus—one-call BM25 or dense retrieval often suffices. The cost of extra retrieval calls (rewrite API, additional index lookups) is non-trivial and should be weighed against expected evidence gain.
We frame this as a one-step offline contextual bandit over retrieval actions, training from logged retrieval outcomes stored as a full-information replay table. The reward combines retrieval-quality signals and a call-cost penalty, so the policy must learn not only which action is best, but also when the gain justifies the extra cost.
Method
Offline routing policy over retrieval actions
Actions
BM25, dense (DPR-style), hybrid BM25+dense, keyword rewrite + BM25. Gemini-backed actions are expensive extensions evaluated separately.
Policy
Direct-method action-value estimators (Linear Thompson Sampling, LinUCB adapted for offline replay). Model selected via 5-fold CV on training reward.
Evaluation
Full-information held-out replay, bootstrap 95% CIs, Doubly Robust OPE, three-seed robustness sweep, and audit-trail artifact checks.
Reward and Constrained Utility
Policy objective
π(s) = argmaxa E[R | s, a]
Reward R
R = Recall@5 + 0.5·MRR
− RewriteCost − ExtraCallCost
Constrained utility (λ sweep)
U(λ) = Recall@5 + 0.5·MRR
− λ · calls/query
State features
Query length · IDF mass · stop-word ratio · BM25 top-1 score · score gap · score variance
Actions (4 core)
BM25 · Dense · Hybrid · Keyword rewrite; Gemini as optional high-cost extension
Training data
Replayable full-information table — all action rewards pre-computed offline from BEIR corpora
Query
Raw user question
→
State s
Lexical + BM25 confidence features
→
Policy π(s)
argmax predicted action value
→
Action a
BM25 / dense / hybrid / rewrite
→
Reward R
Recall@5 + MRR − cost
→
Evidence
CIs, OPE, audit trail
Results
Selective routing improves held-out reward on both datasets.
| Dataset | Method | Obj. | Value | Delta | 95% CI | Calls |
|---|---|---|---|---|---|---|
| SciFact | Train-best fixed | reward | 1.0568 | 0.0000 | — | 2.00 |
| SciFact | Selected-action bandit | reward | 1.0621 | +0.0054 | — | 1.03 |
| SciFact | Selective policy | reward | 1.0905 | +0.0337 | [0.0089, 0.0583] | 1.41 |
| SciFact | Constrained λ=0.03 | utility | 1.0894 | +0.0474 | [0.0143, 0.0822] | 1.27 |
| SciFact | Oracle action | reward | 1.2062 | +0.1495 | — | 1.01 |
| NFCorpus | Train-best fixed | reward | 0.3772 | 0.0000 | — | 1.00 |
| NFCorpus | Selected-action bandit | reward | 0.4040 | +0.0269 | — | 1.14 |
| NFCorpus | Selective policy | reward | 0.4071 | +0.0299 | [0.0054, 0.0548] | 1.00 |
| NFCorpus | Constrained λ=0.03 | utility | 0.4098 | +0.0327 | [0.0106, 0.0584] | 1.00 |
| NFCorpus | Oracle action | reward | 0.4641 | +0.0870 | — | 1.02 |
Reward delta vs. train-best
Orange bars = oracle headroom. Selective policy closes 22.6% of the SciFact oracle gap and 34.4% of the NFCorpus gap. Bootstrap 95% CIs are strictly above zero on both datasets.
Policy behavior
The learned policy adapts its action mix by domain.
SciFact action mix
BM25 35.7%
Dense 23.0%
Hybrid 41.3%
Hybrid and keyword rewrite (49.0% of queries) are frequent—scientific claim queries benefit from multi-index evidence and lexical expansion.
NFCorpus action mix
BM25 44.7%
Dense 55.3%
Hybrid is not selected on held-out NFCorpus. Short biomedical titles are well-served by single-index retrieval; rewrite appears on 31.3% of queries.
Budget sweep (λ)
At λ=0.06, SciFact keeps +0.039 utility delta at only 1.05 calls/query—essentially matching one-call cost. NFCorpus already achieves its best constrained utility at exactly 1.00 call/query.
Cross-Dataset Generalization
Policies trained on one domain transfer to another.
We test zero-shot transfer by applying a policy trained on one BEIR dataset to a completely different target domain—without any retraining. Both transfer directions produce positive reward deltas over the target's own train-best fixed action, suggesting the feature-based policy learns domain-agnostic routing signals.
The SciFact→NFCorpus transfer incurs extra retrieval calls (1.597 vs 1.000), indicating that some SciFact-learned preferences for hybrid retrieval do not pay off in the biomedical domain. The NFCorpus→SciFact direction transfers cleanly at 1.000 calls/query.
| Train on | Eval on | Policy Reward | Target Train-Best | Delta | Calls/Query |
|---|---|---|---|---|---|
| SciFact | SciFact | 1.093 | 1.057 | +0.036 | 1.477 |
| SciFact | NFCorpus | 0.386 | 0.377 | +0.009 | 1.597 |
| NFCorpus | NFCorpus | 0.407 | 0.377 | +0.030 | 1.000 |
| NFCorpus | SciFact | 1.064 | 1.057 | +0.007 | 1.000 |
Where the result gets stronger
Routing matters most on hard, high-disagreement queries.
+0.049
SciFact reward gain on high action-disagreement queries
+0.096
NFCorpus reward gain on high action-disagreement queries
+0.0006
Margin-gate improvement over selective on SciFact — confirming the policy is already well-calibrated
These subgroup analyses were pre-specified and confirm the routing intuition: when fixed strategies disagree most—high intra-action variance—the learned policy gains the most by selecting the right action. The margin-gate baseline adds near-zero improvement, confirming the policy is not easily beaten by simple confidence thresholding.
Diagnostics
Evidence that supports and bounds the claim
Selected-action replay
Linear Thompson Sampling delta vs. train-best: +0.005 on SciFact and +0.027 on NFCorpus under selected-action feedback only.
Replay regret
LinUCB selected-feedback regret is 82.1 vs 67.4 full-information direct on SciFact, and 43.2 vs 29.6 on NFCorpus—confirming that partial-feedback learning is strictly harder than full-replay.
OPE coverage warning
SciFact IPS error rises from 0.127 (uniform logging) to 0.268 (sparse train-best logging). DR-based estimates are used as the primary offline evaluator.
3-seed robustness
Selective win rate is 0.667 on both datasets across three seeds. Mean reward deltas: +0.021 (SciFact) and +0.025 (NFCorpus).
Downstream QA Pilot
Policy-routed Gemini reader: preliminary signal
While our main claim is retrieval-stage only, we ran a bounded Gemini-reader pilot (n=30 per dataset) to verify that retrieval improvements can propagate to answer quality. These are api_pilot results—single seed, small scale, not final QA benchmark evidence.
HotpotQA (n=30)
| Method | EM | F1 |
|---|---|---|
| BM25 + Gemini | 0.733 | 0.823 |
| Train-best fixed + Gemini | 0.767 | 0.852 |
| Learned policy + Gemini | 0.800 | 0.890 |
Promising: policy EM +0.067 over BM25 baseline.
Natural Questions (n=30)
| Method | EM | F1 |
|---|---|---|
| BM25 + Gemini | 0.133 | 0.171 |
| Train-best fixed + Gemini | 0.133 | 0.166 |
| Learned policy + Gemini | 0.133 | 0.170 |
Mixed: no answer-quality separation on NQ — further work needed.
Limitations & Claim Boundary
What this work does and does not claim
Retrieval-stage evidence only
We measure Recall@5 and MRR, not downstream EM/F1 or generation quality. The QA pilot is exploratory and small-scale.
Full-information replay table
Training uses pre-computed rewards for all actions on each query. This is offline supervised learning with a bandit framing—not true online bandit learning from deployed logs.
Simulated OPE
Off-policy evaluation uses simulated behavior logging, not real user interaction data. IPS estimates degrade under sparse logging.
Semantic features analysis-only
BERT/embedding features were explored in ablation but remain outside the main policy. The production policy uses only lexical and BM25-confidence features.
Small corpora
SciFact (5,183 docs) and NFCorpus (3,633 docs) are both small BEIR datasets. Scaling to Wikipedia-scale corpora is future work.
No LLM fine-tuning
Deliberately out of scope: we do not fine-tune any retriever or reader model. All improvements come from routing, not representation learning.
Poster
Final project poster
If the preview is unavailable, open poster.pdf directly.
Citation
Project reference
@misc{tsai2026selective,
title = {Selective and Cost-Aware Retrieval Routing for RAG},
author = {Tsai, Shang-Rong and Lee, Yi-Hsin and Lee, Cheng-Wei},
year = {2026},
note = {Reinforcement Learning Final Project (535514),
National Yang Ming Chiao Tung University}
}