MEMREASONER: A MEMORY-AUGMENTED LANGUAGE MODEL ARCHITECTURE FOR MULTI-HOP REASONING

Authors: Payel Das, Ching-Yun Ko, Sihui Dai, Georgios Kollias, Subhajit Chaudhury, Aurelie Lozano
Primary area: generative models
Keywords: reasoning, multi hop, memory, large language models, generalization

Abstract:
Recent benchmarks suggest that there remains significant room to improve large language models’ ability to robustly reason across facts distributed in extremely long documents. In this work, we propose MemReasoner, a new memory-augmented LLM architecture that is trained to perform temporal reasoning, along with multiple computational steps,  over the context stored in the memory. Experiments show that MemReasoner trained on the core reasoning facts generalizes better, when compared to off-the-shelf large language models and existing recurrent models, on a test distribution  where  the required facts are scattered across long natural text up to 128k tokens. Further, MemReasoner demonstrates robust reasoning performance relative to the baselines, when the answer distribution in test samples  differs from that in the training set.

Decision: reject
Reviewer scores: 3, 6, 3

Reviews and discussion:
<thread>
creation date:
1730578004169

thread_type:
official_review

<thread_object>
creator: reviewer_pgx8

<content>
summary:
This paper proposes MemReasoner, a method that augments language models with memory and is designed to perform reasoning over long contexts. The authors aim to enable models trained on short sequences to generalize and work effectively on longer sequences without retraining. To achieve this, they build on the Larimar model, which augments language models with episodic memory, and enhance it with two key modifications: 1. An iterative mechanism to read information from memory and update the query accordingly. 2. Adding explicit information about temporal ordering of facts within the context (via position embeddings or BiGRU). The approach is evaluated on the bAbI dataset and its long sequence variant BABILong, demonstrating the ability to perform on longer sequences while being trained on short ones, and improving performance over existing methods in this setup.

strengths:
- General idea of iterative reading from memory for multi-hop tasks is clear and reasonable.
- The motivation to address the challenge of generalizing from short to long sequences without extensive retraining is strong and relevant, and addresses a key limitation of current methods.
- Results show that the proposed approach is indeed able to generalize to longer sequences being trained on short only.
- The evaluation setup, using the bAbI dataset and its longer version, is clear and suitable for testing the ability of the method to solve inference tasks over varying sequence lengths.

weaknesses:
- Giving that idea is clear, unfortunately Section 3 that describes the proposed method is hard to read and lacks clarity. I believe clarity of presentation of the proposed method should be largely improved. Details are provided in the suggestions section. 
- Lack of comparison with Transformer models. The authors do not evaluate the performance of transformer-based LMs that support long contexts (e.g., LLama-3.2, Phi-3.5, or Qwen-2.5). These models could be fine-tuned on bAbI and evaluated on BABILong, providing relevant baselines for comparison. Including such evaluations would contextualize the performance of MemReasoner relative to current state-of-the-art models.
- Generalization of the proposed method to tasks other than bAbI is not supported. There are other datasets that require multi-hop reasoning, such as MultiHopQA, MuSiQue, HotPotQA. Their context could be extended by extracting relative paragraphs from e.g. Wikipedia.
- The paper does not provide empirical evaluations of the method’s inference time or memory consumption compared to other methods. However, authors provide theoretical time complexity in Appendix A.2.

questions:
Suggestions on improving presentation:
- Figure 2 is hard to comprehend – it is not clear where to start from. I would suggest making two schemes: 1. conceptual scheme of the method. 2. a detailed one with how the memory module and the query are iteratively updated.
- Similar problem with Section 3 - overloaded preliminary section hindering comprehension. Section 3.1 contains too many details on how memory read/write operations are constructed, making it challenging to grasp the overall picture. Splitting the explanation into two parts – a conceptual overview and detailed descriptions of the memory module and reasoning within it – would help readers understand the main ideas before diving into technical specifics. As the method is based on the Larimar model, a dedicated section describing Larimar would greatly improve clarity.
- L292-296 present formulas for the loss components that are difficult to match with their explanations. Defining each loss component separately and explaining their roles can improve clarity and comprehension.
- How do P_a and P_s look like in L270–271? P_a is later defined in L371–372. Providing these definitions earlier would help improve clarity.
- Typos and Minor Issues
  - L315: "facts from bAbI is distributed" -> are
  - L318–319: missing space "tokens.For"

