TIS 2.0: Token Importance Scoring Now Eliminates Position Bias in RAG

Research
21

Umm… this has gotten a lot better, but I still have this vague feeling that something may or may not be missing here​:sweat_smile::

I spent some more time trying to reproduce the v2.3 result from the public artifacts.

The short version is: I do not think model quality itself is a reason to hold up closing TIS 2.0. The released v2.3 checkpoint looks healthy in an independent rerun. What I still cannot quite close is the exact evaluator/provenance bridge between that checkpoint and the published headline result.

With the dataset, tokenizer/base model, checkpoint, prompt, separator logic, library versions, and scoring path pinned, my clean public-contract rerun lands around MRR 0.55. The most source-faithful run I have now — BF16 head plus position-corrected batched Mistral inference — gave approximately:

Evaluation MRR Hit@1 Hit@3 Hit@5
public-contract rerun 0.5546 169/483 323/483 416/483
published v2.3 result 0.5102 146/483 300/483 390/483

So this does not look to me like “the published checkpoint falls apart when somebody else runs it.” If anything, I am getting a stronger ranking result from the public checkpoint.

The remaining question is why the published evaluator produces a materially different result.

What seems to be the missing link

The published result/log describes an optimized, batch-tokenized evaluation path and reports:

3935 separator hits + 37 midpoint fallbacks

The release material also refers to an optimized evaluator / batch-scoring implementation. But the reproducible paths I can currently connect do not seem to be the same thing:

README rerun path ≠ committed serial v2.3 evaluator ≠ published optimized-result path

I can find the serial v2.3 evaluator, and I can follow the README route, but I have not been able to connect either one exactly to the optimized evaluator that generated the headline 0.5102.

If that exact optimized evaluator exists somewhere I simply missed, then that is probably the highest-value artifact to point at. Otherwise, committing it — or regenerating the headline table with the evaluator that is actually committed — would close a large part of this by itself.

I tried to rule out several simpler explanations before landing there.

The 37 separator fallbacks are not enough by themselves

Using the pinned tokenizer/base and the exact 500-row MS MARCO slice, I get 3972/3972 separator hits, including through the normal batched-tokenizer route.

I then went further and calculated, for every passage, both:

  • the score using the correct Passage: boundary;
  • the score obtained if that passage is deliberately forced through the midpoint fallback.

That makes it possible to solve the counterfactual problem exactly: choose exactly 37 passages to fall back and see whether the published metric vector can be reproduced.

It cannot.

Even choosing the 37 fallbacks that hurt MRR the most only gets down to about:

MRR 0.5163

rather than 0.5102.

More importantly, I cannot simultaneously reproduce the published Hit counts 146 / 300 / 390 with any set of exactly 37 midpoint fallbacks. Even allowing very favorable tie ordering does not make the full metric vector line up.

So I think the safe conclusion is:

the 37 midpoint fallbacks may be part of the difference, but they cannot be the whole difference.

There was a real batched-position issue, but fixing it does not explain 0.5102 either

One useful thing did turn up while checking the batching path.

For a left-padded Mistral batch, directly calling the model with the attention mask but without padding-corrected position_ids gave hidden states that differed from serial inference.

Using position IDs derived from the attention mask:

position_ids = attention_mask.long().cumsum(-1) - 1

and masking the padding positions fixes that. This is also consistent with the way Transformers constructs position IDs for generation.

After that correction, my smoke cases became exactly serial-equivalent:

  • hidden-state max difference: 0
  • BF16 passage-score difference: 0

That is the evaluation route behind the 0.5546 result above.

So the padding/position issue is real and worth making explicit in an optimized evaluator, but it still does not explain the published 0.5102.

BF16 vs FP16 is also too small to explain it

I checked that direction as well.

The pinned prequantized Unsloth Mistral config specifies NF4 4-bit quantization with BF16 compute, and the committed v2.3 importance head is also naturally evaluated in BF16.

Switching between the relevant BF16/FP16 head variants moves the ranking somewhat, but only by a few thousandths of MRR — nowhere near the roughly 0.044 gap to the published result.

So at this point I do not see tokenizer batching, dtype, position IDs, or the 37 fallbacks individually or together as a sufficient explanation. Some additional difference in the optimized scoring/evaluation path still seems necessary.

There is a second, more mechanical cleanup I think would be useful before calling the evaluation side completely finished: regenerate the metric/comparison table from one evaluator contract.

Metric and comparison-table cleanup

Recall@k looks more like Hit@k / Success@k

The MS MARCO v1.1 slice actually contains queries with multiple selected passages.

In the 483 evaluable queries here, 46 have more than one positive passage.

The current metric called Recall@k effectively asks whether at least one designated relevant passage occurs in the top-k. That is a useful metric, but it is normally easier to interpret as Hit@k or Success@k.

I would probably report both:

  • Hit@k: did any relevant passage appear in the top-k?
  • Recall@k: what fraction of all relevant passages appeared in the top-k?

That avoids a small terminology trap for future readers.

NDCG deserves one clean recomputation

The current NDCG path also appears not to implement standard binary NDCG: the relevance/DCG side and ideal denominator do not use the actual set of positive passages in the usual way.

Using the position-corrected public ranking and standard all-positive binary NDCG, I get roughly:

  • NDCG@5 ≈ 0.606
  • NDCG@10 ≈ 0.659

So I would not interpret the very low published NDCG values as evidence that v2.3 ranking quality collapses. They look much more like a metric-definition issue.

I would regenerate the whole v2.2 / v2.3 / BM25 row together

There also seem to be multiple v2.2 evaluation families in the published artifacts.

The +25.6% figure makes sense for approximately:

0.4063 → 0.5102

but not for:

0.471 → 0.5102

which is only around an 8% increase.

