README.md

metadata

license: mit
language:
  - en
tags:
  - kv-cache
  - token-importance
  - llm-compression
  - mistral
  - pytorch
base_model: mistralai/Mistral-7B-v0.3

TIS Stage 3 — Closed-Loop Retrieval Checkpoint

v2 — replaces earlier incorrect upload (June 2026 release used wrong artifact)

This is the Token Importance Scoring (TIS) checkpoint for learned KV cache compression, trained via the closed-loop retrieval objective on Mistral-7B-v0.3.

What Changed in v2

The original tis-stage3-ert upload contained a 128-step prototype training run (stage3_ert_local_fresh) rather than the trained checkpoint. That artifact had near-zero out_proj weights (max 0.0007) and produced non-discriminative scores (std ≈ 1.7, range 45–54 for all tokens), which explains the 0%/12.5% NIAH results reported by external testers.

This v2 replaces it with the correct closed_loop_retrieval_v6 checkpoint (2000 steps, out_proj max 0.028, score std ≈ 22, full 0–100 range).

SHA256 of tis_components.pt: 794cb761d8d840709afb0bea6f0f9b73... (run sha256sum tis_components.pt to verify)

Architecture

• Base model: mistralai/Mistral-7B-v0.3
• TIS components: ImportanceUpdateHead with RMSNorm + cross-attention
• State dict keys: 7 (cross_attn.in_proj_weight, cross_attn.in_proj_bias, cross_attn.out_proj.weight, cross_attn.out_proj.bias, out_proj.weight, out_proj.bias, score_norm.scale)
• Attention hook lambda: 0.0 (hook inactive; scoring via out_proj(hidden) direct path)
• Importance embedding projection: initialized to zero

Benchmark Results

Measured with scripts/eval_niah_hard.py on 50 examples, context 2048 tokens, seed 42. These are hard-NIAH numbers (answer token must appear in top-5 final-position logits).

Budget	Learned	SnapKV proxy	Heuristic	No eviction
10%	4%	2%	4%	52%
25%	22%	22%	14%	52%
50%	74%	24%	28%	52%
75%	78%	32%	40%	52%

Evidence survival at 50% budget: 99.9%
Evidence survival at 75% budget: 100%

Note: The no_eviction ceiling is 52% because this hard evaluator tests whether the answer token appears in the top-5 logits of the final token position. Eviction can actually improve accuracy by removing distractors, which is the mechanism being measured here.

Scorer Path

The eval_niah_hard.py evaluator uses the direct token scorer:

# Direct scorer — what eval_niah_hard.py uses
scores = sigmoid(importance_head.out_proj(hidden)) * 100.0

This is different from the full ImportanceUpdateHead.forward() which uses cross-attention and RMSNorm. The direct scorer applies out_proj token-by-token to final-layer hidden states. See Source Code README for details on all scorer paths.

Training Details

• Training type: closed-loop retrieval
• Steps: 2000
• Learning rate: 0.001
• Loss weights: α_rank=1.0, β_retrieve=2.0, γ_stability=0.05
• Final loss: ~0.79
• Final evidence survival (training metric): ~70%

Usage

# Clone repository
git clone https://github.com/nitroxido/token-importance-scoring.git
cd token-importance-scoring

# Setup
python -m venv .venv && source .venv/bin/activate
pip install -e .

# Download this checkpoint
hf download oldman-dev/tis-stage3-ert --local-dir checkpoints/stage3_ert_learned

# Run hard NIAH evaluation
python scripts/eval_niah_hard.py \
  --learned-checkpoint checkpoints/stage3_ert_learned \
  --budgets 0.25 0.5 0.75 \
  --num-tests 50 \
  --context-tokens 2048 \
  --device cuda \
  --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

Transformers 5 breaks the PatchedCausalLM forward path (attention mask injection conflicts with the SDPA backend). The direct scorer (out_proj(hidden)) works under any version because it does not go through the patched attention path.

Related Checkpoints

• tis-v8b-hard-anchor: Hard-anchor training on top of this checkpoint. Better at low budgets (25%).
• tis-stage1-oracle: Oracle baseline using ground-truth labels.

License

MIT — see repository