Comments:
- One of the results is that RMT being trained on short sequences (single segment) can not generalize to longer ones (multiple segments). RMT inherently can not generalize from 1 segment to multiple segments. To learn to use memory (to pass information sequentially) it needs at least 2 segments. It was shown in RMT paper and generalization to larger lengths is possible with curriculum learning procedure (“Beyond Attention: Breaking the Limits of Transformer Context Length with Recurrent Memory”). So this makes comparison with RMT not so fair.

Questions:
- Table 1 and 2 miss results for RMT on 64k and 128k. Authors mention that it “means unavailable due to out of memory errors or maximal input length constraints” (L426). However, RMT has constant memory consumption as it is required to pass only memory states between segments, and there is no need to all keep intermediate outputs. So, it is not clear why these values are missing? However, I acknowledge that they would be low.
- Is it possible to train MemReasoner on longer sequences and compare it with RMT/Mamba which has also been trained on longer sequences? Will the performance of MemReasoner continue to improve?
- L182 - what does “mimicking Bayesian inference” mean here? How is it motivated and why is it mimicking? It is unclear how this concept is applied and whether it is key to constructing an effective memory module.
- After the last readout operation, the z_i most similar to z_r is selected. Z_i, representing a single fact from bAbI, goes to a decoder. It seems to be sufficient for QA1-2 tasks to get the correct answer from bAbI, but would it be sufficient for general tasks to extract only one fact from memory? Why not to use z_r itself?
- The paper does not clearly explain whether the number of memory read operations correlates with the number of fact hops required in reasoning tasks. For example, in QA2 tasks, a fixed number of two read operations is used. Is this number optimal? Would using only one read operation be insufficient to retrieve the answer? Providing an analysis or justification for the chosen number of read operations would clarify this aspect.

Personally, I like the idea and direction of the work and find it interesting. However, I think it could be improved a lot.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730644232597

thread_type:
official_review

<thread_object>
creator: reviewer_n6d6

<content>
summary:
The paper introduces MemReasoner, a memory-augmented LLM designed to improve temporal and multi-hop reasoning over extremely long contexts. MemReasoner integrates an episodic memory module and introduces a novel mechanism for explicitly learning temporal relationships between events and for iteratively reading and updating queries, aiming to address the limitations of existing LLMs in processing and reasoning across long documents. The proposed model is evaluated on the BABILong benchmark, demonstrating superior generalization compared to existing LLMs.

strengths:
- The authors present an architecture that uses a memory module for capturing temporal relationships and enabling multi-hop reasoning, which adds a valuable improvement over standard LLMs.
- MemReasoner demonstrates a robust ability to reason across contexts of up to 128k tokens, outperforming existing LLMs and memory-augmented architectures on the challenging BABILong benchmark.

weaknesses:
- The experiments are largely based on synthetic datasets like BABILong, which, while controlled, may not fully reflect the complexities of real-world language tasks. Extending evaluations to diverse, natural datasets would strengthen the validity of the model’s utility.
- The architecture involves multiple components like GRU-based temporal encoding and iterative query updates, which might increase computational complexity, potentially making it less scalable for broader applications.

questions:
- What improvements does MemReasoner have compared to Larimar?
- Has the model been tested on other multi-hop reasoning datasets or more general NLP tasks?
- How does MemReasoner compare to traditional, non-LLM memory networks? Are there any experiments comparing it to these earlier methods, and what would the outcomes likely be?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730960486916

thread_type:
official_review

<thread_object>
creator: reviewer_vesz

<content>
summary:
MemReasoner is a novel encoder component of an encoder-decoder transformer architecture. On top of a transformer encoder stack, it uses a recurrent component to write memory and read memory.

The readout is fed into the decoder similar to key-value cache to improve performance on question-answeringe tasks.

The paper benchmarks their method on BABILong. Their method solves the task as shorter contexts but does not seem to length generalize. However, their method does do better than the baselines when symbolically manipulating locations.

