README.md

---
license: mit
base_model: mistralai/Mistral-7B-v0.3
tags:
  - passage-ranking
  - information-retrieval
  - ms-marco
  - query-aware
  - token-importance
language:
  - en
datasets:
  - microsoft/ms_marco
pipeline_tag: text-ranking
---

# TIS v2.2 — Supervised Passage Reranker

**Token Importance Scoring v2.2**: Query-aware passage ranking trained on MS-MARCO relevance labels.

This is the first TIS checkpoint trained with *supervised* relevance labels (MS-MARCO `is_selected`). Earlier checkpoints (Stage3, v8b) used unsupervised ERT objectives. v2.2 resolves the score-direction ambiguity: **high-first (descending) is established by training construction**.

## Performance (500 locked test queries, MS-MARCO v1.1)

| Method | MRR | Recall@1 | Recall@5 | NDCG@5 |
|---|---|---|---|---|
| BM25 (baseline) | 0.432 | 0.205 | 0.532 | — |
| TF-IDF (baseline) | 0.369 | 0.144 | 0.428 | — |
| **TIS v2.2 (this model)** | **0.471** | **0.253** | **0.795** | **0.529** |

**+9.1% MRR over BM25** (0.432 → 0.471, 500 test queries, seed=42).

Release status: **Tier 2 Conditional** — beats BM25, below Tier 1 target (MRR ≥ 0.50). TIS v2.3 in progress.

## Architecture

- **Base model**: Mistral-7B-Instruct-v0.3 (frozen, 4-bit NF4)
- **Importance head**: `QueryAwareImportanceHead` — 4-head cross-attention from passage tokens to mean-pooled query, followed by 3-layer MLP scorer
- **Training**: Pairwise ranking loss (`margin − (score_relevant − score_distractor)`, margin=5.0)
- **Dataset**: MS-MARCO v1.1 passage ranking, 79,704 train queries
- **Steps**: 1000, lr=5e-5, batch=1, gradient accumulation=8
- **Hardware**: RTX 5070 (8 GB VRAM), ~19 min training

## Scoring Contract

```python
# Passage scored with query context (query + passage in same forward pass)
# Token scores aggregated by arithmetic mean (high-first)
# Score space: sigmoid(MLP_output) ∈ [0, 1]
# Direction: descending (high score = more relevant) — established by supervised loss
```

## Checkpoint Structure

```
tis_components.pt:
  importance_head  → QueryAwareImportanceHead state dict
  importance_embedding  → token embedding bias (from base architecture)
  attn_hook_lambda  → attention hook weight
```

## Usage

```python
import torch
from token_importance.model.patched_model import PatchedCausalLM
from token_importance.model.importance_head import QueryAwareImportanceHead

# Load checkpoint
ckpt = torch.load("tis_components.pt", map_location="cpu", weights_only=True)
model.importance_head.load_state_dict(ckpt["importance_head"])

# Score passage given query
def score_passage(model, tokenizer, query, passage, device="cuda"):
    text = f"{query}\n\nPassage: {passage}"
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
    with torch.no_grad():
        out = model._base_model(**inputs, output_hidden_states=True)
        hidden = out.hidden_states[-1]
        # Split query / passage at separator
        sep = inputs["input_ids"][0].tolist().index(28712)  # '\\n\\n' token
        query_h = hidden[:, :sep, :]
        passage_h = hidden[:, sep:, :]
        scores = model.importance_head(doc_hidden=passage_h, query_embeddings=query_h)
    return scores.mean().item()  # mean aggregation, descending = more relevant
```

## Reproduction

```bash
git clone https://github.com/nitroxido/token-importance-scoring.git
cd token-importance-scoring
pip install -e .

# Download this checkpoint
hf download oldman-dev/tis-v2.2-passage-reranker --local-dir checkpoints/v2.2_query_aware_mean

# Evaluate (requires data/msmarco_relevance/test.parquet)
python scripts/evaluate_test_set_v2.2.py \
    --checkpoint checkpoints/v2.2_query_aware_mean/final/tis_components.pt \
    --data-path data/msmarco_relevance/test.parquet
```

Full results: [`results/v2.2_test_final_results.json`](https://github.com/nitroxido/token-importance-scoring/blob/main/results/v2.2_test_final_results.json)

## Checkpoint Identity

| Field | Value |
|---|---|
| SHA-256 (tis_components.pt) | `d26012b28d10b22c5f9c7260b3125ae0c001eb1ef701fef10266fbdc60ea576b` |
| Source commit | [fb04cbc](https://github.com/nitroxido/token-importance-scoring/commit/fb04cbc) |
| Base model | unsloth/mistral-7b-instruct-v0.3-bnb-4bit |
| Training objective | Pairwise ranking (is_selected labels, margin=5.0) |

## Related Checkpoints

| Checkpoint | Task | Notes |
|---|---|---|
| [tis-stage3-ert](https://huggingface.co/oldman-dev/tis-stage3-ert) | KV compression + LITM | ERT trained; context-utility signal |
| [tis-v8b-hard-anchor](https://huggingface.co/oldman-dev/tis-v8b-hard-anchor) | NIAH 82% @ 25% budget | Best KV compression |
| [tis-passage-reranker](https://huggingface.co/oldman-dev/tis-passage-reranker) | LITM elimination | TIS 2.0; LITM gap 0.000 |

## License

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