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

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?

We have performed additional experiments to evaluate the performance of MemN2N[1]. The results are summarized in the following tables and demonstrate the lack of generalization ability of MemN2N compared to MemReasoner.

Table 1 (Performance on bAbi Tasks):

Table 2 (BABILong Task 1 Results):

Table 3 (BABILong Task 2 Results):

Table 4 (Performance on bAbi task 2 → BABILong task 1 generalization):

[1] Sukhbaatar, Sainbayar, Jason Weston, and Rob Fergus. "End-to-end memory networks." Advances in neural information processing systems 28 (2015).

We thank the reviewer for their constructive feedback and for recognizing our work as " a valuable improvement over standard LLMs". Please see our responses below:

Extending evaluations to diverse, natural datasets

We thank the reviewer for this specific comment. We would like to clarify why we evaluate MemReasoner on the BABIlong benchmark. First of all, the original bAbi tasks 1 and 2, though simple in nature, provide a good testbed for evaluating LM's logical reasoning abilities, which is even found to be challenging for more powerful GPT-3 model, even when combined with advanced prompting methods like CoTs. The tasks do comprise many basic elements of reasoning, such as order learning among facts and multi-hop reasoning. And, it should be noted that to date, most logical reasoning datasets are synthetic in nature.

Second, BABIlong provides a controlled testbed for testing reasoning generalization over long-context, where core reasoning facts (that are synthetic) are diluted with irrelevant naturally occurring text throughout the context of varying length. This benchmark is found to be challenging to current LLMs and RAG frameworks. With this construct, due to its synthetic nature and careful construction, the interference from LM's parametric knowledge can be kept at a minimum. This is not the case for many real-world natural language reasoning tasks, which are at risk of training data contamination during LLM pretraining phase. As a result, LLM may not use the non-parametric knowledge provided in the context while answering (as shown in Wu, et al. https://arxiv.org/abs/2402.11924v4.). Thus, BABIlong allows us to train the model on classical logical reasoning tasks such as babi and show evaluation on test distributions that are significantly different from the training samples, which is the main goal of this work.

A third reason to be considered is the presence of reasoning shortcuts in the natural language task samples itself, which does not ensure language models actually being required to perform multi-hop processing over the context, as shown for HotpotQA in Jiang, et al. 2019 https://arxiv.org/abs/1906.07132. This is not the case for BABIlong.

To show the generality of MemReasoner across different tasks, we have now added experiments on another task, namely variable tracking, from the RULER benchmark (see the global response section "additional datasets".). The new results on the variable tracking experiments also show MemReasoner outperforming other baselines such as RMT and Mamba.

Computational complexity

Please refer to global response section on "Computational cost of MemReasoner".

What improvements does MemReasoner have compared to Larimar?

We note that the Larimar architecture was designed for fact-editing in a single-hop setting. While MemReasoner changes the Larmar architecture and training procedure toward multi-hop reasoning, where learning the relative order of facts within an episode is crucial. In terms of architectural changes, MemReasoner introduces an iterative read with a query update procedure and a temporal encoding module into the Larimar architecture. The temporal encoding module allows us to capture the relative ordering of sentences within an episode, which allows us to perform reasoning on tasks where the ordering of facts within the context matters (for example, bAbi dataset). Additionally, the iterative read and the query update allow the model to perform multi-hop reasoning, with each read operation retrieving the supporting fact for each hop. In terms of training modifications, we introduce additional loss terms into MemReasoner (reconstruction of supporting facts and ordering loss) to ensure that the read operations retrieve the correct supporting fact.

RMT training with at least 2 segments

We thank the reviewer for pointing out the importance of training with multiple segments using a curriculum learning procedure. In order to train with more segments while exposing the model to only bAbi data, we reduce the segment size to 64 for task 1 and 128 for task 2. This leads to 2 segments in training for task 1 and 2-4 segments in training for task 2. In order to mimic the curriculum learning process described in the referenced work, we filter the data so that we train with inputs with token length up to the segment size for 10 epochs, up to 2 times the segment size for another 10 epochs, and so on. We present updated values for RMT below, which will be added to the revised manuscript:

BABILong Task 1:

BABILong Task 2:

Results for RMT-.77B on babi + location swap

RMT-.77 Task 2 -> Task 1 generalization

Table 1 and 2 miss results for RMT on 64k and 128k.

Thank you for pointing this out, we realized that in our evaluation we were storing intermediate states which was causing the out of memory error. Please refer to the tables above with training RMT with multiple segments as these tables have values for all BABILong lengths.

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?

Please refer to the global response section on "Finetuning on long data".

182 - 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.

We thank the reviewer for raising this point. We plan to add the following discussion in the paper to clarify this point as follows: MemReasoner follows the earlier works on Kanerva Machine (Wu et al., 2018a), which is inspired by Kanerva’s sparse distributed memory model (Kanerva, 1988), where the memory is viewed as a global latent variable in a generative model. In this framework, the goal is to learn a memory dependent data prior and learnable addresses, where the memory update and read/write are considered as Bayesian inference, i.e., the posterior parameters are updated as new data arrives. Later, Pham et al., 2021, reformulated the encoding of new memories and decoding data from memories from Bayesian updates to an equivalent minimization problem, which essentially amounts to solving a linear system of equations, efficiently done via computing matrix pseudo inverses. MemReasoner follows this approach of reformulation of the Bayesian update as finding the least-squares solution to a linear system for updating the memory parameters (see Algorithm 1).

Yan Wu, Greg Wayne, Alex Graves, and Timothy Lillicrap. "The Kanerva Machine: A Generative Distributed Memory." In International Conference on Learning Representations. 2018.

Pentti Kanerva. Sparse distributed memory. MIT press, 1988.

Kha Pham, Hung Le, Man Ngo, Truyen Tran, Bao Ho, and Svetha Venkatesh. "Generative pseudo-inverse memory." In International Conference on Learning Representations. 2022.

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.

In our experiments, the number of memory read operations corresponds to the number of hops. The reason for this is that in the MemReasoner training pipeline, we enforce that each read operation retrieves the relevant supporting fact for that hop via the reconstruction of supporting facts and ordering losses. To illustrate, consider the context ["Mary is in the kitchen.", "Mary got the apple there."] and the query "Where is the apple?". In training, we ensure that the first read retrieves the sentence "Mary got the apple there.", the readout is then used to update the query so that the query now has some encoded information about "Mary", and then we ensure that a second read with this updated query retrieves "Mary is in the kitchen.", which is then given to the decoder to generate the answer.

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?

First, to clarify, the readout $z_r$ also corresponds to a single fact in bAbI that was written to memory. The difference between $z_r$ and $z_i$ is that $z_i$ is the encoding prior to passing the lines through the GRU, so $z_r$ contains information about the relative ordering of the line in the context while $z_i$ does not. This relative order learning is important, as in the bAbi tasks the agents's location can be updated within one sample (i.e., one episode) See Figure 1, Right as an example, where Sandra's location is updated from bedroom to hallway. Therefore, it is crucial that the decoder's response is built on top of the most relevant (aka, recent) location of Sandra. Since the decoder only needs information about what is present in the most recent retrieved bAbi fact and not the relative order of it, we chose to pass $z_i$ to the decoder in the MemReasoner design. Passing $z_r$ to the decoder can cause the decoder to overfit to the ordering information in $z_r$ which hinders length generalization.

As noted by the reviewer, some reasoning tasks may require outputting information from multiple facts as opposed to just the final supporting fact. For example, for the variable tracking task (see global response "Additional Datasets"), the prompt asks for all variables with a certain value (with the number of variables in the answer corresponding to the number of hops) and each fact specifies only the value of a single variable. In order to handle this case, we can take all readouts $z_{r_1}...z_{r_h}$ where $h$ is the number of hops, map each readout to unordered encodings $z_{i_1}...z_{i_h}$ and then these are individually passed to the decoder and the outputs from the decoder are then aggregated as the final list of $h$ variables.

