Efficient Full-Context Retrieval for Long Documents

Authors: Weili Cao, Jianyou Wang, Youze Zheng, Longtian Bao, Qirui Zheng, Taylor Berg-Kirkpatrick, Ramamohan Paturi, Leon Bergen
Primary area: applications to computer vision, audio, language, and other modalities
Keywords: Long-context retrieval, Efficient language models, Mamba architecture, Synthetic data generation

Abstract:
Long document question answering is challenging due to the quadratic computational cost of transformer-based LLMs. In resource-constrained environments, Retrieval-Augmented Generation (RAG) uses document chunking to maintain linear computational costs but often loses sight of the global context. We introduce the Mamba retriever 130M and Mamba retriever 1.3B retrievers, capable of processing entire long documents in linear time and integrating earlier context to retrieve relevant sentences for answering questions. Mamba retrievers outperform state-of-the-art embedding models across 41 long-document Q\&A benchmarks while maintaining speed and computational efficiency. Their performance is comparable to GPT-4o on long documents over 256k tokens while using significantly fewer tokens. Mamba retrievers are trained on synthetic data generated from our novel link-based method, which enhances the retrievers' ability to leverage long-range document connections. We further demonstrate the effectiveness of our link-based method over baseline synthetic data methods. All code, datasets, and model checkpoints are available at https://github.com/MambaRetriever/MambaRetriever

Decision: reject
Reviewer scores: 6, 5, 6, 6

Reviews and discussion:
<thread>
creation date:
1729183181194

thread_type:
official_review

<thread_object>
creator: reviewer_ws7s

<content>
summary:
The authors introduce a novel approach for long document understanding using the Mamba architecture as an efficient information retriever. Unlike traditional methods that rely on document chunking, this approach processes the entire document, maintaining important contextual dependencies that chunking-based methods may lose. To overcome the challenge of limited labeled datasets for training such a model, the authors employ synthetic data generation strategies—specifically chunk-based, pair-based, and link-based approaches—prompting an LLM to the training examples. The model, trained on both 130M and 1.3B parameter variants, is evaluated on 41 benchmark tasks, spanning domains such as financial documents, government reports, and creative works. The results demonstrate that their Mamba retriever consistently outperforms embedding-based retrievers and exhibits near-GPT-4-level performance on very long documents exceeding 256K tokens. Ablation studies further highlight the importance of using full-document context and the benefits of the link-based synthetic training dataset, which effectively improves retrieval accuracy compared to the other data generation strategies.

strengths:
- The proposed approach effectively removes the need for document chunking, enabling a more comprehensive and uninterrupted representation of document context. This is particularly advantageous for handling long documents where preserving contextual relationships is crucial.
- The Mamba-2 model, available in both 130M and 1.3B parameter variants, demonstrates impressive performance, notably surpassing models significantly larger in size, such as GritLM-7B. This result underscores the efficiency of the model design, making it a valuable solution for environments with limited computational resources.
- To address the lack of annotated datasets for training their long-context models, the authors introduce diverse synthetic data generation methods, including chunk-based, pair-based, and link-based strategies. These approaches help to create high-quality training data that may enable effective model learning.
- The paper includes a well-executed analysis, supported by ablation studies that provide insight into the contributions of different model components and training strategies.

weaknesses:
I found the proposed idea, experiments, and analyses conducted by the authors to be valuable, especially in terms of their potential impact on low-resource scenarios. However, for the paper to fully meet the ICLR standards, there are still areas that need additional work and detail. Below, I outline several key points for improvement. I would be pleased to substantially raise my scores if the authors address these suggestions and enhance the paper accordingly.

**General Feedback**
- I noticed that the title of the paper does not match the one listed on OpenReview.
- The main text should indicate when additional detailed discussions are deferred to the Appendix for better reader guidance.

**Introduction**

- The Introduction lacks foundational references to support key claims. Both the second and third paragraphs would benefit from citations to strengthen the arguments. For instance, the statement: "This method eliminates the need for document chunking, *a common limitation in current retrieval systems that often results in loss of context and reduced accuracy*" needs a supporting citation to substantiate this point.
- The sentence: "Second, to be competitive with embedding approaches, a retrieval language model needs to be small" requires further justification. The authors should include in the paper a complexity analysis comparison discussing time and GPU memory consumption to support this assertion.

**Related Work**

