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

Token Importance Scoring (TIS) 2.0 eliminates position bias in retrieval-augmented generation, reducing the lost-in-the-middle gap to zero and improving exact match accuracy by 5 percentage points over baseline, according to developer oldman-dev. The existing tis-stage3-ert checkpoint, originally trained for KV cache compression, achieves identical passage reordering performance without retraining, demonstrating transferable importance patterns. New features include a query-aware architecture, speculative decoding integration with 12.5% speedup, and open-source checkpoints on HuggingFace.

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 https://discuss.huggingface.co/t/presenting-tis-token-importance-scoring-a-new-way-to-compress-kv-cache/177429 , 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 https://huggingface.co/oldman-dev/tis-stage3-ert : KV compression + passage reordering no additional training needed - tis-passage-reranker oldman-dev/tis-passage-reranker · Hugging Face https://huggingface.co/oldman-dev/tis-passage-reranker new : dedicated passage reordering head - tis-v8b-hard-anchor oldman-dev/tis-v8b-hard-anchor · Hugging Face https://huggingface.co/oldman-dev/tis-v8b-hard-anchor : 82% NIAH @ 25% budget GitHub : GitHub - nitroxido/token-importance-scoring: Public repository for TIS system · GitHub https://github.com/nitroxido/token-importance-scoring 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 bash 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 . Download checkpoint hf download oldman-dev/tis-stage3-ert --local-dir checkpoints/stage3 ert Run 4-pipeline LITM evaluation 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 https://github.com/nitroxido/token-importance-scoring - HuggingFace : oldman-dev Old Man Developer https://huggingface.co/oldman-dev - Reproducibility guide : REPRODUCIBILITY-GUIDE.md token-importance-scoring/REPRODUCIBILITY-GUIDE.md at main · nitroxido/token-importance-scoring · GitHub https://github.com/nitroxido/token-importance-scoring/blob/main/REPRODUCIBILITY-GUIDE.md - Architecture & specs : ARCHITECTURE-TECHNICAL-SPECS.md token-importance-scoring/ARCHITECTURE-TECHNICAL-SPECS.md at main · nitroxido/token-importance-scoring · GitHub https://github.com/nitroxido/token-importance-scoring/blob/main/ARCHITECTURE-TECHNICAL-SPECS.md — 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?