Overflow Prevention Enhances Long-Context Recurrent LLMs

We thank the reviewer for their detailed and valuable feedback. We address the reviewer’s concerns and questions below:

.

OPRM and LLMxMapReduce - Differences and Comparison

Differences
LLM×MapReduce is an agentic pipeline that handles long contexts by first splitting the context into several chunks for LLMs to process and reason over, and then aggregating their intermediate answers using additional LLM calls to produce the final output. In contrast, OPRM is an inference technique for recurrent LLMs that speculatively processes all context chunks and then selects the best one for decoding, all to avoid memory overflows. OPRM requires only a single forward pass.

Comparison
We applied LLMxMapReduce to Falcon3-Mamba-Inst-7B. While it produced reasonable answers for some queries, it often failed to follow the full set of instructions at different stages of the pipeline. Notably, the LLMxMapReduce paper evaluates Transformer models with a minimum size of 70 billion parameters, whereas Falcon3-Mamba-Inst-7B, at 7 billion parameters, is currently the largest available recurrent LLM. We believe this performance gap is largely due to the significant difference in model scale. It would be interesting to evaluate LLMxMapReduce on larger-scale recurrent LLMs once they become available.

Efficiency
We report that OPRM is over an order of magnitude faster than LLMxMapReduce: while the latter requires 225.2 seconds per query, OPRM completes a query in 8.6 seconds.

We will ensure that LLMxMapReduce is properly cited and discuss the differences between the methods, as well as the conditions under which LLMxMapReduce can be applied to recurrent LLMs.

.

multi-hop benchmarks should be evaluated to verify the effectiveness of OPRM

We refer to Section 5.1 (LongBench), where multiple popular multi-hop reasoning benchmarks are evaluated, including: HotpotQA (2-hop), MuSiQue (up to 4-hop), and 2WikiMultihopQA (up to 5-hop). These benchmarks are designed to prevent shortcut-based solutions, and their difficulty has been further increased through the inclusion of distractor passages.
Below, we summarize the performance improvements achieved by OPRM on these benchmarks. We report the average results across all four SoTA recurrent LLMs that were tested:

LB is LongBench and 0-4k, 4-8k, 8k+ are LongBench-E’s length groups. We note that MuSiQue only has a LB dataset. The baseline models are: Falcon3-Mamba-Inst-7B, Falcon-Mamba-Inst-7B, RecurrentGemma-IT-9B, RWKV6-Finch-7B. ‘Improvement’ shows the additional improvement gained by OPRM. The data is from Tables 9-12 in the Appendix.

The results above demonstrate that OPRM substantially enhances multi-hop reasoning performance, with improvements exceeding 100% in certain benchmarks across all evaluated SoTA models. These findings show that memory overflows degrade multi-hop reasoning significantly and that OPRM offers an efficient and effective training-free solution.

To address the concern as thoroughly as possible, we additionally evaluate our method on a large, diverse portion of the proposed InfiniteBench benchmark:

Consistent with our earlier findings, OPRM significantly outperforms the baseline. We note that Code.Run is not solved by either model. As reported by InfiniteBench, this synthetic task - which requires computing a complex composition of many functions - remains unsolved even by leading proprietary models such as Claude, Mistral, and Kimi-Chat. The same is for Math.Calc (all proprietary models score 0 on it).

We conclude that OPRM provides significant improvements in real-world long-context multi-hop tasks, as demonstrated both in our original experiments and in the challenging InfiniteBench evaluation. In the revised version of the paper, we will include these additional results and further emphasize the effectiveness of OPRM for multi-hop reasoning.

.

Improving the clarity of the Method section

We will add a formal description of OPRM using LaTeX’s algorithm environment, and compress some of the language descriptions.

We thank the reviewer for their detailed and valuable feedback.

The reviewer claims that our method is a form of RAG and points out that we have not compared it with RAG baselines, so the study appears unjustified. We respond as follows:

(i) Research focus:
The main objective of our research is to study the fundamental limitations of recurrent LLMs in long-context tasks and to explore ways to improve them. We find that recurrent LLMs, when provided with an appropriately sized input segment, can outperform full-context inference. This result is both important and surprising, given that recurrent LLMs are specifically designed for long-sequence processing. Our proposed solution, OPRM, highlights both the limitations of these models and the substantial benefits of context truncation - offering a simple and efficient approach.