- The sentence "Large Language Models are found to be inefficient processing long-context documents" should be rewritten for clarity, for example: "Large Language Models are inefficient when processing long-context documents."
- The statements "Transformer models suffer from quadratic computation during training and linear computation during inference" and "However, transformer-based models are infeasible to process extremely long documents due to their linear inference time" are incorrect. Transformers, as presented in "Attention is All You Need," scale quadratically in both training and inference.
- The statement regarding State Space Models (SSMs) having "linear scaling during training and constant scaling during inference" is inaccurate. SSMs have linear complexity for both training and inference. The term "constant scaling" implies no dependence on sequence length, which is incorrect.
- The Related Work section is lacking details. The paragraph on long-context language models should provide a more comprehensive overview of existing methods and their limitations, positioning SSMs appropriately. This includes discussing sparse-attention mechanisms [1, 2], segmentation-based approaches [3, 4, 5], memory-enhanced segmentation strategies [6], and recursive methods [7] for handling very long documents.
- Similarly, the paragraph on Retrieval-Augmented Generation should specify how prior works addressed different long document tasks. Examples include successful applications of RAG in long-document summarization [8, 9] and query-focused multi-document summarization [10, 11], which are closely aligned with the present work.

**Figures**
- Figures 1 and 2 are clear but need aesthetic improvements to meet the conference's standard presentation quality.

**Model Architecture**
- The description "a subset of tokens are specially designated, and the classification head is applied to these tokens. In the current work, the classification head is applied to the last token of each sentence, giving sentence-level resolution" is ambiguous. Clarify whether new tokens are added to the sequence or if existing tokens (e.g., periods) are used to represent sentence ends.

**Synthetic Data Generation**

- The "lost in the middle" problem when processing long documents [12] is not explicitly discussed. Have the authors considered the position of chunks during synthetic data generation? Ablation studies varying the position and distance between linked chunks would provide valuable insights into Mamba’s effectiveness in addressing this issue.
- More details are needed regarding the data decontamination pipeline, chunk size, and the relative computational cost of the link-based method versus other strategies.
- The authors claim that synthetic data generation is computationally expensive but provide no supporting quantitative evidence. Information such as time estimates and GPU demand would strengthen this argument and assess feasibility.
- There is no detailed evaluation of the synthetic data’s quality. An analysis of correctness and answer factuality would help validate the impact on retrieval performance beyond benchmark metrics.

**Training**
- This section is too brief. Consider merging it with Section 3, "Model Architecture," for a more cohesive presentation.
- What was the training time for the 130M model?

**Experimental Method**
- Fix minor formatting issues, such as adding a space after the comma in ",LVeval."
- Specify in Table 1 which datasets use free-form versus multiple-choice answers, including the number of answers and average answer lengths.
- Consider experimenting with GPT-4 as a retriever.
- Expand on "The accuracy of freeform answers is judged using GPT-4."
- Elaborate on the validation of the scoring pipeline, particularly regarding "0.942 macro F1." Clarify the data and method used for validation.
- Justify the selection of "50 sentences" for Mamba retrievers and explain chunk creation methods for embedding models. Did the chunks consist of 300 fixed-length segments, or was semantic chunking employed [3, 5]? Sentence-level embedding-based retrieval could be explored to align better with the Mamba setting.
- The assertion that "embedding models were allowed to retrieve more information than Mamba" implies an unfair comparison, but more context can sometimes degrade performance [12].
- Clarify the use of the sliding window approach for documents longer than 128k tokens, especially given the claim that Mamba could process up to 256K tokens directly.

**Results**
- Remove redundancy in Section 7.1.2, such as restating the synthetic data generation strategies.
- Expand the ablation studies to cover different input sequence lengths during training and varying the number of retrieved sentences to explore robustness to configuration changes.
- Highlight that using fewer training examples (500K vs. 1M) achieved comparable accuracy (i.e., 59.4 vs. 60.0, respectively).
- Why not train both the 130M and 1.3B models on a dataset size of 500K examples, but compare using 1M and 400K examples, respectively?

**Limitations**
- The high cost of generating synthetic training data is mentioned but lacks quantification. How computationally expensive is it in terms of time or resources?

**Appendix**
- Note that all figures in Appendices B and C are the same, suggesting an error that needs correcting.

**Missing References**

[1] Longformer: The Long-Document Transformer. arXiv 2020.

[2] LongT5: Efficient Text-To-Text Transformer for Long Sequences. NAACL 2022.

[3] Semantic Self-Segmentation for Abstractive Summarization of Long Documents in Low-Resource Regimes. AAAI 2022.

[4] Summ^n: A Multi-Stage Summarization Framework for Long Input Dialogues and Documents. ACL 2022.

[5] Align-Then-Abstract Representation Learning for Low-Resource Summarization. Neurocomputing 2023.

[6] Efficient Memory-Enhanced Transformer for Long-Document Summarization in Low-Resource Regimes. Sensors 2023.

[7] Recursively Summarizing Books with Human Feedback. arXiv 2021.

[8] DYLE: Dynamic Latent Extraction for Abstractive Long-Input Summarization. ACL 2022.

[9] Towards a Robust Retrieval-Based Summarization System. arXiv 2024.

