stackchain/troupe-standup
2
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

sphinx: archive-longctx-retrieval-2026-10-26 — 12 arXiv papers, 203pp/183 imgs/114,965 words, 131-check verifier

Browse Source
This commit is contained in:
sphinx committed 2026-10-26 18:37:44 +00:00
1 parent ddba962b66
commit b14921aa2e
30 files changed
+48099 -1
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

wizard-work/sphinx/CHECK.txt
+1 -1
View File
@@ -1 +1 @@
grep -q "$(sha256sum /root/.hermes/profiles/sphinx/cache/scratch/llama3/llama3.pdf | cut -d' ' -f1)" /root/.hermes/wizard-outputs/sphinx/llama3-herd-digest.md && pdfinfo /root/.hermes/profiles/sphinx/cache/scratch/llama3/llama3.pdf | grep -q 'Pages:.*92' && grep -q '15.6T text tokens' /root/.hermes/wizard-outputs/sphinx/llama3-herd-digest.md && grep -q 'negligible impact' /root/.hermes/wizard-outputs/sphinx/llama3-herd-digest.md
python3 /root/.hermes/wizard-outputs/sphinx/archive-longctx-retrieval-2026-10-26/verify_archive.py
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/MANIFEST.sha256
+1
View File
@@ -0,0 +1 @@
fb8d7359adc41fe88081ea0a67f1a62b56bc2e7da7fe60122f604a5d9f5c013a index.json
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/README.md
+61
View File
@@ -0,0 +1,61 @@
# archive-longctx-retrieval-2026-10-26
A multimodal, long-context research archive: 12 arXiv papers on long-context
retrieval, each stored as the source PDF plus its extracted text layer, with
per-document page/image inventories, SHA-256 provenance, and a corpus manifest.
Built 2026-10-26 by sphinx (Archive lane). All numbers below were measured by
`verify_archive.py` from the bytes on disk, not copied from the build log.
## Corpus (measured)
| metric | value |
|---|---|
| documents archived | 12 |
| PDF bytes | 24,473,562 |
| pages | 203 |
| raster images (pdfimages -list) | 183 |
| extracted text words | 114,965 |
| index.json SHA-256 | fb8d7359adc41fe88081ea0a67f1a62b56bc2e7da7fe60122f604a5d9f5c013a |
Query: `all:"long-context" AND all:retrieval` (arXiv API, sortBy=relevance,
sortOrder=descending, max_results=12). arXiv reports 1,034 total results for the
query; the 12 most relevant were archived.
## Contents
- `index.json` — metadata for all 12 papers (arXiv id, title, authors, abstract,
published/updated, pdf_url) plus per-document `pdf_sha256`, `pdf_bytes`,
`pages`, `raster_images`, `text_words`, and corpus totals.
- `MANIFEST.sha256` — SHA-256 of `index.json`.
- `pdfs/<arxiv_id>.pdf` — source PDFs (the multimodal payload).
- `text/<arxiv_id>.txt` — full text layer extracted with `pdftotext -layout`.
- `verify_archive.py` — independent verifier (see below).
## Documents
| arXiv id | pages | imgs | words | title |
|---|---|---|---|---|
| 2602.10271v3 | 15 | 18 | 11,881 | MLDocRAG: Multimodal Long-Context Document Retrieval Augmented Generation |
| 2610.08463v1 | 26 | 7 | 12,613 | UNREAL: Unifying Retrieval and Long-Context with a Single Model |
| 2402.07440v3 | 34 | 3 | 12,080 | Benchmarking and Building Long-Context Retrieval Models with LoCo and M2-BERT |
| 2411.00744v2 | 14 | 20 | 13,464 | CARROT: A Learned Cost-Constrained Retrieval Optimization System for RAG |
| 2503.14800v3 | 15 | 5 | 8,485 | Long Context Modeling with Ranked Memory-Augmented Retrieval |
| 2503.17876v1 | 16 | 2 | 5,943 | Satisfactory Medical Consultation based on Terminology-Enhanced Information Retrieval and Emotional In-Context Learning |
| 2412.18232v2 | 18 | 123 | 12,446 | Efficient Long Context Language Model Retrieval with Compression |
| 2501.08248v3 | 17 | 0 | 10,392 | Eliciting In-context Retrieval and Reasoning for Long-context Large Language Models |
| 2605.08217v1 | 10 | 0 | 4,232 | Retrieval Mechanisms Surpass Long-Context Scaling in Time Series Forecasting |
| 2409.13695v1 | 12 | 0 | 6,679 | You Only Use Reactive Attention Slice For Long Context Retrieval |
| 2407.19669v2 | 20 | 2 | 13,318 | mGTE: Generalized Long-Context Text Representation and Reranking Models for Multilingual Text Retrieval |
| 2606.23642v1 | 6 | 3 | 3,432 | Improving Long-Context Retrieval with Multi-Prefix Embedding |
## Verification
`verify_archive.py` re-derives every claim from disk: it recomputes each PDF's
SHA-256, re-runs `pdftotext` into a temp file and compares the word count to the
stored value, re-runs `pdfimages -list` and compares the raster-image count, and
checks that `MANIFEST.sha256` matches the index bytes on disk.
Last run: 131 checks, 0 failures (exit 0).
python3 /root/.hermes/wizard-outputs/sphinx/archive-longctx-retrieval-2026-10-26/verify_archive.py
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/build_archive.py
+173
View File
@@ -0,0 +1,173 @@
#!/usr/bin/env python3
"""Build a multimodal long-context research archive from arXiv.
For each paper: fetch metadata (Atom), download the PDF, extract the text layer
(pdftotext), inventory raster images (pdfimages -list), and hash the PDF.
Writes index.json (metadata + per-doc stats + corpus manifest).
"""
import hashlib
import json
import os
import subprocess
import sys
import time
import urllib.request
import xml.etree.ElementTree as ET
OUT = "/root/.hermes/wizard-outputs/sphinx/archive-longctx-retrieval-2026-10-26"
PDFDIR = os.path.join(OUT, "pdfs")
TXTDIR = os.path.join(OUT, "text")
QUERY = 'all:"long-context" AND all:retrieval'
N = 12
UA = "sphinx-archive/1.0 (research archive; contact sphinx@agents.invalid)"
NS = {"a": "http://www.w3.org/2005/Atom", "arxiv": "http://arxiv.org/schemas/atom"}
def fetch(url, dest=None, timeout=120):
req = urllib.request.Request(url, headers={"User-Agent": UA})
with urllib.request.urlopen(req, timeout=timeout) as r:
data = r.read()
if dest:
with open(dest, "wb") as f:
f.write(data)
return data
def sha256(path):
h = hashlib.sha256()
with open(path, "rb") as f:
for chunk in iter(lambda: f.read(1 << 20), b""):
h.update(chunk)
return h.hexdigest()
def run(cmd):
return subprocess.run(cmd, capture_output=True, text=True)
def parse_entries(xml_bytes):
root = ET.fromstring(xml_bytes)
total = root.find("opensearch:totalResults", {"opensearch": "http://a9.com/-/spec/opensearch/1.1/"})
entries = []
for e in root.findall("a:entry", NS):
aid = e.find("a:id", NS).text.strip()
pdf = None
for l in e.findall("a:link", NS):
if l.get("title") == "pdf":
pdf = l.get("href")
if pdf is None:
continue
entries.append({
"arxiv_id": aid.rsplit("/abs/", 1)[-1],
"title": " ".join(e.find("a:title", NS).text.split()),
"authors": [a.find("a:name", NS).text for a in e.findall("a:author", NS)],
"abstract": " ".join(e.find("a:summary", NS).text.split()),
"published": e.find("a:published", NS).text,
"updated": e.find("a:updated", NS).text,
"pdf_url": pdf,
})
return (int(total.text) if total is not None else -1), entries
def main():
os.makedirs(PDFDIR, exist_ok=True)
os.makedirs(TXTDIR, exist_ok=True)
api = ("https://export.arxiv.org/api/query?search_query="
+ urllib.parse.quote(QUERY)
+ "&start=0&max_results=%d&sortBy=relevance&sortOrder=descending" % N)
xml_bytes = fetch(api)
total_results, entries = parse_entries(xml_bytes)
print("arxiv total results for query: %d ; entries with pdf: %d" % (total_results, len(entries)))
docs = []
for i, ent in enumerate(entries, 1):
aid = ent["arxiv_id"]
pdf_path = os.path.join(PDFDIR, aid + ".pdf")
txt_path = os.path.join(TXTDIR, aid + ".txt")
ok = True
err = ""
for attempt in (1, 2):
try:
fetch(ent["pdf_url"], pdf_path)
break
except Exception as ex:
err = str(ex)
time.sleep(2 * attempt)
else:
ok = False
print("[%2d/%d] %-14s PDF FETCH FAILED: %s" % (i, len(entries), aid, err))
continue
info = run(["pdfinfo", pdf_path])
pages = -1
if info.returncode == 0:
for line in info.stdout.splitlines():
if line.startswith("Pages:"):
pages = int(line.split(":", 1)[1].strip())
run(["pdftotext", "-layout", pdf_path, txt_path])
words = 0
if os.path.exists(txt_path):
with open(txt_path, encoding="utf-8", errors="replace") as f:
words = len(f.read().split())
imgs = run(["pdfimages", "-list", pdf_path])
raster = 0
if imgs.returncode == 0:
for line in imgs.stdout.splitlines():
parts = line.split()
if len(parts) >= 3 and parts[2] == "image":
raster += 1
docs.append({
"arxiv_id": aid,
"title": ent["title"],
"authors": ent["authors"],
"abstract": ent["abstract"],
"published": ent["published"],
"updated": ent["updated"],
"pdf_url": ent["pdf_url"],
"pdf_path": os.path.relpath(pdf_path, OUT),
"text_path": os.path.relpath(txt_path, OUT),
"pdf_sha256": sha256(pdf_path),
"pdf_bytes": os.path.getsize(pdf_path),
"pages": pages,
"raster_images": raster,
"text_words": words,
})
print("[%2d/%d] %-14s pages=%-4d imgs=%-4d words=%-6d sha=%s" % (
i, len(entries), aid, pages, raster, words, docs[-1]["pdf_sha256"][:12]))
time.sleep(1)
index = {
"archive": "archive-longctx-retrieval-2026-10-26",
"version": 1,
"built_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
"query": QUERY,
"query_api": api,
"arxiv_total_results": total_results,
"documents_requested": N,
"documents_archived": len(docs),
"corpus": {
"total_pdf_bytes": sum(d["pdf_bytes"] for d in docs),
"total_pages": sum(d["pages"] for d in docs),
"total_raster_images": sum(d["raster_images"] for d in docs),
"total_text_words": sum(d["text_words"] for d in docs),
},
"documents": docs,
}
with open(os.path.join(OUT, "index.json"), "w") as f:
json.dump(index, f, indent=2, sort_keys=True)
with open(os.path.join(OUT, "index.json"), "rb") as f:
digest = hashlib.sha256(f.read()).hexdigest()
with open(os.path.join(OUT, "MANIFEST.sha256"), "w") as f:
f.write("%s index.json\n" % digest)
print("wrote index.json (%d docs) sha256=%s" % (len(docs), digest))
return 0 if len(docs) >= N else 1
if __name__ == "__main__":
sys.exit(main())
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/index.json
+291
View File
@@ -0,0 +1,291 @@
{
"archive": "archive-longctx-retrieval-2026-10-26",
"arxiv_total_results": 1034,
"built_utc": "2026-10-26T18:34:56Z",
"corpus": {
"total_pages": 203,
"total_pdf_bytes": 24473562,
"total_raster_images": 183,
"total_text_words": 114965
},
"documents": [
{
"abstract": "Understanding multimodal long-context documents that comprise multimodal chunks such as paragraphs, figures, and tables is challenging due to (1) cross-modal heterogeneity to localize relevant information across modalities, (2) cross-page reasoning to aggregate dispersed evidence across pages. To address these challenges, we are motivated to adopt a query-centric formulation that projects cross-modal and cross-page information into a unified query representation space, with queries acting as abstract semantic surrogates for heterogeneous multimodal content. In this paper, we propose a Multimodal Long-Context Document Retrieval Augmented Generation (MLDocRAG) framework that leverages a Multimodal Chunk-Query Graph (MCQG) to organize multimodal document content around semantically rich, answerable queries. MCQG is constructed via a multimodal document expansion process that generates fine-grained queries from heterogeneous document chunks and links them to their corresponding content across modalities and pages. This graph-based structure enables selective, query-centric retrieval and structured evidence aggregation, thereby enhancing grounding and coherence in multimodal long-context question answering. Experiments on datasets MMLongBench-Doc and LongDocURL demonstrate that MLDocRAG consistently improves retrieval quality and answer accuracy, demonstrating its effectiveness for multimodal long-context understanding.",
"arxiv_id": "2602.10271v3",
"authors": [
"Yongyue Zhang",
"Yaxiong Wu"
],
"pages": 15,
"pdf_bytes": 4948611,
"pdf_path": "pdfs/2602.10271v3.pdf",
"pdf_sha256": "890500145ba9a423ecb31dcbf84d425cbded6048cc4d617bb1463ac8ea36fb52",
"pdf_url": "https://arxiv.org/pdf/2602.10271v3",
"published": "2026-02-10T20:29:10Z",
"raster_images": 18,
"text_path": "text/2602.10271v3.txt",
"text_words": 11881,
"title": "MLDocRAG: Multimodal Long-Context Document Retrieval Augmented Generation",
"updated": "2026-02-13T03:25:41Z"
},
{
"abstract": "Long-context inference and Retrieval-Augmented Generation (RAG) handle evidence selection at vastly different scales, from a single long prompt to an entire corpus. We ask whether a single model-internal mechanism can select evidence across this range. We introduce UNifying REtrieval And Long-Context with a Single Model (UNREAL), a model-native evidence selection framework to span corpus retrieval and long-context inference. UNREAL encodes chunks and derives retrieval queries directly from the frozen LLM's internal representations. It adds fewer than 500K trainable parameters and leaves the backbone unchanged. On a 3B-token, 21M-chunk Wikipedia index, all four dense and hybrid UNREAL backbones outperform state-of-the-art retriever-reranker systems. The best model raises recall from 49.1% to 73.2% on HotpotQA, from 31.7% to 60.1% on 2WikiMultiHopQA, and from 8.8% to 14.4% on MuSiQue. Applied to long-context tasks, the same selection mechanism removes distractors before generation, raising NoLiMa accuracy from 1.0% to 24.83% at its maximum context length of 128K tokens, and LV-Eval's F1 score from 49.97% to 54.66% at 256K. UNREAL also reduces FLOPs and time-to-first-token relative to full-context inference from roughly 32K tokens onward, with larger gains as context grows. Together, these results establish model-internal evidence selection as a common foundation for corpus retrieval and evidence-sparse long-context inference.",
"arxiv_id": "2610.08463v1",
"authors": [
"Edan Kinderman",
"Elad Hoffer",
"Yochai Blau",
"Brian Chmiel",
"Ron Banner",
"Daniel Soudry",
"Boris Ginsburg"
],
"pages": 26,
"pdf_bytes": 1097365,
"pdf_path": "pdfs/2610.08463v1.pdf",
"pdf_sha256": "c75963a9db70fd2afbf0e7d6d6e06fb8cd9501db784e36185024b57c5cf4952b",
"pdf_url": "https://arxiv.org/pdf/2610.08463v1",
"published": "2026-10-06T14:48:53Z",
"raster_images": 7,
"text_path": "text/2610.08463v1.txt",
"text_words": 12613,
"title": "UNREAL: Unifying Retrieval and Long-Context with a Single Model",
"updated": "2026-10-06T14:48:53Z"
},
{
"abstract": "Retrieval pipelines-an integral component of many machine learning systems-perform poorly in domains where documents are long (e.g., 10K tokens or more) and where identifying the relevant document requires synthesizing information across the entire text. Developing long-context retrieval encoders suitable for these domains raises three challenges: (1) how to evaluate long-context retrieval performance, (2) how to pretrain a base language model to represent both short contexts (corresponding to queries) and long contexts (corresponding to documents), and (3) how to fine-tune this model for retrieval under the batch size limitations imposed by GPU memory constraints. To address these challenges, we first introduce LoCoV1, a novel 12 task benchmark constructed to measure long-context retrieval where chunking is not possible or not effective. We next present the M2-BERT retrieval encoder, an 80M parameter state-space encoder model built from the Monarch Mixer architecture, capable of scaling to documents up to 32K tokens long. We describe a pretraining data mixture which allows this encoder to process both short and long context sequences, and a finetuning approach that adapts this base model to retrieval with only single-sample batches. Finally, we validate the M2-BERT retrieval encoder on LoCoV1, finding that it outperforms competitive Transformer-based models by at least 23.3 points, despite containing upwards of 90x fewer parameters.",
"arxiv_id": "2402.07440v3",
"authors": [
"Jon Saad-Falcon",
"Daniel Y. Fu",
"Simran Arora",
"Neel Guha",
"Christopher R\u00e9"
],
"pages": 34,
"pdf_bytes": 3443083,
"pdf_path": "pdfs/2402.07440v3.pdf",
"pdf_sha256": "214e9c848e3aa062a51a9c697f4e6ef712b00e11d259a9ed0acb8ca46120d502",
"pdf_url": "https://arxiv.org/pdf/2402.07440v3",
"published": "2024-02-12T06:43:52Z",
"raster_images": 3,
"text_path": "text/2402.07440v3.txt",
"text_words": 12080,
"title": "Benchmarking and Building Long-Context Retrieval Models with LoCo and M2-BERT",
"updated": "2024-11-18T01:52:49Z"
},
{
"abstract": "Large Language Models (LLMs) have demonstrated impressive ability in generation and reasoning tasks but struggle with handling up-to-date knowledge, leading to inaccuracies or hallucinations. Retrieval-Augmented Generation (RAG) mitigates this by retrieving and incorporating external knowledge into input prompts. In particular, due to LLMs' context window limitations and long-context hallucinations, only the most relevant \"chunks\" are retrieved. However, current RAG systems face three key challenges: (1) chunks are often retrieved independently without considering their relationships, such as redundancy and ordering; (2) the utility of chunks is non-monotonic, as adding more chunks can degrade quality; and (3) retrieval strategies fail to adapt to the unique characteristics of different queries. To overcome these challenges, we design a cost-constrained retrieval optimization framework for RAG. We adopt a Monte Carlo Tree Search (MCTS) based strategy to find the optimal chunk combination order, which considers the chunks' correlations. In addition, to address the non-monotonicity of chunk utility, instead of treating budget exhaustion as the termination condition, we design a utility computation strategy to identify the optimal chunk combination without necessarily exhausting the budget. Furthermore, we propose a configuration agent that predicts optimal configurations for each query domain, improving our framework's adaptability and efficiency. Experimental results demonstrate up to a 30% improvement over baseline models, highlighting the framework's effectiveness, scalability, and suitability. Our source code has been released at https://github.com/wang0702/CARROT.",
"arxiv_id": "2411.00744v2",
"authors": [
"Ziting Wang",
"Haitao Yuan",
"Wei Dong",
"Gao Cong",
"Feifei Li"
],
"pages": 14,
"pdf_bytes": 1510082,
"pdf_path": "pdfs/2411.00744v2.pdf",
"pdf_sha256": "42b36d58d5e1b0a26014ca3a4032c75f4e6ea12aa5e6ceb7c1a68ca74b798325",
"pdf_url": "https://arxiv.org/pdf/2411.00744v2",
"published": "2024-11-01T17:11:16Z",
"raster_images": 20,
"text_path": "text/2411.00744v2.txt",
"text_words": 13464,
"title": "CARROT: A Learned Cost-Constrained Retrieval Optimization System for RAG",
"updated": "2026-03-12T11:55:58Z"
},
{
"abstract": "Effective long-term memory management is crucial for language models handling extended contexts. We introduce the Enhanced Ranked Memory Augmented Retrieval (ERMAR) framework, which dynamically ranks memory entries based on relevance. Unlike prior models, ERMAR employs a novel relevance scoring mechanism and a pointwise re-ranking model for key-value embeddings, inspired by learning-to-rank techniques in information retrieval. By integrating historical usage patterns and adaptive retrieval, ERMAR achieves state-of-the-art results on standard benchmarks, demonstrating superior scalability and performance in long-context tasks.",
"arxiv_id": "2503.14800v3",
"authors": [
"Ghadir Alselwi",
"Hao Xue",
"Shoaib Jameel",
"Basem Suleiman",
"Flora D. Salim",
"Imran Razzak"
],
"pages": 15,
"pdf_bytes": 1269291,
"pdf_path": "pdfs/2503.14800v3.pdf",
"pdf_sha256": "badc1996f829b43c15f47acd976cc4da58fa19315e5d21c3b00b4141e5a9c568",
"pdf_url": "https://arxiv.org/pdf/2503.14800v3",
"published": "2025-03-19T00:24:01Z",
"raster_images": 5,
"text_path": "text/2503.14800v3.txt",
"text_words": 8485,
"title": "Long Context Modeling with Ranked Memory-Augmented Retrieval",
"updated": "2026-05-18T08:00:08Z"
},
{
"abstract": "Recent advancements in Large Language Models (LLMs) have marked significant progress in understanding and responding to medical inquiries. However, their performance still falls short of the standards set by professional consultations. This paper introduces a novel framework for medical consultation, comprising two main modules: Terminology-Enhanced Information Retrieval (TEIR) and Emotional In-Context Learning (EICL). TEIR ensures implicit reasoning through the utilization of inductive knowledge and key terminology retrieval, overcoming the limitations of restricted domain knowledge in public databases. Additionally, this module features capabilities for processing long context. The EICL module aids in generating sentences with high attribute relevance by memorizing semantic and attribute information from unlabelled corpora and applying controlled retrieval for the required information. Furthermore, a dataset comprising 803,564 consultation records was compiled in China, significantly enhancing the model's capability for complex dialogues and proactive inquiry initiation. Comprehensive experiments demonstrate the proposed method's effectiveness in extending the context window length of existing LLMs. The experimental outcomes and extensive data validate the framework's superiority over five baseline models in terms of BLEU and ROUGE performance metrics, with substantial leads in certain capabilities. Notably, ablation studies confirm the significance of the TEIR and EICL components. In addition, our new framework has the potential to significantly improve patient satisfaction in real clinical consulting situations.",
"arxiv_id": "2503.17876v1",
"authors": [
"Kaiwen Zuo",
"Jing Tang",
"Hanbing Qin",
"Binli Luo",
"Ligang He",
"Shiyan Tang"
],
"pages": 16,
"pdf_bytes": 5842685,
"pdf_path": "pdfs/2503.17876v1.pdf",
"pdf_sha256": "ac2307d34388bd9246c0b3e6820348f199bf62db4215b7b7d3928312c5cff697",
"pdf_url": "https://arxiv.org/pdf/2503.17876v1",
"published": "2025-03-22T23:01:07Z",
"raster_images": 2,
"text_path": "text/2503.17876v1.txt",
"text_words": 5943,
"title": "Satisfactory Medical Consultation based on Terminology-Enhanced Information Retrieval and Emotional In-Context Learning",
"updated": "2025-03-22T23:01:07Z"
},
{
"abstract": "Long Context Language Models (LCLMs) have emerged as a new paradigm to perform Information Retrieval (IR), which enables the direct ingestion and retrieval of information by processing an entire corpus in their single context, showcasing the potential to surpass traditional sparse and dense retrieval methods. However, processing a large number of passages within in-context for retrieval is computationally expensive, and handling their representations during inference further exacerbates the processing time; thus, we aim to make LCLM retrieval more efficient and potentially more effective with passage compression. Specifically, we propose a new compression approach tailored for LCLM retrieval, which is trained to maximize the retrieval performance while minimizing the length of the compressed passages. To accomplish this, we generate the synthetic data, where compressed passages are automatically created and labeled as chosen or rejected according to their retrieval success for a given query, and we train the proposed Compression model for Long context Retrieval (CoLoR) with this data via preference optimization while adding the length regularization loss on top of it to enforce brevity. Through extensive experiments on 9 datasets, we show that CoLoR improves the retrieval performance by 6% while compressing the in-context size by a factor of 1.91. Our code is available at: https://github.com/going-doer/CoLoR.",
"arxiv_id": "2412.18232v2",
"authors": [
"Minju Seo",
"Jinheon Baek",
"Seongyun Lee",
"Sung Ju Hwang"
],
"pages": 18,
"pdf_bytes": 3483911,
"pdf_path": "pdfs/2412.18232v2.pdf",
"pdf_sha256": "2ad2603c6b95d55240754716f21c9b4610e9765af3f1b9e27e7471b87775f1bb",
"pdf_url": "https://arxiv.org/pdf/2412.18232v2",
"published": "2024-12-24T07:30:55Z",
"raster_images": 123,
"text_path": "text/2412.18232v2.txt",
"text_words": 12446,
"title": "Efficient Long Context Language Model Retrieval with Compression",
"updated": "2025-05-28T12:52:57Z"
},
{
"abstract": "Recent advancements in long-context language models (LCLMs) promise to transform Retrieval-Augmented Generation (RAG) by simplifying pipelines. With their expanded context windows, LCLMs can process entire knowledge bases and perform retrieval and reasoning directly -- a capability we define as In-Context Retrieval and Reasoning (ICR^2). However, existing benchmarks like LOFT often overestimate LCLM performance by providing overly simplified contexts. To address this, we introduce ICR^2, a benchmark that evaluates LCLMs in more realistic scenarios by including confounding passages retrieved with strong retrievers. We then propose three methods to enhance LCLM performance: (1) retrieve-then-generate fine-tuning, (2) retrieval-attention-probing, which uses attention heads to filter and de-noise long contexts during decoding, and (3) joint retrieval head training alongside the generation head. Our evaluation of five well-known LCLMs on LOFT and ICR^2 demonstrates significant gains with our best approach applied to Mistral-7B: +17 and +15 points by Exact Match on LOFT, and +13 and +2 points on ICR^2, compared to vanilla RAG and supervised fine-tuning, respectively. It even outperforms GPT-4-Turbo on most tasks despite being a much smaller model.",
"arxiv_id": "2501.08248v3",
"authors": [
"Yifu Qiu",
"Varun Embar",
"Yizhe Zhang",
"Navdeep Jaitly",
"Shay B. Cohen",
"Benjamin Han"
],
"pages": 17,
"pdf_bytes": 731879,
"pdf_path": "pdfs/2501.08248v3.pdf",
"pdf_sha256": "533189f5421edd53b77daf8f5ea31366a9ca0d4050368349c699efaa3010ecae",
"pdf_url": "https://arxiv.org/pdf/2501.08248v3",
"published": "2025-01-14T16:38:33Z",
"raster_images": 0,
"text_path": "text/2501.08248v3.txt",
"text_words": 10392,
"title": "Eliciting In-context Retrieval and Reasoning for Long-context Large Language Models",
"updated": "2025-06-09T07:37:32Z"
},
{
"abstract": "Time Series Foundation Models (TSFMs) have borrowed the long context paradigm from natural language processing under the premise that feeding more history into the model improves forecast quality. But in stochastic domains, distant history is often just high-frequency noise, not signal. Hence, the proposed work tests whether this premise actually holds by running continuous context architectures (PatchTST included) through the ETTh1 benchmark. The obtained results contradict the premise: an inverse scaling law shows up clearly, with forecasting error rising as context gets longer. A 3,000-step window causes performance to drop by over 68%, evidence that attention mechanisms are poor at ignoring irrelevant historical volatility. Retrieval-Augmented Forecasting (RAFT) is evaluated as an alternative. RAFT achieves a mean squared error (MSE) of 0.379 with a fixed 720-step window and selective retrieval, outperforming both long-context configurations and zero-shot foundation models (Chronos, Moirai) despite requiring far less computation. In addition, the retrieval step injects only the most relevant historical segments as dynamic exogenous variables, which gives the model a context-informed inductive bias it cannot build on its own from raw sequences. Therefore, foundation models going forward need to shift architecturally toward selective retrieval.",
"arxiv_id": "2605.08217v1",
"authors": [
"Rishi Ahuja",
"Kumar Prateek",
"Simranjit Singh",
"Vijay Kumar"
],
"pages": 10,
"pdf_bytes": 292570,
"pdf_path": "pdfs/2605.08217v1.pdf",
"pdf_sha256": "f8b2b0e1dc1ed119c721b9edd74b81b1d559a4ff4742b5b86a62f8a59c72b487",
"pdf_url": "https://arxiv.org/pdf/2605.08217v1",
"published": "2026-05-06T12:15:05Z",
"raster_images": 0,
"text_path": "text/2605.08217v1.txt",
"text_words": 4232,
"title": "Retrieval Mechanisms Surpass Long-Context Scaling in Time Series Forecasting",
"updated": "2026-05-06T12:15:05Z"
},
{
"abstract": "Supporting longer context for Large Language Models (LLM) is a promising direction to advance LLMs. As training a model for a longer context window is computationally expensive, many alternative solutions, such as Retrieval Augmented Generation (RAG), have been used. However, most existing RAG methods adopt embedding-based retrieval that falls short on long contexts. To address such challenges, we propose an attention-based retrieval technique, You Only Use Reactive Attention slice (YOURA). YOURA leverages a novel retrieval heuristic called reaction score to rank the relevance of each sentence in the input context with the query sentence. Intuitively, we measure how the per-token attention score \"reacts\" to the query and greedily retrieves the most reactive sentences. Internally, YOURA generates a token-indexed vector (called reaction vector) for the whole input context. To map each sentence to the token-indexed vector, we propose an Embedding-Agnostic Sentence Yield (EASY), a best-effort token wiggling algorithm. We evaluate our retrieval technique on three open-source pre-trained LLM models across six LongBench QA datasets. Our technique achieves up to 30% vLLM inference throughput improvement for serving long-context queries with a nearly identical quality score to the simple yet effective truncate-middle approach.",
"arxiv_id": "2409.13695v1",
"authors": [
"Yun Joon Soh",
"Hanxian Huang",
"Yuandong Tian",
"Jishen Zhao"
],
"pages": 12,
"pdf_bytes": 476786,
"pdf_path": "pdfs/2409.13695v1.pdf",
"pdf_sha256": "a0fac2a32a6d6d384c15353fbf8e5e601002bf491ff0e725e00199abcdaf6a26",
"pdf_url": "https://arxiv.org/pdf/2409.13695v1",
"published": "2024-09-03T15:30:57Z",
"raster_images": 0,
"text_path": "text/2409.13695v1.txt",
"text_words": 6679,
"title": "You Only Use Reactive Attention Slice For Long Context Retrieval",
"updated": "2024-09-03T15:30:57Z"
},
{
"abstract": "We present systematic efforts in building long-context multilingual text representation model (TRM) and reranker from scratch for text retrieval. We first introduce a text encoder (base size) enhanced with RoPE and unpadding, pre-trained in a native 8192-token context (longer than 512 of previous multilingual encoders). Then we construct a hybrid TRM and a cross-encoder reranker by contrastive learning. Evaluations show that our text encoder outperforms the same-sized previous state-of-the-art XLM-R. Meanwhile, our TRM and reranker match the performance of large-sized state-of-the-art BGE-M3 models and achieve better results on long-context retrieval benchmarks. Further analysis demonstrate that our proposed models exhibit higher efficiency during both training and inference. We believe their efficiency and effectiveness could benefit various researches and industrial applications.",
"arxiv_id": "2407.19669v2",
"authors": [
"Xin Zhang",
"Yanzhao Zhang",
"Dingkun Long",
"Wen Xie",
"Ziqi Dai",
"Jialong Tang",
"Huan Lin",
"Baosong Yang",
"Pengjun Xie",
"Fei Huang",
"Meishan Zhang",
"Wenjie Li",
"Min Zhang"
],
"pages": 20,
"pdf_bytes": 667716,
"pdf_path": "pdfs/2407.19669v2.pdf",
"pdf_sha256": "ffcd09d1e626fe372847f18c9f51b5b5c521efffbabfa3983807cfcfb24a634f",
"pdf_url": "https://arxiv.org/pdf/2407.19669v2",
"published": "2024-07-29T03:12:28Z",
"raster_images": 2,
"text_path": "text/2407.19669v2.txt",
"text_words": 13318,
"title": "mGTE: Generalized Long-Context Text Representation and Reranking Models for Multilingual Text Retrieval",
"updated": "2024-10-14T12:19:44Z"
},
{
"abstract": "Long-context retrieval exposes a tension: single-vector embeddings lose fine-grained detail, while token-level multi-vector methods incur prohibitive storage. We propose Multi-Prefix Embedding (MPE), which partitions a document into chunks separated by EOS tokens, encodes the full sequence in a single causal forward pass, and extracts one embedding at each prefix boundary. MPE retains cross-chunk context, enables chunk-level MaxSim matching, and trains with only document-level relevance labels. Experiments on MLDR-en, BrowseComp-Plus, and LongEmbed show that MPE is competitive with or outperforms single-vector, independent-chunk, and multi-vector baselines, while providing a natural source attribution mechanism for locating evidence chunks.",
"arxiv_id": "2606.23642v1",
"authors": [
"Zhenglin Yu",
"Xueguang Ma",
"Shengyao Zhuang",
"Zhichao Xu",
"Luyu Gao",
"Crystina Zhang",
"Jimmy Lin"
],
"pages": 6,
"pdf_bytes": 709583,
"pdf_path": "pdfs/2606.23642v1.pdf",
"pdf_sha256": "df7f59fdfecfc88acb530af64d1b6574d2e08536f66e1ab7b0f5ce6cad6e97b9",
"pdf_url": "https://arxiv.org/pdf/2606.23642v1",
"published": "2026-06-22T17:31:02Z",
"raster_images": 3,
"text_path": "text/2606.23642v1.txt",
"text_words": 3432,
"title": "Improving Long-Context Retrieval with Multi-Prefix Embedding",
"updated": "2026-06-22T17:31:02Z"
}
],
"documents_archived": 12,
"documents_requested": 12,
"query": "all:\"long-context\" AND all:retrieval",
"query_api": "https://export.arxiv.org/api/query?search_query=all%3A%22long-context%22%20AND%20all%3Aretrieval&start=0&max_results=12&sortBy=relevance&sortOrder=descending",
"version": 1
}
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/pdfs/2402.07440v3.pdf
BIN
View File
Binary file not shown.
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/pdfs/2407.19669v2.pdf
BIN
View File
Binary file not shown.
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/pdfs/2409.13695v1.pdf
BIN
View File
Binary file not shown.
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/pdfs/2411.00744v2.pdf
BIN
View File
Binary file not shown.
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/pdfs/2412.18232v2.pdf
BIN
View File
Binary file not shown.
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/pdfs/2501.08248v3.pdf
+8576
View File
File diff suppressed because it is too large. Load diff
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/pdfs/2503.14800v3.pdf
BIN
View File
Binary file not shown.
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/pdfs/2503.17876v1.pdf
BIN
View File
Binary file not shown.
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/pdfs/2602.10271v3.pdf
+20545
View File
File diff suppressed because it is too large. Load diff
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/pdfs/2605.08217v1.pdf
BIN
View File
Binary file not shown.
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/pdfs/2606.23642v1.pdf
+4875
View File
File diff suppressed because it is too large. Load diff
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/pdfs/2610.08463v1.pdf
BIN
View File
Binary file not shown.
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/text/2402.07440v3.txt
+1741
View File
File diff suppressed because it is too large. Load diff
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/text/2407.19669v2.txt
+1393
View File
File diff suppressed because it is too large. Load diff
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/text/2409.13695v1.txt
+780
View File
@@ -0,0 +1,780 @@
Y OU O NLY U SE R EACTIVE A TTENTION S LICE
FOR L ONG C ONTEXT R ETRIEVAL
Yun Joon Soh 1 Hanxian Huang 1 Yuandong Tian 2 Jishen Zhao 1
A BSTRACT
Supporting longer context for Large Language Models (LLM) is a promising direction to advance LLMs. As
arXiv:2409.13695v1 [cs.CL] 3 Sep 2024
training a model for a longer context window is computationally expensive, many alternative solutions, such
as Retrieval Augmented Generation (RAG), have been used. However, most existing RAG methods adopt
embedding-based retrieval that falls short on long contexts.
To address such challenges, we propose an attention-based retrieval technique, You Only Use Reactive Attention
slice (YOURA). YOURA leverages a novel retrieval heuristic called reaction score to rank the relevance of each
sentence in the input context with the query sentence. Intuitively, we measure how the per-token attention score
“reacts” to the query and greedily retrieves the most reactive sentences. Internally, YOURA generates a token-
indexed vector (called reaction vector) for the whole input context. To map each sentence to the token-indexed
vector, we propose an Embedding-Agnostic Sentence Yield (EASY), a best-effort token wiggling algorithm.
We evaluate our retrieval technique on three open-source pre-trained LLM models across six LongBench QA
datasets. Our technique achieves up to 30% vLLM inference throughput improvement for serving long-context
queries — with a nearly identical quality score to the simple yet effective truncate-middle approach (Bai et al.,
2023).
1 I NTRODUCTION a. Input Context & Query c. Our proposal vs. embeddig-distance based
C1: A woman is wearing a red dress. Current Our Proposal
Pre-trained Large Language Models (LLM) are widely ex- S2: She is walking on t he st reet .
Emb. Distance Reaction Score
plored for various Natural Language Processing (NLP) tasks. S3: Her name is Mary.
S1 S2 S3 S1 S2 S3
To enhance the model for more complex tasks, researchers
Q: Where is Mary?
seek ways to extend the LLM’s context window size via Retrieve Retrieve
Semantically Close Most reactive
fine-tuning the pre-trained model (Chen et al., 2023b; Man- b. Performance Teaser
grulkar et al., 2022), proposing advanced attention mecha- Augmented Augmented
nisms (Xiao et al., 2024; Zhang et al., 2023), and improving Quality Her name is Mary.
Where is Mary?
She is walking on t he st reet .
Where is Mary?
+3% / +9% / +11%
Retrieval Augmented Generation (RAG). However, fine- Llama2 Llama3 Mistral
tuning requires additional computational power for each Generated Generated
Inference Mar y i s i n t he
pre-trained model and is often targeted for specific tasks Throughput l i v i ng r oom.
On t he St r eet
or datasets. Advanced attention mechanisms conceptually +2% / 13% / 31%
Llama2 Llama3 Mistral
Context unaware,
generates incorrect
Correct response
with a better
extend the context window by selectively choosing a subse- responses. context-awareness
quence of tokens for attention but are susceptible to details
due to the nature of dropping information at the token level.
RAG splits the contextual text into smaller chunks and re- Figure 1. YOURA improves the retrieval quality and inference
trieves a subset of chunks before inferencing based on the throughput by retrieving only the “reactive” sentences.
semantical distance to the query in the embedding vector
1
Department of Computer Science and Engineering, Univer-
sity of California, San Diego, San Diego, USA 2 Meta AI (FAIR), space. RAG’s retrieval nature has many benefits; it reduces
Menlo Park, California, USA. Meta-affiliated author served in
an advisory role and all data processing occurred at the Univer-
the input size to process, which reduces the computation dur-
sity of California, San Diego. Correspondence to: Jishen Zhao ing inference (Luo et al., 2024) and improves the accuracy
<jzhao@ucsd.edu>. by eliminating distracting information (Catav, 2023).
A high-quality retrieval is crucial for RAG but often suf-
fers from lower accuracy because (1) the common words
in both the query and text chunk put the semantic distance returns token indices for each sentence.
closer and (2) the semantically identical, but alphabetically
We evaluate (1) how much YOURA improves the answer
different words distance the query and text chunk in the
quality for the 6 LongBench QA datasets (Bai et al., 2023),
embedding vector space. For example, the similarity dis-
(2) how YOURA reduces the inferencing computation cost
tance between a query (“Where is Mary?”) and a sentence
by retrieving fewer, yet sufficient tokens for proper question-
(“Her name is Mary”) may be closer to a sentence (“She is
answering, and (3) how well our EASY algorithm maps
at home”) resulting in suboptimal retrieval. 1
raw texts to the token-indices. Without any pre-training or
We propose a fine-tuning-free, attention-based retrieval strat- additional embeddings, YOURA shows a higher retrieval
egy, which can be applied to various off-the-shelf pre-trained ratio (context token count / retrieved token count) for open-
LLMs to achieve improved text generation quality. Our key source pre-trained models (Llama2, Llama3, Mistral) by up
insight is that the attention in a pre-trained model already to ×1.3 and still achieves a higher answer quality than prior
contains information about which tokens are important for embedding-based retrieval approaches. With the smaller
the given query. Figure 2 illustrates the difference between number of retrieved tokens, YOURA improves up to 30%
the embedding-based retrieval and our proposal; (a) is a typi- vLLM serving throughput. Our EASY algorithm shows an
cal embedding-based retrieval where embedding vectors for average of 93.7% sentence-to-token-index mapping accu-
both the query and context chunk list are generated through racy.
an embedding layer. Figure 2 (b) is another simple yet effec-
We contribute the following.
tive alternative to the embedding distance-based approach,
truncate-middle. Since the introduction and conclusion of • We propose You Only Use the Reactive Attention
most documents contain the most important information, slice (YOURA) algorithm, a fine-tuning-free, attention-
the truncation approach naively takes out the middle of the based retrieval algorithm. YOURA improves QA qual-
context so that the remaining text fits within the LLM’s ity by up to 10% while improving inference throughput
context window size. Unlike embedding-based retrieval or by up to 30%2 .
truncation, we propose an attention-based retrieval.
• We propose an Embedding-Agnostic Sentence Yield
Our proposed technique, You Only Use Reactive Attention (EASY) algorithm, which can map individual sen-
slice (YOURA) Figure 2 (c), is a novel attention-based tences to their vector slice at 93% accuracy.
retrieval technique. Anonlogically speaking, we use the
attention vector generated without query as the baseline • We provide a detailed analysis of YOURA’s QA quality
and measure how each token “reacts” to the query. We improvement on LongBench QA datasets.
refer to the “reactiveness” as Reaction Vector, an absolute
difference in attention scores calculated with and without the 2 R ELATED W ORK
query. The per-sentence average of the reaction vector called
reaction score, is a good ranking metric when retrieving Long-context Language Models. Training LLMs on se-
relevant sentences, as we show in our evaluation. quences with a maximum length while still ensuring they
infer well on sequences longer than the maximum is chal-
Properly slicing the reaction vector into sentences is cru-
lenging. Previous studies proposed techniques such as posi-
cial because YOURA splits the token sequence into chunks
tional interpolation (Chen et al., 2023a; Li et al., 2023), posi-
(as opposed to prior works, which split the raw text into
tional extrapolation (Press et al., 2021), external emory (Wu
chunks). Such a difference sparks a new challenge: map-
et al., 2020; Martins et al., 2021), memory-retrieval augmen-
ping sentences to a token sequence. Merely associating
tation strategies (Mohtashami & Jaggi, 2023; Wang et al.,
specific token values (e.g., punctuation token) to a sentence
2024) to push the limit of the context window size to 2
boundary is insufficient because different tokenizer uses
million and even longer (Ding et al., 2024; Bertsch et al.,
different vocabulary. For example, a tokenizer may encode
2024). Although these works improved model accuracy on
.The (a sentence boundary with a missing white space) as
various tasks, the large context window size implies more
two tokens: [ .T , he ]. In such case, .T token becomes computational and memory costs.
part of the two consecutive sentences. Treating the token be-
fore or after .T token as a sentence boundary would easily Retrieval-Augmented Generation. Retrieval-
garble the remaining sentence splits. For this challenge, we Augmented Generation (RAG) retrieves relevant document
propose the Embedding-Agnostic Sentence Yield (EASY) chunks from the external knowledge base. Prior works
— a best-effort token wiggling algorithm — that takes a list integrate RAG with language models for question an-
of raw sentences and a single token sequence as input and swering with long documents (Stelmakh et al., 2022)
1
Measured with Sentence-Transformer “All-MiniLM-L6-V2”, 2
The code has been open-sourced and is available at
scipy.spatial.distance. https://github.com/yjsoh/youra.
2
Context Query Context Query Context Query
Chunk ... Middle Truncation Transformer
Embedding Reaction Vec. Emb.
Head Truncated Tail Agnostic
Sentence
EmbChunk ... EmbQ
Chunks ... Yield
Less
important
Top-K w/ Cos. Dist. Top-K w/ React. Score
LLM Context Window Size LLM Context Window Size LLM Context Window Size
a. Embedding-based Retrieval b. Truncate Middle Approach c. Our Proposal: YOURA
Figure 2. Comparison of a typical embedding-based retriever (cosine similarity distance), truncation-based approach, and our attention-
based retriever (reaction score). In our evaluation, we show that reaction score is a better retrieval heuristic in terms of quality and
performance. Unlike the embedding-based approach, which splits the raw string context into chunks (e.g., sentences) before retrieval,
YOURA splits the token-indexed vector, reaction-vector in (c), into chunks. Such an approach requires mapping the raw sentences to
token indices. We propose Embedding Agnostic Sentence Yield (EASY) algorithm for the challenge.
and in open-domain (Giorgi et al., 2022). Furthermore, served that the small model is capable of identifying the
language models adopted retrieval at various stages such key relevant information, and devised a way to mitigate the
as pretraining (Wang et al., 2023; Izacard et al., 2023; “positional bias” due to the positional embeddings by shuf-
Borgeaud et al., 2022), fine-tuning (Jiang et al., 2022; fling the long context. Unlike their approach, which relies
Zhang et al., 2024) and inference stage (Khandelwal et al., on a relatively complex perplexity calculation, our approach
2019). merely simply takes the mean within the attention matrix.
Attention Mechanism. Researchers proposed self- 3 YOU O NLY U SE R EACTIVE ATTENTION
attention alternatives to mitigate its quadratic complexity, (YOURA)
which becomes a computational bottleneck for a long
context. These include sparse attention mechanisms with Limitation of Embedding Vectors — Context-Unaware
pre-defined sparsity patterns (Zhu et al., 2021; Liu et al., The limitation of embedding-based retrieval is that the input
2023), recurrence-based method (Bulatov et al., 2022), texts are the sole variable when generating the embedding
low-rank projection attention (Zhao et al., 2024), and vectors. For example, consider the following two paragraphs
memory-based mechanisms (Lou et al., 2024). However, each with two sentences: “Her name is Mary. She is at
these approximate methods introduce inductive bias (e.g., home.” and ”Her name is Alice. She is at home.”. In both
pre=defined sparsity) that can fit well for specific domains, cases, ”She is at home” is embedded into the identical vector,
but may reduce model quality in general LLM training. resulting in equal distance to any sentence.
Fine-tuning for Long Context. Recent works (Ding et al.,
2024; Chen et al., 2023b; Peng et al., 2023) show that a
pre-trained LLM context window can be extended by fine- Attention-Based Retrieval — Context-Aware We pro-
tuning on longer texts. However, fine-tuning approaches pose an attention-based heuristic for context-aware retrieval.
still require multiple self-attention computations demanding At a high level, our proposal is measuring how the query sen-
both high-quality training datasets and computational costs. tence impacts the language model’s attention distribution.
Our approach does not require fine-tuning. We observe that attention values of relevant information
change more drastically than irrelevant information. From
Context Compression. LongLLMLingua (Jiang et al., this observation, we introduce reaction vector to quantify
2023) is one of the closest approaches to ours; they ob- the attention shift caused by the query sentence.
3
Design Overview Figure 3 illustrates the overview of our Reaction Score Given the reaction vector, we define the
proposed design. The left portion of the figure depicts how reaction score (rs) as follows for a vector slice represented
a reaction vector is calculated and the right portion depicts with left-right index, [l, r):
how a reaction vector is used for the retrieval and how the
language model leverages the retrieved sentences to answer ReactionScore(ReactionVector, l, r)
the query.
= avg(ReactionVector[l : r]) (6)
3.1 Problem Statement & Definitions
Annotations We annotate ZC ∈ Rd×c , ZQ ∈ Rd×q as Problem Statement We define the retrieval task as finding
the continuous representation of the context and question, a set of non-overlapping chunks S = {(li , ri ) | 0 ≤ li ≤
respectively, and d is the pre-trained model’s hidden dimen- ri ≤ |Zctxt |, ri ≤ li+1 }, such that
sion, c is the context token count and q is the question token
count.
argmaxS AnsQuality(LLMInference(concat(S , q))) (7)
Reaction Vector We define the reaction vector as follows:
where:
C Q
ReactionVector(Z , Z )
= abs AttnVec(ZC ) − AttnVec(concat(ZC , ZQ ))1:c