That is not necessarily a problem if 0.4063 is the intentionally chosen common-evaluator baseline — in fact using a common evaluator is preferable — but then I would make that lineage explicit in the table.

BM25 MRR is comparatively easy to close: with the same MS MARCO slice and the apparent rank_bm25 contract I can independently reproduce approximately:

BM25 MRR = 0.4315

The secondary BM25 columns are less clean, though, so again I think rebuilding the entire comparison table from one evaluator is safer than correcting cells individually.

A small generated table carrying an evaluator/run ID would make this almost self-auditing:

system evaluator MRR Hit@1 Hit@3 Hit@5 NDCG@5 NDCG@10
BM25 same frozen evaluator … … … … … …
TIS v2.2 same frozen evaluator … … … … … …
TIS v2.3 same frozen evaluator … … … … … …

Then all percentage improvements can be generated from that same table rather than maintained separately.

One other wording point: I would probably describe the 0.5137 validation → 0.5102 test relationship as encouraging, rather than lean too hard on the very small “generalization gap.”

As I understand the training path, the best checkpoint was selected from repeated checks on the same fixed 50-query validation subset. That is useful for checkpoint selection, but it means 0.5137 is itself a selected value rather than an untouched estimate.

Likewise, the 500-query set is already useful as a stable regression benchmark across the TIS versions, which is valuable, but that makes it somewhat different from a completely fresh holdout opened once at the end.

This does not require retraining anything. The low-cost options seem to be either:

  • just soften the generalization wording; or
  • evaluate once on a previously unused fixed slice if a stronger generalization claim is useful.

What I would do before closing TIS 2.0

I do not think another major training cycle is necessary.

The shortest path I can see is something like:

  1. Make one evaluator authoritative.
    Either commit the exact optimized evaluator that generated 0.5102, or rerun the release numbers with the committed evaluator.

  2. Have the README call that exact evaluator.

  3. Regenerate BM25 / v2.2 / v2.3 from the same metric implementation.

  4. Rename or split Hit@k vs literal Recall@k, and use standard NDCG.

  5. When separator fallback happens, log the qid/passage index and detected boundary, not only the total fallback count.

  6. Emit a tiny run manifest.
    Something like source commit, checkpoint revision/hash, base-model revision, tokenizer revision, dataset revision/slice, Python/Transformers/tokenizers/bitsandbytes versions, and the effective quantization configuration.

That would turn most of the remaining uncertainty into things a future reader can check in a few minutes.

So my current impression is:

I would not block closing TIS 2.0 on the model. The v2.3 checkpoint itself looks quite good in an independently pinned evaluation. The remaining loose end is mostly the exact evaluation/provenance path between that checkpoint and the published headline table.

If the missing optimized evaluator explains that path, I think most of the remaining puzzle snaps into place — and then moving on to TIS 3.0 / TIS+PIN seems much cleaner.

22

As always, thanks for the detailed feedback! I will need to stop relying so much in the “semi-automatic” packager and try to make the packaging of the public versions myself or micro-manage them!

As soon as I have some time, I will start preparing the correct upload, because the results are correct, but maybe the attribution is wrong! I used the `–seed 42` parameter during training / testing so it kept the same system, and fixed the use of a different model for the training and testing steps.

23

Well! I have uploaded to GitHub the missing parts and documentation! I hope everything is good now, but I liked your MRR results, higher than mine during my testing/training!

24

Thanks. I think I may have reproduced it (maybe):

The evaluator you added in #23 helped a lot. I froze the newly published f3ed5d5 scoring path together with the public data/model/tokenizer/checkpoint revisions, then ran the full 483 usable-query slice.

I get this:

Evaluation path MRR Hit@1 Hit@3 Hit@5
Historical artifact 0.5101523218 146 300 390
A — published f3ed5d5 canonical semantics 0.5096108975 146 300 389
B — same base hidden states, but remove left padding before query pooling 0.5538573400 169 323 417
C — B + left-padding-corrected position IDs 0.5545729403 169 323 416

So I think the large ~0.51 vs ~0.554 discrepancy is now mostly explained.

The main difference does not seem to be checkpoint quality. It appears to be the optimized evaluator’s treatment of left padding when constructing the query representation.

In the published evaluator, the whole padded prefix is passed to the importance head:

query_hidden = seq_hidden[:, :query_end, :]

and the head then computes the query representation with an unmasked mean:

query_rep = query_embeddings.mean(dim=1)

So for a left-padded batch the effective query representation is approximately:

[left padding | real Question ... | Passage:]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
             mean-pooled

If I keep the same base-model hidden states and remove only the left-padding positions before that pooling step, MRR moves from 0.5096109 to 0.5538573.

A few full-slice sanity checks make this look fairly direct rather than just another correlated evaluator change:

  • 3972 passages total
  • 496 passages have no left padding
    • A and B are identical on all 496/496
  • 3476 passages have left padding
    • their scores can move substantially
  • mean absolute A/B score difference on padded passages is about 4.74 score points
  • correlation between padding fraction and |A-B| is about
    • 0.731 over all passages
    • 0.633 considering padded passages only
  • A/B ranking order changes on 469 / 483 queries

The subsequent position-ID correction is real too, but it is much smaller in aggregate MRR:

A: padding included in query mean
   0.50961090
        |
        | remove left-padding states before query pooling
        v
B: 0.55385734
        |
        | correct positions for left-padded real tokens
        v
C: 0.55457294

So if I were packaging this going forward, I would probably keep two explicit evaluator contracts rather than trying to make one result replace the other:

  1. Historical/legacy evaluator — preserve the current optimized semantics so the historical result remains reproducible.
  2. Corrected evaluator — mask/remove left padding before query pooling, use padding-correct positions, and define the ranking/metric semantics explicitly.

That way the old number does not need to be rewritten; it just gets a precise historical contract.

