README.md

---
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**:

```python
# 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](https://github.com/nitroxido/token-importance-scoring/blob/main/SOURCE-CODE-README.md)
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

```bash
# 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](https://huggingface.co/oldman-dev/tis-v8b-hard-anchor): Hard-anchor
  training on top of this checkpoint. Better at low budgets (25%).
- [tis-stage1-oracle](https://huggingface.co/oldman-dev/tis-stage1-oracle): Oracle baseline
  using ground-truth labels.

## License

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