(1) • concat(S , q) denotes the concatenation of the chunks
in S with the query q.
where
• LLMInference is the function that processes the con-
AttnVec(Z) = Meancol (AttnMatrix(Q, K)) (2) catenated chunks and the query using a language
model.
AttnMatrix is the dense attention function before multiply-
ing the value projection and Q, K are the query and key • AnsQuality measures the quality (e.g., F1 Score) of
projection of Z, respectively, similar to the description in the answer produced by the inference function.
the “Attention is all you need” paper (Vaswani et al., 2023):
QK T 3.2 Embedding-Agnostic Sentence Yield (EASY)
 
AttnMatrix(Q, K) = softmax √ Algorithm
d
(3)
Q = Z × WQ Challenge Given a token sequence, finding the token in-
K =Z×W K dex of sentence boundary is non-trivial for two reasons: (1)
the embedding model for sentence split task (e.g., stanza) is
different from the model’s embedding, and (2) embedding
Meancol takes the attention matrix and returns a per-column
models with huge dictionary may result in an impossible
mean vector (note the subscription 1 : c to indicate trunca-
token sequence to achieve a perfect sentence split. For ex-
tion from 1 to c-th column):
ample, widely known sentence-splitting algorithms (e.g.,
Stanza (Qi et al., 2020)) use a different dictionary from the
d d d
!
1 X X X
Meancol ([aij ]d×n )1:c = ai1 , ai2 , ..., aic LLM models (e.g., Llama3-8B-Instruct). Thus, the total
d i=1 i=1 i=1 token count varies across the embedding models (Table 1).
d
!c
1 X Another challenge is that some models prefer huge dictio-
= aij
d i=1 naries resulting in a token sequence with an impossible case
j=1
(4) to articulate a sentence boundary. For example, llama-3
treats .T as a single token. Although it may be useful in
For multi-head attention, the reaction vector is a mean across code generation, a missing whitespace between a punctua-
heads: tion period and a subsequent sentence that starts with the
H d
!c word “The” would result in an impossible case. In other
h
 1 X 1X h
