tis-v2.2-passage-reranker

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)

+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

# 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

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

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

Checkpoint Identity

Related Checkpoints

License

MIT — see repository.