There is one more interesting detail: the tiny remaining difference between my canonical result and the historical artifact is now small enough to live entirely inside exact BF16 score ties. I would therefore still stop just short of saying I have forensically reproduced the exact historical runtime, but I think the historical evaluation regime is now reconstructed very closely.

Reproduction details, diagnostics, and a few cleanup ideas

1. Frozen reproduction contract

For these runs I pinned the evaluation inputs rather than following mutable main/latest artifacts.

The relevant public pieces were:

The public evaluation contract I froze was:

data:
  microsoft/ms_marco
  config: v1.1
  immutable dataset revision
  validation parquet
  numeric query_id sort
  rows 500:1000

slice:
  500 rows
  17 with no positive passage
  483 usable queries
  3972 passages

prompt:
  Question: {query}

  Passage: {passage}

  Answer:

marker IDs:
  [7031, 1233, 29515]

max_length:
  2048

canonical tokenizer:
  left padding
  pad token explicitly replaced with EOS
  pad_token_id = 2

head:
  public v2.3 checkpoint
  BF16

ranking:
  descending passage score

The Hub supports pinning downloads to a branch/tag/commit through the revision argument, so most of this can be represented directly in a small release manifest rather than depending on whichever cached snapshot happens to be available later: Hugging Face Hub download/revision docs.

One scope note: this is a reconstruction of the published scoring semantics on the independently pinned public validation slice. I would distinguish that from claiming the current shell command/data-loader path is itself already a literal end-to-end reproduction of the historical run.

2. Why the padding behavior changes the score

The distinction I tested was deliberately narrow.

A — current canonical path

batch passages belonging to one query
        |
        v
left pad shorter prompt sequences
        |
        v
Mistral forward with attention_mask
        |
        v
take hidden states through query_end
        |
        v
mean the whole query prefix
        |
        v
TIS passage scoring

For a short candidate in a query batch:

[PAD PAD PAD PAD Question ... Passage: | passage tokens ...]
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
               query mean

B — padding-trimmed control

The base forward is unchanged.

Only after obtaining the same hidden states:

attention_mask == 1
        |
        v
drop left-padding hidden positions
        |
        v
find the same structural boundary
        |
        v
mean real query-prefix states only

That is why A/B is useful as a control: it does not require a second checkpoint or a different base-model forward.

Hugging Face’s Mistral documentation treats attention_mask and position_ids as model inputs. The attention mask prevents attention to masked padding positions inside the model, but a later custom pooling operation still has to apply its own mask if it wants padding-invariant pooling.

The tokenizer API likewise treats padding_side as an explicit tokenizer property: Transformers tokenizer docs.

So I think either of these would be reasonable for a corrected path:

# Option 1: trim before calling the importance head
real = attention_mask.bool()
hidden_real = hidden[real]

or make the query pooling itself mask-aware:

masked = query_hidden * query_mask.unsqueeze(-1)
query_rep = masked.sum(dim=1) / query_mask.sum(dim=1, keepdim=True)

I slightly prefer trimming before the head here because it preserves the current head API and makes the intended semantic sequence explicit.

There is also a training/evaluation consistency reason I lean in this direction: the v2.3 training/early-validation path scores individual passages rather than constructing this same left-padded multi-passage query batch. So the padding-inclusive mean looks specific to the optimized test path rather than an essential property of the learned head.

3. Position IDs look like a second, much smaller factor

The padding issue above and position handling are separate.

For decoder models, an attention mask and token position are different pieces of state. The Mistral API exposes both attention_mask and position_ids.

In my controls:

B:
same left-padded base batch
padding removed before query pooling
uncorrected base positions

MRR = 0.5538573400

and:

C:
padding removed before query pooling
left-padding-corrected positions

MRR = 0.5545729403

So position correction contributes another roughly +0.000716 MRR here.

That is measurable, but it is much too small to explain the original ~0.044 MRR gap by itself.

I would therefore treat these as two regression tests:

1. pooling should be invariant to irrelevant left padding
2. real-token positions should be invariant to irrelevant left padding

rather than bundling them into one generic “batching” issue.

4. The reported 3935 / 37 separator split appears to be a counter/sentinel collision

This one became much clearer once the newly published evaluator was available.

Conceptually, the helper does:

if marker_found:
    return marker_end

return len(input_ids) // 2

But the caller then effectively infers whether the marker was found from:

if query_end == len(input_ids) // 2:
    separator_not_found += 1

Those are not equivalent predicates.

There is a third case:

marker really was found
AND
marker_end happens to equal len(input_ids) // 2

On the full pinned slice I get:

true marker found:     3972
true marker missing:      0

while running the legacy counter condition gives:

reported found:        3935
reported not_found:      37

Exactly the historical fingerprint.

So I do not think those 37 entries represent 37 real midpoint fallback events. They appear to be valid marker boundaries whose numeric index happens to collide with the sentinel value used to indicate fallback.

That also explains why treating “exactly 37 passages use the midpoint boundary” as a counterfactual did not reproduce the historical ranking vector.

A small API change would make this unambiguous:

query_end, marker_found = find_query_end(...)

if not marker_found:
    separator_not_found += 1

or just fail closed on an actual missing structural boundary.

That would let the manifest say something like:

3972 markers found
0 true fallback uses

while still preserving the historical 3935/37 counter output if you want it documented as a legacy diagnostic artifact.

5. The remaining historical difference is now at exact-tie scale

This is the part that makes me reluctant to overstate “exact historical reproduction”, but it also makes the remaining discrepancy much less mysterious.

Full canonical A:

MRR   0.5096108974990962
Hit@1 146
Hit@3 300
Hit@5 389

Historical artifact:

MRR   0.5101523217982845
Hit@1 146
Hit@3 300
Hit@5 390

The MRR gap is only about 0.0005414.