Meancol,heads [aij ]d×n 1:c = a words, Stanza would create a token sequence ( . - The ),
H d i=1 ij
h=1 j=1 whereas Llama3 tokenizer would output ( .T - he ). In
H d
!c such a case, an exact sentence boundary for Llama3 cannot
1 X 1 X
= ah be expressed as a token index. For these reasons, finding
H d i=1 ij the sentence boundary token index across the embedding
h=1 j=1
(5) model is non-trivial.
4
Our Proposal - You Only Use Reactive Attention (YOURA) Section 3.3 Embedding-Agnostic Sentence Yield (EASY)
: Maps raw sentence text to index of token sequence
Context Only Attention Matrix Context + Query Attention Matrix Retrieval
I Token Reaction Score
Sentence
am Index : mean of reaction vector slice
(higher = relevant to query)
an I am an
amazing [ 0, 7 ) Mean ( )
amazing
researcher researcher. Reaction Ignore
.
Vector Slice
I I like
like [ 7, 11 ) Mean ( )
LLM.
LLM
. Retrieved! (? higher reaction score)
Who
am Augmented
I Query: "Who am I?"
? Prompt: "Answer the question. {Retrieved Context} {Q}"
Augmented Prompt
Reaction Vector Generation
Ignore
= abs ( - ) Language Models
Llama2, Llama3, Mistral
Answer:
Amazing researcher!
Attention Vec w/o Query Attention Vec w/ Query
= Meancol (Attention Matrix) = Meancol (Attention Matrix)
Figure 3. Overview of YOURA and where it is used in the Retrieval Augmented Generation (RAG) with an example (example context: ”I
am an amazing researcher. I like LLM.”, example query: ”Who am I?”). The first step is calculating the reaction vector, an absolute
difference between the attention vector with and without the query (left side of the figure). The highlighted cells in the attention matrix
indicate that the token pair exhibits a relatively high value (e.g., Who vs. I). Once the reaction vector has been calculated, each sentence is
assigned a reaction score, the mean of a corresponding reaction vector slice. To map each sentence to a token sequence, we propose
Embedding-Agnostic Sentence Yield (EASY) algorithm (Section 3.2). The retriever passes on the sentences with high reaction scores to
the augmentor. The pre-trained LLM models generate answers using the augmented text, which includes the task-specific prompt, the
retrieved context, and the question.
researcher
like
LLM
amazing
I
am
.
I
.
an
Embedding-Agnostic Sentence Yield (EASY) To re- 3.3 Retrieval using Reaction Score — End-to-end
solve the problem, we propose the Embedding-Agnostic
The retrieval algorithm starts with calculating the reaction
researcher
like
LLM
amazing
?
Who
I
am
.
I
.
am
I
an
Sentence Yield (EASY) algorithm. The input to the algo-
vector. If the input context is longer than the model’s context
rithm is a single token sequence and a list of sentence strings
window size, we calculate the attention vector for each
acquired from the LLM tokenizer and the sentence splitting
chunk of context window size and concatenate them before
NLP models (e.g., Stanza (Qi et al., 2020)), respectively.
returning. At this point, we have the reaction vector whose
The EASY algorithm is a repetition of guessing the sen- length equals the input context’s token sequence.
tence boundary and making best efforts adjustments. First,
To slice the reaction vector for each sentence, we apply
EASY estimates a candidate index as the encoded length of
the Embedding-Agnostic Sentence Yield (EASY) algorithm
sentences that are already processed (Algorithm 1 line 7).
(Section 3.2). Then the retriever slices the reaction vector
EASY decodes the token subsequence between the most-
and calculates the geometric mean of the reaction vector
recently confirmed boundary index (m in Algorithm 1)
slice using the output of the EASY algorithm, a list of
and the current candidate and compares it against the
sentence boundary indices. At this point, we have a list of
target sentence (Algorithm 1 line 16). Based on the compar-
sentences and their associated reaction scores, which will
ison outcome (MATCH, INCLUDED, NO-MATCH), EASY
be used as a retrieval heuristic.
wiggles (+1 or -1) the candidate index until the decoded
subsequence matches the target or a termination condition For the actual retrieval, we greedily add sentences starting
is met (lines 17-23). Lastly, EASY returns the list of inte- with the highest reaction score until the retrieval token bud-
gers, where each value is a token index corresponding to a get is depleted or the sentence count threshold is reached
sentence boundary. If for some reason, the best effort fails, (80% of total sentence count). The retrieved sentences are
EASY restores the initial candidate index and appends to the reordered according to their original position as the final
boundary list. This retains the one-to-one mapping between cleanup step.
the target sentence and the boundary token index.
5
Dataset ID Source Avg # Tokens Metric # Data
Llama2 Llama3 Mistral
Single-Document QA
NarrativeQA 1-1 Literature, Film 35937 29777 35937 F1 200
Qasper 1-2 Science 5603 4922 5517 F1 200
MultiFieldQA-en 1-3 Multi-field 8046 6889 7877 F1 150
Multi-Document QA
HotpotQA 2-1 Wikipedia 15244 12780 14890 F1 200
2WikiMultihopQA 2-2 Wikipedia 8381 7096 8281 F1 200
MuSiQue 2-3 Wikipedia 18459 15544 18069 F1 200
Table 1. Evaluated dataset details.
4 E XPERIMENTS chunks. We measure the cosine similarity between
each chunk’s embedding vector generated by the
We evaluate how the YOURA improves the answering per- “paraphrase-MiniLM-L6-v2” model3 .
formance of 5 pre-trained models on LongBench (Bai et al.,
2023) single-doc and multi-doc QA datasets. • with Truncation We use LongBench’s prompt along
with its context truncation approach: the whole context,
Datasets We evaluate YOURA on the single and multi- if it fits within the pre-trained model’s context window
document QA (Table 1). size, or the concatenation of the context head and tail
with an equal context window budget, otherwise.
Models We used three different open-sourced LLM mod-
els: LLama2-7B, Llama3-8B, and Mistralv0.2-7B, where • with YOURA Our proposed system as described in
each model has a context window size of 4K, 8K, and 32K, Section 3.
respectively.
4.1 Longbench QA Result
Inference Setup For inference, we used the vLLM (Kwon
et al., 2023). We used the LongBench (Bai et al., 2023) code How does YOURA impact the output quality? From Ta-
to evaluate the generated answers. ble 2, which shows the answer quality using three different
open-sourced models (Llama2, Llama3, and Mistralv0.2),
Machine Setup We measured the vLLM (Kwon et al., we make the following observations. (1) We observe that
2023) inference performance by slightly modifying the YOURA showed better overall answer quality than all of the
vLLM’s throughput benchmark (v0.5.3) code on a single retrieval approaches (BM25 and embedding-based, except
NVIDIA H100-80G GPU. We used the default vLLM set- for Llama2 w/ B50 70) at a higher retrieval ratio — i.e.,
tings except for changing the ignore eos to false . retrieving less information from the long context. (2) When
compared to the truncate-middle approach, YOURA shows
Experiment Setup We compare the following scenarios better overall quality for Llama2 and Llama3, and slightly
to evaluate how the YOURA improves the output quality less (-0.15) for Mistral.
or performance. For each setup, as done with the Long- (2-1) When broken down into dataset types, we observe
Bench (Bai et al., 2023), we equalize the retrieval budget to that YOURA is better at multi-document QA than single-
3500, 7500, and 31500 for models with 4K, 8K, and 32K document QA for Llama2 and Llama3. We claim that Trunc.
context window sizes, respectively. truncates the key information to correctly answer multi-
document QAs, but the head and tail of a single document
• Model without context No context is passed to the include key information in many cases. As for YOURA,
LLM. the key information is retrieved properly for both the single-
doc and multi-doc QAs resulting in a comparable answer
• with BM25 Retrieval Bx y indicates chunking at quality. For YOURA’s lower single-document score, we
roughly x tokens and retrieving the top y chunks. conjecture that the retrieval budget (4k, and 8k for Llama2
• with Embedding Retrieval Ex y indicates chunk- 3
https://huggingface.co/sentence-transformers/paraphrase-
ing at roughly x tokens and retrieving the top y MiniLM-L6-v2
6
Algorithm 1 EASY Algorithm Setup Retr. Ratio SDOC MDoc Overall
Require: seq - Encoded token sequence Avg Avg Avg Avg
Require: T S - Target Sentences generated by NLP models
Require: tol ← 30 - Max wiggling tries LLaMA2 INF 13.65 19.30 16.47
1: P ← [] {Processed sentences} w/ B500 7 3.28 22.82 20.17 21.49
2: B ← [] {Sentence boundary indices} w/ B50 70 3.28 23.86 23.25 23.56
3: i, m ← 0 {Initialize indices} w/ E500 7 3.28 17.13 19.57 18.35
4: while P ̸= T S do w/ E50 70 3.28 24.15 22.87 23.51
5: APPEND T S[i] to P w/ Trunc. 3.28 25.74 22.34 24.04
6: c ← Length(Encode(Concat(P ))) {Candidate} w/ YOURA 3.39 23.22 25.17 24.20
7: saved ← c
LLaMA3 INF 11.59 23.16 17.37
8: visited ← ∅
w/ B500 15 1.57 34.07 33.27 33.73
9: while true do
w/ B50 150 1.57 32.92 35.83 34.17
10: if c ∈ visited or c > |seq| or |c − saved| > tol
w/ E500 15 1.57 33.33 30.82 32.08
then
w/ E50 150 1.57 31.50 35.70 33.60
11: c ← saved
w/ Trunc. 1.57 36.83 34.72 35.93
12: break
w/ YOURA 1.76 36.57 36.44 36.50
13: end if
14: ADD c to visited Mistralv0.2 INF 7.02 11.47 9.25
15: s ← Decode(seq[m : c]) w/ B500 63 1.04 32.29 24.66 29.02
16: if s == T S[i] then w/ B50 630 1.04 29.14 23.28 26.63
17: break w/ E500 63 1.04 26.79 22.06 24.76
18: else if s is substring of T S[i] then w/ E50 630 1.04 29.14 23.28 26.63
19: c←c+1 w/ Trunc. 1.04 32.65 26.62 29.64
20: else w/ YOURA 1.33 33.14 25.84 29.49
21: c←c−1
22: end if
Table 2. Answer quality comparison for LongBench’s single-
23: end while
document and multi-document QA datasets. YOURA shows better
24: APPEND c to B
quality with a higher retrieval ratio — i.e., retrieved fewer tokens.
25: m←c All retrieval scenarios use the maximal budget (4K, 8K, and 32K
26: i←i+1 tokens depending on the model) except for without-context setup
27: end while and “YOURA”. The retrieval ratio (second column) is the average
28: return B of the total token count divided by the retrieved token count.
and Llama3, respectively) is insufficient to retrieve all of the
relevant information. As a supporting fact, the quality differ-
From the observations, we conclude that the balance be-
ence between Trunc. and YOURA decreases significantly
tween the pre-trained model’s context window size and the
from -2.52 to -0.26. Experiment results for a pre-trained
retrieval ratio is important for long context QA.
model with a 16k context window size would be useful for
proving/disproving the conjecture but no solid 16k model is How does YOURA’s higher retrieval ratio improve in-
available at the time of the publication. ference performance? Table 3 shows the average vllm
inference throughput (requests per second) assuming offline
(2-2) For Mistral, which has a large context window size
truncation/retrieval. We make the following observations.
of 32k, truncation rarely happens (as shown in the small
First, YOURA improves the overall inference throughput,
retrieval ratio). Without truncation, distractful information
especially for models with larger context window sizes. This
is included as part of the augmented context resulting in
is because the relative retrieval ratio is larger for models
worse single-document answer quality than YOURA. 4 Our
with larger context window sizes. Second, the retrieval ra-
explanation for multi-document QA using Mistral, where
tio is a good relative throughput estimator. The relative
YOURA showed a slightly worse quality (-0.78), is sim-
throughput improvements (+2%, +13%, +31%) are on par
ilar to the explanation in (2-1); YOURA’s relatively high
with the relative retrieval ratio (+3%, +12%, +28%). We
retrieval ratio resulted in leaving out some key information.
conclude that the performance of LLM serving platforms
4
The impact of distractful information on output quality is such as vLLM is sensitive to the retrieved context size and
observed in other work (Catav, 2023) as well. underscores the importance of good information retrieval.
7
Mod Retrieval SDOC MDOC Overall Avg Mod 1-1 1-2 1-3 2-1 2-2 2-3 Avg
L2 Trunc. 5.25 5.76 5.50 L2 94.1 83.1 93.5 98.4 98.8 97.8 94.3
L2 YOURA 5.43 5.78 5.60 (+2%) L3 99.9 99.6 96.5 86.5 86.9 85.6 92.5
L3 Trunc. 3.23 2.79 3.01 M 94.0 83.1 93.5 98.4 98.8 97.8 94.3
L3 YOURA 3.76 3.03 3.39 (+13%) Avg 96.0 88.6 94.5 94.4 94.8 93.7 93.7
M Trunc. 1.97 1.36 1.66 L2: Llama2-7B, L3: Llama3-8B, M: Mistralv0.2
M YOURA 2.54 1.82 2.18 (+31%)
Table 4. Evaluation of EASY algorithm’s average sentence match
Table 3. vLLM Serving performance (request per second, higher rate (exact match / total # of sentences) for each model and dataset
the better) for LongBench’s single-document and multi-document (Table 1) pairs.
QA datasets with respect to Truncation and YOURA. YOURA
retrieved fewer tokens resulting in up to 30% higher inference
throughput. Mod 1-1 1-2 1-3 2-1 2-2 2-3 Avg
L2 3.26 3.48 2.54 2.61 2.36 2.67 2.82
NZ 5.40 4.80 4.35 4.22 4.13 4.31 4.54
4.2 EASY Sentence Splitter Accuracy
L3 1.71 3.45 2.77 3.16 2.80 3.23 2.85
How well can the EASY algorithm split tokens into sen- NZ 3.70 4.76 4.52 4.53 4.29 4.64 4.41
tences? Table 4 shows the quality of our EASY sentence M 3.54 3.52 2.56 2.65 2.38 2.71 2.89
splitting algorithm. We used Stanza (Qi et al., 2020), an NZ 5.52 4.85 4.40 4.29 4.17 4.38 4.60
open-sourced NLP toolkit, to split raw text into sentences.
Then we ran the EASY algorithm for each tokenizer model L2: Llama2-7B, L3: Llama3-8B, M: Mistralv0.2
to identify the sentence boundary within a token sequence. NZ* (Non-Zero): average of non-zero distances.
After decoding each token segment (delimited by the sen-
tence boundary index), it is a match if the Stanza’s sentence Table 5. Average Levenshtein distance between the target sentence
(namely target sentence) matches exactly after the cleanup and EASY algorithm’s sentence split. A smaller number indicates
(e.g., trailing whitespace). We report the average match rate fewer character edits (insert, remove, replace) to transform the
EASY output sentence to the target sentence, and thus a better
for each model and dataset pair.
quality split.
We make the following observations. Overall the EASY
algorithm can identify the sentence boundary for models
with varying dictionary sizes: Llama2, Llama3, and Mis- The NZ row in Table 5 is the average Levenshtein distance
tralv0.2 tokenizers with 32000, 128256, and 32000 words, of unmatched sentences. From these rows, we conclude
respectively. Llama2 and Mistral showed almost identical that the EASY algorithm is resilient to incorrect split and
match ratios, as both models’ tokenizers are based on simi- its cascading effect.
lar vocabulary. Llama3 uses roughly ×4 more vocabulary
resulting in fewer total token counts as shown in Table 1 5 C ONCLUSION
(e.g., .T into a single token). Due to the rich vocabulary, a
typo easily blurred sentence boundaries causing many im- We propose a simple, yet effective attention-based retrieval
perfect matches and thus the smallest overall match quality algorithm. We show that attention vectors are useful re-
as in the 2-1, 2-2, and 2-3 datasets. For 1-1, 1-2, and 1-3, trieval and propose novel retrieval heuristics called reaction
datasets, the rich vocabulary worked in favor. Many atypical scores. Our fine-tuning approach improves the output qual-
char sequences (e.g., <b>, math formula) were present in ity or the inference performance of off-the-shelf models for
these datasets, which were properly represented with tokens long-context QA tasks.
using Llama3.
How resilient is the EASY algorithm from edge cases?
To quantify the resilience despite typos and overly represen-
tative dictionaries, we measured the Levenshtein distance
between the target sentence and the sentence yielded by
our EASY algorithm (Table 5). On average, one needs less
than 3 character edits (insert, remove, replace) to the EASY
output sentences to match the target sentences.
8
R EFERENCES Jiang, Z., Gao, L., Araki, J., Ding, H., Wang, Z., Callan,
J., and Neubig, G. Retrieval as attention: End-to-end
Bai, Y., Lv, X., Zhang, J., Lyu, H., Tang, J., Huang, Z., Du,
learning of retrieval and reading within a single trans-
Z., Liu, X., Zeng, A., Hou, L., Dong, Y., Tang, J., and Li,
former, 2022. URL https://arxiv.org/abs/
J. Longbench: A bilingual, multitask benchmark for long
2212.02027.
context understanding, 2023.
Khandelwal, U., Levy, O., Jurafsky, D., Zettlemoyer, L.,
Bertsch, A., Alon, U., Neubig, G., and Gormley, M. Unlim- and Lewis, M. Generalization through memorization:
iformer: Long-range transformers with unlimited length Nearest neighbor language models. arXiv preprint
input. Advances in Neural Information Processing Sys- arXiv:1911.00172, 2019.
tems, 36, 2024.
Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu,
Borgeaud, S., Mensch, A., Hoffmann, J., Cai, T., Rutherford, C. H., Gonzalez, J. E., Zhang, H., and Stoica, I. Efficient
E., Millican, K., Van Den Driessche, G. B., Lespiau, J.-B., memory management for large language model serving
Damoc, B., Clark, A., et al. Improving language models with pagedattention. In Proceedings of the ACM SIGOPS
by retrieving from trillions of tokens. In International 29th Symposium on Operating Systems Principles, 2023.
conference on machine learning, pp. 2206–2240. PMLR,
2022. Li, S., You, C., Guruganesh, G., Ainslie, J., Ontanon, S.,
Zaheer, M., Sanghai, S., Yang, Y., Kumar, S., and Bho-
Bulatov, A., Kuratov, Y., and Burtsev, M. Recurrent memory janapalli, S. Functional interpolation for relative posi-
transformer. Advances in Neural Information Processing tions improves long context transformers. arXiv preprint
Systems, 35:11079–11091, 2022. arXiv:2310.04418, 2023.
Catav, A. Less is more: Why use retrieval in- Liu, Z., Wang, J., Dao, T., Zhou, T., Yuan, B., Song, Z.,
stead of larger context windows, Jul 2023. Shrivastava, A., Zhang, C., Tian, Y., Re, C., et al. Deja
URL https://www.pinecone.io/blog/ vu: Contextual sparsity for efficient llms at inference time.
why-use-retrieval-instead-of-larger-context/. In International Conference on Machine Learning, pp.
22137–22176. PMLR, 2023.
Chen, S., Wong, S., Chen, L., and Tian, Y. Extending Lou, C., Jia, Z., Zheng, Z., and Tu, K. Sparser is faster and
context window of large language models via positional less is more: Efficient sparse attention for long-range
interpolation. arXiv preprint arXiv:2306.15595, 2023a. transformers, 2024. URL https://arxiv.org/
abs/2406.16747.
Chen, Y., Qian, S., Tang, H., Lai, X., Liu, Z., Han, S., and
Jia, J. Longlora: Efficient fine-tuning of long-context Luo, K., Liu, Z., Xiao, S., and Liu, K. Bge landmark embed-
large language models. arXiv preprint arXiv:2309.12307, ding: A chunking-free embedding method for retrieval
2023b. augmented long-context large language models, 2024.
Ding, Y., Zhang, L. L., Zhang, C., Xu, Y., Shang, N., Xu, Mangrulkar, S., Gugger, S., Debut, L., Belkada, Y., Paul,
J., Yang, F., and Yang, M. Longrope: Extending llm S., and Bossan, B. Peft: State-of-the-art parameter-
context window beyond 2 million tokens. arXiv preprint efficient fine-tuning methods. https://github.
arXiv:2402.13753, 2024. com/huggingface/peft, 2022.
Giorgi, J., Soldaini, L., Wang, B., Bader, G., Lo, K., Martins, P. H., Marinho, Z., and Martins, A. F. ∞-
Wang, L. L., and Cohan, A. Exploring the challenges former: Infinite memory transformer. arXiv preprint
of open domain multi-document summarization. ArXiv, arXiv:2109.00301, 2021.
abs/2212.10526, 2022.
Mohtashami, A. and Jaggi, M. Landmark attention:
Random-access infinite context length for transformers.
Izacard, G., Lewis, P., Lomeli, M., Hosseini, L., Petroni,
arXiv preprint arXiv:2305.16300, 2023.
F., Schick, T., Dwivedi-Yu, J., Joulin, A., Riedel, S., and
Grave, E. Atlas: Few-shot learning with retrieval aug- Peng, B., Quesnelle, J., Fan, H., and Shippole, E. Yarn:
mented language models. Journal of Machine Learning Efficient context window extension of large language
Research, 24(251):1–43, 2023. models. arXiv preprint arXiv:2309.00071, 2023.
Jiang, H., Wu, Q., Luo, X., Li, D., Lin, C.-Y., Yang, Y., and Press, O., Smith, N. A., and Lewis, M. Train short, test
Qiu, L. Longllmlingua: Accelerating and enhancing llms long: Attention with linear biases enables input length
in long context scenarios via prompt compression, 2023. extrapolation. arXiv preprint arXiv:2108.12409, 2021.
9
Qi, P., Zhang, Y., Zhang, Y., Bolton, J., and Manning, C. D.
Stanza: A python natural language processing toolkit for
many human languages, 2020.
Stelmakh, I., Luan, Y., Dhingra, B., and Chang, M.-W.
Asqa: Factoid questions meet long-form answers. arXiv
preprint arXiv:2204.06092, 2022.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones,
L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention
is all you need, 2023.
Wang, B., Ping, W., Xu, P., McAfee, L., Liu, Z., Shoeybi,
M., Dong, Y., Kuchaiev, O., Li, B., Xiao, C., et al.
Shall we pretrain autoregressive language models with
retrieval? a comprehensive study. arXiv preprint
arXiv:2304.06762, 2023.
Wang, W., Dong, L., Cheng, H., Liu, X., Yan, X., Gao, J.,
and Wei, F. Augmenting language models with long-term
memory. Advances in Neural Information Processing
Systems, 36, 2024.
Wu, Q., Lan, Z., Qian, K., Gu, J., Geramifard, A., and Yu,
Z. Memformer: A memory-augmented transformer for
sequence modeling. arXiv preprint arXiv:2010.06891,
2020.
Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Ef-
ficient streaming language models with attention sinks,
2024.
Zhang, T., Patil, S. G., Jain, N., Shen, S., Zaharia, M., Stoica,
I., and Gonzalez, J. E. Raft: Adapting language model
to domain specific rag, 2024. URL https://arxiv.
org/abs/2403.10131.
Zhang, Z., Sheng, Y., Zhou, T., Chen, T., Zheng, L., Cai,
R., Song, Z., Tian, Y., Ré, C., Barrett, C., Wang, Z., and
Chen, B. H2 o: Heavy-hitter oracle for efficient generative
inference of large language models, 2023.
Zhao, J., Zhang, Z., Chen, B., Wang, Z., Anandkumar,
A., and Tian, Y. Galore: Memory-efficient llm train-
ing by gradient low-rank projection. arXiv preprint
arXiv:2403.03507, 2024.
Zhu, C., Ping, W., Xiao, C., Shoeybi, M., Goldstein, T.,
Anandkumar, A., and Catanzaro, B. Long-short trans-
former: Efficient transformers for language and vision.
Advances in neural information processing systems, 34:
17723–17736, 2021.
10
A D ETAILED R ESULTS
Table 6, Table 7 are the detailed version of Table 2 and
Table 3, respectively.
11
Setup Retr. Rate Single-Document QA SDoc Multi-Document QA MDoc Overall
Avg 1-1 1-2 1-3 Avg 2-1 2-2 2-3 Avg Avg
LLaMA2 INF 8.89 13.87 18.19 13.65 23.05 25.67 9.17 19.30 16.47
w/ B500 7 3.28 12.59 21.44 34.42 22.82 25.92 27.47 7.11 20.17 21.49
w/ B50 70 3.28 16.37 21.51 33.71 23.86 27.48 28.62 13.64 23.25 23.56
w/ E500 7 3.28 12.85 16.08 22.47 17.13 25.51 24.84 8.35 19.57 18.35
w/ E50 70 3.28 17.15 20.12 35.18 24.15 28.80 27.15 12.66 22.87 23.51
w/ Trunc. 3.28 18.84 21.71 36.68 25.74 27.55 31.08 8.40 22.34 24.04
w/ YOURA 3.39 15.48 23.41 30.77 23.22 29.55 31.11 14.86 25.17 24.20
LLaMA3 INF 9.62 14.38 10.78 11.59 28.43 30.46 10.60 23.16 17.37
w/ B500 1.57 18.27 40.21 43.72 34.07 41.75 41.18 16.88 33.27 33.73
w/ B50 1.57 19.86 36.51 42.39 32.92 47.04 40.58 19.86 35.83 34.17
w/ E500 1.57 16.38 39.85 43.76 33.33 42.28 34.93 15.25 30.82 32.08
w/ E50 1.57 19.05 35.44 40.02 31.50 48.05 36.85 22.19 35.70 33.60
w/ Trunc. 1.57 21.54 44.21 44.75 36.83 46.43 36.23 21.50 34.72 35.93
w/ YOURA 1.76 21.47 42.39 45.85 36.57 50.78 34.24 24.29 36.44 36.50
Mistralv0.2 INF 7.24 7.69 6.13 7.02 16.41 11.87 6.13 11.47 9.25
w/ B500 63 1.04 20.74 29.22 46.91 32.29 34.26 22.26 17.47 24.66 29.02
w/ B50 630 1.04 18.84 25.08 43.50 29.14 32.85 18.66 18.32 23.28 26.63
w/ E500 63 1.04 14.02 24.14 42.22 26.79 32.81 17.55 15.82 22.06 24.76
w/ E50 630 1.04 18.84 25.08 43.50 29.14 32.85 18.66 18.32 23.28 26.63
w/ Trunc. 1.04 20.80 29.23 47.92 32.65 37.28 21.72 20.86 26.62 29.64
w/ YOURA 1.33 22.37 30.17 46.88 33.14 36.11 21.45 19.95 25.84 29.49
The bold score indicates the best score for each pre-trained model.
Table 6. Detailed experiment results for LongBench’s single-document and multi-document QA datasets. YOURA shows better quality
with a higher retrieval ratio — i.e., retrieved fewer tokens. All retrieval scenarios use the maximal budget (4K, 8K, and 32K tokens
depending on model) except for without-context setup and “YOURA”. Retrieval ratio (second column) is an average of the total token
count divided by the retrieved token count.
Model Retrieval Single-Document QA Multi-Document QA Overall Avg
Method 1-1 1-2 1-3 2-1 2-2 2-3 Req./Sec
LLaMA2 Trunc. 5.58 4.68 5.50 5.71 5.87 5.68 5.50
LLaMA2 YOURA 5.70 4.88 5.72 5.70 6.00 5.63 5.60 (+2%)
LLaMA3 Trunc. 2.50 3.95 3.23 2.63 3.23 2.52 3.01
LLaMA3 YOURA 2.60 4.87 3.82 2.72 3.83 2.53 3.39 (+13%)
Mistralv0.2 Trunc. 0.55 3.17 2.19 1.08 2.14 0.86 1.66
Mistralv0.2 YOURA 0.64 4.02 2.97 1.45 2.85 1.15 2.18 (+31%)
Table 7. vLLM Serving performance (request per second, higher the better) for LongBench’s single-document and multi-document QA
datasets with respect to Truncation and YOURA. YOURA retrieved fewer tokens resulting in upto 30% higher inference throughput.
12
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/text/2411.00744v2.txt
+1284
View File
File diff suppressed because it is too large. Load diff
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/text/2412.18232v2.txt
+1255
View File
File diff suppressed because it is too large. Load diff
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/text/2501.08248v3.txt
+1296
View File
File diff suppressed because it is too large. Load diff
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/text/2503.14800v3.txt
+940
View File
@@ -0,0 +1,940 @@
Long Context Modeling with Ranked Memory-Augmented Retrieval
Ghadir Alselwi1 , Hao Xue1,2,3 , Shoaib Jameel4 , Basem Suleiman1 ,
Flora D. Salim1 , Imran Razzak5,1
1
University of New South Wales, Sydney, NSW, Australia
2
Hong Kong University of Science and Technology (Guangzhou)
3
The Hong Kong University of Science and Technology
4
University of Southampton, Southampton, UK
5
Mohamed Bin Zayed University of Artificial Intelligence, Abu Dhabi, UAE
{g.alselwi,b.suleiman,flora.salim}@unsw.edu.au
haoxue@hkust-gz.edu.cn, M.S.Jameel@southampton.ac.uk, imran.razzak@mbzuai.ac.ae
Abstract ory bank model is a large, non-trainable store
of past context representations. Instead of re-
Effective long-term memory management is
arXiv:2503.14800v3 [cs.IR] 18 May 2026
computing representations for all past tokens every
crucial for language models handling extended
contexts. We introduce the Enhanced Ranked time, these representations are pre-computed and
Memory Augmented Retrieval (ERMAR) stored. Given the current context, MemLong re-
framework, which dynamically ranks mem- trieves relevant segments from the memory bank.
ory entries based on relevance. Unlike prior It uses a dot product similarity search to find the
models, ERMAR employs a novel relevance memory entries most related to the current context.
scoring mechanism and a pointwise re-ranking This allows the model to focus only on the most
model for key-value embeddings, inspired pertinent past information, rather than processing
by learning-to-rank techniques in informa-
the entire history. However, its treatment of all
tion retrieval. By integrating historical us-
age patterns and adaptive retrieval, ERMAR key-value (K-V) pairs with equal weight, regard-
achieves state-of-the-art results on standard less of their contextual relevance, often leads to
benchmarks, demonstrating superior scalabil- information overload and reduced retrieval preci-
ity and performance in long-context tasks. sion. This limitation becomes particularly evident
Code will be released at https://github. in scenarios requiring context management.
com/cruiseresearchgroup/ERMAR.
We have developed a novel model that addresses
1 Introduction the aforementioned limitations by building upon
Memlong (Liu et al., 2024), a publicly available
Large Language Models (LLMs) face a fundamen- baseline on GitHub1 . Our Enhanced Ranked
tal limitation in processing long-context scenarios Memory Augmented Retrieval (ERMAR) model
due to the quadratic complexity of attention mecha- has a novel relevance scoring mechanism that fun-
nisms and increasing memory demands during gen- damentally improves context retrieval and utiliza-
eration (Vaswani, 2017; Tworkowski et al., 2024). tion for K-V embeddings. Unlike MemLong, ER-
Consider a scenario in an automated customer ser- MAR employs multiplication (Cao et al., 2007)
vice system: A customer reports an issue with their to compute relevance scores, enabling a more nu-
printer, referencing a setup process from a previous anced and context-aware assessment of semantic
conversation that occurred two hours ago. After alignment between queries and stored memory. ER-
50 messages of troubleshooting, the customer men- MAR also incorporates a re-ranking mechanism
tions that the same error from the beginning has that dynamically reorders K-V embeddings based
resurfaced. Traditional LLMs, constrained by their on their relevance scores, ensuring that the most
context window, would struggle to access the cru- pertinent information is prioritized during retrieval.
cial earlier context about the initial setup process, This re-ranking process, combined with an adaptive
leading to inconsistent or incomplete responses, retrieval system that integrates historical usage pat-
Figure 1. It is well known that handling extended terns, allows ERMAR to capture subtle contextual
contexts remains a significant challenge, particu- relationships better and refine memory prioritiza-
larly in applications requiring document analysis tion. As shown in Figure 1, ERMAR processes
and sustained dialogue interactions. incoming queries and long-context conversations
The recent MemLong (Liu et al., 2024) architec- through a novel ranking architecture, employing
ture stores and accesses historical context through
1
basic chunk-level memory operations. The mem- https://github.com/Bui1dMySea/MemLong
K-V pairs ranking (K0 -V0 , K1 -V1 ,...,Ki -Vi ) and cor- Longformer (Beltagy et al., 2020) and BigBird (Za-
responding embeddings to perform semantic search heer et al., 2020) reduce computational complexity
and ranking of relevant historical information. while maintaining model capabilities through selec-
Our novel ERMAR model introduces three key tive attention patterns. Position encoding adapta-
improvements: (i) A semantic similarity metric tions like RoPE (Su et al., 2024) and ALiBi (Press
to measure contextual alignment between query et al., 2022) have enhanced models’ ability to han-
embeddings and key-value pairs; (ii) A weighted dle longer sequences. YARN (Peng et al., 2023)
scoring function that considers content similarity further advanced this through dynamic position em-
and contextual relevance; and (iii) Integration of beddings, demonstrating reliable generalization up
historical usage patterns to refine relevance assess- to 128k tokens.
ment.
2.3 Learning-to-Rank and Information
2 Related Work Retrieval
Our ERMAR framework draws inspiration from
The challenge of enabling language models to ef- learning-to-rank techniques in information retrieval.
fectively process extended contexts has driven re- Traditional ranking approaches like BM25 struggle
search across multiple domains. We organize re- with semantic similarity, leading to neural ranking
lated work into key areas that inform our ERMAR models using learned representations. Pointwise
framework. ranking approaches (Cao et al., 2007) predict rel-
evance scores for query-document pairs, directly
2.1 Memory-Augmented Neural Networks
inspiring our relevance scoring mechanism. Dense
Early architectures like Neural Turing Ma- retrieval methods like DPR (Karpukhin et al., 2020)
chines (Graves et al., 2014) and Differentiable Neu- and ColBERT (Khattab and Zaharia, 2020) demon-
ral Computers (Graves et al., 2016) introduced ex- strate the effectiveness of learned dense represen-
ternal memory beyond model parameters. Recent tations for ranking through similarity scores—a
models for long-context language modeling include paradigm we adapt for memory entry ranking.
RetroMAE (Xiao et al., 2022), which struggled BERT-based re-ranking models (Nogueira and Cho,
with semantic coherence, and Memorizing Trans- 2019) show that sophisticated re-ranking signifi-
formers (Wu et al., 2022), which improved cross- cantly improves retrieval quality. This two-stage
sequence information storage but faced scalability retrieve-then-rerank paradigm directly influences
limits. MemLong (Liu et al., 2024) stores histor- our design, where we first retrieve candidate mem-
ical context via chunk-level memory operations, ory entries and then apply learned re-ranking.
but its uniform treatment of key-value pairs can
cause information overload and reduced retrieval 2.4 Current Limitations and Gaps
precision. Despite significant progress, current approaches
face limitations that motivate ERMAR:
2.2 Retrieval-Augmented Generation and Static Memory Management: Most ap-
Long-Context Modeling proaches use fixed memory structures that do not
Retrieval-augmented generation (RAG)(Lewis adapt to content importance or usage patterns, lead-
et al., 2020) pioneered integrating dense passage ing to inefficient memory utilization.
retrieval with generative models, while Fusion-in- Uniform Memory Treatment: Existing meth-
Decoder (FiD)(Izacard and Grave, 2020) improved ods treat all memory entries equally, lacking mech-
efficiency through independent passage process- anisms to prioritize more relevant information.
ing. REALM (Guu et al., 2020) introduced end- Limited Semantic Understanding: Memory
to-end learning of retrieval and generation, and systems often rely on simple similarity metrics
DPR (Karpukhin et al., 2020) established dense without considering contextual nuance or temporal
passage retrieval standards. However, these meth- relevance.
ods typically focus on static corpora rather than Scalability Trade-offs: Current approaches face
dynamic context-aware memory management. The difficult trade-offs between memory capacity, re-
quadratic complexity of attention mechanisms has trieval accuracy, and computational efficiency.
driven research into efficient long-context archi- Our ERMAR framework addresses these limita-
tectures. Sparse attention mechanisms such as tions through dynamic relevance scoring, adaptive
Figure 1: Motivating example of the ERMAR retrieval workflow. Left: A long customer support conversation where
the current query references information from much earlier context. Right: The four-stage workflow — ⃝ 1 Long
Memory Retrieval, ⃝ 2 Search, ⃝3 Reranking, ⃝ 4 Memory Fusion Generation. The ERMAR output (top) correctly
references earlier context, whereas the output without reranking (bottom) produces a generic response.
memory management, and sophisticated re-ranking the model’s output on the most relevant historical
mechanisms inspired by information retrieval tech- context.
niques, providing a more principled approach to This design directly addresses the information
long-context memory management. dilution problem in long contexts: without rerank-
ing (bottom of Figure 1), the model produces a
3 Our Novel ERMAR Model generic response that fails to reference the earlier
Figure 1 presents a concrete motivating example of printer setup context; with reranking (ERMAR out-
how ERMAR’s contextual ranking mechanism op- put, top), the model correctly retrieves and incor-
erates in practice, while Figure 2 illustrates the porates the relevant earlier information.
overall architecture of the proposed framework. As shown in Figure 2, the ERMAR architec-
As shown in Figure 1, ERMAR processes a long- ture consists of three interacting components: a
context query through four sequential stages. First, frozen lower transformer block, a trainable up-
the Long Memory Retrieval module encodes the per block with retrieval-augmented attention, and
full input history into chunk-level key–value (KV) a memory retrieval module with reranking. Input
pairs stored in a memory bank, alongside embed- text is processed by the frozen lower block (Causal
dings that capture semantic representations of each Attention and Feed Forward layers), whose hid-
chunk. Second, when a new query arrives, its den states form the key–value pairs {(Kj , Vj )} ∈
embedding is used to search the memory bank Rdmodel stored in the Memory Bank. In parallel,
for candidate KV pairs via similarity matching. the Dense Embedder independently encodes the
Third, the retrieved candidates pass through the same input into chunk-level retrieval embeddings
Reranking module — the core novelty over Mem- E1 , . . . , Ei ∈ Rdret , which serve as the semantic
Long — which assigns relevance scores to each KV index for similarity-based retrieval. The Rerank-
pair and re-orders them by contextual importance, ing module then scores all stored embeddings and
elevating high-relevance entries (K6 -V6 , ranked) KV pairs using Equation 2, re-ordering them by
above lower-relevance ones (K0 -V0 , unranked). Fi- relevance to the current query and isolating the
nally, only the top-ranked KV pairs are passed to top-ranked subset (right stacks in Figure 2). The
Memory Fusion Generation, which conditions Retrieval module performs top-K selection over
with a relevance scoring mechanism,
M : |Rd{z }×R
model dmodel
| {z } × R
dret
|{z} → S , (1)
keys values embeddings
where dmodel denotes the transformer hidden di-
mension, dret the retrieval embedding dimension,
and S the space of scored memory items (see Ap-
pendix A.8 for embedding dimensionality analy-
sis).
Relevance Score. Given a query embedding
q ∈ Rdret and a matrix of key embeddings K =
[k1 , . . . , km ] ∈ Rm×dret , where each row ki ∈ Rdret
corresponds to a key embedding, the relevance
Figure 2: Architecture of the ERMAR model. The score is defined as:
frozen lower block (blue dashed, snowflake) produces
qK⊤
 
hidden states stored as key–value pairs (dmodel ) in the α(q, K) = softmax √ , (2)
dret
Memory Bank, while the Dense Embedder encodes the
input into chunk-level retrieval embeddings (dret ). The √
where dret normalizes the similarity scores to pre-
Reranking module scores and ranks all stored entries,
and the Retrieval module injects the top-K ranked pairs
vent excessively large values.
into the trainable upper block via Retrieval Causal The resulting vector α(q, K) represents a prob-
Attention (orange arrow) to condition the final out- ability distribution over the keys, where each en-
put. Frozen and trainable components are marked by try αi denotes the relative importance (attention
snowflake and flame icons respectively. weight) of the i-th key with respect to the query q.
The final
Pm context or weighted output is obtained as
o = i=1 αi Vi , where Vi denotes the correspond-
ing value embedding associated with key ki .
this ranked subset and injects the selected KV pairs
We now introduce the Relevance Scoring with
into the trainable upper block via Retrieval Causal
Adaptive Retrieval (RSAR) mechanism, which
Attention (orange arrow), producing a context-
dynamically ranks memory entries according to
enriched representation from which the final output
their importance to the current query, thereby en-
is generated.
hancing the retrieval process. The relevance score
ERMAR maintains consistency through frozen α(q, K), as defined in Equation 2, is employed to
lower layers and selective parameter updates. rank the memory entries.
Specifically, ERMAR (i) stores important infor- RSAR augments the memory module by intro-
mation from earlier parts of the text, (ii) assigns ducing ranked key–value entries, represented as
relevance scores to stored information based on (Kj , Vj , sj ), where sj denotes the relevance score
its importance to the current context, and (iii) re- for each entry. These scores allow the system to pri-
trieves only the most relevant historical information oritize the most informative content during retrieval
when needed. The relevance scoring mechanism while maintaining computational efficiency. To op-
is analogous to the attention operation, enabling timize memory usage, a pruning strategy is applied
the model to focus on salient parts of the mem- to remove less relevant entries. Specifically, entries
ory. There is also a “loose” pointwise connection: with scores below a predefined threshold are dis-
while relevance-based reranking acts as an auxil- carded, preserving only the most critical contextual
iary mechanism that enhances memory retrieval information.
quality, the primary optimization target remains The enhanced retrieval mechanism is defined as:
next-token prediction likelihood.
RRSAR (tq , s) = TopK {sim(E(tq ), e) · maxj sj | e ∈ s}, (3)
Let V be a finite vocabulary, and x =
(x1 , . . . , xn ) ∈ V n be a token sequence with a where tq denotes the query token, E(tq ) represents
preceding context x<i . The embedding function the encoded query, and the operation identifies the
E : V ∗ → Rdret maps sequences to a retrieval space. top-K relevant entries based on their scores; this
We introduce a memory function M augmented mechanism enables efficient retrieval of the most
relevant information, even in extended contexts 3.1.2 Model Architecture and Configuration
(see Appendix A.6 for the impact of different top- We fine-tuned OpenLLaMA-3B, a pre-trained
K values on performance). LLM with rotational position encoding (Su et al.,
Ranked Key–Value Pairs. Each embedding e 2024), using LoRA (Hu et al., 2021) for parameter-
maintains a ranked set of key–value pairs: efficient adaptation. The model consists of L = 26
Rranked (e) = {(Kj , Vj , sj )}m
j=1 , (4) transformer layers, H = 32 attention heads, and
a hidden dimension of d = 100. The 13th layer
where sj = α(e, Kj ) is the relevance score be-
serves as the memory layer where historical con-
tween the embedding e and the key Kj .
text is stored, while layers [14, 18, 22, 26] are aug-
Given a sequence x, the ERMAR objective func-
mented with retrieval mechanisms to access stored
tion is formulated as maximizing
memories. ERMAR employs a memory capacity
n
L(θ) =
X
pθ (xi | RRSAR (ti , s), x<i ), (5) of 32,768 key–value pairs and utilizes BGE-M3
i=1 embeddings for semantic similarity computation.
subject to Complete training hyperparameters and configura-
tion details are provided in Appendix A.2.
si = M(K1:i−1 , V1:i−1 ; E(t1:i−1 )), (6)
ti = text(c⌈i/τ ⌉ ), (7) 3.1.3 Baseline Models
αi = α(E(ti ), K1:i−1 ), (8)
ERMAR was evaluated against state-of-the-art
where c denotes chunk indices, τ represents the models across two parameter scales to ensure com-
chunk size, αi guides key–value pair selection, and prehensive comparison. The 7B models include
pθ denotes the model’s probability distribution. LLaMA-2-7B (Touvron et al., 2023b) as a standard
For new content (Kn , Vn ), the memory state is transformer baseline, LongLoRA-7B-32k (Chen
updated as: et al., 2023) which employs sparse attention mech-
( anisms for 32k-token contexts, and YARN-128k-
M(Kn , Vn ; E(tn )) if |si | < capacity,
si+1 = (9) 7B (Peng et al., 2023) featuring dynamic position
Mu (si , Kn , Vn , αn ) otherwise,
embeddings that support up to 128k tokens.
where capacity is the maximum allowed number For the 3B parameter scale, we compared against
of memory entries, Mu prunes the least relevant OpenLLaMA-3B (Touvron et al., 2023a) as the
entries based on historical scores by ranking and base architecture, LongLLaMA-3B (Tworkowski
removing those with the lowest relevance relative et al., 2024) evaluated in two retrieval configu-
to the current context, and subscript n denotes the rations (4 and 18 memory entries), MemLong-
new key, value, and relevance score corresponding 3B (Liu et al., 2024) as our direct baseline
to the latest input. with chunk-level memory operations, and Phi3-
3.1 Experimental Setup 128k (Abdin et al., 2024) which demonstrates
strong performance across varying context lengths.
3.1.1 Datasets This diverse benchmark suite encompasses differ-
We fine-tuned ERMAR on the SlimPajama ent long-context strategies including sparse atten-
dataset (Fu et al., 2024), a high-quality, dedupli- tion, position encoding extensions, and memory-
cated corpus designed for long-context tasks. It augmented architectures, ensuring robust evalua-
contains 84.7K training rows, making it a com- tion of ERMAR’s retrieval-based approach against
pact yet effective resource for pre-training and fine- complementary methodologies.
tuning. The dataset was preprocessed with a sliding
window approach using 512-token strides to ensure 3.1.4 Evaluation Metrics
comprehensive coverage of long sequences. We employ perplexity as the primary metric for
Performance evaluation was conducted on three language modeling performance, computed on the
benchmark datasets: WikiText-103 (Merity et al., final 2048 tokens of each sequence to focus on
2016) (4,358 test rows), PG-19 (Rae et al., 2019) long-range dependency modeling. For in-context
(100 test rows), and Proof-Pile (Azerbayev et al., learning tasks, we report accuracy on five natu-
2023) (46.3K test rows). Performance was mea- ral language understanding benchmarks: SST-2,
sured across context lengths from (1k-32k) tokens, MR, Subj, SST-5, and MPQA, evaluated in both
using perplexity on the last 2048 tokens (Yen et al., 4-shot and 20-shot settings. Memory efficiency is
2024) following standard evaluation protocols. assessed through peak GPU memory usage and to-
kens processed per second, while computational these are edge cases at the boundary of the training
overhead is measured via inference latency across distribution (the model was fine-tuned up to 32k
different context lengths. tokens): ERMAR maintains consistent improve-
ments at practical context lengths (1K–4K), where
3.2 Results and Discussion memory pressure is highest, and fully recovers at
3.2.1 Long-Context Language Modeling 32k tokens (Table 2), achieving equal or better per-
Following the experimental strategy adopted in plexity across all three datasets. The 16k degra-
(Liu et al., 2024), Table 1 reports the mean per- dation therefore reflects an expected efficiency–
plexity scores of ERMAR across various sequence benefit trade-off rather than a systematic failure
lengths and datasets. Evaluation was conducted of the proposed enhancements.
on the test splits of three standard benchmarks: These findings confirm that ERMAR’s ranked
WikiText-103 (Merity et al., 2016) (4,358 rows), memory retrieval effectively mitigates informa-
PG-19 (Rae et al., 2019) (100 rows), and Proof- tion dilution across extended sequences, improv-
Pile (Azerbayev et al., 2023) (46.3k rows). ing both accuracy and stability over prior memory-
Among the 7B models, YARN-128k-7B augmented models.
achieves the lowest perplexity on shorter contexts, 3.2.2 Scalability to Extended Contexts
whereas LongLoRA-7B-32k scales effectively up
to 16k-token sequences, albeit with mild perfor- To further assess scalability, we evaluated ER-
mance degradation. This result highlights the trade- MAR at a 32k-token context length (Table 2). ER-
off between scalability and modeling precision MAR maintains consistent performance advantages
commonly observed in long-context transformers. across all datasets, achieving a perplexity of 9.765
For the 3B parameter models, ERMAR demon- compared to 9.858 on PG19 (a 0.90% improve-
strates substantial gains in long-context scenarios. ment) and 7.880 versus 7.938 on WikiText-103 (a
While OpenLLaMA-3B exhibits sharp degradation 0.06% improvement). Both models yield identical
beyond 4k tokens and Phi3-128k maintains moder- results on Proof-Pile with a perplexity of 3.063.
ate stability, ERMAR achieves consistently com- These results demonstrate ERMAR’s robustness
petitive performance across all sequence lengths. and scalability in ultra-long context settings, reflect-
At 2k tokens, ERMAR surpasses MemLong on ing the effectiveness of its ranked memory retrieval
Proof-Pile (2.98 vs. 3.16) and sustains strong re- in preserving contextual coherence over extended
sults across the remaining datasets. Its advantage sequences.
is particularly evident on PG-19, where ERMAR 3.2.3 In-Context Learning Performance
achieves the best performance among 3B mod-
Table 3 summarizes ERMAR’s performance across
els at 1k and 2k tokens (10.32 and 9.75, respec-
five natural language understanding tasks under
tively). Notably, ERMAR shows exceptional sta-
4-shot and 20-shot settings.
bility when scaling to longer contexts, with only
In the 4-shot configuration, ERMAR achieves
a 0.31% increase in perplexity from 4k to 16k to-
superior accuracy across all tasks, outperforming
kens on PG-19 (9.78 → 9.81), indicating superior
OpenLLaMA and other memory-augmented base-
scalability. On WikiText-103, ERMAR consistently
lines. It performs notably well on challenging
outperforms other 3B models at all evaluated con-
benchmarks such as SST-5 and MPQA, demonstrat-
text lengths, further validating the effectiveness of
ing strong generalization even with limited exam-
its relevance-based memory retrieval mechanism
ples. The consistent results across different mem-
for long-context modeling.
ory configurations further indicate its robustness in
Overall comparison against MemLong. Across low-resource scenarios.
all 12 configurations in Table 1 (3 datasets × 4 con- In the 20-shot setting, ERMAR continues to de-
text lengths), ERMAR outperforms MemLong in liver leading results, achieving top accuracy on
10 out of 12 settings (83.3%). The two excep- MPQA and Subj, and establishing a new bench-
tions occur at 16k tokens on PG-19 (9.81 vs. 9.64, mark on SST-5. Although MemLong marginally
+1.7%) and Proof-Pile (3.18 vs. 2.99, +6.4%), surpasses ERMAR in isolated tasks, ERMAR main-
where the overhead of the relevance scoring and tains the highest overall average performance, re-
reranking mechanisms marginally outweighs the flecting its scalability as the number of in-context
retrieval benefit at very long contexts. Crucially, examples increases.
PG19 Proof-pile WikiText-103
Model 1k 2k 4k 16k 1k 2k 4k 16k 1k 2k 4k 16k
7B Model
YARN-128k-7b 7.22 7.47 7.17 - 3.03 3.29 2.98 - 5.71 6.11 5.71 -
LongLoRA-7B-32k 9.76 9.71 10.37 7.62 3.68 3.35 3.23 2.60 7.99 7.83 8.39 5.47
LLaMA-2-7B 10.82 10.06 8.92 - 3.24 3.40 2.72 - 10.82 6.49 5.66 -
3B Model
Phi3-128k 11.31 9.90 9.66 -/9.65 4.25 3.11 2.77 -/3.08 7.54 7.22 7.01 -/7.20
OpenLLaMA-3B 11.60 9.77 > 103 - 2.96 2.70 > 103 - 10.57 8.08 > 103 -
LongLLaMA-3B* 10.59 10.02 > 103 - 3.55 3.15 > 103 - 8.88 8.07 > 103 -
LongLLaMA-3B† 10.59 10.25 9.87 - 3.55 3.22 2.94 - 10.69 8.33 7.84 -
MemLong-3B* 10.66 10.09 > 103 - 3.58 3.18 > 103 - 8.72 7.93 > 103 -
w/ 4K MemLong 10.54 9.95 9.89 9.64 3.53 3.16 3.15 2.99 8.53 7.92 7.87 7.99
w/ 4K ERMAR 10.32 9.75 9.78 9.81 3.24 2.98 3.03 3.18 8.42 7.61 7.62 7.80
Table 1: Perplexity comparison of 7B and 3B models across PG19, Proof-pile, and WikiText-103, using a sliding
window evaluation. "-" denotes Out of Memory (OOM) errors, and "x/y" indicates results from single/dual GPU
setups. Memory-augmented models are tested with varying capacities. All runs use a single GPU.
Dataset ERMAR MemLong Difference to 16.61 GB, corresponding to a 30% improvement
PG19 9.765 9.858 0.90%
WikiText-103 7.880 7.938 0.06% over MemLong. It also maintains lower memory
Proof-pile 3.063 3.063 0 consumption per token across all evaluated lengths,
demonstrating superior scalability and efficiency
Table 2: Perplexity at 32k context length, evaluated on
in memory management. This improvement stems
NVIDIA L40S GPU.
from ERMAR’s ranked retrieval and adaptive mem-
InC SST-2 MR Subj SST-5 MPQA ory allocation, which reduce redundancy in long-
Model Avg.
,InM ACC↑ ACC↑ ACC↑ ACC↑ ACC↑ context storage.
OpenLLaMA 4,N/A 90.7 84.0 58.2 41.0 70.5 68.9
w./ Rag 4,4 90.9 90.5 61.6 39.2 63.2 69.1
LongLLaMA 4,4 90.4 83.9 64.3 40.0 64.2 68.6
MemLong 4,4 91.5 84.5 61.5 41.4 70.2 69.8 Context Model Peak Mem Reserved Mem Mem/Token
ERMAR 4,4 93.6 90.8 65.3 45.8 85.2 76.14 Length (GB) (GB) (MB)
LongLLaMA 4,18 91.4 87.1 59.1 41.0 64.5 68.7 1024 ERMAR 7.97 8.16 7.97
MemLong 4,18 91.0 89.6 61.7 43.5 69.4 71.0 MemLong 8.08 8.49 8.08
ERMAR 4,18 93.6 90.8 65.3 45.9 85.2 76.16 2048 ERMAR 8.45 8.71 4.22
OpenLLaMA 20,N/A 93.6 91.2 55.4 38.2 66.4 69.0 MemLong 8.67 9.38 4.33
w./ Rag 20,18 92.2 91.3 75.8 39.8 57.6 71.3 4096 ERMAR 9.42 9.87 2.35
LongLLaMA 20,18 94.1 90.8 64.2 41.4 72.1 72.7 MemLong 9.72 10.58 2.43
MemLong 20,18 93.5 93.8 65.8 43.3 70.6 73.4 16384 ERMAR 15.20 16.61 0.95
ERMAR 20,18 94.7 91.7 82.8 47 86.5 80.54 MemLong 15.60 23.77 0.97
32768 ERMAR 22.87 25.56 0.71
MemLong 23.27 26.05 0.72
Table 3: 4-shot and 20-shot ICL accuracy [%] on 5 NLU
tasks (SST-2, MR, Subj, SST-5, MPQA). We compare Table 4: Memory efficiency comparison of ERMAR
OpenLLaMA, LongLLaMA, MemLong, and ERMAR. and MemLong across context lengths on WikiText-103.
Note: InC = In-Context, InM = In-Memory. Mem/Token is computed as Peak Memory divided by
context length.
Overall, ERMAR consistently performs well
across varying context lengths and task complexi- Memory Pruning Strategy. ERMAR further
ties, effectively leveraging ranked memory retrieval improves efficiency through a dynamic memory
to enhance representation quality. Its ability to pruning mechanism. The algorithm balances three
scale with more examples and preserve perfor- criteria: retaining the most recent 10% of mem-
mance stability across tasks underscores its poten- ory, prioritizing the middle 80% based on retrieval
tial as a general-purpose framework for language frequency, and discarding the oldest 10% as po-
understanding and long-context modeling. tentially outdated. This hybrid strategy allows
ERMAR to adaptively manage memory content
3.2.4 Memory Efficiency Analysis
according to actual usage patterns, ensuring that
Table 4 compares the peak and reserved memory valuable information is preserved while minimiz-
usage of ERMAR and MemLong across different ing unnecessary memory overhead. Such adaptive
context lengths. pruning complements the ranked retrieval process,
As shown in Table 4 and Figure 3, ERMAR ex- enhancing memory efficiency especially at longer
hibits clear memory optimization advantages, par- context lengths. A systematic validation of these
ticularly beyond 8K tokens. At a 16K context, ratios across 11 pruning strategies and five context
ERMAR reduces reserved memory from 23.77 GB lengths is provided in Appendix A.7.
served on WikiText-103 compared to PG-19 likely
stems from dataset-specific characteristics and dif-
ferences in evaluation methodology.
Seq Perplexity Memory Throughput Latency
Length /Token (tokens /Token
(GB) /sec) (ms)
1K 8.42 7.13 3125 0.32
2K 7.61 3.81 2904 0.35
Figure 3: Memory usage comparison between ERMAR 4K 7.62 2.14 2109 0.47
8K 7.76 1.31 1836 0.54
and MemLong across context lengths (1K–32K tokens), 16K 7.80 0.90 1727 0.58
showing (left) peak memory and (right) reserved mem-
ory. ERMAR demonstrates 7–30% lower reserved mem- Table 6: ERMAR’s standalone performance scaling on
ory requirements at longer contexts. Color coding: blue WikiText-103, showing memory efficiency and through-
bars represent MemLong performance, while orange put across sequence lengths.
bars show ERMAR improvements.
3.2.6 Ablation Studies
3.2.5 Latency and Throughput Analysis Ablation on Embedders. To investigate the in-
Table 5 compares the latency and throughput of fluence of embedding architectures on ERMAR’s
ERMAR and MemLong on the PG-19 dataset. performance, we compare BGE and LLM embed-
ders across five NLU tasks (SST-2, Subj, SST-5,
Latency Throughput
Context Model Latency /Token (tokens/ MPQA, and MR) under various configurations. Ex-
(ms) sec) periments examine the effects of Flash versus Ea-
1024 ERMAR 207.97 ± 53.01 0.203 5135
MemLong 190.68 ± 154.55 0.186 6803 ger attention modes, different context lengths (0
2048 ERMAR 423.79 ± 52.56 0.206 4896
MemLong 323.49 ± 204.89 0.157 7446 and 2048), and varying In-C/In-M demonstration
4096 ERMAR 1358.36 ± 55.58 0.331 3020 settings. Results are presented in Tables 7 and 8.
MemLong 1184.16 ± 170.56 0.289 3511
16384 ERMAR 6862.98 ± 45.65 0.418 2387
MemLong 6496.00 ± 213.56 0.396 2524 Att CLen InC InM SST2 Subj SST5 MPQA MR Avg
32768 ERMAR 13679.03 ± 58.76 0.417 2395 Flash 0 4 4 88.07 57.30 41.42 80.04 83.84 70.13
4 18 88.07 57.30 41.42 80.33 83.86 70.20
MemLong 13449.90 ± 56.22 0.410 2436 20 18 93.92 50.35 47.05 73.63 91.50 71.69
Eager 0 4 4 88.19 51.55 41.42 79.92 83.89 69.59
4 18 88.19 51.55 41.51 80.29 83.92 69.89
Table 5: Latency and throughput comparison on the PG- 20 18 94.04 82.85 46.96 73.63 91.54 77.80
19 dataset. Latency is reported with standard deviation, Flash 2048 4 4 88.07 51.55 41.42 80.04 83.84 68.98
4 18 88.07 51.25 41.42 80.33 83.86 68.99
and throughput is measured in tokens per second. 20 18 93.92 82.75 47.05 73.63 91.50 77.77
Eager 2048 4 4 88.19 51.55 41.51 79.86 83.89 68.96
4 18 88.19 51.55 41.51 79.60 83.92 68.95
20 18 94.04 82.85 46.96 73.63 91.34 77.76
ERMAR exhibits slightly higher latency (9–35%
relative increase at 1K–4K tokens, narrowing to Table 7: Performance comparison of the BGE embedder
1.7–5.5% at 16K–32K) while maintaining compet- across configurations on five NLU tasks. Accuracy is
itive throughput of approximately 2.4K tokens/s reported in percentage. Abbreviations: Att = Attention,
at longer contexts. Importantly, ERMAR demon- CLen = Context length, Avg = Average.
strates much lower latency variance (±45–59 ms)
compared to MemLong (±154–214 ms), reflecting Att CLen InC InM SST2 Subj SST5 MPQA MR Avg
Flash 0 4 4 88.07 57.30 41.42 80.04 83.84 70.13
greater runtime stability in long-context settings. 4 18 88.07 57.30 41.42 80.33 83.86 70.20
20 18 93.92 50.35 47.05 73.63 91.54 71.70
This reduced variance aligns with ERMAR’s adap- Eager 0 4 4 88.19 57.30 41.51 79.86 83.89 70.15
4 18 88.19 57.30 41.51 79.60 83.92 70.10
tive ranked retrieval strategy, which streamlines 20 18 94.04 50.35 46.96 73.63 91.54 71.70
key–value access and minimizes redundant mem- Flash 2048 4
4
4
18
88.07
88.07
51.55
51.55
41.42
41.42
80.04 83.84
80.33 83.86
68.98
69.05
ory operations. Eager 2048
20
4
18
4
93.92
88.19
82.85
51.25
47.05
41.51
73.63 91.50
79.92 83.89
77.79
68.95
4 18 88.19 51.25 41.51 80.29 83.92 69.03
Single-Model Performance Scaling. ERMAR’s 20 18 94.04 82.75 46.96 73.63 91.34 77.74
standalone performance across sequence lengths on Table 8: Performance comparison of the LLM embedder
the WikiText-103 dataset (16K memory capacity) is across configurations on five NLU tasks. Accuracy is
presented in Table 6. reported in percentage.
These results highlight ERMAR’s scalability,
achieving an eightfold improvement in per-token The results indicate that both BGE and LLM
memory efficiency (7.13 MB → 0.90 MB) while embedders yield comparable overall performance,
maintaining stable perplexity across increasing se- with BGE showing slight advantages under cer-
quence lengths. The slightly higher throughput ob- tain configurations. Notably, in the 20-shot, 18
in-memory setup with Eager attention and zero ity compared to MemLong and maintains stable
context length, BGE achieves a higher average ac- performance at 32K-token contexts. In few-shot
curacy (77.80%) compared to LLM (71.70%). The in-context learning tasks, ERMAR outperforms
largest gap occurs in the Subj task, where BGE baselines by 6.1% on average in 4-shot and 7.1%
reaches 82.85% versus LLM’s 50.35%, suggest- in 20-shot settings, including substantial gains on
ing that BGE better captures fine-grained semantic MPQA (+21%) and Subj (+26%), confirming its
representations. Nevertheless, both embedders re- effectiveness in retrieval precision and generaliza-
main broadly consistent, highlighting ERMAR’s tion.
adaptability to different embedding backbones. ERMAR also exhibits strong computational ef-
ficiency, reducing reserved memory by up to 30%
Ablation on Relevance and Re-ranking. We
at long contexts and sustaining a throughput of ap-
further evaluate the effect of the Relevance+Re-
proximately 2.4K tokens/sec with minimal latency
ranking mechanism on model perplexity across
variance. These results establish ERMAR as an ef-
varying context lengths for WikiText-103 and PG-
ficient and scalable retrieval-augmented framework
19. Results are shown in Table 9.
for long-context modeling. Future work will focus
Dataset Context
Without With on optimizing ERMAR for specialized datasets and
Relevance+Reranker Relevance+Reranker
Length
Perplexity Perplexity expanding its applicability to complex reasoning
1K 8.841 7.919 tasks.
2K 7.984 7.410
Wiki 4K 7.438 7.437
16K 7.267 7.082
32K 7.938 8.008
5 Limitations
1K 11.451 10.322
2K 10.412 9.746 While ERMAR improves retrieval efficiency and
PG19 4K 9.932 9.780
16K 9.910 9.809 context retention, it has certain limitations. Its re-
32K 9.858 9.765
liance on ranked memory structures increases com-
Table 9: Ablation study on the Relevance+Re-ranking putational overhead compared to standard LLMs,
mechanism showing perplexity differences for WikiText- particularly for large-scale retrieval, as discussed
103 and PG-19 across context lengths. in Subsections 3.2.4 and 3.2.5, as well as Ap-
pendix A.5. Additionally, performance variations
The relevance scoring mechanism provides across different task domains indicate the need for
the greatest improvements at shorter context further tuning. The framework’s effectiveness in
lengths—10.4% at 1K and 7.2% at 2K tokens on real-world, noisy environments also requires fur-
WikiText-103. The gains taper off beyond 4K to- ther validation.
kens and become marginally negative at 32K, sug-
gesting that the computational overhead of ranking Acknowledgments
may outweigh its benefits when sufficient mem- Shoaib Jameel is supported by the Turing’s De-
ory capacity is available. In contrast, PG-19 ex- fence and Security programme through a partner-
hibits consistent improvements (1.0–9.9%) across ship with the UK government in accordance with
all context lengths, indicating that narrative text the framework agreement between GCHQ & The
benefits more from semantic ranking, where long- Alan Turing Institute. This research is partially sup-
range dependencies and non-sequential references ported by the Technology Innovation Institute, Abu
are frequent. Dhabi, UAE. Computational experiments were con-
ducted using the Wolfpack cluster, supported by
4 Conclusion the School of Computer Science and Engineering
This paper introduced ERMAR, a Ranked at UNSW Sydney.
Memory-Augmented Retrieval framework that en-
hances long-context language modeling through dy-
namic relevance scoring and adaptive memory man- References
agement. By prioritizing semantically relevant in- Marah Abdin, Jyoti Aneja, Hany Awadalla, Ahmed
formation during retrieval, ERMAR improves con- Awadallah, Ammar Ahmad Awan, Nguyen Bach,
Amit Bahree, Arash Bakhtiari, Jianmin Bao, Harkirat
textual reasoning and scalability over extended se- Behl, et al. 2024. Phi-3 technical report: A highly ca-
quences. Experimental results demonstrate that ER- pable language model locally on your phone. arXiv
MAR achieves up to a 3.2% reduction in perplex- preprint arXiv:2404.14219.
Zhangir Azerbayev, Edward Ayers, and Bartosz Pi- Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio
otrowski. 2023. Proofpile: A pre-training dataset Petroni, Vladimir Karpukhin, Naman Goyal, Hein-
of mathematical texts. rich Küttler, Mike Lewis, Wen-tau Yih, Tim Rock-
täschel, et al. 2020. Retrieval-augmented generation
Iz Beltagy, Matthew E Peters, and Arman Cohan. 2020. for knowledge-intensive nlp tasks. Advances in neu-
Longformer: The long-document transformer. arXiv ral information processing systems, 33:9459–9474.
preprint arXiv:2004.05150.
Weijie Liu, Zecheng Tang, Juntao Li, Kehai Chen, and
Zhe Cao, Tao Qin, Tie-Yan Liu, Ming-Feng Tsai, and Min Zhang. 2024. Memlong: Memory-augmented
Hang Li. 2007. Learning to rank: from pairwise retrieval for long text modeling. arXiv preprint
approach to listwise approach. In Proceedings of the arXiv:2408.16967.
24th international conference on Machine learning,
pages 129–136. Stephen Merity, Caiming Xiong, James Bradbury, and
Richard Socher. 2016. Pointer sentinel mixture mod-
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, els. arXiv preprint arXiv:1609.07843.
Zhijian Liu, Song Han, and Jiaya Jia. 2023. Longlora:
Rodrigo Nogueira and Kyunghyun Cho. 2019. Pas-
Efficient fine-tuning of long-context large language
sage re-ranking with bert. arXiv preprint
models. arXiv preprint arXiv:2309.12307.
arXiv:1901.04085.
Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Han- Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and En-
naneh Hajishirzi, Yoon Kim, and Hao Peng. 2024. rico Shippole. 2023. Yarn: Efficient context window
Data engineering for scaling language models to 128k extension of large language models. arXiv preprint
context. arXiv preprint arXiv:2402.10171. arXiv:2309.00071.
Alex Graves, Greg Wayne, and Ivo Danihelka. Ofir Press, Noah A. Smith, and Mike Lewis. 2022.
2014. Neural turing machines. arXiv preprint Train short, test long: Attention with linear bi-
arXiv:1410.5401. ases enables input length extrapolation. Preprint,
arXiv:2108.12409.
Alex Graves, Greg Wayne, Malcolm Reynolds,
Tim Harley, Ivo Danihelka, Agnieszka Grabska- Jack W Rae, Anna Potapenko, Siddhant M Jayakumar,
Barwińska, Sergio Gómez Colmenarejo, Edward and Timothy P Lillicrap. 2019. Compressive trans-
Grefenstette, Tiago Ramalho, John Agapiou, et al. formers for long-range sequence modelling. arXiv
2016. Hybrid computing using a neural network with preprint arXiv:1911.05507.
dynamic external memory. Nature, 538(7626):471–
476. Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan,
Wen Bo, and Yunfeng Liu. 2024. Roformer: En-
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasu- hanced transformer with rotary position embedding.
pat, and Mingwei Chang. 2020. Retrieval augmented Neurocomputing, 568:127063.
language model pre-training. In International confer-
ence on machine learning, pages 3929–3938. PMLR. Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier
Martinet, Marie-Anne Lachaux, Timothée Lacroix,
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal
Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Azhar, et al. 2023a. Llama: Open and effi-
and Weizhu Chen. 2021. Lora: Low-rank adap- cient foundation language models. arXiv preprint
tation of large language models. arXiv preprint arXiv:2302.13971.
arXiv:2106.09685.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Al-
bert, Amjad Almahairi, Yasmine Babaei, Nikolay
Gautier Izacard and Edouard Grave. 2020. Leverag-
Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti
ing passage retrieval with generative models for
Bhosale, et al. 2023b. Llama 2: Open founda-
open domain question answering. arXiv preprint
tion and fine-tuned chat models. arXiv preprint
arXiv:2007.01282.
arXiv:2307.09288.
Vladimir Karpukhin, Barlas Oğuz, Sewon Min, Patrick Szymon Tworkowski, Konrad Staniszewski, Mikołaj
Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Pacek, Yuhuai Wu, Henryk Michalewski, and Piotr
Wen-tau Yih. 2020. Dense passage retrieval for Miłoś. 2024. Focused transformer: Contrastive train-
open-domain question answering. arXiv preprint ing for context scaling. Advances in Neural Informa-
arXiv:2004.04906. tion Processing Systems, 36.
Omar Khattab and Matei Zaharia. 2020. Colbert: Effi- A Vaswani. 2017. Attention is all you need. Advances
cient and effective passage search via contextualized in Neural Information Processing Systems.
late interaction over bert. In Proceedings of the 43rd
International ACM SIGIR conference on research Yuhuai Wu, Markus N Rabe, DeLesley Hutchins, and
and development in Information Retrieval, pages 39– Christian Szegedy. 2022. Memorizing transformers.
48. arXiv preprint arXiv:2203.08913.
Shitao Xiao, Zheng Liu, Yingxia Shao, and Zhao Cao. • Retrieval attention layers: [14, 18, 22, 26]
2022. Retromae: Pre-training retrieval-oriented
language models via masked auto-encoder. arXiv • Memory capacity: 32,768 key-value pairs
preprint arXiv:2205.12035.
• Memory group size: 128 tokens per memory
Howard Yen, Tianyu Gao, and Danqi Chen. 2024. Long-
context language modeling with parallel context en-
group
coding. arXiv preprint arXiv:2402.16617.
• Retrieval group size: 8 (TopK retrieval)
Manzil Zaheer, Guru Guruganesh, Kumar Avinava
Dubey, Joshua Ainslie, Chris Alberti, Santiago On- • Gate mechanism: Disabled (use_gate=False)
tanon, Philip Pham, Anirudh Ravula, Qifan Wang,
Li Yang, et al. 2020. Big bird: Transformers for Training Hyperparameters:
longer sequences. Advances in neural information
processing systems, 33:17283–17297. • Learning rate: 5 × 10−5 with 1,000 warmup
steps
A Appendix
• Weight decay: 1 × 10−4
This appendix supplements the main text with de-
tailed implementation and evaluation specifics for • Batch size: 1 per device (with gradient accu-
the Enhanced Ranked Memory Augmented Re- mulation)
trieval (ERMAR) framework. It includes compre-
hensive descriptions of dataset preprocessing, train- • Sequence length: 1,024 tokens
ing configurations, hyperparameters, and key ter- • Last context length: 1,024 tokens
minology, along with extended analyses of context-
length performance and memory efficiency. The • Training epochs: 1 epoch on SlimPajama 0.5B
provided details ensure reproducibility and offer subset
deeper insights into ERMAR’s architectural and
• Training mode: LoRA-freeze (partial parame-
operational nuances.
ter updates)
A.1 Dataset Preprocessing
LoRA Configuration:
The SlimPajama dataset, used for fine-tuning ER-
MAR, underwent several preprocessing steps to • Target modules: q_proj, k_proj, v_proj,
ensure compatibility with long-context tasks. The o_proj
dataset was tokenized using the OpenLLaMA tok-
• Trainable parameters: Layer normalization
enizer, with a maximum sequence length of 32768
and embeddings
tokens. Duplicate sequences were removed using a
hash-based deduplication algorithm, reducing the • Frozen layers: Layers 0-13 (up to memory
dataset to 84.7K unique training rows. To handle layer)
variable context lengths, we applied a sliding win-
dow approach with a stride of 512 tokens, ensuring • Position encoding: Zero position type for ex-
that the model could process contexts ranging from tended contexts
1024 to 32768 tokens. Special tokens were added Embedder Setup:
to denote document boundaries, and padding was
applied to align sequences to the nearest multiple • Embedder: BAAI/bge-m3 (BGE embedder)
of 128 tokens for efficient batch processing.
• Embedding dimension: 1,024
A.2 Training Configuration
• GPU-based similarity search for efficient re-
ERMAR was trained using a two-stage fine-tuning trieval
approach with the following configuration:
Model Parameters: Hardware and Infrastructure:
• Base model: OpenLLaMA-3B-v2 with LoRA • Primary training: Single NVIDIA 3090 24GB
adaptations GPU
• Memory layer: Layer 13 for historical context • Extended context (32k): NVIDIA L40S
storage 44.4GB GPU
• Distributed training: ZeRO-2 optimization Figure 4 presents detailed perplexity measurements
with Accelerate across incremental context lengths from 1000 to
16000 tokens, providing fine-grained insights into
• Memory optimization: Sequential batching ERMAR’s scaling behavior. The fine-grained anal-
and continual fine-tuning ysis reveals that WikiText-103 achieves optimal
performance around 13K-16K tokens (perplexity
A.3 Glossary of Terms
7.36), while PG-19 maintains consistent perfor-
To aid understanding, we provide definitions for mance (9.74-9.97) across all context lengths. This
key terms used in the ERMAR framework: validates ERMAR’s robustness and suggests that
factual content benefits more from extended con-
• Relevance Score (α): A normalized score text than narrative text.
computed via softmax over the dot product of
query and key embeddings, representing the A.5 Relevance+Reranker Visual Analysis
contextual importance of a memory entry (see
Equation 2).
• Key-Value Pair: A tuple (Kj , Vj ) storing
contextual information, where Kj is the key
embedding and Vj is the corresponding value
embedding in the memory bank.
• RSAR (Relevance Scoring with Adaptive
Retrieval): The mechanism that dynamically Figure 5: Visualization of Table 9, showing the relative
ranks key-value pairs based on their relevance impact of the Relevance+Reranker mechanism across
to the query, incorporating historical usage context lengths. Color coding highlights: blue bars
patterns. represent baseline performance, while orange bars show
improvements with our full mechanism(ERMAR).
• Memory Bank: A non-trainable storage of
pre-computed key-value embeddings, used to Figure 5 provides complementary visual evidence
retain historical context without recomputa- for the patterns discussed in Section 3.2.6. The
tion. stronger improvements at shorter contexts are vi-
sually apparent through the larger differentials be-
• TopK Retrieval: The process of selecting the
tween the orange and blue bars, while differences
top-K most relevant memory entries based on
in mechanism effectiveness across datasets are
their relevance scores for use in the current
immediately noticeable in the side-by-side com-
context.
parisons. Additionally, the 32K edge case on
A.4 Fine-Grained Context Length Analysis WikiText, where the mechanism underperforms,
is clearly highlighted in the visualization, under-
scoring context-length sensitivities of the Rele-
vance+Reranker mechanism.
A.6 Impact of Top-K on Performance
We conducted experiments to examine how varying
the Top-K retrieval parameter affects system per-
formance across different sequence lengths on the
WikiText-103 dataset. The Top-K values (1, 4, 8, 16,
32) were tested for sequence lengths of 1K–16K
tokens, measuring latency, throughput, memory
Figure 4: ERMAR perplexity performance across fine- usage, and perplexity.
grained context lengths for PG-19 and WikiText-103
datasets. The analysis reveals dataset-specific scaling Results. Table 10 summarizes latency and mem-
patterns and validates performance stability across ex- ory results across all configurations. Increasing
tended contexts. Top-K consistently reduces latency per token and
improves throughput, particularly between K = 1 Embedder In-C,
In-M
Dim MPQA SST2 SST5 SUBJ MR
and K = 4, with diminishing gains beyond K = 8. 512 88.28 93.00 45.96 94.90 97.15
Memory consumption (both peak and reserved) 4,4 1024
2048
87.93
87.99
91.51
91.51
46.32
46.14
94.65
94.35
97.09
97.08
remains constant for a given sequence length, con- BGE 512 88.02 92.89 46.14 94.90 97.11
firming that Top-K selection does not affect model 4,18 1024
2048
88.11
88.17
91.40
91.40
46.32
46.14
94.65
94.35
97.09
97.08
footprint. 512 91.23 94.50 52.59 95.55 97.83
20,18 1024 91.26 93.35 49.41 96.60 97.44
2048 91.29 93.23 49.68 96.75 97.40
Seq TopK Latency/ Latency Throughput Peak Reserved
Len Token (s) (tokens/sec) Memory Memory 512 88.65 93.00 46.32 89.05 97.62
(ms) (GB) (GB) 4,4 1024 88.76 92.20 46.32 88.15 96.65
2048 88.76 92.20 46.32 88.15 96.65
1 285.38 ± 54.98 0.279 3684.15 7.98 8.19
4 277.15 ± 60.41 0.271 3810.19 7.98 8.21 LLM 512 88.59 93.00 46.32 89.25 97.62
1024 8 270.90 ± 57.72 0.265 3893.17 7.98 8.21 4,18 1024 88.72 92.20 46.32 88.15 96.65
16 268.60 ± 57.54 0.262 3928.15 7.98 8.21 2048 88.72 92.20 46.32 88.15 96.65
32 269.26 ± 56.84 0.263 3915.67 7.98 8.21
512 90.35 92.78 50.32 95.15 98.08
1 630.68 ± 49.07 0.308 3264.37 8.46 8.83 20,18 1024 91.01 93.35 48.14 94.85 98.04
4 556.76 ± 51.17 0.272 3704.12 8.46 8.81 2048 91.01 93.35 48.14 94.85 98.04
2048 8 555.35 ± 51.52 0.271 3714.39 8.46 8.81
16 549.25 ± 50.81 0.268 3755.02 8.46 8.81
32 549.17 ± 50.93 0.268 3756.11 8.46 8.81 Table 12: Accuracy results for embedding dimensional-
1 1873.54 ± 54.43 0.457 2187.96 9.43 9.89 ity experiments across datasets and settings.
4 1804.90 ± 56.60 0.441 2271.53 9.43 9.91
4096 8 1759.28 ± 53.03 0.430 2330.24 9.43 9.89
16 1738.93 ± 44.91 0.425 2356.96 9.43 9.90
32 1719.64 ± 50.91 0.420 2383.85 9.43 9.90
1 4264.45 ± 63.90 0.521 1921.41 11.36 12.16 A.7 Pruning Strategy Validation
4 4116.54 ± 71.41 0.503 1990.60 11.36 12.18
8192 8 4094.69 ± 79.01 0.500 2001.35 11.36 12.16
16 4107.60 ± 94.04 0.501 1995.37 11.36 12.16
The proportional settings (10% recent, 80% middle,
32 4098.25 ± 96.92 0.500 2000.00 11.36 12.16 10% oldest) were initially heuristic. We present sys-
1 9136.22 ± 176.62 0.558 1793.95 15.21 16.69
4 8965.27 ± 115.76 0.547 1827.79 15.21 16.71 tematic experiments validating these ratios across
16348 8 8961.31 ± 55.20 0.547 1828.37 15.21 16.69
16 8945.37 ± 42.26 0.546 1831.60 15.21 16.70
context lengths and text genres.
32 8843.33 ± 48.01 0.540 1852.75 15.21 16.70
Setup. We evaluated 11 pruning strategies on
Table 10: Performance metrics across Top-K values and WikiText-103 at five context lengths (1K–32K
sequence lengths on WikiText-103.
tokens). Baselines: Uniform (equal priority),
Recency-only, and Frequency-only. Hybrid
Constant Metrics. Table 11 reports memory-per- strategies follow the X-Y -Z format (recent%–
token and perplexity, which remain stable across middle%–oldest%): 0-90-10, 5-80-15, 5-85-10, 10-
Top-K variations. This indicates that increasing re- 80-10 (proposed), 10-85-5, 10-90-0, 15-75-10, and
trieval scope affects efficiency but not model qual- 20-70-10. All strategies yield identical perplexity
ity. per context length (1K: 8.42; 2K: 7.61; 4K: 7.62;
16K: 7.80; 32K: 7.86), so results below reflect effi-
Sequence Memory/Token Val Perplexity ciency differences only.
Length (MB) Perplexity
1024 7.98 2.13 8.42 Results. Tables 13 and 14 report efficiency and
2048 4.23 2.03 7.61
4096 2.36 2.03 7.62 latency results respectively for all strategies across
8192 1.42 2.05 7.76 all context lengths. Table 13 shows that hybrid
16348 0.95 2.05 7.80 strategies consistently achieve lower latency per
token and higher throughput than baselines, while
Table 11: Constant metrics across Top-K values for each
Table 14 reveals that hybrids also exhibit tighter
sequence length.
standard deviations, indicating more stable infer-
ence. Table 15 summarizes the best-performing
Findings. Increasing Top-K from 1→4 yields the strategy per context length regime.
largest latency reduction (e.g., 285.4→277.1 ms at
1K; 9136.2→8965.3 ms at 16K), while further in- Findings. (1) Hybrid strategies outperform
creases provide minimal additional gains. Through- baselines, yielding 10–45% throughput gains
put improvements are proportional, especially for across all lengths (e.g., +45.4% vs. Uniform at
longer sequences (e.g., 2188→2384 tokens/s at 1K; +17.0% at 32K). (2) Optimal ratios shift
4K). Memory and perplexity remain unaffected. with context length: at short contexts (1K–2K),
Thus, K = 4–8 offers the best trade-off between 5-85-10 dominates by minimising the recent slot
efficiency and computational cost. and maximising the middle; at 4K our proposed
1K 2K 4K 16K 32K
Strategy ms/tok tok/s ms/tok tok/s ms/tok tok/s ms/tok tok/s ms/tok tok/s
0-90-10 0.285 3619 0.266 3798 0.374 2676 0.444 2253 0.465 2153
5-80-15 0.259 4023 0.259 3906 0.347 2884 0.441 2269 0.458 2184
5-85-10 0.206 5069 0.205 4946 0.336 2981 0.416 2404 0.414 2413
10-80-10 (proposed) 0.262 3981 0.241 4202 0.329 3044 0.453 2209 0.424 2359
10-85-5 0.259 3979 0.228 4428 0.344 2913 0.424 2358 0.446 2243
10-90-0 0.308 3512 0.264 3899 0.363 2760 0.452 2213 0.473 2116
15-75-10 0.283 3875 0.224 4631 0.368 2736 0.417 2400 0.435 2297
20-70-10 0.254 4067 0.242 4197 0.346 2894 0.432 2316 0.452 2212
Frequency-only 0.278 3737 0.267 3790 0.372 2688 0.466 2145 0.482 2076
Recency-only 0.264 3921 0.263 3850 0.360 2779 0.453 2206 0.478 2094
Uniform 0.295 3486 0.283 3566 0.375 2673 0.469 2133 0.485 2062
Table 13: Latency per token (ms/tok) and throughput (tok/s) for all pruning strategies across context lengths on
WikiText-103. All strategies yield identical perplexity at each context length. Bold denotes the best throughput per
context length. Hybrid strategies are separated from baselines by a horizontal rule.
Strategy 1K (s) 2K (s) 4K (s) 16K (s) 32K (s)
0-90-10 1.39±59.71 544.75±61.74 1533.22±63.16 7273.54±65.57 15222.90±76.59
5-80-15 265.24±67.89 529.87±60.99 1422.65±58.80 7222.95±134.59 15003.15±76.86
5-85-10 210.65±54.01 420.44±58.35 1376.97±64.84 6814.59±52.75 13580.93±41.21
10-80-10 (prop.) 267.94±68.49 492.74±57.10 1348.11±60.15 7419.04±81.91 13894.24±98.75
10-85-5 265.54±57.85 467.92±57.57 1408.47±57.39 6947.99±53.99 14608.47±81.73
10-90-0 315.16±121.40 540.34±110.69 1487.03±66.01 7405.31±50.24 15488.75±77.02
15-75-10 289.29±123.63 458.33±109.72 1506.06±125.70 6826.80±60.86 14264.39±89.72
20-70-10 260.48±59.69 495.15±69.11 1417.23±57.05 7073.92±47.16 14811.33±64.79
Frequency-only 284.33±67.71 545.86±61.15 1524.97±44.72 7641.64±127.35 15787.38±72.64
Recency-only 270.18±62.35 538.19±66.10 1476.49±67.39 7428.95±49.05 15648.46±61.91
Uniform 302.29±63.34 579.82±62.77 1535.08±71.37 7683.59±87.02 15888.88±85.68
Table 14: Total latency (mean ± std, in seconds) for all pruning strategies across context lengths on WikiText-103.
Lower values and tighter standard deviations indicate more stable inference. Hybrid strategies are separated from
baselines by a horizontal rule.
Context Best Strategy Throughput (tok/s) A.8 Sensitivity to Embedding Dimensionality
1K / 2K 5-85-10 5069 / 4946 in ICL Retrieval
4K 10-80-10 3044
16K / 32K 5-85-10 2404 / 2413
We analyzed how embedding dimensionality influ-
ences in-context retrieval performance using both
Table 15: Best-performing strategy per context length BGE and LLM embedders. Experiments were con-
regime. ducted on MPQA, SST-2, SST-5, SUBJ, and MR
datasets with in-context/in-memory settings of (4,
4), (4, 18), and (20, 18). Embedding dimensions of
512, 1024, and 2048 were evaluated.
Results. Table 12 reports results for both em-
bedders. Overall, reducing the dimension to 512
(via PCA) maintains or slightly improves accuracy
10-80-10 is best; at long contexts (≥16K), 5-85-
across most datasets compared to higher dimen-
10 reasserts via aggressive middle-biased pruning.
sions, suggesting redundancy in high-dimensional
(3) Quality is unaffected: perplexity is identical
embeddings.
across all strategies at each context length, confirm-
ing that pruning ratios govern efficiency only. (4) Findings. Reducing dimensionality to 512 pre-
Hybrids are most stable: hybrid configurations serves retrieval accuracy while slightly improving
achieve latency variance of ±40–70 ms vs. ±45– performance on SST-5 (+3.2 pp) and MPQA (+0.3
127 ms for baselines, making them more reliable pp). This indicates that lower dimensions sup-
in deployment. press noise and redundancy. Dimensional padding
to 2048 introduces negligible changes (<0.5 pp).
Across both embedders, SUBJ shows minor degra-
dation at high dimensionality, whereas MR remains
stable. Overall, 512-dimensional embeddings offer
the best trade-off between accuracy and efficiency.
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/text/2503.17876v1.txt
+753
View File
@@ -0,0 +1,753 @@
Satisfactory Medical Consultation based on
Terminology-Enhanced Information Retrieval
and Emotional In-Context Learning⋆
Jing Tang1,∗ , Kaiwen Zuo2,∗ , Hanbing Qin3 , Binli Luo1 , Ligang He2 , and
Shiyan Tang4,∗∗
arXiv:2503.17876v1 [cs.CL] 22 Mar 2025
1
School of Mechanical Science and Engineering, Huazhong University of Science and
Technology, Wuhan 430074, Hubei, China, focusers@163.com
2
School of Computer Science, University of Warwick, Coventry, CV4 7AL, UK,
Kaiwen.Zuo@warwick.ac.uk,ligang.he@warwick.ac.uk
3
School of Software and Microelectronics, Peking University, Beijing 100871, China,
qinhanbing@stu.pku.edu.cn
4
MRC Weatherall Institute of Molecular Medicine, University of Oxford, Oxford,
OX1 2JD, UK, shiyan.tang@ndcls.ox.ac.uk
Abstract. Recent advancements in Large Language Models (LLMs)
have marked significant progress in understanding and responding to
medical inquiries. However, their performance still falls short of the
standards set by professional consultations. This paper introduces a
novel framework for medical consultation, comprising two main mod-
ules: Terminology-Enhanced Information Retrieval (TEIR) and Emo-
tional In-Context Learning (EICL). TEIR ensures implicit reasoning
through the utilization of inductive knowledge and key terminology re-
trieval, overcoming the limitations of restricted domain knowledge in
public databases. Additionally, this module features capabilities for pro-
cessing long context. The EICL module aids in generating sentences with
high attribute relevance by memorizing semantic and attribute informa-
tion from unlabelled corpora and applying controlled retrieval for the
required information. Furthermore, a dataset comprising 803,564 consul-
tation records was compiled in China, significantly enhancing the model’s
capability for complex dialogues and proactive inquiry initiation. Com-
prehensive experiments demonstrate the proposed method’s effectiveness
in extending the context window length of existing LLMs. The experi-
mental outcomes and extensive data validate the framework’s superiority
over five baseline models in terms of BLEU and ROUGE performance
metrics, with substantial leads in certain capabilities. Notably, ablation
studies confirm the significance of the TEIR and EICL components. In
addition, our new framework has the potential to significantly improve
patient satisfaction in real clinical consulting situations.
Keywords: Large Language Models (LLMs) · Medical Consultation ·
Terminology-Enhanced Information Retrieval (TEIR) · Emotional In-
Context Learning (EICL)
⋆
These authors contributed to the work equally and should be regarded as co-first
authors.
** Corresponding author.
2 J. Tang et al.
1 Introduction
Medical consultation has always been a cornerstone of healthcare services, pro-
viding necessary guidance and information to individuals seeking medical atten-
tion [1]. Traditionally, these consultations occurred face-to-face within clinical
settings. However, the uneven distribution of medical resources and an underde-
veloped referral system in China have led to prolonged waiting times and limited
consultation periods for many patients [2]. Additionally, patients often encounter
fear and anxiety when seeking medical help. Consequently, the language em-
ployed in communication during medical consultations holds great significance
for both healthcare providers and patients [2]. Should doctors fall short of meet-
ing patients’ unrealistically high expectations during medical consultations [1],
it can result in dissatisfaction among patients. Therefore, patients’ emotions and
satisfaction during these consultations should never be neglected.
Nowadays, the rapid evolution of digital technology has ushered in a new
era of medical service delivery. A significant proportion of medical consultations
transpire online [3–5], offering unprecedented accessibility and convenience to
patients. This digital transformation has not only redefined the interaction be-
tween healthcare providers and patients but has also generated a massive influx
of data [6]. This article delves into the realm of medical consultation in China
by meticulously analyzing 803,564 instances of such data, collected through ad-
vanced web crawlers designed to navigate and retrieve information from a myriad
of online consultation platforms. An example of medical dialogue, as illustrated
in Figure 1, emphasizes that by engaging in proactive keywords and emotional
inquiries, a more nuanced understanding of the patient’s experience and poten-
tial concerns can be obtained, thereby contributing to a more comprehensive
medical diagnosis.
The motivation for embarking on this research journey is rooted in the need
to enhance the existing knowledge-enhanced information retrieval models preva-
lent in the field. While these models have significantly contributed to the ad-
vancement of medical consultations, they are not devoid of limitations. One of
the primary concerns is the accuracy and relevance of the information provided
during the consultations [7]. Misinformation or irrelevant advice can lead to pa-
tient dissatisfaction, mistrust, and in severe cases, adverse health outcomes [8].
Additionally, the impersonal nature of online consultations often overlooks the
emotional and psychological nuances of patient-provider interactions, further
complicating the effectiveness of digital medical advisories.
Another pertinent issue is the efficiency of information retrieval during con-
sultations. Patients seeking medical advice are often met with overwhelming
amounts of information [9], much of which may be irrelevant or too complex
to understand. This not only prolongs the consultation process but also adds
to the cognitive load of the patients, potentially exacerbating their stress and
anxiety [10]. The existing systems need to be more adept at quickly providing
precise, understandable [11], and relevant information tailored to the individual
needs of each patient.
In light of these challenges, this article introduces a novel framework aimed
at significantly enhancing the quality and effectiveness of online medical consul-
Title Suppressed Due to Excessive Length 3
tations. The proposed framework is the culmination of extensive research and
development, leveraging the latest advancements in information technology, data
analytics, and LLMs. The contributions of this study are manifold, addressing
the core issues identified in the current systems and setting new benchmarks for
what digital medical consultations can achieve. The enhancements are specifi-
cally designed to improve the accuracy, efficiency, and user satisfaction of online
medical consultations, making healthcare advice more accessible and reliable for
everyone. Our contributions are detailed as follows:
1. Generating Enhancement Based on Terminology-Enhanced Infor-
mation Retrieval (TEIR): Traditional terminology match methods are
often limited by their literal approach to query understanding, frequently
overlooking the contextual nuances and user intent. Our approach revolu-
tionizes this mechanism by incorporating advanced linguistic models and
terminology-aware algorithms. These enhancements enable the system to
obtain the underlying intent and specifics of user queries, leading to a more
nuanced and accurate retrieval of information. This not only improves the
relevance of the consultation content but also reduces the time users spend
sifting through unnecessary information.
2. Generate Enhancements Based on Emotional In-Context Learning
(ECIL): ECIL ensures that the information provided is coherent, logically
structured, and easy to navigate. At the same time, emotional annotation
adds a layer of emotional intelligence to the system, enabling it to recognize
and respond to the emotional cues of the users. This not only makes the con-
sultation more personalized and empathetic but also significantly enhances
user engagement and satisfaction.
Retrieval 🎙️Input Emotion
📖 Q: How to take off the fixed retainer?
Endodontics
Rehabilitation medicine
A：Fixed braces need to go to a regular hospital, under the guidance of the doctor to use
professional instruments to remove and cancel, fixed braces will cause damage to the teeth 😭
📰
can not be used again, with braces is mainly the use of professional glue and teeth together to
protect teeth, and then remove the fixed braces, you need to re-do braces for installation. Be
Negative
The association of sure to pay attention to oral hygiene during wearing braces.
orthodontic treatment
and fixed retainers with
gingival health < fixed retainer > 𝐹𝑙𝑎𝑔 = 1
Fig. 1. An example of model responses, which involves retrieving information about
dental treatments, processing inquiries about these treatments, and identifying the
emotional tone of the information presented.
2 Related Work
2.1 Keyword Retrieval Generation
Recently, many studies have provided a diverse look into the application of key-
word retrieval generation in online medical consultations [12]. These cover areas
4 J. Tang et al.
from the training of large language models (LLMs) for medical chats to the
systematic extraction of information from electronic medical records and the
use of metadata and knowledge graphs for improved dialogue generation [13].
These insights are crucial for developing a new framework aimed at improving
the accuracy, efficiency, and patient satisfaction of online medical consultations.
In the realm of enhancing online medical consultation, recent studies have
introduced various innovative approaches. ChatDoctor [14] details the training
of a generic conversation model tailored for medical consultations, emphasiz-
ing advanced keyword retrieval techniques. Concurrently, Yu et al. [15] propose
a taxonomy for the utility of ChatGPT in healthcare, highlighting the role of
Large Language Models in keyword retrieval and information generation. Fur-
ther, Ford et al. [16] examine the use of US hospital-based EMRs, employing
keyword searches and rule-based algorithms to enhance medical document re-
trieval, and Malet et al. [17] discuss the use of Medical Core Metadata to improve
keyword-based retrieval in medical consultations. IAG [18] combines the advan-
tages of retrieval and prompt methods, generates knowledge through keywords
inductive enhancement, and inputs it into the generator along with the retrieved
documents, leading to strong performance in open-domain question-answering
tasks. Lastly, Nature introduces a method for knowledge graph enhanced med-
ical dialogue generation, leveraging the Unified Medical Language System for
more accurate and relevant medical dialogues [19].
Our proposed framework encompasses a multifaceted approach to enhance
online medical consultations. It includes an advanced keyword retrieval gen-
eration system designed to improve the accuracy and relevance of information
retrieval. Additionally, the framework integrates Template Engineering and Sen-
timent Annotation Generation, a feature that analyzes and responds to the emo-
tional tone of patient inquiries, ensuring a more empathetic and personalized
consultation experience. Further details of each component will be introduced
and explored in the following sections.
2.2 Template Engineering and Sentiment Annotation Generation
In the quickly emerging field of online medical consultation, the integration of
various methods such as Template Engineering and Sentiment Annotation shows
immense scope for the enhancement of accuracy, efficiency, and patient satisfac-
tion [20, 21]. This approach, as the proposed paper plans to further explore,
primarily employs techniques previously explored in different fields, the most
notable among them being sentiment analysis and multidimensional trust.
Wouter van Atteveldt et al. demonstrated the effectiveness of manual anno-
tation and crowd coding in sentiment analysis [20], a component that plays a
pivotal role in gaining insights from patients’ sentiments during online medical
consultation. This historical emphasis on human-embedded sentiment analysis
tools might be supplemented by the AraSenCorpus model introduced by Al-laith
et al., which has redefined the scale of sentiment analysis with proficient senti-
ment annotation on a large-scale text corpus [22]. GRACE memorizes semantic
and attribute information from unmarked corpora and obtains the required in-
formation through controllable re-retrieval [23].
Title Suppressed Due to Excessive Length 5
Online medical consultation heavily relies on a patient’s trust, a significantly
multidimensional element [21]. Trust bears an impact on patients’ decision to
continue the use of online services, and as shown by Yang et al., the incorporation
of both technological and interpersonal aspects improves the chances of service
continuity [21]. Further, analysis by Wan et al. revealed other influencing factors
such as doctors’ soft skills and overall service experience that enhance the volume
of doctor consultation on online platforms [24].
In addition to trust, efficiency and patient satisfaction have been recognized
as substantial factors contributing to the success of an online medical consulta-
tion framework. Using the employment of scribes as an example, Gidwani et al.
demonstrated their positive impact on physician satisfaction and charting effi-
ciency, without impinging upon patient satisfaction [25]. The shifting dynamics
were further reflected in Jiang et al.’s study, pointing out service-related features
like service delivery quality and patient involvement to be more significant than
physician reputation, in patients’ payment decisions [26].
Overall, combining these aspects, the proposed method is well-supported
by earlier research, and it is thoughtfully aimed at bringing a comforting user
experience by improving the efficiency of medical consultation [20–22, 24–26].
3 Methodology
Our framework introduces an innovative approach to enhance the reliability of
the generated content: (1) Terminology-Enhanced Information Retrieval
(TEIR) constitutes the first part of our framework, focusing on the extraction
and utilization of contextually significant information. The overall pipeline is
shown in 3.1. (2) Emotional In-Context Learning (EICL) is the second part
of our methodology, dedicated to understanding and integrating the emotional
dimensions of the text. It leverages the generative capabilities of the framework,
enriched with a nuanced understanding of emotional attributes.
3.1 Terminology-Enhanced Information Retrieval (TEIR)
qt Terminology Retriever
Terminology Detector
Retrieval
Enhanced
Tt1 T t2 ... T tk Sentence R1,...,RM
Tt1 Tt1->R1 ... Tt1->RM Tt2 Tt2->R1 ... Tt2->RM Ttk Ttk->R1 ... Ttk->RM
Terminology
Memory
Memory1 Memory2 Memory3
Fig. 2. The overall pipeline of Terminology-Enhanced Information Retrieval (TEIR).
6 J. Tang et al.
The first part of our framework, Terminology-Enhanced Information Re-
trieval, focuses on extracting and utilizing contextually significant information
through a synergy of three core components in the enhancement:
Terminology Detector (TD): Utilize a Universal Information Extraction
(UIE)-based model D to realize terminology detection, ensuring the retrieval
process is sharply focused and aligned with the specified attributes.
Terminology Memory (TM): After TD process we can get N different
terminologies T = {Tn }N n=1 . However, in the actual medical knowledge base,
different terminologies may correspond to the same meaning, so we have to
do a simple filtering and alignment. We save the processed terms to the TM
module, which ensures that subsequent questions can also be indexed to relevant
terminologies and in turn allows the model to be progressively tuned as the
dialogue progresses.
It is worth noting that in this process we link the terminology in the TM
one-to-one with the documents in the knowledge database, which will be used
to speed up the subsequent retrieval process.
Enhanced Sentence Generation (ESG): A more accurate question spe-
cially generated for information retrieval can be achieved through terminology
retrieval and combination using TM, reducing the effect of domain-specific vo-
cabulary and facilitating robust information retrieval.
After that, we perform Information Retrieval based on Enhanced Sentence.
First, we link documents based on terminology T = {Tn }kn=1 in Enhanced Sen-
tence, and the linked document is represented by R = {Rm }M m=1 . After that,
only the document in the link needs to be retrieved.
Considering the specialized nature of medical databases and the similarity of
some documents, here we additionally use the following mapping to amplify the
gap between the confidence probabilities for better document retriever.
exp(pn )
p̂n = map(pn ) = PN , n = 1, ..., N (1)
i=1 exp(pi )
3.2 Emotional In-Context Learning
Emotional In-Context Learning, the second part of our methodology, is dedicated
to the integration and understanding of emotional dimensions:
1. Emotionally Attuned Generation: The generative component G is tasked
with sentence formation while ensuring that the text resonates with the de-
sired emotional undertones, aligning with the target attributes.
2. Gradient-Guided Refinement: Employs a gradient-guided retrieval-generation
mechanism to adjust the generated text at each step, ensuring continual
alignment with the target emotional attributes for contextually and emo-
tionally coherent output.
As shown in Figure 3, we can also understand the EICL in terms of a prob-
(i)
ability space. The regions Ie , Iz , and Id are interpreted as states within the
Title Suppressed Due to Excessive Length 7
�
�� �
�
�� (�)
�� �
�
�
Fig. 3. The geometric diagram illustrates the emotion-influenced iterative generation
process. Representations for the event schema, the document, and the I-th discrete
(i)
demonstration are symbolized by Ie , Ix , and Id correspondingly.
probability space. Ie represents the initial emotional context, Iz the influence
(i)
of the document context, and Id the iterative demonstration towards the true
value. The overlapping areas and the vector R symbolize the adaptive process of
the learning model.In each iteration, the system enhances the fidelity of learn-
(i)
ing outcome Id , bringing it nearer to the true value.Using such an effective
guide, we can make the final generated content closer to the patient’s space,
which refers to a concept or psychological space that summarizes the patient’s
emotional and health background. It plays a significant role in predicting and
aligning with the emotional and informational needs of patients in the model.
To realize this, we adopt a suffix-tuning method to transform continuous tasks
into discrete data usable for training while keeping the LM (Language Model)
parameters unchanged. Mathematically speaking, this is the search for the opti-
mal parameters ϕ that maximize the log-likelihood function given a training set
D and fixed model parameters θ.
X
max P (y|x; θ; ϕ) = max P (yi |h>i ; θ; ϕ) (2)
ϕ ϕ
yi
In equation 2, h>i represents the set of all hidden states after the i-th position.
If simplified, it can be considered as a state of the LM. h>i is obtained through
the model Mϕ , which represents the state of the LM given the model parameters
ϕ. It can also be understood by the following example.
Train: Patient: I have a fever today; Doctor: Drink more hot water. We
predict patient feedback: Negative.
Test: I have a fever today;
This dialogue serves as a training instance for the model to understand and
predict emotional responses based on conversational context. Here, the simplic-
ity of the doctor’s advice in response to a possibly serious symptom might not
address the patient’s concerns adequately, leading to negative feedback. Emo-
tional In-Context Learning aims to use such instances to teach the model the
8 J. Tang et al.
Table 1. Case study of regenerated method
Input: I have a fever today;
Answer: Drink more hot water.
Sentiment: Negative.
Input: I have a fever today. Please do not say, ”Drink more hot water.”;
Answer: Drink adequate amounts of fluids.
Water, juices, clear soups, or hot lemonade are all good choices.
Avoid caffeine and alcohol; these ingredients increase fluid loss.
Sentiment: Positive.
subtleties of emotional undercurrents in conversations, helping it predict and,
in real-world applications, provide responses that are not only factually correct
but also emotionally attuned to the patient’s needs.
More specifically, here we construct a regenerated method using ECIL that
utilizes sentiment to achieve better quality response content generation. In the
table 1 we show a real case study.
Considering the specialized nature of medical databases, some of the re-
sponses are difficult to be well understood by the consultants, and our approach
can achieve increased consultant satisfaction through such context-based self-
supervised feedback.
4 Experiment
In this section, we first introduce the dataset, experimental setup, and base-
line methods, followed by an exposition of the evaluation metrics and overall
results, which include 803,564 medical consultation data entries. To ensure the
efficacy of our model evaluation, this study also employs two metrics: BLEU-5
and ROUGE-2/L for assessment.
4.1 Experiment setup
Datasets In order to develop and enhance the model’s performance, our dataset
utilizes a powerful web crawler system to gather a large dataset containing
803,564 medical consultation records. This dataset spans 12 distinct medical
departments and includes more than 10 varied medical Q&A scenarios, it encom-
passes patient inquiries about symptoms and treatments, health consultations,
medical advice, and general medical knowledge. Furthermore, the data has been
meticulously processed to ensure the removal of any personal identifiers, uphold-
ing patient confidentiality and adhering to strict privacy standards. These data
are beneficial for training the model. Utilizing 30% of the dataset for testing
purposes and allocating the remaining 70% for training the model significantly
enhances the model’s reliability. This distribution of data ensures a comprehen-
sive learning process while maintaining a robust framework for validation and
assessment of the model’s performance.
Title Suppressed Due to Excessive Length 9
Evaluation Metrics. Our results are presented using a range of sophisticated
evaluation metrics to ensure a comprehensive analysis of the model’s perfor-
mance in various aspects of language understanding and generation.
– BLEU-5 (Bilingual Evaluation Understudy): This metric is an ex-
tension of the standard BLEU score, which is widely used to evaluate the
quality of text that has been machine-translated from one language to an-
other. BLEU-5 specifically refers to using the average precision of up to 5
grams (sequences of five consecutive words) to assess translation quality,
providing a more nuanced understanding of the model’s linguistic accuracy
and fluency.
– ROUGE-2/L (Recall-Oriented Understudy for Gisting Evaluation):
ROUGE is primarily used to evaluate automatic summarization and machine
translation. ROUGE-2 refers to the overlap of bigrams between the gener-
ated text and a set of reference texts, effectively measuring the quality of
content reproduction. ROUGE-L, on the other hand, focuses on the longest
common subsequence and is particularly adept at assessing the sentence-level
structure similarity between the generated text and reference texts.
In employing these advanced metrics, we aim to thoroughly assess the model’s
ability to generate coherent, contextually relevant, and linguistically accurate
content. By focusing on both the micro-structure of language (as captured by
BLEU-5) and the effective conveyance of meaning and structure (as captured
by ROUGE-2/L), we provide a holistic evaluation of the model’s performance
in critical tasks of language processing and generation.
Baselines. For comprehensive evaluation and fair comparison, we introduce
five baseline models that are prominent in the field of language generation and
attribute-based modeling.DoctorGLM [27], a specialized model for medical
inquiries; ChatGPT-3.5 [28], an iteration of the GPT series with 3.5 billion
parameters; ChatGLM-6B [29], a 6 billion parameter model designed for con-
versational tasks; Ziya-LLaMA-13B [30], optimized for linguistic tasks with
13 billion parameters; and HuatuoGPT [31], our proprietary model before fine-
tuning. Additionally, we showcase our model T5, which is HuatuoGPT after
extensive fine-tuning to enhance its performance on a wide array of tasks.
Experimental settings. Our experimental setup is built on a robust hardware
environment to facilitate efficient and reliable model training and evaluation.
Specifically, we utilize NVIDIA-V100-32GB GPUs, known for their powerful
computation capabilities, particularly in handling large-scale machine learning
tasks. We initialize our models with the pre-trained T5 model, available in the
HuggingFace Transformers library. We consider two model sizes, base and large,
containing respectively 220M and 770M parameters. Due to GPU memory lim-
itations, we use different batch sizes for different models: 8 for T5-large and 16
for T5-base.
10 J. Tang et al.
Table 2. AMT study results for medical consultant
Model Patient satisfaction
DoctorGLM 5.67% (34/600)
ChatGPT-3.5 7.17% (43/600)
ChatGLM-6B 2.83% (17/600)
Ziya-LLaMA-13B 2.17% (13/600)
HuatuoGPT 12.50% (75/600)
Ours(w/o TEIR) 17.50% (105/600)
Ours(w/o EICL) 15.33% (92/600)
Ours 36.83% (221/600)
4.2 Experiment result
Table 1 presents the results of a study measuring patient satisfaction across var-
ious models. The first column lists the names of the models evaluated, including
DoctorGLM, ChatGPT-3.5, ChatGLM-6B, Ziya-LLaMA-13B, HuatuoGPT, and
variations of Ours. Each model represents a different approach or version in the
medical consultation simulation or relevant language model application fields.
The second column, labeled ”Patient satisfaction,” displays the percentage of
patient satisfaction along with the raw number of satisfied instances out of 600.
The satisfaction rates are expressed as percentages and are derived from a pos-
sibly larger-scale study involving 600 individual assessments or interactions for
each model.
To ensure the integrity of the Amazon Mechanical Turk (AMT) study assess-
ing patient satisfaction with medical consultation responses, participants were
provided with unequivocal instructions. Satisfaction was gauged using a 5-point
Likert scale, ranging from 1 (not satisfied ) to 5 (extremely satisfied ). Quality
control was maintained through embedded attention checks to validate partici-
pant engagement. Consistency in responses was assured by interspersing identi-
cal questions in varying forms throughout the survey, enabling cross-verification
of data reliability. These methodologies combined to form a robust framework,
optimizing the authenticity of the satisfaction metrics obtained.
Notably, the patient satisfaction percentages progressively increase for the
variants of Ours - Ours (w/o TEIR), Ours (w/o EICL), and Ours, with
Ours achieving the highest satisfaction rate at 36.83% (221 out of 600). This is
attributed to the positive impact of the TEIR and EICL features on the Ours
model’s ability to attract and satisfy patients. The inclusion and optimization of
these features in the final iteration of Ours evidently contribute to its superior
performance in patient satisfaction, underscoring the significance of continuous
model enhancement and feature integration in the development of medical con-
sultation models.
In the evaluation of various language models, as displayed in Table 2, we have
witnessed significant advancements in model performance on the Chinese medical
QA dataset. The table outlines the results for a series of models, including T5
Title Suppressed Due to Excessive Length 11
(fine-tuned),the previously five model, and Ours model. It’s noteworthy that
Ours, which is an enhanced iteration based on the fundamental approach of
HuatuoGPT through the Integrated Argument Generation (IAG) enhancement,
demonstrates marked improvements over its counterparts, particularly in metrics
such as BLEU-1, GLEU, ROUGE-1, and Distinct-2—critical indicators of the
model’s linguistic precision and contextual understanding, as shown in Table 2.
The results indicate a significant improvement of Ours over the state-of-the-art
(SOTA) models (24.68 to 26.12 for BLEU-1 and 8.07 to 8.23 for GLEU), and a
comparable leap over another SOTA improvement (27.93 to 29.07 for ROUGE-1
and 0.93 to 0.96 for Distinct-2).
Compared to ChatGLM variants, our model shows a significant advantage;
however, DoctorGLM performs poorly on BLEU and GLEU metrics. This could
be due to the model’s potential difficulty in handling the subtleties of context
within medical texts. Medical language is typically laden with jargon and can
be quite nuanced. If DoctorGLM is not fine-tuned on medical-specific data, it
may fail to capture the complexity required for higher BLEU and GLEU scores.
Nonetheless, it performs well on the Distinct metrics. This is because Doctor-
GLM might employ a generation strategy that favors diversity over precision.
Such an approach might manifest in a model that is less likely to repeat the
same phrases and more inclined to explore the breadth of its training data when
generating responses. In summary, our model outperforms nearly all evaluated
metrics when compared to others.
(a) GPT-4 Evaluation (b) Doctor Evaluation
Fig. 4. 1v1 battle between different models and ours
In order to make a more direct comparison of different models, we used
GPT-4 and a professional doctor to evaluate our model and other models one-
on-one. It can be seen that our model has a significant advantage over the others;
doctorGLM does not perform well on some of the quantitative metrics, but
when viewed by a professional doctor, it has a very good performance, probably
because the content it generates is more in line with the human mind. Overall,
our model won almost all 1v1 battles.
12 J. Tang et al.
Table 3. Test Result on Chinese medical QA dataset
Model BLEU-1 BLEU-2 BLEU-3 BLEU-4 GLEU
T5 (fine-tuned) 21.27 13.65 9.89 7.21 9.07
DoctorGLM 9.89 5.36 2.56 1.45 4.43
ChatGPT-3.5 17.87 6.54 2.37 0.99 4.87
ChatGLM-6B 23.53 12.07 6.54 3.65 8.07
Ziya-LLaMA-13B 22.01 11.34 6.35 4.01 6.99
HuatuoGPT 24.68 13.25 7.88 4.34 7.30
Ours(w/o TEIR) 25.54 13.45 8.01 4.64 7.56
Ours(w/o EICL) 25.73 14.01 8.16 4.97 8.12
Ours 26.12 14.34 8.43 5.01 8.23
Model ROUGE-1 ROUGE-2 ROUGE-L Distinct-1 Distinct-2
T5 (fine-tuned) 30.08 13.76 24.43 0.39 0.50
DoctorGLM 22.79 5.43 11.73 0.82 0.93
ChatGPT-3.5 19.79 2.46 12.76 0.67 0.89
ChatGLM-6B 27.69 6.17 17.98 0.65 0.87
Ziya-LLaMA-13B 27.41 6.90 13.24 0.75 0.91
HuatuoGPT 27.93 7.28 15.20 0.74 0.93
Ours(w/o TEIR) 29.13 8.10 14.01 0.81 0.93
Ours(w/o EICL) 28.23 7.69 13.32 0.76 0.90
Ours 29.07 8.10 14.50 0.84 0.96
4.3 Ablation Study
The ablation study focuses on two components of the Ours model: TEIR and
EICL. By analyzing the performance impact of each component’s removal, we
gain insights into their contributions.
Ablation of TEIR Without TEIR, Ours (w/o TEIR)exhibits a decrease
in BLEU-1 and BLEU-4 scores, indicating the importance of TEIR in captur-
ing accurate unigrams and coherent longer n-gram sequences. The reduction in
the GLEU score suggests that TEIR enhances sentence-level alignment with ref-
erence translations. Lower ROUGE scores point to TEIR’s role in replicating
essential content and maintaining structural integrity.
Ablation of EICL The absence of EICL in Ours (w/o EICL) leads to lowered
performance across ROUGE and Distinct metrics. This suggests that EICL may
be instrumental in capturing critical details and enhancing the diversity and
uniqueness of the generated text.
Combined Impact The results of the ablation study indicate that both TEIR
and EICL significantly contribute to the Ours model’s ability to handle complex
tasks in the Chinese medical QA dataset. Their combined contributions are
pivotal in enhancing the model’s precision and capacity to generate contextually
rich and diverse language.
Title Suppressed Due to Excessive Length 13
5 Conclusion
In conclusion, our method is a retrieval-augmented generation framework. It
innovatively integrates Terminology-Enhanced Information Retrieval and Emo-
tional In-Context Learning, significantly enhancing its ability to understand
complex medical dialogues and emotional contexts. Following rigorous evalua-
tion with 803,564 medical consultations, our method has demonstrated superior
performance compared to existing models. This advancement represents a sub-
stantial contribution to the field of artificial intelligence, particularly in retrieval-
augmented methodologies. It holds immense potential to revolutionize medical
information retrieval and generation, setting a new trajectory for the application
of artificial intelligence in healthcare. This research highlights the transforma-
tive implications of our method, not just as a technological innovation but as a
pivotal step towards more sophisticated, responsive, and empathetic healthcare
solutions. Given the limited time doctors spend with each patient and ongoing
patient complaints, satisfaction with medical consultations is not consistently
achieved. Our new framework has the potential to create a more satisfying med-
ical consultation environment for patients by addressing their emotions. Ad-
ditionally, our framework promises to contribute to future medical education,
offering insights on establishing effective patient relationships, which could lead
to higher patient satisfaction rates.
References
[1] Ilan S Schwartz, Katherine E Link, Roxana Daneshjou, and Nicolás Cortés-
Penfield. Black box warning: large language models and the future of in-
fectious diseases consultation. Clinical Infectious Diseases, page ciad633,
2023.
[2] Yuxin Wang and Shunda Du. Time to rebuild the doctor-patient relation-
ship in china. Hepatobiliary Surgery and Nutrition, 12(2):235, 2023.
[3] Hussain A Younis, Taiseer Abdalla Elfadil Eisa, Maged Nasser,
Thaeer Mueen Sahib, Ameen A Noor, Osamah Mohammed Alyasiri, Sani
Salisu, Israa M Hayder, and Hameed AbdulKareem Younis. A system-
atic review and meta-analysis of artificial intelligence tools in medicine and
healthcare: Applications, considerations, limitations, motivation and chal-
lenges. Diagnostics, 14(1):109, 2024.
[4] Per Engelseth, BE White, Ingunn Mundal, Trude Fløystad Eines, and
Duangpun Kritchanchai. Systems modelling to support the complex na-
ture of healthcare services. Health and Technology, 11:193–209, 2021.
[5] Ian L Alberts, Lorenzo Mercolli, Thomas Pyka, George Prenosil, Kuangyu
Shi, Axel Rominger, and Ali Afshar-Oromieh. Large language models (llm)
and chatgpt: what will the impact on nuclear medicine be? European journal
of nuclear medicine and molecular imaging, 50(6):1549–1552, 2023.
[6] Abdulmotaleb El Saddik and Sara Ghaboura. The integration of chatgpt
with the metaverse for medical consultations. IEEE Consumer Electronics
Magazine, 2023.
[7] Maksut Senbekov, Timur Saliev, Zhanar Bukeyeva, Aigul Almabayeva,
Marina Zhanaliyeva, Nazym Aitenova, Yerzhan Toishibekov, Ildar
Fakhradiyev, et al. The recent progress and applications of digital tech-
nologies in healthcare: a review. International journal of telemedicine and
applications, 2020, 2020.
[8] Changyu Wang, Siru Liu, Hao Yang, Jiulin Guo, Yuxuan Wu, and Jialin
Liu. Ethical considerations of using chatgpt in health care. Journal of
Medical Internet Research, 25:e48009, 2023.
[9] Jonah Kenei. Supporting Information Retrieval From Clinical Narrative
Texts Using Text Classification and Visualization Techniques. PhD thesis,
University of Nairobi, 2022.
[10] Tin Kam Ho, Yen-Fu Luo, and Rodrigo Capobianco Guido. Explainability
of methods for critical information extraction from clinical documents: A
survey of representative works. IEEE Signal Processing Magazine, 39(4):96–
106, 2022.
[11] Piet van der Keylen, Johanna Tomandl, Katharina Wollmann, Ralph
Moehler, Mario Sofroniou, Andy Maun, Sebastian Voigt-Radloff, and Luca
Frank. The online health information needs of family physicians: systematic
review of qualitative and quantitative studies. Journal of medical Internet
research, 22(12):e18816, 2020.
Title Suppressed Due to Excessive Length 15
[12] Yubing Ren, Yanan Cao, Ping Guo, Fang Fang, Wei Ma, and Zheng Lin.
Retrieve-and-sample: Document-level event argument extraction via hybrid
retrieval augmentation. In Annual Meeting of the Association for Compu-
tational Linguistics, 2023.
[13] Xinke Jiang, Ruizhe Zhang, Yongxin Xu, Rihong Qiu, Yue Fang, Zhiyuan
Wang, Jinyi Tang, Hongxin Ding, Xu Chu, Junfeng Zhao, and Yasha Wang.
Think and retrieval: A hypothesis knowledge graph enhanced medical large
language models, 2023.
[14] Yunxiang Li, Zihan Li, Kai Zhang, Ruilong Dan, Steve Jiang, and You
Zhang. Chatdoctor: A medical chat model fine-tuned on a large language
model meta-ai (llama) using medical domain knowledge, 2023.
[15] Ping Yu, Hua Xu, Xia Hu, and Chao Deng. Leveraging generative ai and
large language models: A comprehensive roadmap for healthcare integra-
tion. In Healthcare, volume 11, page 2776. MDPI, 2023.
[16] Elizabeth Ford, John A Carroll, Helen E Smith, Donia Scott, and Jackie A
Cassell. Extracting information from the text of electronic medical records
to improve case detection: a systematic review. Journal of the American
Medical Informatics Association, 23(5):1007–1015, 2016.
[17] Gary Malet, Felix Munoz, Richard Appleyard, and William Hersh. A model
for enhancing internet medical document retrieval with “medical core meta-
data”. Journal of the American Medical Informatics Association, 6(2):163–
172, 1999.
[18] Zhebin Zhang, Xinyu Zhang, Yuanhang Ren, Saijiang Shi, Meng Han,
Yongkang Wu, Ruofei Lai, and Zhao Cao. Iag: Induction-augmented gen-
eration framework for answering reasoning questions. In Proceedings of the
2023 Conference on Empirical Methods in Natural Language Processing,
pages 1–14, 2023.
[19] Deeksha Varshney, Aizan Zafar, Niranshu Kumar Behera, and Asif Ekbal.
Knowledge grounded medical dialogue generation using augmented graphs.
Scientific Reports, 13(1):3310, 2023.
[20] Wouter van Atteveldt, Mariken A.C.G. van der Velden, and Mark Boukes.
The validity of sentiment analysis: Comparing manual annotation, crowd-
coding, dictionary approaches, and machine learning algorithms. Commu-
nication Methods and Measures, 15:121 – 140, 2021.
[21] Ming Yang, Jinglu Jiang, Melody Y. Kiang, and Fangyun Yuan. Re-
examining the impact of multidimensional trust on patients’ online medical
consultation service continuance decision. Information Systems Frontiers,
24:983 – 1007, 2021.
[22] Ali Al-laith, Muhammad Shahbaz, Hind Alaskar, and Asim Rehmat.
Arasencorpus: A semi-supervised approach for sentiment annotation of a
large arabic text corpus. Applied Sciences, 2021.
[23] Zhihua Wen, Zhiliang Tian, Zhen Huang, Yuxin Yang, Zexin Jian,
Changjian Wang, and Dongsheng li. Grace: Gradient-guided controllable
retrieval for augmenting attribute-based text generation. pages 8377–8398,
01 2023.
16 J. Tang et al.
[24] Yan Wan, Ziqing Peng, Yalu Wang, Yifan Zhang, Jinping Gao, and Baojun
Ma. Influencing factors and mechanism of doctor consultation volume on
online medical consultation platforms based on physician review analysis.
Internet Res., 31:2055–2075, 2021.
[25] Risha Gidwani, Cathina T. Nguyen, Alexis Kofoed, Catherine Carragee,
Tracy A Rydel, Ian Nelligan, Amelia L Sattler, Megan Mahoney, and Steven
Lin. Impact of scribes on physician satisfaction, patient satisfaction, and
charting efficiency: A randomized controlled trial. The Annals of Family
Medicine, 15:427 – 433, 2017.
[26] Jinglu Jiang, Ann Frances Cameron, and Ming Yang. Analysis of massive
online medical consultation service data to understand physicians’ economic
return: Observational data mining study. JMIR Medical Informatics, 8,
2020.
[27] Honglin Xiong, Sheng Wang, Yitao Zhu, Zihao Zhao, Yuxiao Liu, Linlin
Huang, Qian Wang, and Dinggang Shen. Doctorglm: Fine-tuning your chi-
nese doctor is not a herculean task. ArXiv, abs/2304.01097, 2023.
[28] Mustafa Said Yıldız. Comparing response performances of chatgpt-3.5,
chatgpt-4 and bard to health-related questions: Comprehensiveness, accu-
racy and being up-to-date. SSRN Electronic Journal, 2023.
[29] Bang Yang, Asif Raza, Yuexian Zou, and Tong Zhang. Customizing general-
purpose foundation models for medical report generation, 2023.
[30] Junyu Lu, Di Zhang, Xiaojun Wu, Xinyu Gao, Ruyi Gan, Jiaxing Zhang,
Yan Song, and Pingjian Zhang. Ziya-visual: Bilingual large vision-language
model via multi-task instruction tuning. ArXiv, abs/2310.08166, 2023.
[31] Hongbo Zhang, Junying Chen, Feng Jiang, Fei Yu, Zhihong Chen, Jianquan
Li, Guimin Chen, Xiangbo Wu, Zhiyi Zhang, Qingying Xiao, Xiang Wan,
Benyou Wang, and Haizhou Li. Huatuogpt, towards taming language model
to be a doctor. In Conference on Empirical Methods in Natural Language
Processing, 2023.
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/text/2602.10271v3.txt
+1034
View File
File diff suppressed because it is too large. Load diff
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/text/2605.08217v1.txt
+596
View File
@@ -0,0 +1,596 @@
ICLR 2026 Workshop on Time Series in the Age of Large Models (TSALM)
R ETRIEVAL M ECHANISMS S URPASS L ONG -C ONTEXT
S CALING IN T IME S ERIES F ORECASTING
Rishi Ahuja† Kumar Prateek† Simranjit Singh† ⊠ Vijay Kumar†
†
Department of Information Technology,
Dr. B.R. Ambedkar National Institute of Technology Jalandhar,
Punjab, 144008, India.
{rishia.it.24, kumarprateek, ⊠ singhsimranjit, vijayk}@nitj.ac.in
arXiv:2605.08217v1 [cs.LG] 6 May 2026
A BSTRACT
Time Series Foundation Models (TSFMs) have borrowed the long context
paradigm from natural language processing under the premise that feeding more
history into the model improves forecast quality. But in stochastic domains, dis-
tant history is often just high-frequency noise, not signal. Hence, the proposed
work tests whether this premise actually holds by running continuous context ar-
chitectures (PatchTST included) through the ETTh1 benchmark. The obtained
results contradict the premise: an inverse scaling law shows up clearly, with fore-
casting error rising as context gets longer. A 3,000-step window causes perfor-
mance to drop by over 68%, evidence that attention mechanisms are poor at ig-
noring irrelevant historical volatility. Retrieval-Augmented Forecasting (RAFT) is
evaluated as an alternative. RAFT achieves a mean squared error (MSE) of 0.379
with a fixed 720-step window and selective retrieval, outperforming both long-
context configurations and zero-shot foundation models (Chronos, Moirai) de-
spite requiring far less computation. In addition, the retrieval step injects only the
most relevant historical segments as dynamic exogenous variables, which gives
the model a context-informed inductive bias it cannot build on its own from raw
sequences. Therefore, foundation models going forward need to shift architec-
turally toward selective retrieval.
Track: Research
1 I NTRODUCTION
Time Series Foundation Models (TSFMs) owe much of their recent popularity to the “scaling hy-
pothesis,” the idea that larger models with longer input windows will produce monotonically better
forecasts. Moirai (Woo et al., 2024) and Chronos (Ansari et al., 2024) have given weight to that
idea: both achieve strong zero-shot results across multiple forecasting domains by utilizing scale.
But a harder question remains unresolved. Can the “Long Context” logic of Large Language Mod-
els (LLMs) be transplanted directly into the stochastic world of time series? There are good reasons
to doubt it. Language tokens retain semantic content even when they appear thousands of positions
earlier in a document. Numerical time series do not behave this way. Old values in a stochastic
process are frequently just noise, and their statistical relationship to future values weakens sharply
over time. Therefore, the proposed work advances the following hypothesis: when the context win-
dow in time series forecasting is expanded without discrimination, a failure mode called Stochastic
Noise Accumulation takes over. Irrelevant historical values flood the attention computation, the
mechanism loses the ability to weight important recent patterns appropriately, and forecasting accu-
racy suffers. Williams et al. (2025) show that short textual descriptions alongside numerical inputs
yield large accuracy gains, highlighting the poverty of raw numerical context. Pineda Arango et al.
(2025) finds that exogenous variables are critical for downstream task adaptation. This work con-
tends that selective retrieval, not blind window expansion, is the correct mechanism for injecting
useful context.
⊠
Corresponding Author.
1
ICLR 2026 Workshop on Time Series in the Age of Large Models (TSALM)
The work uses ETTh1 benchmark for empirical validation and compares three model types: vanilla
Transformers, PatchTST (Nie et al., 2023) (the current state-of-the-art), and a Retrieval-Augmented
Forecasting (RAFT) baseline. The experiments produce an Inverse Scaling Law for stochastic time
series: a) Long Context Fails: PatchTST accuracy drops by 68% when context grows from 720
to 3,000 tokens. Even advanced patching cannot screen out the noise in long historical windows.
b) Retrieval Wins: RAFT (Han et al., 2025) uses a 720-step window plus selective retrieval and
scores the best MSE (0.379), outperforming all long-context and zero-shot foundation model base-
lines. Retrieved segments function as dynamic exogenous variables that provide context-informed
inductive bias, and they do so without the noise penalties attached to long continuous windows.
2 R ELATED W ORK
2.1 T IME S ERIES F OUNDATION M ODELS AND C ONTEXT S CALING
Moirai (Woo et al., 2024) and Chronos (Ansari et al., 2024) achieve strong zero-shot transfer across
multiple forecasting domains by training on large heterogeneous corpora, demonstrating that univer-
sal temporal patterns can be captured at scale. PatchTST (Nie et al., 2023) advances supervised long-
context forecasting through channel-independent patching, reducing quadratic attention overhead
while enabling very long input sequences. However, a premise shared across these architectures,
namely that longer lookback windows translate into higher accuracy, does not hold in stochastic set-
tings. The stochastic portion of the input sequence is dominated by high-frequency volatility with
little predictive value. Blindly extending context therefore creates a noise accumulation bottleneck,
a failure mode absent from the text-based tasks where these designs originated.
2.2 C ONTEXT-I NFORMED F ORECASTING AND R ETRIEVAL
Auxiliary information has proven valuable for forecast robustness. Williams et al. (2025) show that
textual grounding of forecasts reduces ambiguity in foundation models substantially and calls con-
text “key.” Furthermore, Pineda Arango et al. (2025) demonstrates a complementary point through
ChronosX: exogenous variables enable pretrained models to specialize for particular downstream
tasks. The shared lesson from both is that predictions benefit when models receive curated context
rather than raw data. Yet autoregressive models lack a built-in mechanism for relevance-based fil-
tering of their input history. RAFT fills that gap. It treats the historical record as a pool from which
only the most similar segments are retrieved via cosine similarity. Stochastic noise is rejected at the
retrieval stage, and what enters the model has high information density. The net effect is a dynamic
inductive bias rooted in the same philosophy as exogenous signal injection but implemented through
a retrieval mechanism.
3 M ETHODOLOGY AND A NALYSIS
3.1 T HE E XPERIMENTAL F RAMEWORK
The proposed work chooses the ETTh1 (Electricity Transformer Temperature) benchmark (Zhou
et al., 2021) because its high-frequency stochastic volatility makes it challenging for any model that
relies on long historical context. Separating genuine signals from random fluctuations is difficult on
this dataset, making it a natural stress test for the “Long Context” hypothesis. A prediction horizon
of H = 96 steps is fixed throughout. The lookback window L is varied across {720, 1440, 3000}
so that the impact of context scaling can be measured in isolation. Mean Squared Error (MSE) and
Mean Absolute Error (MAE) are used as the evaluation metrics. The experiment pits Continuous
Context models against a Retrieved Context model: a) Vanilla Transformer (Baseline): It uses
Encoder-decoder architecture with dmodel = 512 and nheads = 8 to observe how standard attention
behaves as context grows. b) PatchTST (SOTA Long-Context): It uses channel independence and
patching (P = 16, S = 8) per Nie et al. (2023). It represents the scaling hypothesis at maximum
strength: long context, low computational cost. c) RAFT (Retrieval Baseline): It keeps the look-
back short (L = 720) and pulls in relevant history through top-k cosine similarity retrieval (Han
et al., 2025). This tests whether high-quality context in a small window outperforms noisy context
in a large one.
2
ICLR 2026 Workshop on Time Series in the Age of Large Models (TSALM)
3.2 T HE I NVERSE S CALING L AW
The experiments yield a result that runs opposite to NLP scaling laws. Figure A1 shows a clear
Inverse Scaling Law: longer context does not reduce error. MSE rises at every step as the window
gets larger. The law holds at prediction horizons beyond H = 96; (see Appendix H for results at
H = 336 and H = 720). The Vanilla Transformer breaks down most severely. Going from 720
to 3,000 steps pushes MSE up by about 200%. Self-attention, given that many positions to attend
over, distributes weight too broadly and cannot recover the relevant signal. PatchTST handles the
problem better but still fails. At 3,000 steps, its MSE increases by 68% (from 0.385 to 0.647).
Patching, though designed for long sequences, fails to reject noise effectively over extended histo-
ries. Therefore, practitioners assuming patching eliminates noise should reconsider these results.
3.3 G ENERALIZATION ACROSS D OMAINS
The work tests two more benchmarks ETTh2 (Zhou et al., 2021) (electricity domain) and Exchange
Rate (financial domain) to verify that the inverse scaling law is not just a quirk of ETTh1. Table 1
Table 1: Cross-Domain Validation: Degradation of Long Context
Dataset Model MSE ↓ Degradation
PatchTST (720) 0.307 —
ETTh2
PatchTST (3000) 0.533 +73.4%
PatchTST (720) 0.093 —
Exchange
PatchTST (3000) 0.350 +276.8%
shows severe degradation across both additional datasets. On Exchange Rate data, where financial
volatility is high, the long-context configuration degrades by 276%. Noise accumulation is not
dataset-specific; it is a structural problem for stochastic time series, and it worsens proportionally
with the volatility of the underlying process.
3.4 M ECHANISM A NALYSIS : S TOCHASTIC N OISE ACCUMULATION
The proposed work attributes the observed degradation to Stochastic Noise Accumulation. NLP
tokens far back in the sequence frequently carry high information density. However, ETTh1 time se-
ries do not share the same property, as distant points are mostly uncorrelated stochastic fluctuations
without predictive utility. Consequently, the softmax denominator grows with L, spreading proba-
bility mass across positions that contribute nothing useful, resulting in a diminishing share allocated
to genuinely informative recent timesteps. Internal representations are therefore over-smoothed: the
model loses the high-frequency detail required for accurate prediction. The observation is consistent
with Williams et al. (2025), who find that context quality matters far more than context volume when
signal-to-noise ratio is low.
Quantitative measurements of attention entropy further confirm the mechanism. As context grows
from L = 336 (42 patches) to L = 3000 (375 patches), the normalized attention entropy of
PatchTST rises from 0.952 to 0.989, approaching the theoretical maximum of 1.0, which corre-
sponds to a uniform distribution. At L = 3000, the effective attention rank collapses to 0.1: no
single patch receives meaningfully differentiated weight. In practice, the attention mechanism de-
generates into uniform averaging over hundreds of irrelevant positions. The full per-layer entropy
measurements appear in Appendix G.
3.5 T HE S UPERIORITY OF S ELECTIVE R ETRIEVAL
RAFT outperforms the scaling-based models because its context-informed inductive bias pre-filters
history before attention ever operates on it. Figure 1 shows RAFT at an MSE of 0.379 with only 720
steps of direct context. PatchTST at 720 steps scores 0.385, so RAFT is slightly better even without
the scaling advantage. At 3000 steps PatchTST degrades to 0.647, making RAFT the clear winner.
Moreover, RAFT stabilizes its forecasts and avoids the noise penalty that scaling imposes by inject-
3
ICLR 2026 Workshop on Time Series in the Age of Large Models (TSALM)
1.4 RAFT Baseline 1.323
1.2
1.0
Test MSE
0.8
0.647
0.6 0.484
0.385 0.416 0.441
0.4 0.379
0.2
0.0
RAFT PatchTST Vanilla TF Vanilla TF Vanilla TF PatchTST Vanilla TF
(Ret-720) (720) (Best, 720) (720) (1440) (3000) (3000)
Figure 1: Performance degradation trend. Error rate (Y-axis) rises monotonically as Context Length
(X-axis) increases from 720 to 3000, illustrating Stochastic Noise Accumulation.
ing retrieved segments as dynamic exogenous variables (Pineda Arango et al., 2025). Therefore,
retrieval is a more effective noise rejection strategy than window extension for stochastic data.
Zero-shot foundation models, including Chronos-T5-Small (Ansari et al., 2024) (MSE 0.483) and
Moirai-1.1-R-Small (Woo et al., 2024) (MSE 0.553), both fall substantially short of RAFT’s 0.379,
despite pretraining on far larger corpora revealing that the advantage extends beyond long-context
scaling. Such a finding motivates the proposed work to evaluate Chronos-T5-Small and Moirai-1.1-
R-Small in Zero-shot settings on the ETTh1 test set to assess whether pretrained foundation models
can bypass the inverse scaling law. Specifically, Chronos operates univariately with autoregressive
rolling (chunk size 64), while Moirai is evaluated channel-independently for a fair comparison with
PatchTST. Table 2 details the comparison of zero-shot foundation models to trained models on
ETTh1 at H = 96.
Table 2: Foundation Model Zero-Shot vs. Trained Models on ETTh1 (H = 96).
Model Context MSE ↓ MAE ↓
RAFT 720 0.379 —
PatchTST 720 0.385 —
Chronos-T5-Small 336 0.497 0.415
Chronos-T5-Small 720 0.483 0.414
Moirai-1.1-R-Small 336 0.575 0.489
Moirai-1.1-R-Small 720 0.553 0.487
Neither foundation model matches RAFT or even PatchTST at L = 720. Chronos at its best (L =
720, MSE 0.483) is 27% above RAFT, and Moirai (L = 720, MSE 0.553) is 46% above. Hence,
Zero-shot generality does not compensate for targeted retrieval on stochastic data.
4 D ISCUSSION
4.1 T HE S TOCHASTIC C ONTEXT D ILEMMA
NLP scaling laws predict that more context produces better outputs. The inverse scaling law docu-
mented here contradicts that prediction; the reason lies in a basic difference between the two data
types. Language tokens and factual claims retain informativeness regardless of how far back they
appear. ETTh1 measurements from 2,000 steps ago are essentially uncorrelated noise with respect to
the prediction target. Furthermore, over 3,000 positions, the attention Softmax inevitably distributes
probability mass over inputs carrying no forecasting information. The proposed work characterizes
the effect as Attention Entropy inflation. Specifically, measured normalized entropy rises from 0.952
(L = 336) to 0.989 (L = 3000), confirming near-complete entropy saturation at extreme context
lengths (Section 3.4). Relevant recent signals are drowned out. More history helps only when it
carries high semantic density; stochastic volatility violates that condition by definition.
4
ICLR 2026 Workshop on Time Series in the Age of Large Models (TSALM)
4.2 T HE E FFICIENCY-ACCURACY T RADE - OFF
Training cost makes the case against long context even stronger: a) Long Context (Inefficient):
Vanilla Transformer at 3,000 steps required 83.47 minutes of training and produced MSE 1.323,
the worst result in the study. b) Retrieval (Efficient): RAFT trained in 2.13 minutes (roughly
40× faster) and produced MSE 0.379, the best result. The expectation behind NLP scaling laws is
that spending more compute yields better model quality. In time series forecasting, spending more
compute on longer context yields worse quality and wastes resources in the process.
4.3 R ETRIEVAL AS A D ENOISING I NDUCTIVE B IAS
RAFT works well because its top-k retrieval stage operates as a “hard attention” gate before the
standard “soft attention” layer. Two stages run in sequence, forming a noise-rejection pipeline: a)
Filter: Retrieval discards segments with low similarity. The surviving segments have high fidelity
to the current forecasting situation. b) Process: The Transformer receives only these filtered inputs,
treating them as dynamic exogenous variables (Pineda Arango et al., 2025) rather than undifferen-
tiated history. Moreover, long-range dependencies remain exploitable through retrieval, without the
noise accumulation that accompanies brute-force window scaling.
5 L IMITATIONS
The three benchmarks tested (ETTh1, ETTh2, Exchange Rate) all exhibit stochastic, mean-reverting
dynamics; i.e., whether the inverse scaling law holds on trend-dominated or strongly seasonal se-
ries remains open. The main experiments fix the value of H=96, and foundation models are not
re-evaluated at longer horizons except at H=336, 720 (detailed in Appendix H). Furthermore, hy-
perparameter sensitivity (detailed in Appendix I) and single-GPU timing measurements may not
generalize across hardware.
6 C ONCLUSION AND F UTURE W ORK
The proposed work has tested the “Long Context” hypothesis for time series forecasting and found
it wanting. On ETTh1, PatchTST can lose up to 68% of its accuracy when context extends from
720 to 3,000 steps, an Inverse Scaling Law driven by noise accumulation in stochastic data. RAFT
(Han et al., 2025), with selective retrieval and a constrained 720-step window, reaches MSE 0.379,
well below the expensive long-context alternatives. Context quality, not context volume, deter-
mines forecast accuracy. Additionally, the inverse scaling law holds at longer prediction horizons
(H = 336, H = 720) and the degradation mechanism is corroborated by attention entropy mea-
surements where normalized entropy reaches 0.989 at L = 3000. Hence, the next frontier is to
integrate retrieval directly into pretraining. Foundation models that use dynamic retrieval heads and
are trained end-to-end to identify the most useful historical segments could achieve robust zero-shot
performance without the noise penalty of static windows. Such an architecture would put into prac-
tice the “Context is Key” philosophy advocated by Williams et al. (2025) and make stochastic time
series forecasting substantially more reliable.
R EFERENCES
Abdul Fatir Ansari, Lorenzo Stella, Caner Turkmen, Xiyuan Zhang, Pedro Mercado, Huibin Shen,
Oleksandr Shchur, Syama Sundar Rangapuram, Sebastian Pineda Arango, Shubham Kapoor,
Jasper Zschiegner, Danielle C. Maddix, Michael W. Mahoney, Kari Torkkola, Andrew Gordon
Wilson, Michael Bohlke-Schneider, and Yuyang Wang. Chronos: Learning the language of time
series. Transactions on Machine Learning Research, 2024.
Sungwon Han, Seungeon Lee, Meeyoung Cha, Sercan O. Arik, and Jinsung Yoon. Retrieval aug-
mented time series forecasting. In Proceedings of the 42nd International Conference on Machine
Learning, volume 267 of Proceedings of Machine Learning Research, pp. 21774–21797. PMLR,
2025.
5
ICLR 2026 Workshop on Time Series in the Age of Large Models (TSALM)
Yuqi Nie, Nam H. Nguyen, Phanwadee Sinthong, and Jayant Kalagnanam. A time series is worth
64 words: Long-term forecasting with transformers. In International Conference on Learning
Representations, 2023.
Sebastian Pineda Arango, Pedro Mercado, Shubham Kapoor, Abdul Fatir Ansari, Lorenzo Stella,
Huibin Shen, Hugo Henri Joseph Senetaire, Ali Caner Turkmen, Oleksandr Shchur, Danielle C.
Maddix, Michael Bohlke-Schneider, Bernie Wang, and Syama Syndar Rangapuram. Chronosx:
Adapting pretrained time series models with exogenous variables. In Proceedings of The 28th
International Conference on Artificial Intelligence and Statistics, volume 258 of Proceedings of
Machine Learning Research, pp. 2242–2250. PMLR, 2025.
Andrew Robert Williams, Arjun Ashok, Étienne Marcotte, Valentina Zantedeschi, Jithendaraa Sub-
ramanian, Roland Riachi, James Requeima, Alexandre Lacoste, Irina Rish, Nicolas Chapados,
and Alexandre Drouin. Context is key: A benchmark for forecasting with essential textual infor-
mation. In Proceedings of the 42nd International Conference on Machine Learning, volume 267
of Proceedings of Machine Learning Research, pp. 66887–66944. PMLR, 2025.
Gerald Woo, Chenghao Liu, Akshat Kumar, Caiming Xiong, Silvio Savarese, and Doyen Sahoo.
Unified training of universal time series forecasting transformers. In Forty-first International
Conference on Machine Learning, 2024.
Haoyi Zhou, Shanghang Zhang, Jieqi Peng, Shuai Zhang, Jianxin Li, Hui Xiong, and Wancai Zhang.
Informer: Beyond efficient transformer for long sequence time-series forecasting. In AAAI Con-
ference on Artificial Intelligence, volume 35, pp. 11106–11115, 2021.
6
ICLR 2026 Workshop on Time Series in the Age of Large Models (TSALM)
A PPENDIX
A AI U SE D ISCLOSURE
Large language models were used to assist with grammar refinement and manuscript editing. All
scientific content, experimental results, and conclusions are the sole work of the authors.
B E XPERIMENTAL C ONFIGURATION
Hyperparameter settings and training configurations are reported here for reproducibility. Figure A1
shows a clear Inverse Scaling Law already discussed in 3.2.
Vanilla Transformer (Long Ctx)
PatchTST (SOTA)
RAFT (Retrieval @ 720)
1.2
Test MSE (Lower is Better)
1.0
Performance degrades
0.8 as context increases
0.6
0.4
720 1440 3000
Context Window Size (L)
Figure A1: The Inverse Scaling Law on ETTh1. RAFT (Green) holds error low (MSE 0.379) with
a fixed window. PatchTST (Orange), when extended to 3000 steps, can see its performance degrade
by up to 68% (MSE 0.647). The Vanilla Transformer (Red) collapses entirely.
B.1 M ODEL H YPERPARAMETERS
Table A1 gives the full configuration of each model. The same optimization settings apply across
all experiments to keep the comparison controlled.
B.2 T RAINING P ROTOCOL
Adam with a learning rate of 0.0001 was used for all models. If validation MSE failed to improve
for 3 consecutive epochs, training was stopped early. The random seed for weight initialization was
2021 throughout, which makes results deterministic.
C DATASET D ETAILS
The inverse scaling law hypothesis is tested on three benchmarks that span different domains and
noise profiles.
C.1 ETT H 1 DATASET
ETTh1 (Electricity Transformer Temperature) contains hourly transformer readings collected over
600 days (14,400 timesteps total). The measured features are HUFL, HULL, MUFL, MULL, LUFL,
7
ICLR 2026 Workshop on Time Series in the Age of Large Models (TSALM)
Table A1: Complete Hyperparameter Configurations
Parameter Vanilla TF PatchTST RAFT
seq len 720/1440/3000 720/3000 720
pred len 96 96 96
label len 48 48 48
dmodel 512 128 128
nheads 8 16 8
elayers 2 3 2
dlayers 1 1 1
dff 2048 256 256
dropout 0.1 0.1 0.1
patch len — 16 —
stride — 8 —
top-k retrieval — — 5
learning rate 0.0001 0.0001 0.0001
optimizer Adam Adam Adam
batch size 32 32 32
epochs 10 10 10
patience 3 3 3
loss MSE MSE MSE
LULL, and OT (Oil Temperature), giving seven channels. A 60/20/20 split was used for training,
validation, and test partitions. Each feature was normalized independently using StandardScaler.
C.2 C ROSS -D OMAIN DATASETS
ETTh2 has the same format as ETTh1 but comes from a different transformer station. The Exchange
Rate dataset contains daily rates for 8 currencies over 7,588 trading days. Both datasets were in-
cluded because their noise characteristics differ from each other, allowing the proposed work to
check whether the inverse scaling law generalizes beyond a single domain.
D C OMPLETE R ESULTS
Table A2 reports MSE and training time for every configuration tested on ETTh1.
Table A2: Complete Results on ETTh1
Model Seq Len MSE ↓ Time (min)
RAFT 720 0.379 2.13
PatchTST 720 0.385 2.90
PatchTST 3000 0.647 18.33
Vanilla TF 720 0.441 7.42
Vanilla TF 1440 0.484 22.00
Vanilla TF 3000 1.323 83.47
D.1 D EGRADATION A NALYSIS
Table A3 reports the MSE increase when context is extended from 720 to 3000 steps, broken down
by model and dataset.
Table A3: Performance Degradation Across Contexts
Model Dataset 720 → 3000 Degradation
Vanilla TF ETTh1 0.441 → 1.323 +200%
PatchTST ETTh1 0.385 → 0.647 +68%
PatchTST ETTh2 0.307 → 0.533 +73%
PatchTST Exchange 0.093 → 0.350 +277%
8
ICLR 2026 Workshop on Time Series in the Age of Large Models (TSALM)
Degradation is consistent across datasets. Financial data shows the largest increase, as expected
given that currency exchange rates exhibit particularly high stochastic volatility.
E C OMPUTATIONAL E NVIRONMENT
Experiments were run on cloud GPU infrastructure (16GB RAM, 8 CPU cores). Python 3.9, Py-
Torch 2.0.1, and CUDA 11.8 formed the core of the software stack; numpy 1.24, pandas 2.0, and
scikit-learn 1.3 were also used. Training time grows quadratically with context length. At 3000
steps the Vanilla Transformer needed 83.47 minutes per run; RAFT needed 2.13 minutes for the
same prediction task, about 40× faster, while also producing a better MSE.
F R EPRODUCIBILITY
All datasets are publicly available from the Informer repository. The train-validation-test split
is 60/20/20 in every experiment, and StandardScaler normalization is applied per feature. Ran-
dom seeds are fixed at 2021. Early stopping halts training if validation MSE does not drop
for 3 epochs. Replication requires downloading the datasets, applying the hyperparameters in
Table A1, and running training with the Adam optimizer at the specified learning rate. Stan-
dard PyTorch was used throughout; no custom modifications were made. The complete source
code, training scripts, and evaluation pipelines are publicly available at https://github.com/
RishiAhuja/ahuja2026retrieval.
G ATTENTION E NTROPY M EASUREMENTS
Normalized attention entropy was computed per layer P across all attention heads in PatchTST at
each context length. Entropy is defined as H = − i pi log pi , normalized by log N where N is
the number of patches. An entropy of 1.0 corresponds to a perfectly uniform (random) attention
distribution and an entropy of 0.0 corresponds to a sparse single-position spike.
At L = 336 (42 patches), mean normalized entropy is 0.952, indicating that attention is already
fairly diffuse but retains some discriminative structure. Similarly, at L = 720 (90 patches), entropy
rises to 0.971. Furthermore, at L = 3000 (375 patches), entropy reaches 0.989, within 1.1% of the
theoretical maximum. Additionally, the effective attention rank (computed as exp(H), the exponen-
tial of the entropy) collapses from 8.3 at L = 336 to 0.1 at L = 3000, confirming that the attention
mechanism has degenerated into near-uniform averaging at extreme context lengths.
H M ULTI -H ORIZON E VALUATION
The inverse scaling law is tested at longer prediction horizons (H ∈ {336, 720}) to verify that the
phenomenon generalizes beyond H = 96. Table A4 details the multi-horizon results on ETTh1.
Table A4: Multi-Horizon Results on ETTh1 (MSE ↓).
Model Seq Len H = 336 H = 720
RAFT 720 0.438 0.474
PatchTST 720 0.503 0.517
PatchTST 3000 0.800 1.513
Vanilla TF 720 0.586 0.755
Vanilla TF 3000 0.763 0.882
As evident from Table A4, the inverse scaling law holds at all tested horizons. PatchTST degradation
worsens at longer horizons: +59% at H = 336 and +193% at H = 720 when the context extends
from 720 to 3000. Additionally, RAFT maintains the lowest MSE at every horizon without any
context extension.
9
ICLR 2026 Workshop on Time Series in the Age of Large Models (TSALM)
I T RAINING C ONFIGURATION S ENSITIVITY
The magnitude of long-context degradation depends on the training protocol. Under the schedule
used in the main experiments (10 epochs, step-decay learning rate), PatchTST reaches MSE 0.647
at L = 3000 (+68% over L = 720). The same architecture, when trained with an extended schedule
(100 epochs, cosine annealing learning rate), achieves MSE 0.426 at L = 3000, narrowing the gap to
approximately 8% over L = 720 (MSE 0.395). Importantly, RAFT (MSE 0.379) remains superior
under both training configurations, and the need for careful hyperparameter tuning to mitigate long-
context degradation is itself a practical limitation of the continuous-context paradigm.
10
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/text/2606.23642v1.txt
+364
View File
@@ -0,0 +1,364 @@
Improving Long-Context Retrieval with Multi-Prefix Embedding
Zhenglin Yu1 , Xueguang Ma1 , Shengyao Zhuang2 , Zhichao Xu3 ,
Luyu Gao4 , Crystina Zhang1 , Jimmy Lin1
1
University of Waterloo, 2 University of Queensland, 3 University of Utah, 4 Carnegie Mellon University
Correspondence: x93ma@uwaterloo.ca
Abstract
Long-context retrieval exposes a tension:
single-vector embeddings lose fine-grained de-
tail, while token-level multi-vector methods
incur prohibitive storage. We propose Multi-
arXiv:2606.23642v1 [cs.IR] 22 Jun 2026
Prefix Embedding (MPE), which partitions
a document into chunks separated by EOS Figure 1: Multi-Prefix Embedding (MPE). A long docu-
tokens, encodes the full sequence in a sin- ment is split into chunks separated by EOS tokens and
gle causal forward pass, and extracts one encoded in one causal forward pass. Prefix embeddings
embedding at each prefix boundary. MPE are extracted at EOS1 , . . . , EOSK , and query–document
retains cross-chunk context, enables chunk- relevance is computed via MaxSim.
level MaxSim matching, and trains with only
document-level relevance labels. Experi-
ments on MLDR-en, BrowseComp-Plus, and in a single causal forward pass. The hidden state
LongEmbed show that MPE is competitive
at each EOS position summarizes the preceding
with or outperforms single-vector, independent-
chunk, and multi-vector baselines, while pro-
content, yielding one prefix embedding per chunk
viding a natural source attribution mechanism boundary without architectural changes or chunk-
for locating evidence chunks. level annotations. Query–document similarity is
computed with MaxSim over the document’s prefix
1 Introduction embeddings, and the model is trained using only
document-level relevance labels.
Long-context retrieval requires models to iden-
tify localized evidence within documents that may MPE occupies a middle ground between single-
span thousands of tokens. Dense retrieval encodes vector, independent-chunk, and token-level multi-
queries and documents as fixed-dimensional vec- vector retrieval. It provides multiple matching
tors for efficient inner-product search (Karpukhin points within a long document, preserves preced-
et al., 2020; Izacard et al., 2022; Qwen Team, ing cross-chunk context through causal attention,
2025; Xu et al., 2025), but compressing a long and stores only one embedding per chunk bound-
document into a single vector can obscure fine- ary rather than one embedding per token. The
grained evidence (Zhu et al., 2024). Token-level MaxSim-selected prefix also identifies the docu-
multi-vector methods such as ColBERT (Khattab ment region most responsible for a match, provid-
and Zaharia, 2020; Santhanam et al., 2022) pre- ing a lightweight mechanism for surfacing candi-
serve finer granularity, but incur substantially larger date evidence.
indexes and higher retrieval cost. Chunk-based We evaluate MPE on MLDR-en (Chen et al.,
retrieval is cheaper, but independently encoded 2024), BrowseComp-Plus (Chen et al., 2025), and
chunks lose cross-chunk context and require post- LongEmbed (Zhu et al., 2024). Fine-tuning Qwen3-
hoc document-level aggregation. This raises a cen- Embedding-0.6B (Qwen Team, 2025) on MLDR-
tral question: can we retain fine-grained matching en, MPE matches or exceeds single-vector and
while preserving context and keeping the represen- independent-chunk baselines on most benchmarks,
tation compact? while random prefix-length augmentation improves
We propose Multi-Prefix Embedding (MPE), a robustness under granularity mismatch, especially
compact multi-vector representation for causal em- on out-of-domain BrowseComp-Plus.
bedding models. As shown in Figure 1, MPE splits Our main contributions are:
a long document into chunks, inserts EOS tokens 1. We propose MPE, a compact multi-vector repre-
between them, and encodes the resulting sequence sentation for long-context retrieval that balances
fine-grained matching, context preservation, and mentation. Figure 1 illustrates the overall frame-
indexing efficiency. work.
2. We show that MaxSim-selected prefixes often Given a document p = [t1 , . . . , tL ], we split it
align with annotated answer locations, suggest- into K consecutive chunks of m content tokens
ing a lightweight mechanism for candidate evi- and append an EOS token after each chunk:
dence surfacing.
3. We present experiments and ablations across x = [ t1 . . . tm , <eos>, tm+1 . . . t2m , <eos> . . . ] (1)
representation methods, attention mechanisms,
and granularity settings. The resulting sequence is encoded in a single for-
ward pass through a pretrained causal language
2 Related Work
model. Under causal masking, the hidden state
Dense retrieval encodes queries and documents at each EOS position attends to the current chunk
into fixed-dimensional vectors for inner-product and all preceding chunks, yielding a contextual-
search (Karpukhin et al., 2020; Izacard et al., ized prefix representation. Unlike independently
2022). DPR uses BERT encoders with [CLS] pool- encoded chunks, each prefix embedding therefore
ing (Karpukhin et al., 2020), while recent decoder- incorporates preceding cross-chunk context. We
only embedding models such as RepLLaMA (Ma L2-normalize the hidden states at EOS positions to
et al., 2024), E5-Mistral (Wang et al., 2024), and obtain document prefix embeddings {c1 , . . . , cK },
Qwen3-Embedding (Qwen Team, 2025) use last- where ck ∈ Rd . When K = 1, MPE reduces to
token pooling. Although efficient, single-vector standard last-token pooling.
embeddings can obscure fine-grained evidence in
long documents. 3.1 Scoring, Training, and Retrieval
Multi-vector methods improve fine-grained Given a query embedding q, MPE scores a docu-
matching by representing each document with mul- ment by MaxSim over its prefix embeddings:
tiple embeddings. ColBERT (Khattab and Zaharia,
2020; Santhanam et al., 2022) uses late interac- s(q, p) = MaxSim(q, p) = max q ⊤ ck (2)
tion over token-level representations, but its stor- k∈{1,...,K}
age and retrieval cost grows with document length.
MaxP (Dai and Callan, 2019) splits documents into This allows a query to match the most relevant
fixed-size chunks and takes the maximum indepen- document region rather than requiring the whole
dently encoded chunk score, but loses cross-chunk document to be compressed into one vector. We
context. train MPE with contrastive loss over document-
MPE is most related to compact and contex- level relevance labels using cross-device in-batch
tual multi-vector representations. BGE Land- negatives.
mark Embedding (Luo et al., 2024) also extracts At retrieval time, all prefix embeddings are in-
hidden states from special positions in a causal dexed in a FAISS (Johnson et al., 2019) flat inner-
LLM, but targets context selection for retrieval- product index together with their document IDs.
augmented long-context LLMs using landmark Prefix-level hits are aggregated by document ID us-
tokens and multi-stage supervision. In contrast, ing the maximum prefix score to produce the final
MPE uses EOS-separated chunk boundaries for document ranking.
document-level retrieval and trains with a single
contrastive objective over document-level labels. 3.2 Random Prefix-Length Augmentation
Late chunking (Günther et al., 2024) pools chunk
Fixed chunk sizes can make the model sensitive to
representations after bidirectional document encod-
a specific granularity. To improve robustness, we
ing, while token pooling (Clavié et al., 2024) com-
sample the content chunk length for each training
presses ColBERT-style token embeddings for short-
passage, m ∼ U[mmin , mmax ], and construct EOS-
passage retrieval.
separated prefixes using that length. All chunks
3 Method within a passage share the same sampled m. This
exposes the model to diverse prefix boundaries dur-
MPE consists of EOS-based prefix construction, ing training and helps a single MPE model general-
MaxSim scoring, and random prefix-length aug- ize across inference chunk sizes without retraining.
4 Experiments
4.1 Setup
Datasets. We evaluate on three long-context re-
trieval benchmarks. MLDR-en (Chen et al., 2024)
contains long-document retrieval queries with doc-
uments up to 8,192 tokens; we use its 800-query
test split. BrowseComp-Plus (Chen et al., 2025)
contains 830 difficult web queries with human-
verified supporting documents and mined hard
negatives. LongEmbed (Zhu et al., 2024) evalu-
ates long-context retrieval ability; we report re-
sults on the real-document subsets NarrativeQA, Figure 2: Granularity mismatch on MLDR-en. Fixed-
2WikiMQA, SummScreen, and QMSum. size MPE degrades under mismatched granularities,
while MPE-Rand tracks the upper envelope with a single
Model and training. We use Qwen3-
model. Star symbols denote matched train–eval sizes;
Embedding-0.6B (Qwen Team, 2025) as the base circle symbols denote mismatched sizes.
model and fine-tune it on the MLDR-en training
split for one epoch using Tevatron (Ma et al., 2025).
Training uses LoRA adapters with rank 16, α = 64, samples the training chunk size from [a, b]. Unless
and dropout 0.1 on all attention and MLP projec- specified, chunk-based and MPE methods use infer-
tions (q/k/v/o/gate/up/down_proj). We ence chunk size 64; LongEmbed and BrowseComp-
use learning rate 1×10−4 , temperature τ = 0.03, Plus are zero-shot after MLDR-en training.
per-GPU batch size 2, group size 4, query length
4.2 Main Results
512, maximum document length 8,192, right-side
padding, and gradient checkpointing. Experiments Table 2 reports nDCG@10 on in-domain and zero-
were run on 8 NVIDIA RTX 5090 GPUs. The to- shot long-context retrieval benchmarks. MPE
tal compute budget for the reported experiments is Fixed-64 achieves the best dense result on MLDR-
approximately 120 GPU-hours, including training, en, while MPE-Rand[32, 1024] is the only multi-
ablations, and inference. vector variant to outperform Single-vector on
For zero-shot evaluation, we reuse the MLDR- BrowseComp-Plus, highlighting the importance of
en fine-tuned model without further training. granularity robustness for out-of-domain retrieval.
We use maximum document length 4,096 for On LongEmbed, MPE gives the best dense result
BrowseComp-Plus and 8,192 for LongEmbed. on 2WikiMQA and QMSum and remains competi-
Configurations. We compare five configura- tive on NarrativeQA; SummScreen is nearly satu-
tions that isolate the effects of multi-vector infer- rated by BM25 and Single-vector retrieval.
ence, MaxSim training, cross-chunk attention, and MPE improves over independent chunking
random prefix-length augmentation (Table 1). through contextual prefix embeddings and gran-
ularity robustness. Under the same inference chunk
Setting MaxSim Train Cross-Chunk Attn. Rand. Size size, MPE Fixed-64 outperforms MaxP-Train on
Single-vector × – × MLDR-en, suggesting that preceding cross-chunk
MaxP × × × context helps. MPE-Rand sacrifices some in-
MaxP-Train ✓ × × domain peak performance on MLDR-en but im-
MPE Fixed-N ✓ ✓ ×
MPE-Rand[a, b] ✓ ✓ ✓ proves BrowseComp-Plus and QMSum.
Granularity robustness. Figure 2 varies the in-
Table 1: Comparison of configuration settings.
ference chunk size from 32 to 1024 on MLDR-
en. Fixed-size MPE models peak near their train-
Single-vector uses last-token pooling without
ing granularity but degrade under mismatch, while
chunking. MaxP (Dai and Callan, 2019) encodes
MPE-Rand[32, 1024] tracks their upper envelope
chunks independently at inference and takes the
with a single checkpoint.
maximum chunk score. MaxP-Train adds MaxSim
training. MPE Fixed-N trains and infers with fixed Causal vs. bidirectional attention. As a diag-
EOS-separated prefixes, while MPE-Rand[a, b] nostic, we replace the causal self-attention mask
Method MLDR Train MLDR-en BrowseComp-Plus NarrativeQA 2WikiMQA SummScreen QMSum
BM25 × 0.679 0.016 0.715 0.965 0.976 0.813
jina-v2-base-en (Günther et al., 2023) × 0.370 – 0.394 0.740 0.935 0.389
OpenAI-Ada-002 – 0.387 – 0.411 0.801 0.918 0.400
E5-Mistral-7B (Wang et al., 2024) × 0.433 – 0.449 – – –
BGE-M3 (Chen et al., 2024) ✓ 0.489 – 0.487 0.780 0.940 0.355
BGE-M3mv (Chen et al., 2024) ✓ 0.558 – 0.554 – – –
Single-vector ✓ 0.548 0.132 0.543 0.863 0.981 0.565
MaxP ✓ 0.758 0.103 0.604 0.944 0.906 0.550
MaxP-Train ✓ 0.776 0.104 0.557 0.945 0.903 0.570
MPE Fixed-64 ✓ 0.783 0.122 0.580 0.948 0.928 0.652
MPE-Rand[32, 1024] ✓ 0.760 0.153 0.579 0.945 0.949 0.705
Table 2: nDCG@10 on long-context retrieval benchmarks. Baseline results are from original papers (Günther
et al., 2023; Zhu et al., 2024; Chen et al., 2024; Wang et al., 2024; Nussbaum et al., 2024). BGE-M3mv denotes
ColBERT-style multi-vector retrieval. Methods below the midrule use Qwen3-Embedding-0.6B; chunk-based and
MPE methods use inference chunk size 64 unless otherwise specified. Bold marks the best dense result.
Retriever Acc. (%) Recall (%) Avg. Calls Cal. Err. (%)1
BM25 34.94 37.78 18.85 51.17
Single-vector 42.29 48.00 17.40 46.99
MPE-Rand 51.45 57.70 15.78 39.74
Table 3: End-to-end BrowseComp-Plus with a Gemini
3 Flash Preview search agent. Accuracy is judged by
GPT-4o-mini. MPE-Rand is trained on [32, 1024] and
Figure 3: MaxSim-selected chunk positions vs. Gemini-
evaluated at chunk size 64.
annotated answer positions on MLDR-en. Left: overlaid
position distributions. Right: per-passage scatter plot
showing substantial rank correlation (Spearman ρ =
setting. Replacing the single-vector retriever
0.77).
with MPE-Rand improves answer accuracy from
42.29% to 51.45%, supporting-document recall
with a bidirectional mask during fine-tuning and from 48.00% to 57.70%, and reduces retrieval calls.
inference. This modification reduces MLDR-en
nDCG@10 from 0.783 to 0.694, suggesting that 5 Conclusion
MPE benefits from the pretrained causal attention We presented Multi-Prefix Embedding (MPE), a
structure rather than simply from exposing each compact chunk-level multi-vector method for long-
EOS position to more tokens. document retrieval. By inserting EOS tokens be-
Storage overhead. MPE stores K = T /C vec- tween chunks, MPE extracts prefix embeddings
tors per document, where T is document length in a single causal LM forward pass, preserving
and C is chunk size. On MLDR-en (T ≤ 8,192), preceding cross-chunk context while enabling fine-
this gives K ≈ 128 at C = 64 and K ≈ 16 at grained MaxSim matching without chunk-level su-
C = 512. For N documents, MPE stores O(KN ) pervision or architectural changes.
vectors, between single-vector retrieval O(N ) and Experiments on MLDR-en, BrowseComp-Plus,
token-level late interaction O(T N ), with K ≪ T . and LongEmbed show that MPE matches or outper-
forms single-vector and independent-chunk base-
Source attribution. Using Gemini-annotated an- lines on most benchmarks. Random prefix-length
swer spans on 714 MLDR-en passages, MaxSim augmentation improves robustness across chunk
selects a chunk within ±1 position of the anno- sizes, with particularly strong gains under granular-
tated chunk for 65.7% of passages, with Spearman ity mismatch and out-of-domain retrieval. MaxSim-
correlation ρ = 0.77 (Figure 3). Although the selected prefixes align with Gemini annotated an-
annotations are LLM-derived, this suggests that swer locations.
MPE can surface candidate evidence in addition to Code and model checkpoints will be released
document-level rankings. upon publication.
End-to-end search agent. Table 3 evaluates 1
Calibration error measures how closely predicted confi-
MPE-Rand in a BrowseComp-Plus search-agent dence matches actual answer accuracy (Chen et al., 2025).
Limitations text embeddings through self-knowledge distillation.
arXiv preprint arXiv:2402.03216.
Our experiments are limited to a single 0.6B-
parameter causal embedding model and English- Zijian Chen, Xueguang Ma, Shengyao Zhuang, Ping
Nie, Kai Zou, Andrew Liu, Joshua Green, Kshama
language benchmarks. Scaling MPE to larger Patel, Ruoxi Meng, Mingyi Su, Sahel Shari-
embedding models, multilingual retrieval, and fymoghaddam, Yanxi Li, Haoran Hong, Xinyu
stronger instruction-tuned retrievers remains future Shi, Xuye Liu, Nandan Thakur, Crystina Zhang,
work. Luyu Gao, Wenhu Chen, and Jimmy Lin. 2025.
BrowseComp-Plus: A more fair and transparent eval-
MPE introduces additional storage and retrieval uation benchmark of deep-research agent. arXiv
cost compared with single-vector embeddings be- preprint arXiv:2508.06600.
cause each document is represented by multiple
Benjamin Clavié, Antoine Chaffin, and Griffin Adams.
prefix embeddings. While this cost is substantially 2024. Reducing the footprint of multi-vector retrieval
smaller than token-level late interaction, practi- with minimal performance impact via token pooling.
cal deployment may require product quantization, arXiv:2409.14683.
residual compression, pruning, or adaptive prefix
Zhuyun Dai and Jamie Callan. 2019. Deeper text under-
selection. standing for ir with contextual neural language mod-
Our comparisons against Landmark Embed- eling. In Proceedings of the 42nd International ACM
ding (Luo et al., 2024) and Late Chunking (Günther SIGIR Conference on Research and Development
et al., 2024) are based on reported results and con- in Information Retrieval, SIGIR’19, page 985–988,
New York, NY, USA. Association for Computing
ceptual differences rather than a fully controlled Machinery.
head-to-head evaluation. A matched comparison
under the same backbone, training data, context Michael Günther, Isabelle Mohr, Daniel J Williams,
Bo Wang, and Han Xiao Rodríguez. 2024. Late
length, and indexing budget would better isolate chunking: Contextual chunk embeddings using
the effect of MPE’s prefix representation. long-context embedding models. arXiv preprint
Finally, our source-attribution analysis relies on arXiv:2409.04701.
Gemini-generated answer-span annotations. Al- Michael Günther, Jackmin Ong, Isabelle Mohr, Alaed-
though the observed correlation suggests that dine Abdessalem, Tanguy Abel, Mohammad Afzal
MaxSim often selects meaningful evidence loca- Akber, Susana Gusber, Georgios Mastrapas, Robin
tions, these labels are not human gold annotations Milliken, and Bo Wang. 2023. Jina embeddings 2:
8192-token general-purpose text embeddings for long
and may contain noise.
documents. arXiv preprint arXiv:2310.19923.
Ethics Descriptions Gautier Izacard, Mathilde Caron, Lucas Hosseini, Sebas-
tian Riedel, Piotr Bojanowski, Armand Joulin, and
This work uses public retrieval benchmarks and Edouard Grave. 2022. Unsupervised dense informa-
does not involve human-subject data collection. tion retrieval with contrastive learning. Transactions
on Machine Learning Research.
Like other retrieval systems, MPE may amplify bi-
ases or misinformation present in training corpora Jeff Johnson, Matthijs Douze, and Hervé Jégou. 2019.
or retrieved documents, especially when used in Billion-scale similarity search with GPUs. IEEE
downstream search or question-answering agents. Transactions on Big Data, 7(3):535–547.
Our BrowseComp-Plus agent evaluation uses LLM- Vladimir Karpukhin, Barlas Oğuz, Sewon Min, Patrick
based answer judging, which can introduce noise Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and
or evaluator bias; we therefore treat it as comple- Wen-tau Yih. 2020. Dense passage retrieval for open-
domain question answering. In Proceedings of the
mentary to retrieval metrics rather than definitive 2020 Conference on Empirical Methods in Natural
human evaluation. We disclose the main compute Language Processing (EMNLP), pages 6769–6781.
setup in Section 4.1; practical deployments should
consider the storage, latency, and energy costs of Omar Khattab and Matei Zaharia. 2020. ColBERT: Effi-
cient and effective passage search via contextualized
multi-vector indexing. late interaction over BERT. In Proceedings of the
43rd International ACM SIGIR Conference on Re-
search and Development in Information Retrieval,
References pages 39–48.
Jianlv Chen, Shitao Xiao, Peitian Zhang, Kun Luo, Defu Kun Luo, Zheng Liu, Shitao Xiao, and Kang Liu.
Lian, and Zheng Liu. 2024. BGE M3-embedding: 2024. BGE landmark embedding: A chunking-
Multi-lingual, multi-functionality, multi-granularity free embedding method for retrieval augmented
long-context large language models. arXiv preprint and follow the terms and licenses of the original
arXiv:2402.11573. component datasets included in LongEmbed.
Our use of these artifacts is consistent with their
Xueguang Ma, Luyu Gao, Shengyao Zhuang, Ji-
aqi Samantha Zhan, Jamie Callan, and Jimmy Lin. intended research purpose: datasets are used for
2025. Tevatron 2.0: Unified document retrieval retrieval benchmarking and evaluation, and the pre-
toolkit across scale, language, and modality. arXiv trained Qwen3-Embedding model is used as a base
preprint arXiv:2505.02466. embedding model for research experiments. We
Xueguang Ma, Liang Wang, Nan Yang, Furu Wei, and do not redistribute the original benchmark data be-
Jimmy Lin. 2024. Fine-tuning llama for multi-stage yond what is permitted by the artifact providers.
text retrieval. In Proceedings of the 47th Interna- Any released code or checkpoints from this work
tional ACM SIGIR Conference on Research and De- are intended for research and reproducibility, and
velopment in Information Retrieval, SIGIR ’24, page
2421–2425, New York, NY, USA. Association for users should comply with the licenses and terms
Computing Machinery. of the underlying datasets, pretrained models, and
software dependencies.
Zach Nussbaum, John X Morris, Brandon Duderstadt,
and Andriy Mulyar. 2024. Nomic embed: Training B LLM Usage
a reproducible long context text embedder. arXiv
preprint arXiv:2402.01613. The authors used generative AI tools for language
polishing, grammar checking, minor editing sug-
Qwen Team. 2025. Qwen3 embedding: Advancing text
embedding and reranking through foundation models. gestions, and assistance with code debugging. The
arXiv preprint arXiv:2506.05176. tools were not used to generate the core technical
ideas, design the experiments, produce the reported
Keshav Santhanam, Omar Khattab, Jon Saad-Falcon,
results, or make final scientific claims. All experi-
Christopher Potts, and Matei Zaharia. 2022. Col-
BERTv2: Effective and efficient retrieval via ments, analyses, results, and final writing decisions
lightweight late interaction. In Proceedings of the were performed and verified by the authors.
2022 Conference of the North American Chapter of
the Association for Computational Linguistics: Hu-
man Language Technologies, pages 3715–3734, Seat-
tle, United States. Association for Computational
Linguistics.
Liang Wang, Nan Yang, Xiaolong Huang, Linjun Yang,
Rangan Majumder, and Furu Wei. 2024. Improving
text embeddings with large language models. arXiv
preprint arXiv:2401.00368.
Zhichao Xu, Fengran Mo, Zhiqi Huang, Crystina Zhang,
Puxuan Yu, Bei Wang, Jimmy Lin, and Vivek Sriku-
mar. 2025. A survey of model architectures in infor-
mation retrieval. arXiv preprint arXiv:2502.14822.
Dawei Zhu, Liang Wang, Nan Yang, Yifan Song, Wen-
hao Wu, Furu Wei, and Sujian Li. 2024. Longembed:
Extending embedding models for long context re-
trieval. arXiv preprint arXiv:2404.12096.
A License
We use the following external artifacts: MLDR-
en, BrowseComp-Plus, LongEmbed, and Qwen3-
Embedding-0.6B. MLDR-en and BrowseComp-
Plus are released under the MIT license. Qwen3-
Embedding-0.6B is released under the Apache-2.0
license. LongEmbed is a public research bench-
mark for long-context retrieval evaluation; its Hug-
ging Face dataset card does not list an explicit li-
cense field, so we use it only for research evaluation
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/text/2610.08463v1.txt
+2005
View File
File diff suppressed because it is too large. Load diff
wizard-work/sphinx/archive-longctx-retrieval-2026-10-26/verify_archive.py
+135
View File
@@ -0,0 +1,135 @@
#!/usr/bin/env python3
"""Independent verifier for archive-longctx-retrieval-2026-10-26.
Re-derives every claim in index.json from the bytes on disk: PDF SHA-256,
text-layer word counts (fresh pdftotext run), page/image counts, corpus totals,
and the MANIFEST hash. Exits 0 only if every check passes.
"""
import hashlib
import json
import os
import subprocess
import sys
import tempfile
OUT = "/root/.hermes/wizard-outputs/sphinx/archive-longctx-retrieval-2026-10-26"
INDEX = os.path.join(OUT, "index.json")
MANIFEST = os.path.join(OUT, "MANIFEST.sha256")
EXPECT_DOCS = 12
MIN_PAGES = 3
MIN_WORDS = 1000
MIN_ABSTRACT = 200
failures = []
checks = 0
def check(cond, label):
global checks
checks += 1
if not cond:
failures.append(label)
print("FAIL %s" % label)
else:
print("ok %s" % label)
def sha256(path):
h = hashlib.sha256()
with open(path, "rb") as f:
for chunk in iter(lambda: f.read(1 << 20), b""):
h.update(chunk)
return h.hexdigest()
def main():
check(os.path.isfile(INDEX), "index.json exists")
if not os.path.isfile(INDEX):
print("cannot continue without index.json")
return 1
with open(INDEX, encoding="utf-8") as f:
idx = json.load(f)
check(idx.get("version") == 1, "version == 1")
check(idx.get("documents_archived") == EXPECT_DOCS, "documents_archived == %d" % EXPECT_DOCS)
check(len(idx.get("documents", [])) == EXPECT_DOCS, "len(documents) == %d" % EXPECT_DOCS)
ids = [d["arxiv_id"] for d in idx["documents"]]
check(len(set(ids)) == len(ids), "arxiv_id values unique")
# MANIFEST must match the index bytes on disk
if os.path.isfile(MANIFEST):
with open(MANIFEST) as f:
line = f.read().strip().split()
check(len(line) == 2 and line[1] == "index.json", "MANIFEST well-formed")
check(line[0] == sha256(INDEX), "MANIFEST sha256 matches index.json on disk")
else:
check(False, "MANIFEST.sha256 exists")
tot_bytes = tot_pages = tot_imgs = tot_words = 0
for d in idx["documents"]:
aid = d["arxiv_id"]
pdf = os.path.join(OUT, d["pdf_path"])
txt = os.path.join(OUT, d["text_path"])
check(os.path.isfile(pdf) and os.path.getsize(pdf) > 10000, "%s pdf present + non-trivial" % aid)
check(os.path.getsize(pdf) == d["pdf_bytes"], "%s pdf_bytes matches disk" % aid)
check(sha256(pdf) == d["pdf_sha256"], "%s pdf sha256 matches" % aid)
check(bool(d["title"].strip()), "%s title non-empty" % aid)
check(len(d["abstract"]) >= MIN_ABSTRACT, "%s abstract >= %d chars" % (aid, MIN_ABSTRACT))
check(d["pages"] >= MIN_PAGES, "%s pages >= %d" % (aid, MIN_PAGES))
# fresh, independent text extraction
if os.path.isfile(txt):
with open(txt, encoding="utf-8", errors="replace") as f:
disk_words = len(f.read().split())
check(disk_words == d["text_words"], "%s text_words matches stored text file" % aid)
check(disk_words >= MIN_WORDS, "%s text_words >= %d" % (aid, MIN_WORDS))
else:
check(False, "%s text file present" % aid)
disk_words = 0
with tempfile.TemporaryDirectory() as td:
fresh = os.path.join(td, "fresh.txt")
subprocess.run(["pdftotext", "-layout", pdf, fresh], capture_output=True)
if os.path.isfile(fresh):
with open(fresh, encoding="utf-8", errors="replace") as f:
fresh_words = len(f.read().split())
check(fresh_words == d["text_words"], "%s fresh pdftotext word count matches" % aid)
else:
check(False, "%s fresh pdftotext produced output" % aid)
imgs = subprocess.run(["pdfimages", "-list", pdf], capture_output=True, text=True)
raster = 0
if imgs.returncode == 0:
for line in imgs.stdout.splitlines():
p = line.split()
if len(p) >= 3 and p[2] == "image":
raster += 1
check(raster == d["raster_images"], "%s raster_images matches pdfimages -list" % aid)
tot_bytes += d["pdf_bytes"]
tot_pages += d["pages"]
tot_imgs += d["raster_images"]
tot_words += d["text_words"]
c = idx["corpus"]
check(c["total_pdf_bytes"] == tot_bytes, "corpus.total_pdf_bytes == sum")
check(c["total_pages"] == tot_pages, "corpus.total_pages == sum")
check(c["total_raster_images"] == tot_imgs, "corpus.total_raster_images == sum")
check(c["total_text_words"] == tot_words, "corpus.total_text_words == sum")
print("\n%d checks, %d failures" % (checks, len(failures)))
if failures:
for fl in failures:
print(" FAILED: %s" % fl)
return 1
print("ARCHIVE VERIFIED: %d docs, %d pages, %d raster images, %d words, %d pdf bytes" % (
len(ids), tot_pages, tot_imgs, tot_words, tot_bytes))
return 0
if __name__ == "__main__":
sys.exit(main())