(ii) Methodological differences:
Although there are some similarities, our approach is not conventional RAG. It does not use an external retrieval model, does not involve additional training, and does not follow the typical RAG pipeline of retrieval, prefill, and decode. Instead, it operates in a single forward pass, like a standard LLM. If we adopt the reviewer’s broader definition, where any chunk-selection procedure qualifies as RAG, then the closest analogy would be zero-shot RAG.

That said, to fully address the reviewer’s concerns, we perform additional experiments comparing our method to two zero-shot RAG baselines:

• DRAGON, which employs a dense retriever designed to generalize in zero-shot settings [1].
• PRP, a retriever-free method that uses additional forward passes to rank chunks [2].

We evaluate these methods across all four Document QA tasks in LongBench_e (HotpotQA, 2WikiMQA, MultiFieldQA, and Qasper), and report the average performance across all benchmarks. As shown in the table below, OPRM outperforms both methods:

For each RAG method, we tested multiple combinations of num_chunks and chunk_size, ensuring a fair comparison by constraining num_chunks x chunk_size = OPRM_chunk_size. It is also worth noting that PRP's sequential ranking process makes it significantly slower than both OPRM and DRAGON: PRP requires several minutes to process each sample, whereas OPRM and DRAGON take only a few seconds.

These results are consistent with prior findings: RAG-based methods generally do not outperform long-context LLMs on long-context tasks [3,4,5]. To the best of our knowledge, there is no definitive explanation, yet some hypotheses suggest that long-context models may be better at uncovering multi-hop relations, as “…long context(s) are beneficial to increase the recall of incorporating all intermediate hops” [5].
Moreover, the recurrent LLMs used in our work share the same computational complexity as RAG-based solutions. This stands in contrast to the Transformer-based long-context LLMs used in [3,4,5], which incur substantial computational overhead. Taken together, these results and other findings in our paper underscore the potential of long-context recurrent LLMs when combined with overflow prevention techniques such as OPRM. We will incorporate these comparisons and discussion in the revised manuscript.

.

The authors need to discuss RAG systems in the related work section.

We will add a subsection to the Related Work section discussing RAG-based systems. This subsection will explicitly cover the RAG pipeline (retrieve, prefill, decode), various types of RAG (external retriever vs. retriever-free, fine-tuning vs. zero-shot, and advanced methods), and the use of RAG in long-context tasks [3,4,5].

.
We would like to thank the reviewer for raising these important points, which provided valuable perspectives and insights that have helped improve our paper.

.
References:
[1] How to Train Your DRAGON: Diverse Augmentation Towards Generalizable Dense Retrieval; Lin et. al; EMNLP Findings 2023
[2] Large Language Models are Effective Text Rankers with Pairwise Ranking Prompting; Qin et. al; NAACL Findings 2024
[3] Retrieval Augmented Generation or Long-Context LLMs? A Comprehensive Study and Hybrid Approach; Li et. al; EMNLP 2024
[4] LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding; Bai et. al; ACL 2024
[5] Retrieval Meets Long Context Large Language Models; Xu et. al; ICLR 2024

We thank the reviewer for the valuable feedback, and address the raised concerns below:

.

The limitations of linear attention in AR are well-known and have been studied in works such as DeciMamba and LongMamba. The new insights in Section 3 are limited.

DeciMamba and LongMamba do not study the Associative Recall (AR) task [1,2]. Instead, they analyze long contexts from the view of length generalization, focusing on explaining performance degradations that occur when the model processes sequences that are longer than those it was trained on. For clarity, all experiments in Section 3 use contexts that are shorter than the training length, hence should not suffer from such issues.

Section 3 Shows That a More Fundamental Problem Exists:
SoTA recurrent LLMs such as Falcon-Mamba-Inst-7B are:
(a) Sufficiently large for good in-context understanding [9]
(b) Match Transformer performance on short sequence tasks [13]
(c) Do not need to length-generalize in some long-context benchmarks (trained on contexts longer than the ones in the benchmark)
(d) Feature a significantly large hidden state - which should correlate with increased recurrent memory capacity [3-6].

Yet, despite all of the above, these models still underperform on real-world long-context tasks.
To investigate this, we turn to the AR task and show that a more fundamental problem exists (L114–115).
Surprisingly, under these “good conditions”, the AR score is much lower than expected: the model cannot reliably store more than 40 key-value pairs (Figure 1, left), and its performance degrades sharply under memory overload. This suggests that while recurrent LLMs trained via next-token prediction may develop some form of recurrent memory, its capacity falls far short of both the demands of long-context tasks and the model’s actual memorization potential.