[10] Discriminative Marginalized Probabilistic Neural Method for Multi-Document Summarization of Medical Literature. ACL 2022.

[11] Retrieve-and-Rank End-to-End Summarization of Biomedical Studies. SISAP 2023.

[12] Lost in the Middle: How Language Models Use Long Contexts. TACL 2024.

questions:
Please address the questions outlined in the weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1729573886370

thread_type:
official_review

<thread_object>
creator: reviewer_psm7

<content>
summary:
### Summary
1. A Mamba architecture based retriever is presented which is trained as a discriminator model
2. The retrievers are much smaller (130M and 1.3B) and outperform much larger open source (and non-finetuned) embedding based models
3. Efficacy of all retrievers is tested by plugging the retrieved chunks (or sentences) into a RAG pipeline


### Contributions
1. The paper showcases how a Mamba architecture model could be finetuned to create an effective retriever for long context documents
2. Methodology and ablation for creating synthetic data is well studied and supported via ablation
3. Synthetic dataset would further research for long context document understanding (and retrieval)

strengths:
1. This work presents a lightweight (small parameter) based retriever for long context documents
2. Complete training code and dataset is available on Github for reproducibility
3. New dataset would further research in the direction
4. Ablations presented help understand the impact of synthetic data generation process (link based) as well as effect of context size on RAG system performance

weaknesses:
1. The baseline retriever (open source embedding model) is weak as it has not seen any synthetic data. It is unclear whether gains seen by Mamba retriever is due to synthetic data or the model's ability to handle long context
2. The paper claims that Mamba is more efficient in handling long document context. No supportive results are given in terms of time taken. It's unclear how efficient the model is when compared to a retriever
3. Embedding based retriever has a different retrieval granularity (chunks) as compared to Mamba (Top 50 sentences)
4. The paper claims there is lack of synthetic data for long context documents, but ends up creating a dataset where maximum context length is 10k. Is my understanding correct?
5. (Minor issue) No qualitative examples are presented in the paper which highlight how Mamba retrieves more relevant context as compared to Embedding based models. Perhaps some examples can be added to the Appendix, if there

questions:
1. (Lines 233 to 239) Why is the retriever setup different for Mamba and embedding based retrievers? Mamba fetches Top 50 sentences whereas embedding based retrievers fetch 5 chunks. How do embedding models perform when you have identical setup?
2. Why was the link based synthetic data not used to finetune an embedding model and compared as baseline? Perhaps the gains seen by Mamba retriever are due to synthetic data. A quick check would be to see how embedding retrievers perform (in terms of recall) on a held out set of the synthetic link based data. Maybe this number could be improved via finetuning and the finetuned embedding model tested as baseline
3. (Lines 198-200) What decontamination pipeline was used to ensure that there is no overlap between synthetic data documents and documents used in the 41 benchmarks. Did you look for exact match of document content? Are the domains completely different?
4. (Lines 230-232) How did you evaluate the quality of an answer using GPT-4o? Did you give some instructions in prompt on how to judge the quality of an answer? Did the model return a binary score (such as relevant or non-relevant)
5. How was BM25 setup? Did it also retriever chunks? What happens when it retrieves sentences (same as Mamba)
6. (Lines 204-207) Which hyper-parameters were optimized? How did you select the best checkpoint? How did you decide when to stop training?
7. (Minor comment): Figures 6-10 in the Appendix seem to have the same exact prompt. Was this intentional or done by mistake?
8. (Minor comment): Lines 89-90 say upto 256 k. Did the authors mean more than 256k?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730106208830

thread_type:
official_review

<thread_object>
creator: reviewer_4h6p

<content>
summary:
The paper presents a retrieval model based on the Mamba architecture to improve long document understanding. By fine-tuning a Mamba checkpoint, the authors aim to handle entire document contexts without chunking, thus enhancing retrieval accuracy. The paper highlights synthetic data generation for training the model, including three methods—chunk-based, pair-based, and link-based. Experimental results indicate that the Mamba retriever performs comparably with existing large language models in various retrieval tasks while being more efficient. The model is evaluated on 41 QA benchmarks across multiple domains, demonstrating that the full-context Mamba retriever outperforms embedding-based retrievers, especially for long documents.

strengths:
1. The paper introduces a novel adaptation of the Mamba architecture, focusing on processing entire documents without chunking, which may provide a more holistic context.
2. The paper addresses the scarcity of long-context retrieval data by developing synthetic data generation methods, specifically the effective link-based generation, which strengthens the retrieval model’s generalization.

