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:
-
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.
-
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.
-
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.
-
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.
-
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:
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
Environment
The comparison was run on:
| Component |
Value |
| Runtime |
Colab Free |
| GPU |
Tesla T4, 16 GB class |
| Python |
3.12.13 |
| PyTorch |
2.11.0 + CUDA 12.8 |
| Transformers |
5.12.1 |
| bitsandbytes |
0.49.2 |
| NumPy |
2.0.2 |
| Base model |
mistralai/Mistral-7B-v0.3 |
| Quantization |
NF4 4-bit, double quantization |
| Base-model compute |
FP16 |
| Direct learned scorer |
FP32 |
| Test examples |
8 |
| Target context size |
512 tokens |
| Budgets |
25%, 50% |
The repository’s Reproducibility Guide pins an older environment:
torch==2.1.2
torchvision==0.16.2
transformers==4.36.0
peft==0.7.0
bitsandbytes==0.41.3
datasets==2.14.5
accelerate==0.25.0
Therefore, this was not a bitwise reproduction of the documented reference environment.
I avoided downgrading the standard Colab environment because replacing Transformers, tokenizers, Hub, PyTorch, and NumPy packages in-place can introduce a second set of environment failures. Instead, this test asks a narrower question:
When Stage 3 and V8b are evaluated under exactly the same current Colab conditions, do they behave like the checkpoint descriptions suggest?
The answer in this small test was: V8b broadly did; Stage 3 did not.
FP32 scorer adaptation
The public hard evaluator requests BF16 for the 4-bit model. T4 has compute capability 7.5 and does not provide native Ampere-style BF16 acceleration, so I used:
- FP16 for the quantized base model;
- FP32 only for the small
4096 → 1 direct scoring projection.
I also compared the FP32 and FP16 versions of that scorer on identical hidden states.
At 50% budget, the selected top-k sets were identical for all eight examples. At 25%, seven of eight were exactly identical, with mean Jaccard overlap around 0.997.
That does not prove equivalence to the original BF16 environment, but it makes scorer precision alone an unlikely explanation for the large Stage 3 / V8b difference.
Important limitations
This test used:
- 8 examples rather than the published 50 or larger benchmark;
- 512-token target contexts rather than the hard evaluator’s current 2048-token default;
- Transformers 5 rather than the documented Transformers 4.36;
- one Colab T4 runtime;
- the permissive answer metric used by the public evaluator.
The results should therefore be read as an artifact-lineage smoke test, not as replacement benchmark numbers.
Exact artifacts and immutable identifiers
The repository appears to be moving quickly, so these identifiers are more useful than referring only to main.
Source snapshot
| Item |
Identifier |
| Source snapshot SHA256 |
96bfdf9071b2acd9668c3cf424f7ed143822fd0fd023f18c6e6f5b24a93d0b35 |
scripts/eval_niah_hard.py SHA256 |
aa536f908bf5926f3cab5bc8050a382e1d0d5f80fcb9ef20dcbc7d7b5f04cc48 |
The public repository is:
TIS GitHub repository
Stage 3
| Item |
Identifier |
| Hub repository |
oldman-dev/tis-stage3-ert |
| Hub revision |
0907e56257e1ea34c194231fad30c0fb0a0b2ed3 |
| Checkpoint filename |
tis_components.pt |
| Checkpoint SHA256 |
878fac63e61ffc213b937165c2ebfe006bd1271f69e638dfddae0a66490cd86c |
| Checkpoint size |
268,789,935 bytes |
V8b
| Item |
Identifier |
| Hub repository |
oldman-dev/tis-v8b-hard-anchor |
| Hub revision |
e05057d9213bbddf51d16649ffff4b23aeaf2c46 |
| Checkpoint filename |
tis_components.pt |
| Checkpoint SHA256 |
779d2362e16d52f98148b3653d3185a09e7c5eacd01d0ead89b6fea2e06917d5 |
Public test notebook
Checkpoint-lineage and artifact-drift probe
The base-model revision was not explicitly pinned in this run. Adding a base-model revision to an evaluation manifest would close that remaining provenance gap.
Evaluation path used
The comparison follows the public hard-NIAH evaluator’s learned-scoring path as closely as practical.
1. Build the synthetic examples
The examples were produced using the repository’s RetrievalDataset, with the same examples reused for both checkpoints.
The eight prompt-token hashes were all distinct, so this was not accidentally the same generated prompt evaluated eight times.
2. Obtain hidden states
The public evaluator creates neutral seed importance scores and calls the patched model:
seed = torch.full(
(seq_len,),
50,
dtype=torch.uint8,
device=device,
)
output = model(
input_ids=input_ids,
importance_scores=seed,
attention_mask=attention_mask,
output_hidden_states=True,
)
hidden = output.hidden_states[-1].float().squeeze(0)
3. Apply the learned direct scorer
The public evaluator describes this as the training path:
raw = model.importance_head.out_proj(
hidden.unsqueeze(0)
)
scores = torch.sigmoid(
raw.squeeze(-1).squeeze(0)
) * 100.0
This is important because it is not the same computation as calling the complete ImportanceUpdateHead.forward() method.
4. Apply the anchor rule
The public evaluator always preserves:
- the first 4 tokens;
- the last 30 tokens.
The remaining budget is allocated to the highest-scoring content positions.
5. Evaluate with a logical attention mask
The full-length input remains present, but non-kept positions are masked:
output = model(
input_ids=input_ids,
importance_scores=scores_uint8,
attention_mask=keep_mask,
)
6. Use the public answer metric
The evaluator inspects the final-token logits and considers the example correct when any token ID from the encoded answer appears in the top 5 predictions.
That is a permissive retrieval proxy rather than exact generated-answer accuracy, but I retained it so the checkpoint comparison would follow the public evaluator contract.
7. Keep all other conditions fixed
For each generated example, the following were identical across Stage 3 and V8b:
- input token IDs;
- evidence mask;
- base model;
- quantization;
- hidden-state path;
- anchor policy;
- budget;
- answer metric;
- runtime.
Only the loaded TIS checkpoint changed.
Detailed same-example results
Stage 3
| Budget |
Correct |
Accuracy |
Mean evidence survival |
| 25% |
0 / 8 |
0% |
about 10.0% |
| 50% |
1 / 8 |
12.5% |
about 26.7% |
V8b
| Budget |
Correct |
Accuracy |
Mean evidence survival |
| 25% |
6 / 8 |
75% |
100% |
| 50% |
6 / 8 |
75% |
100% |
Published values
The Stage 3 Model Card currently reports:
| Budget |
NIAH |
| 25% |
98% |
| 50% |
100% |
| 75% |
100% |
| 100% |
100% |
The V8b Model Card currently reports:
| Budget |
NIAH |
| 10% |
8% |
| 25% |
82% |
| 50% |
68% |
| 75% |
40% |
The repository’s Project Evolution Report also describes more than one experimental lineage:
- ERT results reported as 100% across NIAH budgets;
- later V8/V8b results around 68–78% at 50%;
- separate changes to hard-anchor and stability-loss behavior.
That history is useful, but it makes an immutable checkpoint-to-result manifest especially important.
Scorer relationship
Across the eight examples, Stage 3 and V8b had:
- mean score correlation of approximately
-0.126;
- 0% exact top-k set matches in the recorded comparisons.
This does not identify which artifact is correct. It does show that the result is not merely a tiny floating-point perturbation near a ranking boundary. The two public checkpoints were applying materially different token rankings.
Stage 3 structure and documentation observations
Several details suggest that the current Stage 3 Model Card and the serialized checkpoint may describe different stages of the implementation.
Model Card description
The Stage 3 Model Card describes:
- an
ImportanceUpdateHead with RMSNorm;
- hard-anchor forcing;
- an importance embedding;
- an attention-hook lambda of
0.1;
- NIAH performance of 98% at 25% and 100% at 50%.
It also calls this “the main checkpoint from the v8b publication.”
The separate V8b Model Card calls V8b itself “the publication checkpoint.”
The overlapping naming alone makes it difficult for a future reader to determine which artifact produced which result.
Serialized Stage 3 checkpoint
The downloaded Stage 3 checkpoint contained:
importance_embedding;importance_head;attn_hook_lambda.
But the observed component state was:
| Observation |
Stage 3 checkpoint |
importance_head state keys |
6 |
score_norm.scale |
absent |
serialized attn_hook_lambda |
0.0 |
| importance embedding projection |
all zero |
learned out_proj |
non-zero |
The current ImportanceUpdateHead defines score_norm, and its normal forward path applies RMSNorm after the projection.
The V8b checkpoint contained 7 importance-head keys, including the norm state. Its Model Card explicitly notes “7 keys versus 6 in Stage 3.”
That means the 6-versus-7 distinction is probably known and intentional at some point in the project history. The mismatch is that the current Stage 3 Model Card simultaneously describes Stage 3 as an RMSNorm checkpoint.
A cautious description would be:
The public Stage 3 checkpoint may be an older compatible checkpoint, but it does not appear to match the architecture description currently displayed on its Model Card.
Lambda configuration versus serialized value
The Stage 3 train_args.json records:
{
"lambda_init_stage2": 0.1,
"load_in_4bit": true,
"max_length": 256,
"max_samples": 128,
"epochs": 1
}
The serialized checkpoint, however, contains:
attn_hook_lambda = 0.0
The Model Card describes the checkpoint lambda as 0.1.
One possible explanation is that the Model Card copied the configured or intended initialization value rather than the value stored in the final uploaded checkpoint.
That is a documentation or serialization-provenance issue rather than necessarily an algorithmic error.
Training log provenance
The Stage 3 repository includes a 128-step training log and arguments consistent with a small 128-sample run.
What it does not currently include is an immutable bridge from that training run to the benchmark table, such as:
- source Git commit;
- final checkpoint SHA256;
- evaluation command;
- evaluator SHA256;
- result CSV;
- benchmark seed;
- benchmark input hashes;
- base-model revision.
Therefore, the public files do not currently make it possible to verify that the uploaded tis_components.pt is the exact artifact that generated the Model Card’s 98% / 100% numbers.
What probably does and does not explain the discrepancy
Missing RMSNorm state is evidence of drift, but not the direct cause here
The missing score_norm.scale is a useful architecture-lineage clue.
However, the public hard-NIAH evaluator computes learned scores using:
model.importance_head.out_proj(hidden)
It does not call the full ImportanceUpdateHead.forward() path and therefore bypasses score_norm.
So it would be incorrect to say:
Stage 3 performed poorly because the RMSNorm key was missing.
A more accurate statement is:
The missing norm state shows that the checkpoint and current architecture description are not from exactly the same implementation state, but it does not directly explain this particular direct-scorer result.
Zero importance embedding and lambda do not explain the Stage 3 / V8b direct-scorer difference
The Stage 3 importance embedding projection was zero, and the serialized lambda was zero.
Those facts matter for the normal runtime wrapper because PatchedCausalLM.forward() incorporates importance through embedding deltas and the merged attention mask.
However, both checkpoints were compared through the direct out_proj(hidden) scoring path. That path does not depend on the importance embedding or attention-hook lambda.
The Stage 3 / V8b difference therefore appears to be concentrated in the learned scorer lineage itself rather than only in the optional embedding, norm, or hook components.
Transformers 5 is a compatibility risk, but probably not the whole explanation
The project’s reference environment is Transformers 4.36, while the test used Transformers 5.12.
Transformers 5 includes broad changes to model loading, tokenization, generation, and other internal contracts, as documented in the official Transformers v5 Migration Guide.
That means exact agreement should not be assumed.
However:
- Stage 3 and V8b were evaluated in the same V5 process;
- V8b behaved in the expected direction;
- FP32 and FP16 direct-scorer rankings were almost identical;
- Stage 3 and V8b rankings differed substantially.
So V5 incompatibility remains a valid secondary issue, but it does not look like a sufficient explanation for the relative checkpoint behavior.
Small sample size could move percentages, but not explain the structural mismatch
Eight examples can easily turn 68% into 75%, or 82% into 75%.
It is much less persuasive as an explanation for all of the following together:
- Stage 3 0% / 12.5%;
- low Stage 3 evidence survival;
- absent RMSNorm checkpoint state despite the current Card description;
- serialized lambda
0.0 versus documented 0.1;
- Model Card commands that do not match current script interfaces;
- materially different Stage 3 and V8b rankings.
Separate implementation issues worth tracking independently
These appear separable from the Stage 3 checkpoint mismatch.
1. Benchmark scorer and runtime update head are different paths
The hard-NIAH evaluator uses:
scores = sigmoid(
importance_head.out_proj(hidden)
)
The normal ImportanceUpdateHead.forward() instead:
- attends from the current hidden state to the context;
- produces one attention summary;
- expands that summary across all
T positions;
- applies
out_proj;
- applies RMSNorm.
In the current implementation:
attn_expanded = attn_out.expand(-1, T, -1)
raw_deltas = self.out_proj(attn_expanded)
normalized_deltas = self.score_norm(raw_deltas)
Because the same summary vector is expanded across the sequence, the normal update path can emit the same delta for every token position.
That may be intentional if it is only scaffolding for a later dynamic update mechanism, but it is not the same learned function evaluated by eval_niah_hard.py.
A short documentation note could distinguish:
- benchmark/direct scorer:
out_proj(hidden_t) for each token;
- runtime update scorer: cross-attention summary expanded across positions;
- query-aware future scorer: per-token query-conditioned path.
2. Logical masking is not physical KV eviction
The public hard-NIAH evaluator passes the full input and a keep mask. In the test, cache sequence length remained equal to the full input length for all logical-mask conditions.
That is consistent with the evaluator’s implementation: it measures which content remains visible through the mask, not whether already-created key/value tensors are physically compacted.
Hugging Face’s cache documentation describes the standard cache as growing with processed sequence length.
A separate physical-shortening test confirmed that cache length decreased when the prefill input itself was shortened. That demonstrates upstream token selection / shortened prefill, not post-prefill cache eviction.
It may help to label future tables explicitly:
| Category |
Meaning |
| Importance scoring |
Assign scores to tokens |
| Logical selection |
Hide non-kept tokens through a mask |
| Shortened prefill |
Remove tokens before cache construction |
| Physical KV compaction |
Remove key/value positions from an existing cache |
This would connect cleanly to the upstream-admission versus downstream-retention distinction already discussed in this thread.
3. “SnapKV” in eval_niah_hard.py is a proxy
The hard evaluator calls one condition snapkv, but its implementation is hidden-state L2 norm multiplied by a recency weight.
The script itself describes this as a proxy for SnapKV-style instruction-region attention scoring rather than the original SnapKV implementation.
It would be clearer in result tables to label it:
SnapKV proxy (hidden-state norm × recency)
That avoids readers interpreting it as a direct run of the official SnapKV code.
4. Transformers 5 generation prefix duplication
In one V5 smoke test, the static wrapper generation path returned a sequence longer than input_length + max_new_tokens.
The current wrapper calls the base generator with both:
input_ids=input_ids
inputs_embeds=inputs_embeds
and then performs:
output_ids = torch.cat(
[input_ids, new_ids],
dim=1,
)
In that V5 run, new_ids already contained the input prefix, so the wrapper prepended it a second time.
This is probably a version-contract issue rather than a TIS research issue.
A small regression test could assert:
assert output_ids.shape[1] <= input_length + max_new_tokens
and normalize the returned sequence based on whether the base generator already included the prefix.
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.
)