M+: Extending MemoryLLM with Scalable Long-Term Memory

Claims And Evidence:

We would like to clarify that by “CPU offloading,” we specifically mean offloading the memory vectors present in each layer of the model. In our setup, each layer contains 12,800 memory vectors, and it is unnecessary to keep all of them simultaneously in GPU memory. Instead, we can store them on the CPU and load them into GPU memory only when the corresponding layer is being computed. It is important to note that other models such as LLaMA 3.1–8B do not have memory vectors, so our CPU offloading can only be applied on MemoryLLM and M+. We will add more clarifications.

Supplementary Material: We will publish the code and model upon acceptance and make sure of the reproducibility.

Essential References Not Discussed: Thank you for bringing these up, we will add these works ([1], [2], [3]) into our paper.

Comments, Suggestions and Questions:
Figure 3: We will revise the scales to ensure consistency across subplots.
Limitation Section:: We briefly mentioned our plan to reduce CPU-GPU communication overhead in future work, we will add a section to discuss this.

[Q1]

(1) Evaluations on LongBench: We would like to note that our evaluations on LongBench are included in Section 4.4.

(2) Experimental Results on NIAH: We would like to emphasize two key points:

• The scope of this paper goes beyond NIAH: Our primary goal is to capture and reason over global information, rather than focusing solely on retrieval. Thus we focus on on longbook question answering and introduce a new dataset, Longbook-Event-QA (Section 4.1.1), which requires the model to understand a broader part of the book. In contrast, NIAH primarily evaluates retrieval ability, regardless of whether the model comprehends the broader story.
• Our knowledge retention experiments share similarities with NIAH: In Section 4.5 (Figures 5 and 6), the model is tasked with recalling information from a context it encountered long ago and answering related questions. This setup resembles document-level retrieval and aligns conceptually with the goals of NIAH, suggesting that our work addresses similar challenges from a different perspective.

[Q2]

We add the following two rows to Table 1 to report the GPU memory consumption of MemoryLLM: MemoryLLM-8B, 21176.24MB; MemoryLLM-8B (offload): 17967.47MB. These results show that MemoryLLM-8B has comparable GPU memory usage to M+.

[Q3]

On 1,000 unseen examples from Fineweb-Edu (2048-token limit), M+ achieves a perplexity of 1.9828 vs. 1.9734 for LLaMA-3.1-8B, showing no degradation in base model performance.

[Q4]

Using BM25-based RAG (retrieving up to 4 chunks of 4,096 tokens), we observe limited gains: LLaMA-3.1-8B + BM25 achieves 0.1623 on LongbookQA and 0.2065 on LongBook-Event-QA, slightly improving over LLaMA-3.1-8B-16k (0.1514 / 0.2362) on one task but worse on the other. In contrast, M+ achieves the best scores: 0.1755 / 0.2470. This shows RAG offers no consistent benefit and may hurt performance when global context is needed.

[Q5]

Retrieval latency scales as $\text{latency} \propto d_r \cdot s \cdot L \propto d \cdot s \cdot L$, where $d_r$ is the retriever hidden size ($d/20$), $s$ is the memory size (maximum 150k, independent of model size), and $L$ is the number of layers. Since $s$ is constant, latency simplifies to $\text{latency} \propto d \cdot L$. Given $M \propto d \cdot L$, we have $\text{latency} \propto M$. As for the memory consumption, Memory vectors overhead also scales linearly with $d$ and $L$.

[Q6]

M+ and LLaMA-3.1-8B have similar FLOPs (e.g., 6.92e13 vs. 5.68e13 at 2k, 1.94e15 vs. 1.75e15 at 64k). At 128k, LLaMA-3.1-8B runs out of memory, while M+ runs property and has the total FLOPs as 3.78e15.

[Q7]

Our memory vectors can be viewed as hidden states within the transformer layers, with miner differences being that they may store more compressed information. Thus the information they capture should be similar to the representations seen in the intermediate layers of a transformer when processing text. Across layers, we hypothesize that the memory vectors follow a similar pattern to what has been observed in prior work on transformer interpretability [4, 5]:

• Lower layers tend to encode more surface-level features,
• Higher layers tend to encode more semantic or abstract information.

