File size: 4,225 Bytes
4529d09 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | ---
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)
|