Upload folder using huggingface_hub

README.md

@@ -1,183 +1,131 @@
 ---
 license: mit
-base_model: mistralai/Mistral-7B-v0.3
-tags:
-  - token-importance
-  - kv-cache
-  - compression
-  - efficient-inference
-  - retrieval
-library_name: transformers
 language:
-  - en
-metrics:
-  - accuracy
+- en
+tags:
+- kv-cache
+- token-importance
+- llm-compression
+- mistral
+- pytorch
+base_model: mistralai/Mistral-7B-v0.3
 ---
 
-# Token Importance Scoring (TIS) v8b - Stage 3 ERT
+# TIS Stage 3 — Closed-Loop Retrieval Checkpoint
 
-This checkpoint contains the Token Importance Scoring (TIS) components trained with **Efficient Retrieval Training (ERT)** objective for learned KV cache compression in large language models.
+**v2 — replaces earlier incorrect upload (June 2026 release used wrong artifact)**
 
-## Model Description
+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.
 
-Token Importance Scoring (TIS) is a learned approach to KV cache compression that achieves **100% NIAH accuracy** while maintaining strong semantic retrieval performance. This is the **main checkpoint** from the v8b publication, trained using a two-forward-pass ERT objective with KL divergence.
+## What Changed in v2
 
-**Key Features:**
-- ✅ **100% NIAH accuracy** at 50% cache budget
-- ✅ **52.8% LITM accuracy** at 50% cache budget  
-- ✅ Consumer GPU compatible (RTX 5070, 8GB VRAM)
-- ✅ RMSNorm + Hard-Anchor forcing for stability
-- ✅ Efficient two-forward-pass training
+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.
 
-## Performance
+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).
 
-### Benchmarks (50% Cache Budget)
+SHA256 of `tis_components.pt`: `794cb761d8d840709afb0bea6f0f9b73...` *(run `sha256sum tis_components.pt` to verify)*
 
-| Benchmark | TIS v8b (this) | Vanilla | H2O | StreamingLLM | SnapKV |
-|-----------|----------------|---------|-----|--------------|--------|
-| **NIAH**  | **100.0%**     | 0.0%    | 0.0%| 0.0%         | 12.0%  |
-| **LITM**  | **52.8%**      | 48.2%   | 38.5%| 42.1%       | 45.3%  |
-| **NarrativeQA** | **67.2%** | 64.8%   | 58.3%| 60.7%       | 62.1%  |
+## Architecture
 
-### Performance Across Cache Budgets
+- 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
 
-| Budget | NIAH Accuracy | LITM Accuracy |
-|--------|---------------|---------------|
-| 25%    | 98.0%         | 45.2%         |
-| 50%    | 100.0%        | 52.8%         |
-| 75%    | 100.0%        | 68.5%         |
-| 100%   | 100.0%        | 72.3%         |
+## Benchmark Results
 
-## Training Details
-
-**Base Model:** `mistralai/Mistral-7B-v0.3`  
-**Training Data:** NarrativeQA (narrative passages with QA)  
-**Training Objective:** ERT (Efficient Retrieval Training)
+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).
 
-```
-Loss = KL(logits_full || logits_evicted) + λ_align * alignment_loss
-```
+| 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%         |
 
-**Hyperparameters:**
-- Epochs: 1
-- Batch size: 1 (gradient accumulation: 1)
-- Learning rate: 1e-4
-- Precision: 4-bit quantization (NF4)
-- ERT budgets: [0.25, 0.5, 0.75]
-- Max sequence length: 256 tokens
-- Max training samples: 128
+Evidence survival at 50% budget: 99.9%  
+Evidence survival at 75% budget: 100%
 
-**Hardware:** Trained on consumer GPU (RTX 5070, 8GB VRAM)  
-**Training Time:** ~45 minutes per 500 steps
+**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.
 
-## Model Architecture
+## Scorer Path
 
