tis-v2.3-passage-reranker

TIS v2.3: Tier 1 Passage Ranking

🏆 Tier 1 Achievement — Test MRR 0.5102 (target: ≥0.50) | +18.1% vs BM25 (0.432)

Major milestone: First TIS checkpoint to exceed the 0.50 MRR threshold on MS-MARCO passage ranking.

Quick Start

import torch
from transformers import AutoTokenizer, AutoModel
from token_importance.model.importance_head import QueryAwareImportanceHead

# Load base model and checkpoint
model = AutoModel.from_pretrained("mistralai/Mistral-7B-v0.3", device_map="cuda")
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.3")

# Load v2.3 importance head
head = QueryAwareImportanceHead(hidden_dim=4096)
checkpoint = torch.load("best/tis_components.pt", map_location="cuda")
head.load_state_dict(checkpoint)

# Score a passage given a query
query = "What is the capital of France?"
passage = "Paris is the capital and most populous city of France."

query_ids = tokenizer.encode(query, return_tensors="pt").to("cuda")
passage_ids = tokenizer.encode(passage, return_tensors="pt").to("cuda")

with torch.no_grad():
    query_hidden = model(query_ids, output_hidden_states=True).hidden_states[-1]
    passage_hidden = model(passage_ids, output_hidden_states=True).hidden_states[-1]
    
    # Aggregate to passage-level score [0-100]
    scores = head.direct_score(passage_hidden, query_hidden)
    passage_score = scores.mean().item()
    print(f"Passage relevance score: {passage_score:.2f}")

Performance Metrics

Test Set Results (500 locked queries, MS-MARCO v1.1, seed=42)

Key metrics:

• Validation-Test gap: 0.35% (excellent generalization)
• Separator detection rate: 99.1% (high confidence separators found)
• Training efficiency: 2250 steps to peak validation (out of 3000 planned)

Training Trajectory

Step 250:  Val MRR 0.4479 (initial)
Step 500:  Val MRR 0.4300 (dip)
Step 750:  Val MRR 0.4801 (recovery)
Step 1000: Val MRR 0.4662 (variance)
Step 1250: Val MRR 0.4752 (holding)
Step 1500: Val MRR 0.4911 (improvement)
Step 1750: Val MRR 0.4869 (slight decline)
Step 2000: Val MRR 0.4734 (further decline)
Step 2250: Val MRR 0.5137 🏆 PEAK (early stop reset)
Step 2500: Val MRR 0.5057 (decline starts)
Step 2750: Val MRR 0.4968 (decline continues)
Step 3000: Val MRR 0.5060 (early stop triggered)

Best checkpoint saved at step 2250 (validation MRR 0.5137, test MRR 0.5102).

Architecture

Model: QueryAwareImportanceHead with batch tokenization optimization

• Input: passage tokens + query representation
• Aggregation: mean pooling (all token scores averaged)
• Score direction: high-first (descending, established by construction)
• Score range: [0, 100] (sigmoid(raw) × 100)
• Query interaction: 4-head cross-attention (passage tokens attend to query)
• Scoring head: 3-layer MLP with ReLU activations

Training Configuration

{
  "base_model": "mistralai/Mistral-7B-v0.3",
  "quantization": "4-bit NF4 (bitsandbytes)",
  "learning_rate": 5e-5,
  "optimizer": "Adam",
  "batch_size": 1 (forced by 8GB VRAM),
  "gradient_accumulation": 8 (effective batch 8),
  "mixed_precision": "bfloat16",
  "loss_fn": "pairwise ranking (margin=5.0)",
  "dataset": "MS-MARCO v1.1 (79K+ train, 500 tune, 500 test)",
  "train_steps": 3000,
  "best_step": 2250,
  "early_stopping": "patience=3 (steps 2500, 2750, 3000)",
  "hardware": "RTX 5070 (8GB VRAM, ~5.5GB peak)",
  "training_time": "~1:08 (50:37 to step 2250)"
}

Scoring Contract

Input:

• Passage tokens (must include [7031, 1233, 29515] separator markers)
• Query representation (mean-pooled hidden states of query)

Output:

• Score ∈ [0, 100] per token (direct_score path)
• Aggregated to passage-level via mean pooling
• Direction: high values = more important (descending order is correct)

Score meaning: Learned relevance signal between passage and query. Trained via supervised pairwise ranking on MS-MARCO relevance labels (higher scores for passages with more selected spans).

Release Status

✅ Tier 1 Achieved — MRR 0.5102 exceeds target (≥0.50)

Ready for production. All bugs fixed. Evaluation contract transparent. Results reproducible.

Related Checkpoints

• v2.2 (0.471 MRR, +9.1% BM25): Baseline supervised model
• stage3_ert (KV compression, LITM transfer): ERT-trained context-utility head
• v8b_hard_anchor (NIAH 82%): Specialized for cache compression

Citation

@article{tis2026v23,
  title   = {TIS v2.3: Tier 1 Passage Ranking via Query-Aware Importance Scoring},
  year    = {2026},
  month   = {August},
  version = {v2.3},
  metrics = {Test MRR: 0.5102, +18.1% vs BM25}
}

License

MIT — See repository for details.

---

Full paper & results: https://github.com/nitroxido/token-importance-scoring
All checkpoints: https://huggingface.co/oldman-dev
Baseline comparison: See V2.3-COMPLETE-COMPARISON.json in results/