README.md

---
license: mit
language:
- en
tags:
- kv-cache
- token-importance
- passage-reranking
- position-bias
- litm
- rag
- mistral
- pytorch
base_model: mistralai/Mistral-7B-v0.3
---

# 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`:
```python
{
  "query_proj.weight": ...,
  "query_proj.bias": ...,
  "context_proj.weight": ...,
  "context_proj.bias": ...,
  # + cross-attention and output projection weights
}
```

## Usage

```bash
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](https://huggingface.co/oldman-dev/tis-stage3-ert): KV cache compression
  checkpoint that also eliminates LITM gap via transfer
- [tis-v8b-hard-anchor](https://huggingface.co/oldman-dev/tis-v8b-hard-anchor): 82% NIAH
  evidence survival at 25% cache budget

## License

MIT — see [repository](https://github.com/nitroxido/token-importance-scoring)