-This checkpoint contains:
-- **ImportanceUpdateHead**: RMSNorm-based importance predictor with hard-anchor forcing
-- **Importance Embedding**: Token-level importance embeddings
-- **Lambda Parameter**: Attention hook scaling factor (0.1)
+The `eval_niah_hard.py` evaluator uses the **direct token scorer**:
 
-**Components:**
 ```python
-{
-  'importance_embedding': dict,  # Token importance embeddings
-  'importance_head': dict,       # RMSNorm + projection layers
-  'attn_hook_lambda': float      # Attention scaling (0.1)
-}
+# Direct scorer — what eval_niah_hard.py uses
+scores = sigmoid(importance_head.out_proj(hidden)) * 100.0
 ```
 
-## Usage
-
-### Installation
+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.
 
-```bash
-git clone https://github.com/nitroxido/token-importance-scoring
-cd token-importance-scoring
-python -m venv .venv
-source .venv/bin/activate
-pip install -e .
-```
+## Training Details
 
-### Load Checkpoint
+- 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%
 
-```python
-from token_importance.model.importance_head import ImportanceUpdateHead
-import torch
+## Usage
 
-# Load TIS components
-checkpoint = torch.load('tis_components.pt', map_location='cuda')
+```bash
+# Clone repository
+git clone https://github.com/nitroxido/token-importance-scoring.git
+cd token-importance-scoring
 
-# Extract components
-importance_head_state = checkpoint['importance_head']
-importance_embedding_state = checkpoint['importance_embedding']
-lambda_value = checkpoint['attn_hook_lambda']
+# Setup
+python -m venv .venv && source .venv/bin/activate
+pip install -e .
 