Regarding long-term memory, it is constructed by randomly dropped vectors from the short-term memory and storing them for extended use. Importantly, long-term memory vectors are structurally identical to short-term ones.

References:

[1] Recurrent Memory Transformer
[2] Scaling Transformer to 1M tokens and beyond with RMT
[3] MemoryPrompt: A Light Wrapper to Improve Context Tracking in Pre-trained Language Models
[4] How Many Layers and Why? An Analysis of the Model Depth in Transformers
[5] What does BERT learn about the structure of language?

We are sincerely grateful for Reviewer v8Qd’s recognition of our work. Below, we address the reviewer’s questions in detail:

**[Q1] Direct Comparison between M+ and MemoryLLM: ** Thank you for raising this important point. In developing M+, we incorporated a number of significant improvements over MemoryLLM-7B, including: (1) a multi-LoRA architecture, (2) refined dataset selection (Fineweb-Edu and SlimPajama), and (3) a carefully designed data curriculum. These enhancements led to substantial performance gains, and as a result, we consider M+ to be in a different performance tier compared to MemoryLLM-7B.

That said, to provide a fair and controlled comparison of the architectural differences—specifically the benefits introduced by M+ – we retrained the original MemoryLLM with: (1) The LLaMA3-8B backbone; (2) The same multi-LoRA design; (3) More advanced datasets (Fineweb-Edu and SlimPajama). The retrained versions are included in our ablation study in Section 4.5, where we have two variants:

MemoryLLM-8B: We switch the backbone to Llama-3.1-8B and use multi-LoRA design and Fineweb-Edu dataset to continually train the model. This is exactly Stage 1 mentioned in Section 3.2.4.

MemoryLLM-8B-Long: After obtaining MemoryLLM-8B, we apply Stage 2 with long documents extracted from SlimPajama to enhance the model’s long document understanding abilities.

Although this section is titled “Ablation Study,” it effectively serves as a direct comparison between the MemoryLLM architecture and M+, isolating the effect of long-term memory while controlling for backbone model and training data. The comparison includes the following aspects:

1. Perplexity (Figure 5): M+ demonstrates lower perplexity than both MemoryLLM-8B and MemoryLLM-8B-Long. The performance gap between M+ and MemoryLLM-8B-Long is smaller, reflecting the benefits of the long-input training in both models.

2. Knowledge Retention (Figure 6): M+ shows a significant advantage over MemoryLLM-8B-Long in retention tasks. Notably, we observed that MemoryLLM-8B and MemoryLLM-8B-Long perform similarly on this task, underscoring the importance of the architectural changes in M+.

3. Performance on Relatively Short Documents: To further address your question, we have conducted additional experiments on shorter documents (8k tokens), and we will include these results in the paper. The results are as follows:

   Model	2wikimqa	hotpotqa	qasper	musique	multifieldqa_en	narrativeqa	Avg
   MemoryLLM-8B (8k)	32.30	33.39	23.88	12.37	35.91	21.46	26.55
   MemoryLLM-8B-Long (8k)	32.23	37.86	31.62	20.35	42.16	23.49	31.29
   M+ (8k)	33.12	37.99	29.91	20.68	40.11	24.18	31.00

These results show that while M+ and MemoryLLM-8B-Long perform similarly on shorter documents (as expected), M+ provides significant gains on long-context tasks (as shown in Figure 6). This further supports our claim that the long-term memory architecture in M+ is beneficial primarily for extended contexts. Additionally, the improved performance of MemoryLLM-8B-Long over MemoryLLM-8B on short documents can be attributed to the inclusion of longer training examples in Stage 2 (4k–64k), whereas Fineweb-Edu (used in Stage 1) contains very few examples longer than 4k.

In summary, Section 4.5 provides a direct and controlled comparison between M+ and MemoryLLM, where the primary architectural difference is the inclusion of long-term memory. We will explicitly clarify this point in the paper and include the new short-document comparison results for completeness.

We sincerely thank reviewer f5Uz for their recognition of the value of our work. We address the reviewer's concerns below:

Relation To Broader Scientific Literature:

