Efficient Full-Context Retrieval for Long Documents

Thank you for your insightful responses. We have rewritten our paper. All changes we made are reflected in red color.

I noticed that the title of the paper does not match the one listed on OpenReview.

Thank you for pointing this out. We have revised the title.

The main text should indicate when additional detailed discussions are deferred to the Appendix for better reader guidance.

In the updated paper, we have added additional references to the appendix.

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.

We have rewritten the introduction and added proper citations. We included all the citations you suggested in Related Work.

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.

We compare the Mamba retriever with three state-of-the-art embedding models, including the leader of the Massive Text Embedding Benchmark (MTEB), evaluating their performance in terms of document processing speed, FLOPs, and parameters.

Mamba retrievers achieve better document processing speed and accuracy than the embedding models. The most accurate embedding model is the 7B NV-Embed-v2, which uses many more FLOPs than the 1.3B or 130M Mamba retriever. The 130M Mamba retriever achieves higher accuracy than any of the embedding models, while using a fraction of the FLOPs and achieving much faster document processing time.

We also included this discussion in Section 5.6, Table 2, Line 373-377, thanks to your suggestion.

Related Work [...]

Thank you for your insightful comments. We have adapted the suggestions in our revised paper's related work section including rewriting sentences for clarity, adding more details and adding missing references.

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

Thank you for the suggestion. We have combined and improved the two figures, which are now presented as Figure 1 in the revised paper.

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.

We revisited the entire section to enhance clarity and provide additional details.

No new token is introduced. The last token of a sentence (e.g., a period) is used to represent the end of the sentence. See Line 104-136 in Section 3 Model Architecture of Mamba Retriever.

Dear Reviewer wS7S,

Thank you again for your thorough review. As the rebuttal period ends in a few hours, we wanted to confirm that our revised manuscript has adequately addressed all your concerns. If your concerns have been addressed, we respectfully request that you increase your scores.

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

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

Based on your suggestion, we have performed these baseline experiments.

We fine-tuned two embedding models: Contriever-110M, a robust retriever used in standard RAG pipelines (Xu et al., 2023), and GTE-Qwen2-1.5B, the second best 1.5B parameter model on the massive text embedding benchmark (MTEB). These two are selected because of their open-sourced training details. These models are comparable in size to the Mamba retriever 130M and 1.3B respectively. We fine-tuned them using the same 1 million link-based synthetic data for one epoch. For each query, relevant sentences were treated as positives, while irrelevant ones were treated as negatives. We employed the same contrastive loss and applied the same hyperparameter settings (e.g., scheduler, optimizer, and temperature $\tau$ in the InfoNCE loss) as reported in their original papers. Additionally, we optimized the learning rates, batch size, and training data size using the same validation sets.

We observed no noticeable improvement after fine-tuning on the synthetic data (the data are also included in Table 1 of the revised paper). This suggests that the gains achieved by our Mamba retrievers are not merely artifacts of the training documents. Instead, the Mamba retriever has learned to synthesize information from different parts of the document, which aligns with the intention of our link-based synthetic training data.

We have incorporated these discussions into the paper (Line 295-300, Table 1, Line 406-409)

Dear Reviewer PSm7,

As the rebuttal period ends in a few hours, we urgently request your feedback on our manuscript. If your concerns have been addressed, we respectfully request that you significantly increase your scores.

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.

We have substantially revised the paper to enhance its readability. All important changes are highlighted in red in the updated version.

Many sections make assertions without proper references, undermining the credibility of the stated findings and comparisons.

We have included the previously missing references in the updated version of the paper.

The proposed method essentially fine-tunes Mamba for retrieval tasks without introducing substantial architectural innovation or unique methodological contributions.

We respectfully disagree. The approach required solving two key challenges: finding the right model adaptation and generating effective training data.

Obvious model adaptations fail: training Mamba for generative retrieval achieves only 33.5% accuracy, while direct question answering achieves only 27.6% (Tables 5 and 6). Additionally, straightforward synthetic data generation strategies like chunk-based (using individual document segments) and pair-based (matching similar chunks) produce questions that don't require understanding document-wide connections and lead to worse performance (Table 4 and Line 452-461).