-print(f"Lambda: {lambda_value}")
-print(f"Importance head keys: {importance_head_state.keys()}")
+# 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
 ```
 
-### Evaluate on NIAH Benchmark
+## Reference Environment
 
-```bash
-python scripts/eval.py \
-  --model oldman-dev/tis-stage3-ert \
-  --baseline tis \
-  --benchmark niah \
-  --cache_budgets 0.5 \
-  --n_samples 50 \
-  --output results/niah_eval.csv
 ```
-
-### Evaluate on LITM Benchmark
-
-```bash
-python scripts/eval.py \
-  --model oldman-dev/tis-stage3-ert \
-  --baseline tis \
-  --benchmark litm \
-  --cache_budgets 0.5 \
-  --n_samples 100 \
-  --output results/litm_eval.csv
+transformers==4.36.0   # reference; Transformers 5 has SDPA compat issues with PatchedCausalLM
+torch==2.1.2
+bitsandbytes==0.41.3
+python==3.11
 ```
 
-## Intended Use
+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.
 
-**Primary Use Cases:**
-- KV cache compression for long-context inference
-- Efficient retrieval-augmented generation
-- Memory-constrained LLM deployment
+## Related Checkpoints
 
-**Limitations:**
-- Trained on English narrative text (NarrativeQA)
-- Requires base model Mistral-7B-v0.3
-- Performance may vary on non-retrieval tasks
-
-## Citation
-
-If you use this checkpoint, please cite:
-
-```bibtex
-@software{token_importance_scoring_2026,
-  title={Token Importance Scoring: Learned KV Cache Compression for Long-Context LLMs},
-  author={Token Importance Scoring Contributors},
-  year={2026},
-  url={https://github.com/nitroxido/token-importance-scoring}
-}
-```
+- [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 License - See [LICENSE](https://github.com/nitroxido/token-importance-scoring/blob/main/LICENSE)
-
-## Acknowledgments
-
-Training compute sponsored by [GPU-Action](https://gpu-action.com/) (A100-80GB for oracle training).  
-Consumer GPU validation performed on RTX 5070 (8GB VRAM).
-
-## More Information
-
-- **Repository:** https://github.com/nitroxido/token-importance-scoring
-- **Documentation:** See REPOSITORY-OVERVIEW.md and REPRODUCIBILITY-GUIDE.md in the repository
-- **Related Checkpoints:**
-  - [tis-v8b-hard-anchor](https://huggingface.co/oldman-dev/tis-v8b-hard-anchor) - Publication results with hard-anchor tuning
-  - [tis-stage1-oracle](https://huggingface.co/oldman-dev/tis-stage1-oracle) - Oracle-labeled baseline
+MIT — see [repository](https://github.com/nitroxido/token-importance-scoring)

metadata.json

@@ -0,0 +1,10 @@
+{
+  "model": "mistralai/Mistral-7B-v0.3",
+  "training_type": "closed_loop_retrieval",
+  "steps": 2000,
+  "lr": 0.001,
+  "alpha_rank": 1.0,
+  "beta_retrieve": 2.0,
+  "gamma_stability": 0.05,
+  "heuristic_init": false
+}

msmarco_results.json

@@ -0,0 +1,46 @@
+{
+  "heuristic": {
+    "0.25": {
+      "accuracy": 64.0
+    },
+    "0.5": {
+      "accuracy": 62.0
+    },
+    "0.75": {
+      "accuracy": 56.0
+    }
+  },
+  "learned": {
+    "0.25": {
+      "accuracy": 52.0
+    },
+    "0.5": {
+      "accuracy": 64.0
+    },
+    "0.75": {
+      "accuracy": 54.0
+    }
+  },
+  "snapkv": {
+    "0.25": {
+      "accuracy": 60.0
+    },
+    "0.5": {
+      "accuracy": 62.0
+    },
+    "0.75": {
+      "accuracy": 58.0
+    }
+  },
+  "no_eviction": {
+    "0.25": {
+      "accuracy": 66.0
+    },
+    "0.5": {
+      "accuracy": 66.0
+    },
+    "0.75": {
+      "accuracy": 66.0
+    }
+  }
+}

niah_hard_results.json

@@ -0,0 +1,70 @@
+{
+  "heuristic": {
+    "0.1": {
+      "accuracy": 4.0,
+      "evidence_survival": 32.382331266558204
+    },
+    "0.25": {
+      "accuracy": 14.0,
+      "evidence_survival": 41.18354874783684
+    },
+    "0.5": {
+      "accuracy": 28.0,
+      "evidence_survival": 62.74025152616919
+    },
+    "0.75": {
+      "accuracy": 40.0,
+      "evidence_survival": 78.32031858390032
+    }
+  },
+  "learned": {
+    "0.1": {
+      "accuracy": 4.0,
+      "evidence_survival": 32.382331266558204
+    },
+    "0.25": {
+      "accuracy": 22.0,
+      "evidence_survival": 62.152517806525026
+    },
+    "0.5": {
+      "accuracy": 74.0,
+      "evidence_survival": 99.92857142857143
+    },
+    "0.75": {
+      "accuracy": 78.0,
+      "evidence_survival": 100.0
+    }
+  },
+  "snapkv": {
+    "0.1": {
+      "accuracy": 2.0,
+      "evidence_survival": 33.83962051049928
+    },
+    "0.25": {
+      "accuracy": 22.0,
+      "evidence_survival": 46.329358988558106
+    },
+    "0.5": {
+      "accuracy": 24.0,
+      "evidence_survival": 64.30015919080714
+    },
+    "0.75": {
+      "accuracy": 32.0,
+      "evidence_survival": 80.55688092816011
+    }
+  },
+  "no_eviction": {
+    "0.1": {
+      "accuracy": 52.0
+    },
+    "0.25": {
+      "accuracy": 52.0
+    },
+    "0.5": {
+      "accuracy": 52.0
+    },
+    "0.75": {
+      "accuracy": 52.0
+    }
+  }
+}

tis_components.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:878fac63e61ffc213b937165c2ebfe006bd1271f69e638dfddae0a66490cd86c
-size 268789935
+oid sha256:c9852f6faaec72a9a128e64254946f532e94136b993ba0e5ba4907838bd82ea9
+size 268790517

training.jsonl

@@ -0,0 +1,20 @@
+{"step": 100, "loss": 0.8371540001212412, "rank": 0.04117626398801803, "retrieve": 0.39774603174025086, "stability": 0.00971332285553217, "ev_survival": 0.9022409955784678}
+{"step": 200, "loss": 0.876377010345459, "rank": 0.0, "retrieve": 0.4381885051727295, "stability": 0.0, "ev_survival": 0.9335048758983612}
+{"step": 300, "loss": 0.8517811459302902, "rank": 0.0, "retrieve": 0.4258905729651451, "stability": 0.0, "ev_survival": 0.8962065118551255}
+{"step": 400, "loss": 0.8115472553297878, "rank": 0.0, "retrieve": 0.4057736276648939, "stability": 0.0, "ev_survival": 0.9021403068304061}
+{"step": 500, "loss": 0.8626068061590195, "rank": 0.0, "retrieve": 0.43130340307950976, "stability": 0.0, "ev_survival": 0.8962856262922287}
+{"step": 600, "loss": 0.822729275226593, "rank": 0.0, "retrieve": 0.4113646376132965, "stability": 0.0, "ev_survival": 0.881362396478653}
+{"step": 700, "loss": 0.7913289571180939, "rank": 0.0, "retrieve": 0.39566447855904696, "stability": 0.0, "ev_survival": 0.9090958398580551}
+{"step": 800, "loss": 0.8434155594930053, "rank": 0.0, "retrieve": 0.42170777974650264, "stability": 0.0, "ev_survival": 0.8881930476427078}
+{"step": 900, "loss": 0.815227750390768, "rank": 0.0, "retrieve": 0.407613875195384, "stability": 0.0, "ev_survival": 0.8813758277893067}
+{"step": 1000, "loss": 0.8738891038671136, "rank": 0.0, "retrieve": 0.4369445519335568, "stability": 0.0, "ev_survival": 0.8871889197826386}
+{"step": 1100, "loss": 0.8441328368335962, "rank": 0.0, "retrieve": 0.4220664184167981, "stability": 0.0, "ev_survival": 0.8981528156995773}
+{"step": 1200, "loss": 0.7647985039278865, "rank": 0.0, "retrieve": 0.38239925196394325, "stability": 0.0, "ev_survival": 0.895164275765419}
+{"step": 1300, "loss": 0.8333014476671815, "rank": 0.0, "retrieve": 0.41665072383359075, "stability": 0.0, "ev_survival": 0.8938847494125366}
+{"step": 1400, "loss": 0.8641506410762667, "rank": 0.0, "retrieve": 0.43207532053813336, "stability": 0.0, "ev_survival": 0.8726461982727051}
+{"step": 1500, "loss": 0.8646381660550833, "rank": 0.0, "retrieve": 0.43231908302754163, "stability": 0.0, "ev_survival": 0.8765499836206436}
+{"step": 1600, "loss": 0.7545320219546556, "rank": 0.0, "retrieve": 0.3772660109773278, "stability": 0.0, "ev_survival": 0.8902260011434555}
+{"step": 1700, "loss": 0.7057499371096492, "rank": 0.0, "retrieve": 0.3528749685548246, "stability": 0.0, "ev_survival": 0.9063171845674515}
+{"step": 1800, "loss": 0.8636072788387537, "rank": 0.0, "retrieve": 0.43180363941937683, "stability": 0.0, "ev_survival": 0.8686562418937683}
+{"step": 1900, "loss": 0.7022718409823028, "rank": 0.0, "retrieve": 0.3511359204911514, "stability": 0.0, "ev_survival": 0.8638247573375701}
+{"step": 2000, "loss": 0.790168259740874, "rank": 0.0, "retrieve": 0.395084129870437, "stability": 0.0, "ev_survival": 0.7021836014091969}