The BF16 passage scores contain a fair number of exact ties. On the completed 483-query result, changing no passage score at all, and allowing only the ordering of exactly equal-score candidates to vary, is enough to produce the historical aggregate exactly.

One minimum witness needs eight query-level tie resolutions. Applying only those legal equal-score rank changes gives:

MRR   0.5101523217982844
Hit@1 146
Hit@3 300
Hit@5 390

The last MRR digit differs from the stored value only at ordinary host floating-point representation scale.

I do not interpret that as proof that those exact eight tie resolutions occurred historically. The minimum witness is not unique.

What it shows is narrower and, I think, useful:

the remaining historical aggregate no longer requires another large scoring mechanism, another checkpoint, or 37 real fallback events; it already exists inside the exact-tie surface of the reconstructed BF16 scores.

I also ran those tie-sensitive cases repeatedly on the current runtime.

Three repetitions under the same runtime produced identical passage scores and identical rankings, so ordinary nondeterminism within that environment does not appear to explain the residual.

Two small diagnostics show how easy it is for a different numerical contract to move these ties, though:

Same BF16 token scores, different final reduction precision

The source-faithful passage score uses a BF16 mean().

Keeping the same BF16 token-score vectors but taking the final mean in FP32 moves 5 of the 8 witness queries to historical-compatible ranks.

PyTorch documents that torch.mean(..., dtype=...) casts the input to the requested dtype before reduction, so these are genuinely different numerical contracts.

Attention implementation

Switching the base attention implementation from the current SDPA path to eager also moves another set of 5/8 tie-sensitive cases to historical-compatible ranks.

Transformers explicitly supports different attention backends such as SDPA and eager.

Again, I would not infer that the historical run necessarily used eager or FP32 reduction. These were diagnostic perturbations only.

I would summarize the remaining uncertainty as:

historical bit-level runtime / reduction / candidate-order provenance

rather than:

another unexplained model-quality gap

PyTorch’s own numerical-accuracy documentation is relevant here: mathematically equivalent floating-point paths are not generally guaranteed to be bitwise identical across implementations, batching strategies, devices, or reduction orders.

6. I would probably make tie behavior part of the evaluator contract

This mattered more than I initially expected.

The current ranking uses ordinary Python sorting on score. Python sorts are stable, so if two candidates have exactly equal scores, their original relative input order is preserved.

That means the effective ranking contract is not merely:

score descending

It is:

score descending
then preserve source candidate order on exact ties

If source candidate order is itself not frozen, aggregate metrics can differ even when passage scores are identical.

For a legacy reproduction evaluator, preserving the original stable-source-order behavior is fine — I would just document it.

For a forward-looking evaluator, another option would be a deterministic secondary key:

sorted(
    candidates,
    key=lambda x: (-x.score, x.candidate_id),
)

The exact policy matters less than making it explicit.

Tie handling is also not an exotic concern in ranking evaluation; for example, scikit-learn’s ndcg_score explicitly distinguishes tie-aware evaluation from the optimization that assumes no ties.

7. A small immutable evaluation manifest would probably prevent most of this kind of drift

If useful, I would make the v2.3 reproduction surface something approximately like this:

source:
  repo: nitroxido/token-importance-scoring
  commit: <full Git commit>
  evaluator_file: scripts/evaluate_v2.3_optimized.py
  evaluator_sha256: <sha256>

data:
  repo: microsoft/ms_marco
  config: v1.1
  revision: <full Hub revision>
  file: v1.1/validation/0000.parquet
  sha256: <sha256>
  selection: numeric query_id sort, rows 500:1000
  qid_list_sha256: <sha256>

base_model:
  repo: unsloth/mistral-7b-instruct-v0.3-bnb-4bit
  revision: <full Hub revision>
  weights_sha256: <sha256>

tokenizer:
  revision: <full Hub revision>
  tokenizer_sha256: <sha256>
  padding_side: left
  pad_token_id: 2

tis_checkpoint:
  repo: oldman-dev/tis-v2.3-passage-reranker
  revision: <full Hub revision>
  file: best/tis_components.pt
  sha256: <sha256>

prompt:
  template: |
    Question: {query}

    Passage: {passage}

    Answer:
  marker_ids: [7031, 1233, 29515]
  max_length: 2048

runtime:
  python: <version>
  torch: <version>
  transformers: <version>
  tokenizers: <version>
  bitsandbytes: <version>
  cuda: <version>
  gpu: <model>
  attention_backend: <sdpa/eager/etc>
  base_compute_dtype: <dtype>
  head_dtype: <dtype>
  score_reduction_dtype: <dtype>

ranking:
  direction: descending
  tie_policy: <explicit policy>

metrics:
  no_positive_query_policy: skip
  definitions:
    - MRR
    - Hit@1
    - Hit@3
    - Hit@5
    - literal Recall@k
    - standard binary NDCG@k

outputs:
  aggregate_metrics: results.json
  per_query_scores_and_ranks: per_query_results.parquet

I would especially keep the per-query score/rank artifact, not just the aggregate JSON.

In this case, if the historical result had included qid → passage scores/ranks, the last 0.00054 residual would have been localizable immediately instead of being recoverable only from the aggregate metric fingerprint.

8. Separate from the padding issue: I think the metric names are worth making explicit

This does not explain the 0.51 vs 0.55 result difference, so I would treat it as separate release cleanup.

For multi-positive queries, code of the form:

1 if first_relevant_rank < k else 0

is normally a Hit@k / Success@k / top-k accuracy style metric.

Literal Recall@k is:

number of relevant candidates retrieved in top k
-------------------------------------------------
total number of relevant candidates

Those happen to agree when there is exactly one positive, but not for multi-positive queries.

Libraries such as BEIR keep Recall@k and top-k accuracy as distinct concepts and use pytrec_eval for standard NDCG/Recall/Precision evaluation. pytrec_eval is also a useful cross-check if the goal is to avoid having the custom scorer and custom metric definitions share the same implementation assumptions.

