oldman-dev/tis-v2.5-multihop-reranker

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

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

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

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

Citation

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

License

MIT