Presenting TIS (Token Importance Scoring) - A new way to compress KV cache

Hi. For now, I gave this a lightweight test run on Colab Free. It is not a full reproduction, but I may have run into a small bug or some artifact drift around the Stage 3 release:

---

Short version

I was able to run the public code, load Mistral-7B-v0.3 in 4-bit, load both public TIS component checkpoints, obtain hidden states, run the direct learned scorer used by the public hard-NIAH evaluator, and complete the anchored-mask evaluation path.

The main observation is:

• tis-v8b-hard-anchor behaved broadly in the direction described by its Model Card in a small same-example smoke test.
• tis-stage3-ert did not.
• The downloaded Stage 3 checkpoint also does not appear to match several structural details currently described by its Model Card.

This looks more like a checkpoint / Model Card / evaluator lineage mismatch than evidence that the underlying TIS idea is invalid.

The most likely resolution seems to be one of these branches:

1. The reported Stage 3 result used another local checkpoint. In that case, uploading that checkpoint and publishing its SHA256 would probably resolve most of the discrepancy.

2. The currently published Stage 3 checkpoint is intentional, but its Model Card is stale. In that case, the architecture description, lambda value, evaluation command, and benchmark table may need updating.

3. The checkpoint is correct, but the reported result used a different source revision or evaluator. Pinning the exact Git commit, command, seed, environment, and result CSV would reconnect the artifacts.

4. Transformers 4.36 is the only reference implementation. In that case, it would help to label that environment as the canonical reproduction path and treat Transformers 5 support as a separate compatibility track.

5. Stage 3 and V8b intentionally use different scoring paths. In that case, naming those paths explicitly in each Model Card would prevent future readers from combining incompatible components.

What worked

The public artifacts were not generally unusable. The following paths worked on a Colab Free T4:

• repository import and static compilation;
• Mistral-7B-v0.3 4-bit loading;
• Stage 3 and V8b component checkpoint loading;
• hidden-state extraction;
• direct out_proj(hidden) scoring;
• anchored top-k mask construction;
• final-token top-5 evaluation;
• same-example Stage 3 / V8b comparison.

So this is not a report that “nothing works.” The narrower problem is that the public Stage 3 artifact combination did not line up with its current documentation or reported behavior in this test.

Same-example result summary

This was a deliberately small test: eight distinct synthetic retrieval examples, the same examples for both checkpoints, with 25% and 50% budgets.

Public checkpoint	25% budget	50% budget	Evidence survival at 25% / 50%
Stage 3 ERT	0 / 8	1 / 8	about 10.0% / 26.7%
V8b Hard Anchor	6 / 8	6 / 8	100% / 100%

For comparison, the current Model Cards report:

Public checkpoint	Published 25%	Published 50%
Stage 3 ERT	98%	100%
V8b Hard Anchor	82%	68%

Eight samples are nowhere near enough to re-estimate the published benchmark. The useful signal here is comparative:

• V8b produced 75% / 75% and perfect evidence survival on this small run, which is at least directionally compatible with its published result.
• Stage 3 produced 0% / 12.5% on the same examples and runtime.
• The two learned scorers did not behave like small numerical variations of the same policy: their mean score correlation was approximately -0.126, and their selected top-k sets were never exactly identical in this run.

That makes a simple “Colab or Transformers 5 broke the entire evaluator” explanation less likely, because V8b worked under the same conditions.

The runnable comparison notebook is here:

Stage 3 / V8b checkpoint-lineage probe

A documentation-level mismatch that may be easy to fix

The evaluation commands currently shown in the Model Cards do not appear to match the current script interfaces.

For example, the Stage 3 Model Card currently suggests passing the component repository itself as --model:

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

However, the current scripts/eval.py loads:

• the base causal LM from args.model; and
• tis_components.pt from the directory passed through args.checkpoint.

The Stage 3 Hub repository contains TIS components rather than a complete standalone Mistral model, so the current code path appears to expect the Mistral base model and checkpoint directory separately.

Similarly, the V8b Model Card shows arguments such as --model, --baseline, --cache_budgets, and --n_samples, while the current eval_niah_hard.py exposes:

--learned-checkpoint
--budgets
--num-tests
--context-tokens
--device
--seed

This may simply be documentation drift after the evaluator was reorganized, but fixing the commands would make first-time reproduction much easier.

A current hard-NIAH invocation appears closer to:

python scripts/eval_niah_hard.py \
  --learned-checkpoint /path/to/checkpoint_directory \
  --budgets 0.25 0.5 0.75 \
  --num-tests 50 \
  --context-tokens 2048 \
  --device cuda \
  --seed 42