Similarly, I would keep standard binary NDCG separate from the older custom v2.3 NDCG calculation.

A clear forward-facing result set could simply report:

MRR

Hit@1
Hit@3
Hit@5

literal Recall@1
literal Recall@3
literal Recall@5

standard binary NDCG@5
standard binary NDCG@10

That also makes comparisons by future readers easier.

9. One small end-to-end data/CLI detail

This is also separate from the scoring finding.

The current public MS MARCO v1.1 dataset card describes rows as:

answers
passages:
    is_selected
    passage_text
    url
query
query_id
query_type
wellFormedAnswers

In other words, is_selected is nested under passages, not a top-level raw-parquet column.

So if the reproduction command in the manifest is intended to consume the raw Hub parquet directly, I think it would help to either:

  1. make the loader understand that nested schema, or
  2. document the preparation step that converts it to the evaluator’s expected table.

Likewise, at f3ed5d5 I would double-check that the documented CLI flag and the evaluator’s argparse name are identical (--data-path vs --test-data in the versions I inspected).

Neither affects the scoring experiment above; it is just the last piece needed for a future reader to go from:

fresh clone

to:

same evaluation result

without guessing which intermediate file format is intended.

10. A few very small regression fixtures could make this robust

I do not think this needs a large test suite. Four tiny fixtures would catch most of what showed up here.

Padding-invariance fixture

Take the same real query/passage sequence and represent it with:

0 left pads
16 left pads
32 left pads
64 left pads

For the corrected evaluator, assert that irrelevant padding does not change the query representation, passage score, or rank.

Separator-sentinel fixture

Include a case where:

marker_found == True
marker_end == midpoint

and assert:

fallback_count does not increase

That directly protects the 3935/37 corner case.

Position-ID fixture

Compare:

serial/no-padding input

against:

left-padded batch
+ padding-correct real-token positions

and verify hidden/score parity within the intended numerical tolerance.

Tie-policy fixture

Use two exactly equal scores and verify that the declared secondary ordering produces a deterministic rank and metric value.

Those four tests are small enough to run as ordinary release CI rather than as another benchmark.

11. One terminology/scope note for future readers

The MS MARCO v1.1 object being scored here is a bounded list of candidate passages associated with each query row.

So I would probably call this result something like:

bounded-candidate MS MARCO v1.1 reranking MRR

rather than implicitly equating it with the official full-corpus MS MARCO passage-ranking leaderboard protocol.

That is not a problem with the experiment; it just makes the benchmark endpoint unambiguous.

12. What I would still leave open

The missing evaluator source was the big blocker in my earlier reply, and #23 closes most of that practical gap.

What I still would not claim from this reconstruction is forensic identity of the historical execution environment.

In particular, I do not have independent proof of the exact historical:

  • evaluator file bytes at the time of the original run
  • base-model cache revision
  • PyTorch/CUDA/bitsandbytes combination
  • attention backend
  • passage input ordering
  • BF16/FP32 reduction details
  • tie resolution

That distinction matters because the remaining difference is now small enough for those details to affect a few ranks.

But importantly, this is no longer a large unexplained behavioral difference:

before:
historical ~0.510
public-contract reconstruction ~0.554
large unexplained gap

now:
historical       0.51015232
canonical A      0.50961090
padding-trim B   0.55385734
position-fix C   0.55457294

and the historical - canonical A residual is fully reachable inside exact BF16 tie groups.

So from my side I would classify the remaining question as:

exact historical numerical provenance

rather than:

whether the public checkpoint actually exhibits the claimed reranking regime

The latter now looks pretty solid to me.

So, broadly: the evaluator upload seems to have supplied the missing bridge I needed. I would preserve that path as the historical/legacy contract, and make the mask-aware + position-correct path the clean forward evaluator. That keeps the historical result reproducible while making the current evaluation semantics much easier for someone else to reason about and reproduce.

25

This lines up with something we found the hard way in our own retrieval work. We had a similarity-based ranker sitting downstream of a graph-walk expansion step, and it turned out the ranker was quietly discarding relevant candidates that didn’t lexically resemble the query, even though they were reached correctly upstream. Measured against a real control, the walk’s contribution to recall was zero, purely because of what the ranker threw away at the last step. Reordering/reranking by something other than raw similarity, the way TIS does with learned importance, seems like exactly the right fix. Did you test whether the position-bias elimination also recovers cases where the content, not just the position, made the ranker discount a relevant passage?

26

Okey .. after your concerns, I made some specific testing for you and found these results!

image
image692×607 49.9 KB

So expect a V2.3 upload with code for how to test this using the TIS system for yourself!

You can also test it yourself using these tips:

image
image655×493 41.2 KB

27

The distinction between position invariance and actual ordering quality is very useful here.

One additional test I would be interested in is a controlled multi-hop setting where the answer requires evidence from two separate passages rather than one answer-bearing passage. A scorer can place one highly useful passage first while still missing a dependency that only becomes visible when two distant passages are retained together.

For that setting, it could be useful to report:

  • answer accuracy

  • recall of both required evidence passages

  • a matched random canonical ordering control

  • the same evaluation under fixed context budget

That would help distinguish “the system removes sensitivity to the original slot” from “the importance policy preserves the dependencies needed for harder long-context reasoning.”

The public-artifact replication and the effort to separate these interpretations are especially valuable.

28

Hey, @corechek! That’s a good question, and I will try to set-up a case that answers it! Maybe I need a new training + testing regime, maybe I can just do a test over one of the available checkpoints!

Thanks for the interest, and await for results!

29

I had to find a moment to make the testing, but now it’s done and these are the results:

image
image673×726 60.5 KB