We thank the reviewer for their constructive feedback and for recognizing our idea and the direction of the work as "interesting", "clear and reasonable", and "strong and relevant". Please see our detailed responses below:

Lack of comparison with Transformer models

Please refer to global response section on "Comparison to decoder-only LMs that support long contexts". We have now aded Qwen2.5 baslines as requested by the reviewer, which further establishes that current LLMs with long-context training does not enable robust reasoning generalization.

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.

We thank the reviewer for this specific comment. We would like to clarify that although traditional multihop QA like HotpotQA, MusiQue, and 2WikiMultihopQA can be applied to evaluate the multi-step reasoning ability, existing studies suggest that those datasets often fail short while evaluating LM's logical reasoning abilities. One underlying reason is the interference from the parametric knowledge, likely caused by the training data contamination during LLM pretraining phase, as all of these benchmarks include data sourced from wikipedia. As a result, LLM may not use the non-parametric knowledge provided in the context while answering (as shown in Wu, et al. https://arxiv.org/abs/2402.11924v4. Another important factor to be considered is the presence of reasoning shortcuts in the task samples themselves, which does not ensure language models are actually being required to perform multi-hop processing over the context, as shown for HotpotQA in Jiang, et al. 2019 https://arxiv.org/abs/1906.07132. Since the goal of our study is to propose and evaluate language model architectures that do indeed perform more than one hop to solve a logical reasoning task in the absence of real-world knowledge, we avoid the potential data contamination issue affecting model's reasoning ability evaluation by training and testing on synthetic logical reasoning tasks. Further, we subject the model to an even more difficult scenario, where the core reasoning facts are diffused with distracting task-irrelevant natural text, the datasets suggested by the reviewer do not fit that framework. And to avoid the issue of model taking reasoning shortcuts instead of performing multiple hops, we focus on classical logical reasoning tasks such as babi and show evaluation on test distributions that are significantly different from the training samples.

However, we do provide additional experimental results on variable tracking tasks (both single-hop and two-hop) from the RULER benchmark, since we agree that it is important to show the generality of MemReasoner, when applied to other long-context tasks. Please refer to the global response section on "Additional datasets" for these results.

The paper does not provide empirical evaluations of the method’s inference time or memory consumption compared to other methods.

Please refer to the global response section on "Computational cost of MemReasoner".

Suggestions on improving presentation

We truly appreciate the reviewer's feedback on improving presentation of the MemReasoner framework by modifying Figure 2 and section 3.1. As suggested, we are revising the section 3 and the Figure 2 in manuscript by beginning with a conceptual description and then providing further details of the model and method. We are also adding a dedicated section on Larimar and explicitly elaborated on each loss term.

We thank the reviewer for their feedback. Please see our responses below:

The model does much worse than the models that are finetuned on long context data. We should see if MemReasoner also benefits from this?

We thank the reviewer for their feedback. We want to clarify the misunderstanding/misinterpretation of the statement "Their method solves the task as shorter contexts but does not seem to length generalize." We understand that the reviewer here is referring to RMT and Mamba baselines in Tables 1 and 2 of the paper, that are finetuned on BABIlong samples. However, the goal of this work is to propose a LLM architecture that robustly learns the core reasoning task and does not get confused when evaluated on long context test samples that contain naturally occuring text as distractors together with the core reasoning facts. Therefore, contaminating the training data of MemReasoner with BABIlong samples will defeat the purpose. That is why we train the proposed MemReasoner model on original Babi samples only and test them on BABIlong data, rather than training the model on BABIlong samples directly, as that will beat the purpose of testing model's reasoning ability on an unseen test distribution while the core reasoning task to be solved remains the same. Further, training on longer sequences is time and cost-intensive. Please refer to global response section on "Finetuning on long data". The results shown in Tables 2 and 3, combined with the new results on fine-tuned Qwen models, clearly establish this point that current transformer-based LLMs (including those that support long context) as well as alternative architectures like recurrent memory transformer (RMT) and Mamba, when finetuned on Babi, does not generalize to BABIlong. Whereas MemReasoner outperforms those baslines on both single and two-hop tasks. We further show MemReasoner's robustness to manipulations of answer locations in test samples for both single and two-hop tasks, as well as on generalzation of two-hop Babi-finetuned model on single-hop BABIlong samples. We have also now added a new task of variable tracking (both in single-hop and two-hop setting), where MemReasoner again outperforms RMT and Mamba baselines.

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

