emrekuruu/RetrievalRouter-lambda-l00

1---2license: mit3language:4- en5base_model:6- Qwen/Qwen3-0.6B-Base7library_name: transformers8pipeline_tag: text-classification9tags:10- retrieval11- document-retrieval12- information-retrieval13- routing14- RAG15- query-routing16- late-interaction17- lora18- peft19datasets:20- emrekuruu/FinReport21- emrekuruu/FinSlides22- emrekuruu/FinQA23- emrekuruu/ConvFinQA24- emrekuruu/VQAonBD25- emrekuruu/TATDQA26- emrekuruu/ArxivQA27- emrekuruu/Wiki-ss28- emrekuruu/MP-DocVQA29- emrekuruu/SciQAG30- emrekuruu/DUDE31metrics:32- ndcg33---34 35# RetrievalRouter (λ=0.0)36 37Official checkpoint from **RetrievalRouter: Joint Modality and Architecture Selection for38Document Retrieval** (EMNLP 2026). Given only the **query text**, RetrievalRouter predicts39*which retrieval pipeline* — across **modality** (text vs. multimodal) and **architecture**40(lexical, dense, or late-interaction rerank) — to run for that query.41 42- 📄 Paper: https://arxiv.org/pdf/2608.2562543- 💻 Code: https://github.com/emrekuruu/retrieval-router44- 🤗 Collection (all checkpoints + datasets): https://huggingface.co/collections/emrekuruu/retrieval-router45 46## Motivation47 48Retrieval pipelines differ in **modality** (search over text, or over page images) and49**architecture** (cheap dense search, or expensive late-interaction). The accurate ones are50slow; the fast ones miss evidence on hard documents. And which one fails depends on the query —51a text pipeline can't answer "what's the red curve in Figure 3?", but a multimodal one is52overkill for a plain factoid. Across 11 benchmarks, **no single pipeline wins on everything**.53RetrievalRouter picks the cheapest pipeline that can still answer each query, so easy queries54stay fast and hard ones still get the heavy pipeline.55 56## What this model is for57 58This is a **router, not a retriever**. It takes a query and predicts which of five retrieval59pipelines to run — in about 15 ms, before any search happens. You then run the chosen pipeline60to fetch documents.61 62Use it when you keep several retrieval setups over the same corpus and want to run the expensive63ones only when they help. Pick the checkpoint by **λ**: `0.0` for best accuracy, `1.0` for best64speed, in between to trade off.65 66It doesn't rank or read documents itself, and assumes your indices already exist. Trained on67English financial, scientific, and open-domain documents; other domains and languages are68untested.69 70## This checkpoint71 72**Trained with λ=0.0** — a **quality-only** objective — it picks the pipeline with the highest expected nDCG@5 and ignores latency. This is the most accurate operating point.73 74| Checkpoint | λ | Objective |75|---|---|---|76| [`RetrievalRouter-lambda-l00`](https://huggingface.co/emrekuruu/RetrievalRouter-lambda-l00) | 0.0 | Accuracy only |77| [`RetrievalRouter-lambda-l10`](https://huggingface.co/emrekuruu/RetrievalRouter-lambda-l10) | 0.1 | Accuracy-leaning |78| [`RetrievalRouter-lambda-l30`](https://huggingface.co/emrekuruu/RetrievalRouter-lambda-l30) | 0.3 | Balanced |79| [`RetrievalRouter-lambda-l50`](https://huggingface.co/emrekuruu/RetrievalRouter-lambda-l50) | 0.5 | Balanced |80| [`RetrievalRouter-lambda-l70`](https://huggingface.co/emrekuruu/RetrievalRouter-lambda-l70) | 0.7 | Latency-leaning |81| [`RetrievalRouter-lambda-l100`](https://huggingface.co/emrekuruu/RetrievalRouter-lambda-l100) | 1.0 | Latency only |82 83## Routing arms84 85| Index | Arm (config name) | Paper name | Modality | Architecture |86|---|---|---|---|---|87| 0 | `MULTIMODAL_RERANK` | MM-Rerank | Multimodal | Dense → late-interaction rerank |88| 1 | `MULTIMODAL-SINGLE` | MM-Dense | Multimodal | Single-vector dense |89| 2 | `TEXT_RERANK` | Text-Rerank | Text | Dense → late-interaction rerank |90| 3 | `TEXT-SINGLE` | Text-Dense | Text | Single-vector dense |91| 4 | `BM25` | BM25 | Text | Lexical |92 93The action space is these five arms. Two further pipelines evaluated in the paper (Text-Late,94MM-Late) are reported as static reference baselines but never routed to.95 96## Architecture97 98- **Encoder:** [Qwen/Qwen3-0.6B-Base](https://huggingface.co/Qwen/Qwen3-0.6B-Base) with LoRA99  adapters on the attention and feed-forward projections (merged into these weights).100- **Pooling:** mean-pool over the final hidden states → a 1024-d query representation.101- **Head:** a single linear layer → logits over the five arms; softmax gives the routing policy.102- Custom modeling code ships in the repo and loads via `trust_remote_code=True`.103 104## Usage105 106```python107import torch108from transformers import AutoModel, AutoTokenizer109 110repo = "emrekuruu/RetrievalRouter-lambda-l00"111tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)112model = AutoModel.from_pretrained(repo, trust_remote_code=True).eval()113 114inputs = tokenizer("In figure 3, what does the red dashed curve represent?",115                   return_tensors="pt", truncation=True, max_length=128)116with torch.no_grad():117    logits = model(**inputs)["logits"]              # shape [1, 5]118arm = model.config.strategy_names[logits.softmax(-1).argmax(-1).item()]119print(arm)   # e.g. "MULTIMODAL_RERANK" -> run that pipeline for this query120```121 122The router returns **which retrieval pipeline to run**, not documents. You then execute the123selected pipeline against your own indices.124 125## Training126 127Trained against **soft targets** from a per-query reward vector over the five arms, rather than a128single hard best-pipeline label (pipelines frequently tie on nDCG@5, and hard labels inject129noise). The reward combines accuracy and efficiency,130 131$$ r_i(q) = (1-\lambda)\, s_i(q) + \lambda\,(1 - \ell_i(q)), $$132 133where $s_i(q)$ is the arm's nDCG@5 and $\ell_i(q)$ its per-query normalized latency. The reward134vector becomes a target distribution via a low-temperature softmax (τ=0.1), and the router135minimizes the KL divergence to it. **λ is the only knob** that differs across the checkpoints136above. Training data spans **85,103 queries across 11 benchmarks**.137 138## Results (headline)139 140Against the strongest static pipeline, RetrievalRouter is **+2.5% nDCG@5 and 12.4× faster**.141Against the prior adaptive strategy-selection baseline142([`emrekuruu/RetrievalRouter_Baseline`](https://huggingface.co/emrekuruu/RetrievalRouter_Baseline)), it achieves significantly143higher nDCG@5 in accuracy-oriented settings and matches or numerically beats it on both nDCG@5144and latency in latency-oriented settings. See the paper for full tables and significance tests.145 146## Citation147 148```bibtex149@misc{kuru2026retrievalrouterjointmodalityarchitecture,150      title={RetrievalRouter: Joint Modality and Architecture Selection for Document Retrieval}, 151      author={Emre Kuru and Mehmet Onur Keskin and Reza Farahbakhsh and Noel Crespi},152      year={2026},153      eprint={2608.25625},154      archivePrefix={arXiv},155      primaryClass={cs.IR},156      url={https://arxiv.org/abs/2608.25625}, 157}158```159