weaknesses:
1. The paper’s readability suffers due to numerous grammatical issues, ambiguous statements, and poorly constructed sentences. A thorough rewrite would significantly improve comprehension and presentation.
2. Many sections make assertions without proper references, undermining the credibility of the stated findings and comparisons.
3. The proposed method essentially fine-tunes Mamba for retrieval tasks without introducing substantial architectural innovation or unique methodological contributions.
4. The authors categorize datasets into four types but fail to clarify the rationale or methodology behind this categorization.
5. Given the method’s focus on long-context retrieval, comparisons with state-of-the-art long LLMs (such as Llama 3.1, ChatGLM, etc.) and leading retrievers from the mteb leaderboard would provide a more competitive and comprehensive evaluation.
6. Efficiency is critical in retrieval tasks, yet the paper omits a direct analysis of retrieval speed and computational resources, which would provide essential context on the practical applicability of the Mamba retriever.

questions:
1. What criteria were used to categorize datasets into four types, and how was this categorization verified or validated?
2. How does the efficiency of the Mamba retriever compare to traditional embedding-based models in terms of speed, memory usage, and computational cost?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730698978885

thread_type:
official_review

<thread_object>
creator: reviewer_f8dj

<content>
summary:
This paper proposes to find-tune a long-context Mamba model for the retrieval task, claiming that such a method is better than embedding-based retriever because the awareness of context. To train such a model, a data synthesis method is proposed, which is to generate query given two semantically relevant text chunks. In the experiments, the proposed method is compared with long-context model and standard RAG methods and shows strong performance.

strengths:
- The paper studies an idea of fine-tuning long-context model as retriever. To gain efficiency for training, the Mamba model is used to save cost for training a Transformer model
- The performance shows the long-context discriminative model as retriever brings improvement to the retrieval performance.

weaknesses:
- The efficiency of proposed method is not studied. I can expect that with pre-built vector index, the standard RAG may be more efficient that the proposed method which needs to pass the entire textual data for each query. Although adding context for retrieval brings improvement, I am not sure whether it worth the cost if its non-trivial.
- I am not clear with the data split. What data is used for data synthesis and training and what for testing? The paper only mentions the test set so I am not sure if it is used for both.
- It is not clear to me what if the long-context model is directly trained to be a reader, i.e. generate the answer given a query. The data synthesis part can generate a query so can generate a corresponding answer as well. This should be a baseline and shows the necessity of training the model as a retrieval followed by another reader model.

questions:
See the weaknesses raised above
</content>
</thread_object>

</thread>
<thread>
creation date:
1733293162254

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of discussion period

comment:
We thank the reviewers for their questions and comments. These have allowed us to strengthen the paper through our revisions. All reviewers increased their scores following the revisions, which focused on three main areas.

Our direct comparisons with embedding models now clearly demonstrate efficiency gains - the 130M Mamba retriever achieves higher accuracy while using 68x fewer TFLOPS than the leading 7B parameter baseline. Our larger 1.3B model achieves higher accuracy while still using 6.5x fewer TFLOPS.

We resolved the core question about whether gains came from synthetic data alone by showing that fine-tuning embedding models on our data did not yield comparable improvements. This demonstrates the gains come from our model's ability to process document-wide context.

We addressed evaluation fairness by showing that improvements over embedding models are robust to different experimental settings. We also added Llama-70B comparisons showing our approach outperforms much larger models.

The remaining concerns from other reviewers focus mainly on presentation rather than technical issues. For example, Reviewer 4h6p acknowledged that our responses "resolve most concerns" but maintains reservations about novelty framing, which we have addressed in the revised introduction. Regarding this question of novelty, prior and contemporaneous work on SSMs for retrieval are embedding models, whereas the current work is not.

The revisions have created reviewer consensus around our key contribution: achieving strong long-document understanding with a small, efficient model.
</content>
</thread_object>

</thread>


Metareview:
metareview:
**Summary:**

The paper introduces the Mamba retriever, a model designed for efficient and effective retrieval in long document question answering. Unlike traditional embedding-based methods that require document chunking and can lose global context, the Mamba retriever processes entire documents linearly, leveraging long-range dependencies to retrieve relevant sentences for answering queries. This approach outperforms state-of-the-art models, achieving near-GPT-4o performance on documents exceeding 256k tokens. Despite its computational efficiency and superior performance, the high cost of synthetic data generation presents a limitation for scalability.

**Strength:**
- It explores fine-tuning long-context models as retrievers, demonstrating improved retrieval performance.
- It Eliminates the need for document chunking, preserving contextual relationships in long documents.
- It includes thorough ablation studies to analyze contributions of model components and training strategies.

**Weakness:**
- The method essentially fine-tunes the Mamba model for retrieval without introducing significant architectural or methodological innovation.
- The baseline retriever does not leverage synthetic data, making it unclear whether Mamba’s improvements are due to data or its ability to handle long contexts.
- Enhancements in clarity, detail, and comparative evaluation would significantly improve the paper’s quality and review scores.