oldman-dev/tis-passage-reranker

TIS Passage Reranker — Query-Aware Position Bias Elimination

A lightweight (~8M parameter) query-aware importance head for Mistral-7B-v0.3 that eliminates Lost-in-the-Middle (LITM) position bias in retrieval-augmented generation pipelines.

Part of the Token Importance Scoring (TIS 2.0) release.
GitHub: https://github.com/nitroxido/token-importance-scoring

What This Checkpoint Does

This head scores retrieved passages by relevance and reorders them before generation, ensuring the model sees the most relevant content regardless of the original retrieval rank. The core finding: reordering alone completely eliminates the LITM position gap.

Results (MS-MARCO, 60 queries × 3 positions = 180 test cases)

Pipeline	EM (early)	EM (middle)	EM (end)	LITM gap
Baseline (no reordering)	18.3%	13.3%	18.3%	0.050
Lexical (TF-IDF)	16.7%	18.3%	18.3%	−0.008
TIS passage reranker	21.7%	21.7%	21.7%	0.000
Oracle (gold always first)	18.3%	18.3%	18.3%	0.000

TIS eliminates the LITM gap (0.050 → 0.000) and improves overall EM by +5pp over baseline. Evaluation used EM (exact match) on generated answers, not top-5 logits.

Transfer Learning Finding

The existing tis-stage3-ert checkpoint (trained for KV cache compression, not passage reordering) achieves identical results without any retraining:

Checkpoint	LITM gap	EM (all positions)
tis-stage3-ert (KV eviction, no retraining)	0.000	21.7%
tis-passage-reranker (trained for reordering)	0.000	21.7%

This suggests TIS learns generalizable importance patterns that transfer across tasks. If you already have tis-stage3-ert, you may not need this checkpoint.

Architecture

• Base: mistralai/Mistral-7B-v0.3 (frozen, 4-bit NF4)
• Head: QueryAwareImportanceHead wrapping ImportanceUpdateHead
  • Query-aware cross-attention over passage hidden states
  • RMSNorm output stabilization
  • ~8M parameters, ~8.5 MB
• Training: InfoNCE contrastive loss on MS-MARCO passage relevance labels
  • 500 examples, 10 epochs, lr=1e-4, batch size=2
  • Best validation loss: 0.6157

Checkpoint Contents

p3b_production/
├── best.pt           # Model weights (~8.5 MB)
└── training_summary.json

SHA256 of best.pt: 46af05996d923aafa1f4c74c... (verify with sha256sum best.pt)

Key structure of best.pt:

{
  "query_proj.weight": ...,
  "query_proj.bias": ...,
  "context_proj.weight": ...,
  "context_proj.bias": ...,
  # + cross-attention and output projection weights
}

Usage

git clone https://github.com/nitroxido/token-importance-scoring.git
cd token-importance-scoring
python -m venv .venv && source .venv/bin/activate
pip install -e .

# Download checkpoint
hf download oldman-dev/tis-passage-reranker --local-dir checkpoints/passage_reranker

# Run passage reordering evaluation (4-pipeline comparison)
python scripts/run_litm_with_baselines.py \
    --checkpoint checkpoints/passage_reranker \
    --n-examples 60 \
    --seed 42

Reference Environment

transformers==4.36.0   # reference; Transformers 5 has SDPA compat issues with PatchedCausalLM
torch==2.1.2
bitsandbytes==0.41.3
python==3.11
mistral-7b-v0.3 (base model)

Limitations

• Requires local model access (hidden state extraction from Mistral-7B)
• Trained on MS-MARCO passage QA; transfer to other tasks may require retraining
• Beginning-position tokens (causal attention limit) score 0% with any learned scorer; the tis_key_match oracle policy bypasses this by text parsing

Related Checkpoints

• tis-stage3-ert: KV cache compression checkpoint that also eliminates LITM gap via transfer
• tis-v8b-hard-anchor: 82% NIAH evidence survival at 25% cache budget

License

MIT — see repository