Indeed, AR has been used in prior work. However, to the best of our knowledge, we are the first to apply it to large-scale, long-context recurrent models and use it to identify memory overflows as a fundamental limitation in long-context settings.

.

The proposed method is not limited to Mamba; a broader range of model families are expected.

We refer to Section 5.1 (LongBench, LongBench v2), Figure 1 (Right), and Tables 1&2, where we evaluate a variety of recurrent and hybrid architectures. Notably, on LongBench, OPRM inference improves RWKV6-Finch-7B by 51%. It also improves RecurrentGemma-IT-9B (a hybrid of Sliding Window Attention and RG-LRU) by 50%. These results (via RWKV, Griffin, Falcon-Mamba and Falcon3-Mamba) show that OPRM generalizes well beyond Mamba, offering substantial gains across diverse recurrent and hybrid architectures.

.

It is unclear whether OPRM can handle scenarios where the required information is spread in different chunks. Combining chunks could be better.

To address this concern, we conduct the following experiment:
We use the OPRM pipeline to first identify top-k (smaller) relevant chunks and then evaluate them together with an additional forward pass. We refer to this method as CC (Combined Chunks). We evaluated CC’s performance across all 4 Document QA tasks in LongBench_e and report the average scores:

Falcon-Mamba-Inst-7B is the baseline (vanilla inference). ‘+’ indicates the inference algorithm. The 4 datasets are HotpotQA, 2WikiMQA, MultiFieldQA, and Qasper. The total length of the CC chunks is constrained to equal one OPRM chunk. We experimented with multiple combinations of chunk length and number, and used the best parameters for testing.

We see that while CC is beneficial for short contexts (0-4k), it offers no clear advantage over OPRM in longer contexts (4-8k and 8k+). Moreover, it adds computational and algorithmic overheads.
One possible explanation is that longer contiguous contexts help the model uncover multi-hop relations [10], as some reasoning steps depend on access to previously seen information. In such cases, given the same length budget, selecting longer segments may be more effective than using multiple smaller chunks: when a chunk is small, its relevance is more likely to be biased toward surface-level similarity with the query. In contrast, longer segments include additional context that can improve the relevance estimation.

We conclude that cross-chunk information fusion warrants further investigation (L333-334). For example, a system-oriented approach, such as an overflow-aware variant of Bahdanau Attention [8] or hierarchical state processing, may hold significant potential. We will add the above experiment and discussion to the revision of the paper.

We thank the reviewer for their valuable comments and suggestions.

.

All experiments of AR were conducted under one sequence length with 1,200 tokens. It is recommended to evaluate AR with different sequence lengths.

Following the reviewer’s suggestion, we extend the AR experiment to multiple sequence lengths: 300, 600, 1200, 2400, 4800. The results can be found here: https://bashify.io/i/xsnIwA
Our findings indicate that memory capacity is not particularly sensitive to the overall context length.
Similar to the original 1,200-token setting, performance starts at around 75% accuracy for 10 key-value pairs (facts), and gradually declines as the number of facts increases, eventually converging toward 0% accuracy.
Note that each fact is 6 tokens long, therefore for sequence length L=300 we can only test a maximum of 50 facts and for L=600 a maximum of 100 facts.

.

It is recommended to control the position of the queried KV pair in the AR task to better understand the recurrent memory failure problem.

In response to the reviewer’s recommendation, we analyze the position of successfully retrieved key-value (KV) pairs. The results can be found here: https://bashify.io/i/R2qtbs
Each plot presents a normalized histogram of the locations of successfully retrieved KV pairs (facts). In all measurements we use 10 equally spaced bins to aggregate the samples, and average over 5 different seeds. When the number of facts is small, the distribution appears relatively uniform, which corresponds with the high AR accuracy. Interestingly, as the number of facts increases, we observe a U-shaped recall pattern, resembling the lost-in-the-middle phenomenon reported in [1].
We conclude that while recall performance degrades across all positions, the middle portions of the context suffer the most. Since this phenomenon is also observed in Transformer-based LLMs, we suspect that positional relevance may be influenced by properties of the data itself rather than the model's architecture. This noteworthy finding warrants further investigation in future work.

.
We will add both experiments to the revised manuscript. We thank the reviewer for suggesting these experiments, which offer deeper insight into the behavior of the recurrent memory in recall-intensive tasks, and help improve the paper.

References:
[1] Lost in the Middle: How Language Models Use Long Contexts; Liu et. al; TACL 2023