So these are the takeaways:

  • TIS improves multi-hop recall (2.2x at K=5) but it doesn’t fill the gap completely.
  • The original canonic ordering doesn’t affect much the results.
  • The bridge passage is at 5.8 average rank, while at tight budgets (K=2) it looses 90.5% of the questions.
  • TIS scores passages independently, but due to their inter-relations, it gives them a slightly higher chance to be closer to each other.

This is the current problem:

image
image678×333 33.7 KB

The system will need further refining in order to be helpful in these cases.

If you are interested, I can retrain TIS (and maybe bump it a revision or two) in order for the system to work good even in the multi-hop case. I already have a plan ready, so it’s just a matter of finding some time to run it @corechek !

30

Thanks. I think I may have figured out a few things here:

Yes — I think a multi-hop-oriented retraining run could be worth trying.

But before spending a full training sprint on it, I would probably do one small diagnostic pass first, because your bridge-passage result seems to narrow the problem much more than the 2.2x @ K=5 headline alone does.

What caught my attention most was this combination:

TIS improves multi-hop recall at K=5
bridge mean rank ≈ 5.8
bridge is lost in 90.5% of cases at K=2
original canonical ordering matters relatively little
current passages are scored independently

That looks compatible with a fairly specific boundary:

A passage can be only weakly useful when scored against the original question in isolation, while becoming essential once another piece of evidence has already been selected.

In other words, there may be a difference between:

individual passage importance

    score(q, p)

and something closer to:

conditional evidence utility

    utility(p | q, evidence already selected)

or, at the context-budget level:

set utility

    utility({p1, p2, ...}, q, budget)

I would not take the current result as proof that independent scoring is the root cause — there are a few cheaper alternatives to rule out first — but it looks like a good hypothesis to test.

My default route would therefore be:

before retraining
      |
      v
check where the bridge is actually being lost
      |
      +-- absent from candidate pool?
      |       -> upstream retrieval / expansion problem
      |
      +-- present, but reranker removes it?
      |       -> scoring / supervision / set-composition problem
      |
      +-- retained, but reader still fails?
              -> downstream reasoning / reader problem

Then, only if the middle branch is the important one, decide what kind of retraining TIS actually needs.

One especially cheap experiment may distinguish several of these possibilities without training anything.

A small high-information probe

Keep fixed:

same questions
same candidate pool
same checkpoint
same reader
same context budget

and record, for every question:

required evidence passage IDs
endpoint passage rank
bridge passage rank
whether all required evidence was present in the original pool
whether all required evidence survived @ K
final answer correctness

Then try two counterfactuals.

1. Oracle bridge retention

Force the bridge passage to remain in the selected context and let the normal policy choose the remaining slots.

If answer accuracy recovers strongly, then preserving the bridge is genuinely useful to the downstream reader and improving the selector is probably worth the effort.

If the bridge is retained but answers still fail, the bottleneck may be partly downstream rather than primarily TIS ranking.

For questions requiring more than two pieces of evidence, I would also try an oracle complete-evidence condition whenever the budget permits it, because forcing only one bridge does not guarantee that the rest of the reasoning chain survived.

2. First-hop-conditioned query control

Instead of immediately changing the model architecture, expose the information discovered at hop 1 to the scorer and see whether the bridge suddenly becomes easy to rank.

For example:

original:
    score(original question, candidate)

control:
    score(original question + discovered entity/fact, candidate)

or turn the first-hop result into a simple second-hop sub-question.

If the bridge goes from rank ~5.8 to near the top, that would be a strong clue that the problem is not simply “TIS does not recognize useful passages.” It may instead be that the bridge is not identifiable from the original query state alone.

This kind of failure is already known in multi-hop retrieval. GoldEn Retriever was motivated by cases where the original question does not contain the entity needed to retrieve the next document, so it iteratively generates a new search query from the context already found. IRCoT makes a closely related observation: in multi-step QA, what should be retrieved next depends on what has already been derived or retrieved.

I would treat this as a mechanism probe, not as a production TIS evaluation, because changing the scorer input can itself be distribution-shifting. But it is cheap, and it tells you what kind of next model may be worth building.

A useful comparison would be to apply the same rewritten query to BM25 or a dense retriever too:

rewritten query helps everything
    -> clue exposure / query formulation may be the main issue

rewritten query helps BM25/dense but not TIS
    -> more TIS-specific scoring/training issue

rewritten query also helps TIS
    -> you may not need a full set-wise architecture yet
Why I think the bridge result is especially interesting

There is a growing line of work that separates ranking individually relevant passages from selecting a jointly sufficient evidence set.

The most direct example I found is Shifting from Ranking to Set Selection for Retrieval Augmented Generation (SetR, ACL 2025).

Its motivation is essentially:

top individually relevant passages
          !=
best collectively sufficient passage set

For a simple one-hop query these may often coincide.

For a multi-hop query they do not have to.

Suppose the question requires:

A -> B -> answer

Passage A may look highly relevant to the original query.

Passage B may be a weak lexical/semantic match to the original query but becomes indispensable after A reveals an intermediate entity.

A point-wise sorter can therefore produce something like:

rank 1: strong endpoint-looking passage
rank 2: another directly relevant passage
rank 3: redundant support
rank 4: another direct match
rank 5: bridge

and look reasonable under ordinary relevance scoring while failing completely at a tight evidence budget.

Your reported mean bridge rank of 5.8 looks compatible with that kind of behavior.

There is an even closer paper, Modeling Contextual Passage Utility for Multihop Question Answering, which explicitly distinguishes ordinary relevance from passage utility that depends on other passages, including complementary evidence and passages that form a crucial link.

That seems unusually close to what your bridge test is exposing.

The important part, though, is that this does not automatically imply:

“TIS now needs expensive candidate-to-candidate cross-attention.”

There is a cheaper possibility.

The contextual-utility work uses richer reasoning/context information to construct utility supervision, but then trains a relatively lightweight scorer to predict that utility.