Our key insight was combining a discriminative retriever architecture with a link-based training strategy that discovers real connections within documents. The effectiveness of this specific combination is demonstrated by: (1) the model outperforms embedding models 8-60x larger while using fewer computational resources (Table 1 and 2), (2) fine-tuning these embedding models on our synthetic data shows no improvement (Line 405-409), and (3) we observe strong performance on documents up to length 256k, which is more than 20x longer than any documents observed during training (Line 400-404).

The authors categorize datasets into four types but fail to clarify the rationale or methodology behind this categorization.

Because the method is evaluated on 41 benchmark datasets, it was not possible to include information about all of these datasets in the main text. Results on each individual dataset can be found in Appendix A.4 Table 10,11,12,13,14.

The test data were divided into four categories for expository purposes only. We clarify our rationale in our revised paper (Line 232-242). Furthermore, details of the datasets and their categorization can be found in Appendix A.1 Table 7 and Appendix A.2 Table 8.

Given the method’s focus on long-context retrieval, comparisons with state-of-the-art long LLMs (such as Llama 3.1, ChatGLM, etc.)

We explore the performance of LLMs in two settings. In the first setting (generative retrievers), the LLM is given a question and document and asked to generate relevant passages from the document. Given a complete document, GPT-4o and Llama-3.1-70B are asked to retrieve relevant sentences for up to 2,000 tokens, which provides a fair comparison with the "50 sentences" setup used in Mamba retrievers (Section 7.3). In the second setting (direct answering), the LLM is given the document and asked to directly answer the question (Section 7.4). We also fine-tuned Mamba-2 models for direct question answering on the same 1 million synthetic data.

Generative Retriever

Direct answer generation from full context.

The above two tables are included as Table 5,6 in the revised paper. All generative retrievers perform significantly worse than discriminative Mamba retrievers. In the direct answering setting, only GPT-4o performs better than the discriminative Mamba retrievers. This suggests that retrieval by LLMs is often lossy in long-context settings, further highlighting the advantages of using discriminative retrieval with our method.

Dear Reviewer 4h6p,

Thank you once again for your insightful suggestions throughout this process. As the rebuttal ends in a few hours, we respectfully request your prompt feedback on our revised manuscript. Your timely response is critical to ensuring we address all concerns effectively.

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.

Thank you for your feedback. Our method is specifically designed for scenarios where a document is introduced for the first time and then queried. We compare the Mamba retriever with three state-of-the-art embedding models, evaluating their performance in processing documents. Embedding models were evaluated in two settings: sentence retrieval and chunk retrieval (chunk size was 300 words). Speed and TFLOPS differ in these two settings due to the increased padding required for sentence retrieval. All embedding models were evaluated using their official repositories.

Mamba retrievers achieve better speed and accuracy than the embedding models. The most accurate embedding model is the 7B NV-embed, which uses many more FLOPs than the 1.3B or 130M Mamba retriever. The 130M Mamba retriever achieves higher accuracy than any of the embedding models while using a fraction of the FLOPs and achieving much faster document processing time. We have reflected these changes in Section 5.6, Table 2, and Line 373-377 in the revised paper.

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.

The documents in the training set were collected from novels in Project Gutenberg, government reports (Huang et al 2021), public-domain financial statements published on the US Securities and Exchange Commission website, and legal contracts from Hendrycks et al’s (Section 4.2 line 211-215). These documents were used to generate synthetic questions as described in the paper. The training sets for the 41 benchmark datasets were not used for training the model. Models were evaluated on the test sets for the 41 benchmark datasets. Detailed statistics for each of these datasets are provided in Appendix A.2 Table 8. A decontamination pipeline was applied in order to ensure that the training documents did not accidentally overlap with the testing documents. The pipeline removed any training document that had more than 1% overlap (at the sentence level) with the test set. The procedure is described in Sections 4.3 Line 218-223.