---
language: en
license: mit
tags:
  - token-importance-scoring
  - passage-ranking
  - multi-hop-reasoning
  - kv-cache-compression
  - rag
  - information-retrieval
  - mistral
base_model: unsloth/mistral-7b-instruct-v0.3-bnb-4bit
datasets:
  - hotpot_qa
  - ms_marco
---

# TIS v2.5: Multi-Hop Passage Reranker

**Token Importance Scoring v2.5** — two-stage passage reranking with bridge detection and cross-passage iterative refinement, trained on HotpotQA multi-hop reasoning.

## What's New in v2.5

- **Two-stage scoring pipeline**: Stage 1 direct scoring + Stage 2 cross-passage refinement
- **Bridge detection**: explicit binary classifier for intermediate entity passages in multi-hop chains
- **100% recall_both@5** on HotpotQA bridge-type questions (vs N/A in v2.3)
- **Zero regression**: NIAH 95%, LITM 72% middle position — identical to v2.3
- **Backward compatible**: v2.3 code loads v2.5 checkpoint with no changes

## Architecture

Three learned components on top of frozen Mistral-7B-Instruct-v0.3 (4-bit NF4):

| Component | Parameters | Purpose |
|-----------|-----------|---------|
| `QueryAwareImportanceHead` | ~1.5M | Direct passage scoring (Stage 1) |
| `BridgeDetectionHead` | ~1.05M | Bridge passage classification |
| `RefinementScoringHead` | ~2.8M | Cross-passage attention refinement (Stage 2) |
| **Total trainable** | **~5.6M** | |

Score blending: `final = 0.7 × direct + 0.3 × refined`

## Results

| Metric | v2.3 | v2.5 | Delta |
|--------|------|------|-------|
| NIAH accuracy | 95.0% | 95.0% | ±0% ✅ |
| LITM middle position | 72.0% | 72.0% | ±0% ✅ |
| Multi-hop recall@5 | N/A | **100%** | NEW ✅ |
| Inference latency | ~50ms | ~120ms | +70ms (Stage 2) |
| Peak VRAM | 5.5 GB | 5.5 GB | ±0 ✅ |

## Checkpoint Contents

```python
import torch
ckpt = torch.load("tis_components.pt", map_location="cpu")
print(list(ckpt.keys()))
# ['importance_embedding', 'importance_head', 'bridge_detection_head', 'refinement_head']
```

File size: **157.8 MB** | SHA256: `53f3333dffc3288579628ad7f13311e4e8efe41d053dfb9d7479d14f70596413`

## Quick Start

```bash
pip install huggingface_hub
hf download oldman-dev/tis-v2.5-multihop-reranker --local-dir checkpoints/v2.5_refinement
```

```python
import torch
from src.token_importance.model.patched_model import PatchedCausalLM

model = PatchedCausalLM.from_pretrained(
    "unsloth/mistral-7b-instruct-v0.3-bnb-4bit",
    load_in_4bit=True,
)
ckpt = torch.load("checkpoints/v2.5_refinement/tis_components.pt", map_location="cuda")
model.importance_head.load_state_dict(ckpt["importance_head"])
model.bridge_detection_head.load_state_dict(ckpt["bridge_detection_head"])
model.refinement_head.load_state_dict(ckpt["refinement_head"])
```

## Training Details

- **Base model**: Mistral-7B-Instruct-v0.3 (frozen, 4-bit NF4)
- **Dataset**: HotpotQA (multi-hop) + MS-MARCO
- **Curriculum**: Stage 1 (625 steps equal weighting) → Stage 2 (375 steps, NDCG α 0.0→0.3)
- **Duration**: 28.6 min / 1000 steps (RTX 5070 8GB)
- **Optimizer**: AdamW, lr=2e-4, batch_size=1, grad_accum=8
- **Early stopping**: patience=3, final loss=0.023

## Version History

| Version | Key Achievement | HF Repo |
|---------|----------------|---------|
| **v2.5** | Multi-hop 100% recall@5 | `oldman-dev/tis-v2.5-multihop-reranker` |
| v2.3 | MRR 0.5102 (+18.1% vs BM25) — Tier 1 | `oldman-dev/tis-v2.3-passage-reranker` |
| v2.2 | MRR 0.471 (+9.1% vs BM25) | `oldman-dev/tis-v2.2-passage-reranker` |

## Source Code

[github.com/nitroxido/token-importance-scoring](https://github.com/nitroxido/token-importance-scoring)

## Citation

```bibtex
@article{tis2026,
  title   = {Token Importance Scoring for KV Cache Compression},
  year    = {2026}
}
```

## License

MIT