So there are at least two separate knobs:

A. change what TIS learns

B. change what TIS can condition on at inference time

and I would test them separately before combining them.

One caveat about interpreting K=2

I would also be slightly careful with the 90.5% bridge loss at K=2 number until the exact task composition is attached to it.

If every evaluated question requires exactly two evidence passages, then K=2 is a very clean and interesting stress test:

budget == minimum complete evidence set

But if the test mixes questions requiring 2, 3, 4, … distinct pieces of evidence, then some cases may be structurally impossible at K=2 regardless of ranking quality.

So for future runs I would probably stratify by:

required evidence count
or
hop count

and report something like:

2-evidence questions:
    complete-evidence success @2
    complete-evidence success @5

3-evidence questions:
    complete-evidence success @3
    complete-evidence success @5

...

This would make the bridge result much easier to interpret.

Also, mean bridge rank = 5.8 and 90.5% bridge loss @2 are related statistics rather than two fully independent pieces of evidence: once most bridge ranks sit outside the first two slots, a large loss at K=2 follows mechanically.

I still think both are useful to show, but I would probably add the rank distribution or CDF:

P(bridge rank <= 1)
P(bridge rank <= 2)
P(bridge rank <= 3)
P(bridge rank <= 5)
P(bridge rank <= 10)

That would tell us much more than the mean alone.

If you retrain TIS, I would separate the training target from the architecture

If you do decide to build the multi-hop revision, I think one useful experimental separation would be:

Axis 1 — what should the score mean?

The current style of supervision can be summarized roughly as:

relevant passage > distractor

For multi-hop, possible targets become richer:

direct relevance

bridge / endpoint role

marginal utility:
    how much does adding this passage improve evidence completeness?

contextual utility:
    is this passage useful given what has already been found?

survival objective:
    do all required evidence passages survive the K/token budget?

A particularly interesting low-cost version would be to keep the existing lightweight scoring shape but change the labels.

For example, if a reasoning trace says:

passage A
    reveals entity X

passage B
    connects X to Y

passage C
    supplies final answer

then B can receive strong training utility even if B is not an especially obvious direct match for the original question.

That would test:

Can a point-wise TIS head learn a prior for bridge-like utility?

before introducing a more complicated inference graph.

Axis 2 — what information does the scorer receive?

Only after that, test progressively richer contracts:

1. score(q, p)

2. score(q + first-hop clue, p)

3. score(q, selected evidence, p)

4. score(q, current selected set, candidate)

5. beam/path/set-wise selection

If option 1 with improved supervision fixes most of the problem, that would preserve much of the current TIS simplicity.

If option 2 or 3 is necessary, then a stateful/sequential variant becomes justified.

If only 4/5 works well, then the problem is genuinely closer to set construction than scalar relevance ranking.

This separation also makes the ablation much easier to interpret:

better labels helped
vs.
additional state helped
vs.
both were required

rather than changing the training data, objective, architecture, and inference loop in one revision and not knowing which part mattered.

There are several intermediate designs before a fully set-wise TIS

I found several existing approaches that map nicely onto different points in this design space.

1. Better set composition

SetR explicitly changes the problem from ranking passages individually to selecting a set that collectively satisfies the query’s information requirements.

That is probably the closest “full set selection” reference.

Their public implementation is here:

https://github.com/LGAI-Research/SetR

I would regard this as a relatively far-end architectural option rather than the first thing to copy.

2. Contextual utility supervision

Modeling Contextual Passage Utility for Multihop Question Answering is interesting precisely because it suggests that some of the set/inter-passage information can be pushed into the training signal, while keeping the final utility model comparatively lightweight.

This looks particularly compatible with the TIS philosophy.

Conceptually:

rich teacher / reasoning trace
        |
        v
derive bridge-aware utility targets
        |
        v
train lightweight TIS scorer
        |
        v
cheap inference

If that works, it could give you multi-hop sensitivity without turning every TIS ranking call into an expensive combinatorial search.

3. Iterative retrieval/query rewriting

GoldEn Retriever and IRCoT represent a different solution:

retrieve
   |
read / derive clue
   |
rewrite retrieval state
   |
retrieve again

This is attractive if the bridge is not recognizable until an intermediate entity becomes visible.

A TIS-native version might look something like:

TIS selects first evidence
        |
derive/update query state
        |
TIS rescoring
        |
select next evidence

That would preserve TIS as the scoring primitive while making the overall policy multi-step.

4. Question decomposition

Another cheap baseline is Question Decomposition for Retrieval-Augmented Generation.

Their pipeline is approximately:

complex query
   |
decompose into sub-questions
   |
retrieve per sub-question
   |
merge candidates
   |
rerank

They report improvements on MultiHop-RAG and HotpotQA without task-specific retriever training.

I like this as a control because it answers:

Is the bridge hard because the scorer cannot model passage interactions, or simply because the original query does not expose the bridge’s retrieval clue?

If decomposition fixes the bridge substantially, a huge TIS redesign may not be necessary.

5. Keep multiple partial evidence paths

If greedy sequential selection becomes the problem, End-to-End Beam Retrieval for Multi-Hop Question Answering is another useful reference.

Instead of committing completely to one first-hop choice, Beam Retrieval retains several partial evidence hypotheses.

Conceptually:

hop 1:
    A
    B
    C
   /|\

hop 2:
    A->D
    A->E
    B->F
    ...

This trades more compute for less risk that one early ranking mistake permanently removes the correct chain.

Again, I would not start there unless the cheaper controls show that greedy/state-less selection really is the limiting factor.

A few very cheap controls may tell you which mechanism is involved

Before a large new architecture, I would keep several simple controls around.

Diversity / redundancy control

A relevance scorer can use its top slots on several passages saying nearly the same thing.