Test environment and scope (click for more details) Exact artifacts and immutable identifiers (click for more details) Evaluation path used (click for more details) Detailed same-example results (click for more details) Stage 3 structure and documentation observations (click for more details) What probably does and does not explain the discrepancy (click for more details) Separate implementation issues worth tracking independently (click for more details)

Small fixes that may resolve most of this

1. Add an evaluation manifest to every released checkpoint

For example:

{
  "artifact_schema": 1,
  "checkpoint_name": "tis-stage3-ert",
  "checkpoint_sha256": "...",
  "base_model": "mistralai/Mistral-7B-v0.3",
  "base_model_revision": "...",
  "source_commit": "...",
  "evaluation_script": "scripts/eval_niah_hard.py",
  "evaluation_script_sha256": "...",
  "evaluation_command": "...",
  "python": "...",
  "torch": "...",
  "transformers": "...",
  "bitsandbytes": "...",
  "quantization": "NF4 4-bit",
  "compute_dtype": "bfloat16",
  "seed": 42,
  "context_tokens": 2048,
  "num_samples": 50,
  "budgets": [0.25, 0.5, 0.75],
  "result_file": "niah_hard_eval.csv"
}

This would make it immediately clear which checkpoint, source, and runtime produced each table.

2. Add a checkpoint contract test

Before evaluation, print or validate:

checkpoint SHA256
top-level state keys
importance-head state keys
expected hidden size
score_norm state present/absent
attention-hook lambda present/value
importance-embedding non-zero count
out_proj non-zero count
architecture version

A mismatch could fail with a message such as:

This checkpoint was trained for architecture schema vN,
but the current runtime implements schema vM.

3. Give the scorer paths explicit names

For example:

direct_token_scorer
runtime_cross_attention_update
query_aware_scorer
oracle_labels

Then record the selected path in the manifest and result CSV.

4. Make one checkpoint the canonical reproduction artifact

There are several reasonable options:

• replace the public Stage 3 checkpoint if another local artifact produced the reported 100%;
• update the Stage 3 table if the current file is intentional;
• temporarily recommend V8b as the reproducible hard-NIAH checkpoint;
• publish both historical and corrected Stage 3 artifacts under distinct names.

Keeping the older artifact is useful if it is renamed rather than overwritten, for example:

tis-stage3-ert-legacy
tis-stage3-ert-publication

5. Attach the actual result CSV

A result CSV with per-example fields would help future debugging:

example_id
input_hash
seed
budget
condition
answer
evidence_positions
selected_positions
evidence_survival
top5_token_ids
correct

6. Add two small regression suites

Reference suite

• pinned Transformers 4.36 environment;
• fixed checkpoint;
• fixed 8–16 examples;
• expected selected positions or minimum evidence survival;
• expected aggregate range.

Compatibility suite

• current supported Transformers release;
• model/checkpoint load;
• direct-score shape;
• finite scores;
• generation output length;
• cache and mask shape;
• no duplicated prefix.

The compatibility suite does not need to reproduce benchmark accuracy. Its job is only to detect API drift.

What I am not concluding

To keep the scope clear, this test does not establish that:

• the TIS idea is invalid;
• the published 50-sample or larger benchmark is disproven;
• V8b’s 75% result here is a new benchmark number;
• FP32 scoring is fully equivalent to the original BF16 environment;
• Transformers 5 is an officially supported reference runtime;
• logical attention masking is the same as physical KV compaction;
• the correct Stage 3 checkpoint no longer exists;
• V8b’s synthetic evidence survival proves general relevance detection without task-specific evidence supervision.

The narrower conclusion is:

The public V8b checkpoint behaved broadly consistently with its description in a small same-example smoke test, while the current public Stage 3 checkpoint did not. The Stage 3 checkpoint, Model Card, and current evaluation instructions also contain several observable structural or interface mismatches. An artifact-lineage or release-packaging issue therefore seems more likely than a failure of the underlying learned-scoring idea.

Practical resolution paths

Any one of these would make the public reproduction path much easier to follow:

• Correct checkpoint available: upload it and publish its SHA256.
• Current checkpoint intentional: update the Stage 3 architecture and benchmark description.
• Different evaluator used: publish the exact commit, command, and result CSV.
• V4-only reference: mark Transformers 4.36 as canonical and track V5 separately.
• V8b currently canonical: say so directly and treat Stage 3 as historical.
• Multiple scorer generations intentional: give each architecture and checkpoint an explicit version/schema name.

The core code and V8b behavior suggest that this is probably fixable by reconnecting the right checkpoint, evaluator, and documentation rather than redesigning the whole method.