To the best of our knowledge, following MemoryLLM, the most recent works on parametric memory include Titans [3] and Memory at Scale [4], which scale models to 760M and 1.3B parameters respectively. However, these efforts remain exploratory and have not yet been evaluated on real-world tasks such as long-context question answering, which is the focus of our paper. As a result, there are currently few parametric memory approaches that offer a meaningful comparison to our work.

As for the questions:

[Q1] Evaluation Metric in LongBench:

Thank you for pointing this out. Following LongBench [1], we use the QA-F1 Score as the evaluation metric for all six benchmarks included in our paper. We will add this clarification in the revised version of our manuscript.

[Q2] Related Work and Comparison with [2]:

We appreciate you highlighting [2]—Assessing Episodic Memory in LLMs with Sequence Order Recall Tasks. It is indeed a relevant and important contribution. We note that it is conceptually similar to our newly proposed Longbook-Event-QA task (Section 4.1.1). While [2] evaluates models by asking them to sort two segments after reading an entire book, Longbook-Event-QA requires selecting the next event given a sequence of events. Both tasks emphasize the model’s comprehension of the narrative and its ability to reason about event order. Given their shared focus, we anticipate consistent results between Longbook-Event-QA and [2], and we will consider incorporating [2] as an extended benchmark in future work.

References:

[1] LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding.
[2] Assessing Episodic Memory in LLMs with Sequence Order Recall Tasks.
[3] Titans: Learning to Memorize at Test Time.
[4] Memory Layers at Scale.

Essential References Not Discussed:

Thank you for highlighting these important works. We will incorporate [1], [2], and [3] into our related work section. Specifically, [1] aligns with the core motivation behind incorporating memory into language models. Both [2] and [3] explore architectural modifications to enable memory mechanisms in transformers. However, these approaches remain exploratory, as they have not been scaled beyond small models or applied to real-world tasks such as long-context question answering. In contrast, M+ is implemented at the 8B parameter scale and is designed to be scalable with additional GPU resources. We will update our related work section accordingly and include a discussion of these points in the paper.

Weaknesses:

[W1] Similarity to Attention-Based Retrieval Methods: We acknowledge that our method shares some similarities with prior approaches that use attention to retrieve keys and values. However, there are critical differences that make our approach unique and practically advantageous:

(1) Efficiency: Methods such as SnapKV maintain and retrieve key-value pairs per head, which becomes extremely costly when scaled. In our setting—with 32 layers and 32 attention heads per layer—this requires 1024 retrievals per query, resulting in significant latency (as noted in line 59 of our paper). In contrast, M+ uses a co-trained retriever to retrieve memory tokens, which are compressed hidden states. This results in only 32 total retrievals—one per layer—dramatically reducing both computational cost and latency.

(2) Performance: In Figure 6, the curve labeled MemoryLLM-8B-Attn follows the SnapKV-style approach of retrieving key-value pairs using attention per head. As shown in the figure, it performs substantially worse than M+, highlighting that our co-trained retriever not only improves efficiency but also yields better results in practice compared with attention-based retrievals.

(3) Design: Note that our training setup includes both relevant and irrelevant documents (See details in Appendix D), making it well-suited for contrastive learning. This allows us to effectively train the retriever, which integrates naturally into our overall training framework.

[W2] Representation of Long-Term Memory (Hidden States vs. KV):

We appreciate the reviewer’s insightful comment regarding the form of long-term memory. We advocate for the use of hidden states over key-value (KV) caches based on two key considerations:

Compression Efficiency: As detailed in the paper, we compress each 512-token chunk into 256 memory vectors per layer in a lossless manner. In contrast, KV-based methods often require downsampling—e.g., dropping half the keys and values—to control memory size, resulting in unavoidable information loss.

Retrieval Efficiency and Performance: As described in [W1], hidden states can be effectively retrieved using our co-trained retriever, requiring only 32 retrievals for each query. In contrast, a KV-cache approach would demand up to 1024 retrievals, significantly increasing computational cost. Furthermore, as shown in Figure 6, using hidden states yields better performance compared to using KV caches.

We believe these benefits make hidden states a more efficient and effective choice for long-term memory representation in our system.

[1] Hybrid computing using a neural network with dynamic external memory.
[2] Titans: Learning to Memorize at Test Time.
[3] Scaling Transformer to 1M tokens and beyond with RMT.