A basic diversity-aware selector — even something MMR-like — could test whether the bridge is being lost mainly because the top-K is redundant.

Interpretation:

diversity control strongly improves complete evidence survival
    -> redundancy / coverage is a meaningful part of the problem

little or no improvement
    -> bridge complementarity is probably more specific than novelty

I would not treat MMR as a real model of multi-hop dependencies. It is just an inexpensive negative control.

Random canonical ordering

Keep the matched random canonical order from the original suggestion.

If TIS improves complete-evidence survival over that control, then the signal is doing something useful beyond merely removing position dependence.

BM25 / lexical / dense

I would also keep at least one lexical and one semantic ranking baseline if convenient.

The important comparison is not necessarily:

Who has the best mean MRR?

but:

Which selector keeps the full evidence chain alive under the same budget?

A scorer can have a respectable ordinary ranking metric and still be poor at multi-hop completeness.

Metrics I think would make the next result much easier to read

For the multi-hop experiment I would separate three layers.

1. Candidate-generation coverage

Before reranking:

endpoint present in candidate pool?
bridge present in candidate pool?
ALL required evidence present in candidate pool?

This matters because no reranker can recover evidence that was never retrieved.

2. Selection / ranking

I would report roles separately:

endpoint mean/median rank
bridge mean/median rank

endpoint Hit@K
bridge Hit@K

but the most useful multi-hop metric may be something like:

complete-evidence success @ K

defined as:

1  if every required evidence passage survives in top K
0  otherwise

For a two-evidence question:

endpoint AND bridge present @ K

This answers a different question from ordinary “did I retrieve at least one relevant passage?”

If there can be multiple valid support passages per role, the metric can operate over acceptable evidence sets rather than one hardcoded passage identity.

3. Reader / answer

Then separately report:

answer EM / F1 / accuracy

conditioned on:
    complete evidence present
    bridge absent
    endpoint absent

That decomposition would be extremely informative:

candidate generation
        |
        v
evidence selection
        |
        v
complete context?
        |
        v
reader succeeds?

HotpotQA is a useful precedent for keeping answer quality and supporting-evidence quality distinct: its evaluation reports answer metrics, supporting-fact metrics, and a joint score rather than collapsing everything into one endpoint.

HotpotQA

This distinction seems especially important for TIS because the original thread has already shown that better passage ordering and better final answer generation do not necessarily move in lockstep.

There may also be a connection to the content-bias question from #25

I think @VSdutch’s observation and this bridge result may share a pipeline-level pattern, although I would not assume they have the same cause.

The common structure can be:

upstream stage successfully exposes useful information
          |
          v
downstream query-based scorer discounts it
          |
          v
reader never gets a chance to use it

But there are at least two distinct cases.

Content-bias case

passage is independently useful for q
but its surface/content similarity is weak

A better relevance/utility scorer may directly fix that.

Bridge case

passage looks weak for q alone
but becomes essential after passage A is known

That is a stronger form of dependency.

So the tests you showed after #25 may be useful controls for the multi-hop work, but I would keep the claims separate.

If the same new TIS revision improves both, that would be interesting evidence that the learned “importance” signal is becoming broader than direct lexical/semantic relevance.

If it improves content-bias recovery but not bridge retention, that would also be useful: it would tell us that multi-hop conditionality is a genuinely separate axis rather than just a harder relevance example.

What I would still leave open from the current screenshot

For now I would treat the reported numbers as a very useful behavioral result, but not yet infer too much about the exact mechanism until the corresponding test artifact appears.

The main things that would affect interpretation are:

dataset / split / N

number of hops or required evidence passages per question

absolute values behind the 2.2x ratio

exact definition of "multi-hop recall"

candidate-pool construction

whether all gold bridge passages were already present before TIS ranking

passage-count K vs fixed token budget

reader model / prompt

checkpoint and evaluator identity

zero-shot existing checkpoint vs a modified/retrained checkpoint

per-question passage ranks

I checked the current public repository while thinking about this; the visible main history is still at the v2.3/f3ed5d5-era release, so I do not think the #29-specific evaluator/result artifact is public there yet:

https://github.com/nitroxido/token-importance-scoring

That is not a blocker for discussing the result — the failure mode is already interesting — but once the artifact lands, the candidate-coverage / evidence-count / rank-distribution checks above are the first things I would look at.

So my current answer to the retraining idea is:

yes, I think it is worth pursuing, but I would make the next training target depend on one small diagnostic first.

Something like:

1. Is the bridge in the candidate pool?
2. If retained by oracle, does it actually restore the answer?
3. Does exposing the first-hop clue make the bridge easy to rank?
4. Does simple diversity already fix much of the loss?

Then the result tells you which revision is worth spending compute on:

bridge missing upstream
    -> iterative retrieval / query decomposition

bridge present, oracle helps,
and rewritten query helps
    -> stateful/sequential TIS may be enough

bridge present, oracle helps,
but query rewriting does not
    -> bridge/contextual-utility supervision becomes interesting

simple diversity fixes it
    -> coverage policy may be enough

all evidence survives but reader still fails
    -> improve the reader/reasoning side before making TIS larger

none of the cheap controls work
    -> set-wise / beam/path modeling becomes much more justified

That seems like a fairly nice next step because none of these outcomes make the current result a dead end.

Even the 90.5% @ K=2 failure is useful: it tells you which property a multi-hop-aware TIS revision would actually need to improve, rather than just asking it to produce a higher generic ranking score.

31

As usual, I will address your feedback too, but I already have trained a multi-step version of TIS that:

  • Doesn’t loose the already achieved results in NIAH and LITM.
  • Does get better results at the bridge tasks after some refining on the training system.

Some of the final results I got:

image
image673×486 19.2 KB

Expect checkpoints and a code update next Monday, and then I will get to work on the points raised by @John6666 if it’s needed!