---
license: mit
language:
- en
library_name: transformers
tags:
- rag
- router
- multimodal
- retrieval
- query-routing
- qwen3
pipeline_tag: text-classification
datasets:
- ananoymous/Wiki-ss
- ananoymous/DUDE
- ananoymous/TATDQA
- ananoymous/ArxivQA
- ananoymous/FinQA
- ananoymous/FinReport
- ananoymous/FinSlides
- ananoymous/ConvFinQA
- ananoymous/MP-DocVQA
- ananoymous/SciQAG
- ananoymous/VQAonBD
---

# IRouterLM: Adaptive Query Routing for Multimodal RAG

<p align="center">
  <a href="https://github.com/ananoymous-submission/sigir26">GitHub</a> •
  <a href="https://hf.co/collections/ananoymous/irouterlm">Training Data</a>
</p>

> A lightweight query-aware router that dynamically selects the optimal retrieval modality and architecture per query. IRouterLM achieves **state-of-the-art accuracy (0.76 nDCG@5)** while reducing latency by **90%** compared to the strongest baseline.

## Model Description

IRouterLM is a fine-tuned Qwen3-0.6B model that classifies queries into optimal RAG retrieval strategies. Given a user query, the model predicts which retrieval pipeline will yield the best results while balancing accuracy and latency.

### Supported Strategies

| Strategy ID | Strategy Name | Description |
|-------------|--------------|-------------|
| 0 | `MULTIMODAL_RERANK` | Multimodal dense retrieval + late-interaction reranking |
| 1 | `MULTIMODAL-SINGLE` | Single-stage multimodal dense retrieval |
| 2 | `TEXT_RERANK` | Text dense retrieval + late-interaction reranking |
| 3 | `TEXT-SINGLE` | Single-stage text dense retrieval |

## Quick Start

```python
from transformers import AutoModel, AutoTokenizer
import torch

# Load model and tokenizer
model = AutoModel.from_pretrained("ananoymous/IRouterLM", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("ananoymous/IRouterLM")

# Example query
query = "What was the revenue growth in Q3 2024?"
inputs = tokenizer(query, return_tensors="pt")

# Get prediction
with torch.no_grad():
    outputs = model(**inputs)
    probs = torch.softmax(outputs["logits"], dim=-1)
    prediction = probs.argmax(dim=-1).item()

# Strategy mapping
strategies = ["MULTIMODAL_RERANK", "MULTIMODAL-SINGLE", "TEXT_RERANK", "TEXT-SINGLE"]
print(f"Predicted strategy: {strategies[prediction]}")
print(f"Confidence: {probs[0][prediction]:.2%}")
```

### Using the `predict` Method

```python
result = model.predict(inputs["input_ids"], inputs["attention_mask"])
print(f"Strategy: {result['strategy_names'][0]}")
print(f"Probabilities: {result['probabilities']}")
```

## Architecture

- **Base Model**: Qwen3-0.6B
- **Fine-tuning**: LoRA (rank=16, alpha=32)
- **Target Modules**: q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj
- **Classification Head**: Mean pooling + Linear (1024 → 4)
- **Training Loss**: Weighted KL Divergence with soft labels

```
Query → Qwen3-0.6B (LoRA) → Mean Pooling → Classifier → Strategy Prediction
```

## Training Details

### Dataset

The model was trained on 80,000+ queries from 11 benchmarks:

| Domain | Datasets |
|--------|----------|
| Financial | FinReport, FinSlides, FinQA, ConvFinQA, TAT-DQA |
| Scientific | ArxivQA, SciQAG |
| General | Wiki-SS, MP-DocVQA, DUDE, VQAnBD,  |

### Hyperparameters

| Parameter | Value |
|-----------|-------|
| Learning Rate | 1e-4 |
| Batch Size | 16 |
| Epochs | 2 |
| Weight Decay | 0.01 |
| Warmup Ratio | 0.1 |
| Scheduler | Cosine |
| Precision | bfloat16 |
| λ (trade-off) | 0.0 (accuracy-focused) |

## Intended Use

IRouterLM is designed for:

- **RAG Systems**: Automatically select the optimal retrieval strategy per query
- **Document QA**: Route queries to text-only or multimodal pipelines based on query semantics
- **Cost Optimization**: Reduce computational costs by avoiding expensive pipelines when simpler ones suffice

### Limitations

- Trained on English queries only
- Optimized for document retrieval tasks (financial, scientific, general domains)

## License

MIT License

## Acknowledgments

This work builds on:
- [Qwen3](https://huggingface.co/Qwen/Qwen3-0.6B-Base) for the base model
- [ColPali](https://github.com/illuin-tech/colpali) for multimodal late-interaction retrieval
- [PEFT](https://github.com/huggingface/peft) for efficient fine-tuning