strengths:
The hypothesis and method are well motivated. The experiments are well done and the model well in short contexts and is robust to symbolic manipulation of the locations.

weaknesses:
I don't think I entirely understand the comparisons. The model does much worse than the models that are finetuned on long context data. We should see if MemReasoner also benefits from this?

questions:
The encoder and memory components have some cost and overhead. What are they?

How does it compare to just give the decoder more parameters?

Can MemReasoner augment existing decoder-only LLMs?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732307188284

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response

comment:
We thank the reviewers for their feedback.  We are encouraged that the reviewers find that the problem and the proposed method of length generalization for multi-hop reasoning are well motivated [vesZ, PgX8], and that our experimental results support that MemReasoner has a robust ability to generalize to long contexts while only trained on short contexts [N6d6, PgX8] and is robust to symbolic manipulation of the locations names [vesZ].  We now address some common questions from reviewers below and will be adding these points into the updated manuscript:

**Computational cost of MemReasoner [vesZ, N6d6, PgX8]**
We provide the inference cost measured in seconds per input for evaluating with Babilong in comparison to the base decoder (gpt2-large). We note that gpt2-large does not support context lengths longer than 1024 tokens.  Overall, we observe that the increase in inference time for MemReasoner is very small for 0k and MemReasoner is more efficient for 1k context length. This is because of utilizing the latent encodings of context, performing one-shot write to the memory, and executing multiple-hops over that memory in latent space.

|  Babilong | 0k | 1k | 2k | 4k | 8k | 16k | 32k | 64k | 128k |  
| -------- | -------- | -------- |-------- | -------- |-------- | -------- |-------- | -------- |-------- |
| gpt2-large | 0.28 | 1.13 | - |- |- |- |- |- |-
| MemReasoner     |  0.30 | 0.33 | 0.40 | 0.61 | 0.98 | 1.94 | 3.26 | 11.25 | 13.77 |

**Finetuning on long data [vesZ, PgX8]**

We want to clarify that the goal of this work is to propose an LLM-based architecture that learns to reason over long-context in an efficient, robust, and generalizable manner. As such, the evaluation framework corresponds to a set-up where the core reasoning facts are diluted in the presence of irrelevant natural text distractors distributed over the context. This setup allows one to test how consistently language model can solve the same reasoning task across different input lengths. This is inspired by the recent research showing that current LLMs’ reasoning performance degrade at much shorter input lengths than their technical maximum (Levy, et al. https://arxiv.org/abs/2402.14848v1, Kuratov, et al. https://arxiv.org/abs/2406.10149). As a solution, we propose MemReasoner, which shows robust reasoning on unseen test distribution with longer length and/or distinct characteristics. At the same time, finetuning on longer sequences presents several practical challenges: (1) The longer sequences with proper (human or machine) annotation should be available during training -- which is typically expensive and is difficult to scale in real-world; (2) Expanding the context window usually incurs a quadratic increase in computational and memory cost for transformer-based LLMs. For example, the  training setup used in (Fu et al. https://arxiv.org/abs/2402.10171) shows that extending the Llama-2 7B model’s context window from 4k to 80k requires 8 A100 GPUs (80G each) for five days. The costs of resources and time further increase significantly for larger models, for longer context length, and for more extended training periods. (3) The test distribution is still expected to match to longer sequences seen during training -- which is not realistic. As a result, straightforward continual pre-raining or fine-tuning on longer sequences may still not solve the fundamental problem of learning to reason over (long) context in a robust and generalizable manner.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper presents a memory-augmented LLM architecture for improved temporal and multi-hop reasoning over long contexts. The key innovation is integrating an episodic memory module with temporal encoding and iterative query updates. While results show strong performance on BABILong compared to baselines (99% vs 36% accuracy for 1k context), there are significant concerns. The method requires explicit supervision for reasoning steps and predefined hop counts, limiting generalizability. The evaluation is restricted to synthetic datasets, and there's insufficient analysis of computational costs and scaling properties. The presentation lacks clarity, particularly in the methodology section and architectural diagrams. Though the core idea of memory-augmented reasoning is valuable, these limitations make the current work insufficient for acceptance.