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

An update to the Token Importance Scoring series — new passage reordering capability, transfer learning findings, and open-source checkpoints.

Following the initial [TIS release](Presenting TIS (Token Importance Scoring) - A new way to compress KV cache), we’ve significantly expanded TIS capabilities. The biggest finding: a learned importance head trained for KV cache compression also eliminates position bias in retrieval-augmented generation without any retraining — suggesting TIS learns generalizable importance patterns.

## What’s New in TIS 2.0

### 1. Position Bias Elimination (Lost-in-the-Middle)

The core discovery: reordering retrieved passages by TIS importance scores completely eliminates the LITM position gap while improving overall accuracy.

| Pipeline | EM (early) | EM (middle) | EM (end) | LITM Gap |

|—|—|—|—|—|

| Baseline (no reordering) | 18.3% | 13.3% | 18.3% | 0.050 |

| Lexical (TF-IDF) | 16.7% | 18.3% | 18.3% | −0.008 |

| TIS (passage reranker) | 21.7% | 21.7% | 21.7% | 0.000 |

| Oracle | 18.3% | 18.3% | 18.3% | 0.000 |

Result: TIS reduces the LITM gap to zero (matching Oracle) while improving EM by +5pp over baseline. Evaluation on 60 MS-MARCO queries × 3 positions = 180 test cases.

### 2. Zero-Shot Transfer Finding

Even more surprising: the existing tis-stage3-ert checkpoint (trained for KV cache compression on synthetic data) achieves identical passage reordering performance:

| Checkpoint | LITM Gap | EM (all positions) | Training |

|—|—|—|—|

| tis-stage3-ert (KV eviction) | 0.000 | 21.7% | None (existing) |

| tis-passage-reranker (contrastive) | 0.000 | 21.7% | MS-MARCO passage pairs |

This suggests TIS learns abstract “importance” patterns that generalize across task granularities (tokens ↔ passages) and objectives (compression ↔ ranking).

### 3. Query-Aware Architecture

QueryAwareImportanceHead extends the base ImportanceUpdateHead with:

• Cross-attention query-context interaction

• Per-passage relevance scoring

• RMSNorm output stabilization

• InfoNCE contrastive training on passage pairs

### 4. Speculative Decoding Integration

TIS importance bias applied to drafter predictions improves acceptance:

| Condition | Accept Length | Speedup | Notes |

|—|—|—|—|

| No TIS | 5.80 / 8 | 0.644 | Baseline |

| TIS depth-scaled | 6.57 / 8 | 0.730 | +12.5% |

(LLaMA-3.1-8B target + LLaMA-3.2-1B drafter, n=30)

### 5. Improved Benchmarking

New evaluation framework with 4 concurrent pipelines:

• baseline: original retrieval order

• lexical: TF-IDF cosine similarity reordering

• tis: learned passage importance

• oracle: gold passage always first

Enables direct comparison of ranking quality vs. baseline and oracle bounds.

## Checkpoints

HuggingFace:

• [tis-stage3-ert](oldman-dev/tis-stage3-ert · Hugging Face): KV compression + passage reordering (no additional training needed)

• **[tis-passage-reranker**](oldman-dev/tis-passage-reranker · Hugging Face) (new): dedicated passage reordering head

• [tis-v8b-hard-anchor](oldman-dev/tis-v8b-hard-anchor · Hugging Face): 82% NIAH @ 25% budget

GitHub: GitHub - nitroxido/token-importance-scoring: Public repository for TIS system · GitHub

## Key Code Changes

New modules:

• QueryAwareImportanceHead — contrastive query-aware scoring

• CrossAttentionImportanceScorer — per-passage relevance

• TISCompositeLoss — 5-component objective (task + KL + budget + churn + saliency)

• TISDrafter — speculative decoding adapter

• Evaluation: run_litm_with_baselines.py (4-pipeline comparison)

Updated:

• ImportanceUpdateHead.direct_score() — explicit named scorer path (evaluation-only)

• PatchedCausalLM — Transformers 5 SDPA compatibility warning + generation prefix fix

• eval/benchmarks.py — full LITM pipeline with scoring policies

## Technical Insights

1. Generalizable importance: Token-level and passage-level importance share underlying patterns. A head trained for KV compression can score entire passage relevance.

2. Position bias root cause: The position effect is strongest at positions where passage content is sparse (middle positions lose context) — not just a positional encoding issue. Reordering addresses this directly.

3. Trade-offs: Improved passage ordering does not guarantee proportional answer quality gains. ~76% of test cases remain incorrect after reordering, indicating answer extraction (not ranking) is the bottleneck in this domain.

## Reproducibility

git clone [GitHub - nitroxido/token-importance-scoring: Public repository for TIS system · GitHub](https://github.com/nitroxido/token-importance-scoring.git)

cd token-importance-scoring

python -m venv .venv && source .venv/bin/activate

pip install -e .


hf download oldman-dev/tis-stage3-ert --local-dir checkpoints/stage3_ert


python scripts/run_litm_with_baselines.py \

--checkpoint checkpoints/stage3_ert \

--n-examples 60 \

--seed 42 Expected: baseline LITM gap 0.050 → TIS gap 0.000, baseline EM 16.7% → TIS EM 21.7%

## Open Questions for the Community

1. Does this transfer to other domains? We validated on MS-MARCO. Does passage reordering work on e.g., NaturalQuestions, HotpotQA, or multi-document summarization?

2. Beginning-position ceiling: All learned scorers (including ours) achieve 0% on beginning-position tokens due to causal masking. Our oracle key_match scorer (text parsing) reaches 67%. Can we close this gap with a non-causal architecture?

3. Cross-model transfer: Can a TIS head trained on Mistral-7B transfer to Llama-3-8B or Qwen-72B without retraining?

4. Inference cost: TIS adds ~276ms per query (score computation). Is this acceptable for real-time RAG, or should we optimize further?

## Links

• GitHub: GitHub - nitroxido/token-importance-scoring: Public repository for TIS system · GitHub

• HuggingFace: oldman-dev (Old Man Developer)

• Reproducibility guide: [REPRODUCIBILITY-GUIDE.md](token-importance-scoring/REPRODUCIBILITY-GUIDE.md at main · nitroxido/token-importance-scoring · GitHub)

• Architecture & specs: [ARCHITECTURE-TECHNICAL-SPECS.md](token-importance-scoring/ARCHITECTURE-TECHNICAL-SPECS.md at main · nitroxido/token-importance-scoring · GitHub)

—

We’d love to hear from the community: Have you tested TIS on your domain? Do the transfer findings hold for other model families? What would make passage reordering more useful for your RAG pipeline?