Please refer to global response section on "Computational cost of MemReasoner".

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

Please refer to global response section on "Comparison to decoder-only LMs that support long contexts".

Can MemReasoner augment existing decoder-only LLMs?

MemReasoner is a model-agnostic way to augment current decoder-only llms with a dynamically updatable memory. The MemReasoner architecture adds an encoder, a memory module, and a temporal encoding module on top of the decoder. The encodings from the encoder are used by the decoder via the past key values of the decoder layers. Via end-to-end training, the architecture learns to write the latent encodings in a fixed-size memory, order them in their order of apperance in the context, and perform multiple hop over that context and update the latent query accordingly. The decoder learns a diffrentiated attention mechanism to the final readout from the memory, in order to accurately generate the final answer and supporting facts (intermediate hops). This ability allows MemReasoner to perform robust multi-hop reasoning on unseen test samples.

Our provided experimental results in the main paper can be thought of as augmenting GPT2-L. To further demonstrate the general applicability of MemReasoner as a way to augment another existing decoder-only LLM of a different scale, we augmented GPT-J-6B with MemReasoner-like architecture and training, and showed the results in Table R1.

Table R1 (Performance on BABILong task 1 and 2 using MemReasoner-6B (bAbi) with GPT-J decoder):

Additional datasets [N6d6, PgX8]

We agree that it is important to show generalization on another dataset and provide experimental results for MemReasoner for the variable tracking (VT) from RULER [1]. In the VT task, the model is given context with lines with information about variable value assignment such as "VAR AAAAA = 16438" or "VAR BBBBB = AAAAA" and the model is prompted to obtain all variables with a specific value. Variable names have the format of 5 repeating letters randomly sampled from the alphabet. We train and evaluate with chains of length 2, 4, 6, 8 , and 10 and return the average accuracy over all chain lengths for the 1 hop and 2 hop VT tasks. In order to pad the context for lengths 1k, 4k, and 16k, we follow the approach taken from RULER of padding with the sentence "The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again.\n" until the context reaches the desired length. This noise is not present during training and the 0k data follow the same distribution as the training data.

Since VT asks for all variables with a specific value, for MemReasoner, we take all unordered readouts of the model and pass them individually to the decoder to get the variables from each reasoning hop, and then concatenate these variables in order to obtain the final answer. For RMT, we train with 2 segments, with segment size set to the median length on the train dataset. We observed that it is difficult to train RMT with 2 segments for the 2-hop VT task, RMT can easily learn a shortcut and have high accuracy on the training set, but does not generalize well to the test set at 0k length and performance degrades further at longer context length.

Single hop:

Two hop:

[1] Hsieh, C. P., Sun, S., Kriman, S., Acharya, S., Rekesh, D., Jia, F., ... & Ginsburg, B. (2024). RULER: What's the Real Context Size of Your Long-Context Language Models?. arXiv preprint arXiv:2404.06654.

Comparison to decoder-only LMs that support long contexts [vesZ, PgX8]

We experiment with Qwen2.5-0.5B and Qwen2.5-1.5B models both of which support long context windows (up to 128k tokens) and provide results that are to be compared with those in the respective Tables 1-5 in the manuscript. Overall, we find that MemReasoner is able to achieve better length generalization to long contexts compared to Qwen2.5. For location changes, we find that MemReasoner outperforms both Qwen2.5-0.5B and Qwen2.5-1.5B for Task 1, but for Task 2 Qwen2.5-1.5B outperforms MemReasoner.

Table 1 (Performance on bAbi tasks):

Table 2 (BABILong Task 1 Results):

Table 3 (BABILong Task 2 Results):

Table 4 (Robustness to location changes in bAbi test set):

Table 5 (Performance on bAbi task 2 → BABILong task 1 generalization):