M+: Extending MemoryLLM with Scalable Long-Term Memory
Authors: Yu Wang
Dmitry Krotov
Yuanzhe Hu
Yifan Gao
Wangchunshu Zhou
Julian McAuley
Dan Gutfreund
Rogerio Feris
Zexue He
Venue: ICML 2025 poster (2025)
Keywords: memory
long-term memory
long context
Primary area: deep_learning->foundation_models
TL;DR: Equipping a long-term memory on MemoryLLM

Abstract: Equipping large language models (LLMs) with latent-space memory has attracted increasing attention as they can extend the context window of existing language models. However, retaining information from the distant past remains a challenge. For example, MemoryLLM (Wang et al., 2024a), as a representative work with latent-space memory, compresses past information into hidden states across all layers, forming a memory pool of 1B parameters. While effective for sequence lengths up to 16k tokens, it struggles to retain knowledge beyond 20k tokens. In this work, we address this limitation by introducing M+, a memory-augmented model based on MemoryLLM that significantly enhances long-term information retention. M+ integrates a long-term memory mechanism with a co-trained retriever, dynamically retrieving relevant information during text generation. We evaluate M+ on diverse benchmarks, including long-context understanding and knowledge retention tasks. Experimental results show that M+ significantly outperforms MemoryLLM and recent strong baselines, extending knowledge retention from under 20k to over 160k tokens with similar GPU memory overhead.

M+: Extending MemoryLLM with Scalable Long- Long-Term Memory
Yu Wang' Dmitry Krotov Yuanzhe Hu'
Yifan Gao
Wangchunshu Zhou Zexue He 123
Julian McAuley
Dan Gutfreund 23
Rogerio Feris? 23
Abstract
abling direct retrieval and manipulation of information at the token level; and (2) Latent-space memory, where memory is stored as high-dimensional vectors in the hidden space, offering a more abstract and compact representation of in- formation. Token-level memory provides adaptability (the base model can be easily replaced) and interpretability (text- based format is easy to understand for humans). However, such bext- memory could be redundant as text format may not be the most compressed method for representing information (Bellard, 2021; Belcak & Wattenhofer, 2024; Rahman et al., 2024), and resolving conflicting informa- tion in text-based memory can be challenging (Pham et al., 2024). Meanwhile, as noted by Fedorenko al. (2024); Hao al. (2024), human reasoning often transcends the token level, leveraging deeper, integrated representations akin to latent spaces.
Equipping large language models (LLMs) with latent-space memory has attracted increasing at- tention they can extend the context window of existing language models. However, retain- ing information from the distant past remains a challenge. For example, MemoryLLM (Wang et al., 2024a), as a representative work with latent- space memory, compresses past information into hidden states across all layers, forming mem- ory pool of IB parameters. While effective for sequence lengths up to 16k tokens, it struggles to retain knowledge beyond 20k tokens. In this work, we address this limitation by introducing M+, a memory- auggented model based on Mem- LryLLM that significantly enhances long-tenn terg-termmm in- formation retention. M+ integrates a long-term memory mechanism with co-trained retriever, dynamically retrieving relevant information dur- ing text generation. We evaluate M+ on diverse benchmarks, including long-context understand- ing and knowledge retention tasks. Experimental results show that M+ significantly outperforms MemoryLLM and recent strong baselines, extend- ing knowledge retention from under 20k to over 160k tokens with similar GPU memory overhead. We open- source our code at https: //tp::::://// / thub. com/ haag/wangy--ooe wwwpaangyuustce///oeee Maattt//MMMMLMMMM.
In contrast, Latent-Space Space Memory offers unique advantages: (1) Efficient Compression: Information is compressed into hidden states (Wang et al., 2024a), internalized into model parameters (Wang et al., 2024c), or stored in a more com- pact latent space (Das al., 2024). These methods reduce storage overhead, with some approaches even embedding knowledge directly into model parameters, eliminating the need for external storage (Wang et al., 2024c). (2) End- to-End Training: Latent-space memory can be involved in gradient-based optimization, allowing it to be updated and refined during training. This enables the integration of memory into the training loop (Yin al., 2024; Wang et al., 2023; Ge et al., 2024). (3) Similarity to Human Memory: As suggested by Fedorenko et al. (2024) and Hao et al. (2024), human reasoning relies on integrated repre- sentations beyond discrete tokens, akin to latent spaces. By encoding knowledge in latent representations, the methods with latent-space memory can more closely mimic the mech- anisms of human memory, which store information within neural activations.
1. Introduction
The integration of memory modules into large language models (LL.Ms) has gained increasing attention (Wang et al.,
2024b). Existing approaches for constructing memory mod- ules can be broadly divided into two main categories: (1) Token-level memory (Packer al., 2023; Modarressi al., 2024), where memory is represented as structured text,
In this paper, we focus Latent-Space Memory. Memo- ryLLM (Wang et al., 2024a), as representative work in this category, enhances transformer- language model by incorporating a large number of memory tokens into each layer, creating a memory pool with billion parameters. This framework introduces a carefully designed update and generate process, achieving superior performance compared
Work done during the internship at MIT-IBM Waston Lab. 'UC. San Diego "MIT-IBM Waston Lab 'IBM Research 'Amazon *OPPO. Correspondence to: Yu Wang <yuw164@ udu..uuuss Zexue He Micxue.H@@ibm.coom
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

M+: Extending McmoryLLM with Scalable Long-Term Memory
to the backbone model Llama-2-7B and other long-context methods. However, MemoryI.LM faces limitations in recall- ing information injected beyond 20k tokens, restricting its long-term retention capabilities. To address this limitation, we propose M+, a novel model incorporating a long-term memory mechanism alongside MLM. Unlike pre- vious approaches such as H2O (Zhang et al., 2023) and SnapKV (Li et al., 2024). which store keys and values from past contexts and perform retrieval separately for each query head and layer-leading to high latency-M+ optimizes retrieval in the space of hidden states via co-training the re- triever and the language model. This allows M+ to retrieve only once per layer for all query heads, significantly improv- ing efficiency. Furthermore, as the long-term memory is stored on the CPU, M+ significantly extends long-term re- tention capabilities without increasing GPU memory usage.
poses treating context and memory as analogous tradi- tional memory in operating systems, enabling more flexible and organized memory structures. These approaches typi- cally rely on text embeddings for memory retrieval, where queries can originate from either the current conversation turn (Zhong et al., 2023; Zhou et al., 2023) or queries gener- ated by the language model itself (Packer et al., 2023). In contrast, ChatDB (Hu et al., 2023) stores knowledge in a database and performs retrieval using SQL queries, while MemLLMMModdareesi et al., 2024) fine-tunes the model to generate function calls that initiate searches within knowl- edge graph, referred to as "Triple Memory" by Modarressi et al. (2024). These methods generally offer benefits such as modularity (with the exception of MemLLM, which re- quires fine-tuning) and interpretability (Yin et al., 2024), allowing for potential integration with external systems (Wu et al., 2022a). However, these approaches have limitations. Some require storing the raw text, which is not the most compressed method to store information (Rahman et al., 2024; Bellard, 2021; Belcak & Wattenhofer, 2024). Oth- ers store knowledge in the form of triplets, which may be unsuitable for representing complex conversations that are difficult to convert into triplets (Wang al., 2024d).
We evaluate M+ across a diverse set of benchmarks, in- cluding tasks such as long-book understanding, knowledge retention, and question answering on relatively short docu- ments. Experimental results demonstrate that M+ achieves significant performance improvements in all long bench- marks compared to previous memory- -baee methods while operating within the same or even smaller inference memory budget. In summary, our contributions are as follows:
2.2. Latent-Space Memory
We enhance MemoryLLM by incorporating a long- term memory mechanism and introducing co-trained retriever for efficient and effective memory retrieval.
Latent-space memory stores information in a compressed format, embedding knowledge into soft prompts (Rako- tonirina & Baroni, 2024). hidden states (Khandelwal et al., 2019; Bulatov et al., 2022: 2023; Wang et al., 2024a), model parameters (Wang al., 2024c), or an external latent space (Das et al., 2024), among other methods. Some ap- proaches use memory slots to encode information (AI Adel & Burtsev, 2021), while others rely on key-value caches stored in memory pools for future retrieval (Wu et al., 2022b; Wang et al., 2023; He et al., 2024: Park & Bak, 2024). Notably, CamelLoT (He et al., 2024) and Memoria (Park & Bak, 2024) incorporate forgetting mechanisms to better emulate human memory. Similarly, MemoryLLM (Wang et al., 2024a) compresses knowledge into hidden states and employs random dropping to prevent unbounded memory growth. The M3 method (Yang et al., 2024) also stores memory in the hidden-state space, archiving a vast pre- training dataset comprising 1.1 x 1 text chunks. Dis- tinct from methods that utilize hidden states or key-value caches, Larimar (Das et al., 2024) introduces a memory matrix that supports read and write operations, demonstrat- ing effectiveness in mhowledge-editingg tasks. Furthermore, SEL.--PARAM (Wang et al., 2024c) explores embedding knowledge directly into model parameters without degrad- ing the model's capabilities or requiring additional param- eters. These latent-space memory techniques have shown promising results across various downstream tasks. By saving information in compressed format and leveraging retrieval during generation, they enable substantial expan-
We design specialized data curriculum for long- context training, enhancing the long-context modeling ability of M+.
Through extensive experiments on multiple bench- marks, we demonstrate that M+ significantly outper- forms the baselines while maintaining a similar or re- duced GPU memory footprint.
2. Related Work
We classify memory-based methods into two categories: Token-Level Memory and Latent-Space Memory, which is similar to the categorizations in Yin al. (2024) where they classify methods into implicit memory and explicit memory.
2.1. Token-level Memory
Teken-level memory refers to memory structures repre- sented in textual forms, which can include raw context, summaries (Zhong et al., 2023; Zhou et al., 2023). knowl- edge graphs (Packer et al., 2023: Gutiérrez et al., 2024), organized text with hierarchical or graph structures (Packer et al., 2023; Chen et al., 2024), or databases (Hu et al., 2023). Methods such as MemoryBank (Zhong et al., 2023), RecurrentGPT (Zhou et al., 2023) incorporate multiple components of memory, including both raw conversational data and summaries. MemGPT (Packer et al., 2023) pro-
2

---

M+: Extending McmoryLLM with Scalable Long-Term Memory
New K tokens
New tokens
Transformer Layer
Transformer Layer Φt
K tokens
Hidden states from the new chunk
K tokens
Hidden states from the new chunk
Memory Pool 0 (N tokens)
Memory Pool 0 (N tokens)
K' tokens (dropped)
(N-K) tokens (remaining)
New K tokens
K tokens
(N-K) tokens (remaining)
New tokens
(dropped)
Long-Term Memory 9 (M tokens)
New Memory Pool 8 (N tokens)
Update
New Memory Pool Bi (N takens)
Query Hidden States (Next Layer)
Generation
Query Hidden States (Next Layer)
Transformer Layer Φ
Extracted LTM (Kg tokens)
Memory Poal 9, (N tokens)
Query Hidden States
Transformer Layer Φ1
Memory Pool (N tokens)
Query Hidden States
Long-Term Memory 0 (M tokens)
(Extracted with the Co- ooairiined Retriever)
MemoryLLM
M+
Figure The left side shows the Update Generation Process of MemoryLLM (Wang et al., 2024a). process chunk with de to obtain new K tokens during the update process, which is perceived by ø using cross-attention during the generation process. The right side shows the Update and Generation Process of M+. For layer during Update, the old memory pool 0; is split into two parts: K dropped tokens and N K remaining tokens. The dropped tokens stored in the long- long-term memory θ while the remaining tokens and new tokens are combined to obtain the new memory pool 8/. Then during generation, we use our co-tra co-trained retriever to retrieve tokens from O, which is fed into the transformer layer along with the short-term memory O and the query hidden states. The major difference between MemoryLLM and M+ is the introduction of Long- Term Memory O,.
sions of the context window without incurring excessive GPU memory costs. Despite the ad advantages and potential of Latent-Space Memory, existing methods within this cate- gory typically fall short when dealing with extremely long input (Das al., 2024; Wang et al., 2024c;a; et al., 2024). In contrast, M+ can have much longer retention compared to existing methods.
Figure 1). Merging is achieved by randomly dropping K tokens from 9 and appending the new К tokens to the end. During generation, the memory pool 0₁ is perceived using cross-attentic
3.2. Equipping MemoryLLM with Long-Term Memory
In this section, explain how we instantiate the long-term memory and how it integrates with the language model and the original memory pool in MemoryLLM. In this paper, we term the original memory pool as short-term memory to distinguish from the new long-term memory.
3. Methodology
3.1. Preliminaries
We introduce the structure of MemoryLLM (Wang et al., 2024a), which serves the base structure of M+. Memo- ryLLM comprises two main components: (the memory pool) and & (a transformer-based decoder-only language model). The memory pool 0 consists of layers: {0:) where the number layers the transformer Φ). For every layer, Ot has N memory tokens, where each token is vector in Rdd with d representing the hidden size of the language model. During the update process, the last K tokens from the /-th layer's 's memory pool, 0, are extracted and combined with the chunk to be injected. The resulting new K tokens are then merged back into 6 (illustrated in
3.2.1. MEMORY STRUCTURES
We denote the long-term memory as θ. Similarly, has layers [0,) Each layer has long-term memory pool where the size flexible. We specify a maximum size for the long-term memory. The maximum size of the long-term memory is denoted as M and the size of long term memory is flexible. In practice, we choose M to be 150k. Then we introduce the update and generate process of M+:
Update Process During the update process, note that in the original MemoryLLM, K tokens are dropped from during updates and are permanently discarded. In M+, the
3

---

M+: Extending Mc with Scalable Long-Term Memory
dropped K tokens are instead stored in the long-term mem- ory O ensuring their retention for extended durations (as illustrated in Figure 1). We assign each token the variable "age" so that after retrieving tokens from @ we can sort these tokens according to age, ensuring that the tokens are chronologically ordered. As for the new K tokens, they are obtained with the same process as in L.M, described in Figure 1. When the memory tokens in the long- term memory reach the maximum capacity, i.e., M tokens, we would drop the tokens with the largest ages.
are denoted as 0_. After that, we run a forward pass on 2n to obtain the hidden states hn (Note that this is general notation for all layers). For the hidden states hn in each layer, we train the retriever using the following objective:
min - log(p+)) - log(1 - p-); fe.Je
where p+ = fa
p = (l(()))) fx<(0-));
ie., we are maximizing the distance between hn and f_ while minimizing the distance between I. and 0+ after applying fg and fa on hn and 0 respectively.
Generation Process During generation, at each layer, we extract Kg tokens from the long-term memory O using a retrieval mechanism described below, sort them by their ages, and concatenate them with the short-term memory 0,. This allows the query hidden states to access both the extracted long-term memory and the short-term memory using cross-attention, enabling the query to retrieve relevant information from the memory.
3.2.3, TRAINING DETAILS
Setting Configurations We build M+ on top of Llama- 3.1-8B (Dubey al., 2024) and train it using eight A100 GPUs. We tried FSDP, deepspeed and deepspeed-stage-3 and we finally choose seepspedd-stageee 2 due to resource limitation and li- brary incompatibility (see the details in Appendix A). Specif- ically, we set = 256, N = 10240 (N is the number of tokens in the short-term memory, see Section 3.1). and the number of tokens of extracted LTM in Figure 1 is set to 2,560. The generation window (i.e., the maximum length of generation) is set to be 2,048. Thus maximally we have the attention matrix of shape (12, 800+2, 048) by 2, 048, which is fit into eight A100 GPUs using peresspeed-stagee Although Llama-3.1-88 can practically handle 128k con- text window, it went through much more extensive training that we cannot afford. Should we have more GPUs and more budget, M+ could also be scaled to 128k level. Given the constraint of GPU resources, we have scaled to 12,800 memory tokens and 2,048 generation context window, re- lying on 1amma-3.1-BB ' capability on a context window of 12, 800 + 2, 048 = 14, 848 tokens. Thus, for fair com- parisons, we mainly focus on Llama-3.1-8B-16k as the baseline.
2,
Multi-LoRA Design In our training, we use two sets of LoRA weights, one is activated during the update process, and the other is activated during the generation process (as shown in Figure 1). Intuitively, the update process com- presses the information (similar to writing) while the gener- ating process loads the information (similar to reading), thus having two LoRA weights could potentially make learning easier for our model. This is similar to the intuition in T5 where they find sharing the weights of encoder and decoder leads to slightly inferior performances (See Table in Raffel et al. (2020)).
3.2.2. RETRIEVER DESIGN AND TRAINING
Retriever Design The retriever has two projectors: query projector fg and key projector ft, which are all instantiated with a two-layer perceptron. The output dimension of both projectors, denoted as dproj. is set to be a small number. In our experiments, we set dproj to be d/20 where d is the hidden size of the language model Φ. When dropping tokens from Og into O (as shown in Figure 1), we apply fk on top of the dropped memory tokens, thus we need an additional pool storing all the key vectors corresponding to the memory tokens in oj Note that the key vectors are the output from Ss. and are of dimension dproj. requiring little additional memory footprint compared to the long- term memory. During generation, given the hidden states from the query, we apply fq on the hidden states to get query vectors and use them to retrieve tokens from O according to the dot product between query vectors and key vectors. Training the Retriever To train the retriever, we first split document into chunks F2, , II, and we inject (x1, , 1 into the short-term memory. Then we can track the embeddings in the short-term memory that are related to 31 - which we denote as 0... Then the memory tokens that are not related to "1,'"' 1nn-1 i.c., the tokens that were there before injecting "1,"" In-11,
3.2.4. DATA CURRICULUM
The training process consists of three stages:
Continual Training of MemoryLLM (Stage 1) Differ- ent from (Wang et al., 2024a) which starts from the back- bone model Llama-2-7B. we start with the backbone model Llama-3.1 1-BBB which serves & as shown in Figure 1. We equip & with N = 12, 800 memory tokens in each layer and set the generation context window as 2,048. We first continually train & equipped with on the dataset fi fineweb-edi (Penedo et al., 2024) for 1,200,000 steps over four weeks, establishing a strong foundation for han- dling short documents. This training stage consists of three key sub-tasks outlined in MemoryLLM (Wang et al., 2024a) (see details in Appendix D)).
Long-Context Modeling with Long Documents (Stage 2) Since most of the fineweb -ddu dataset (used in Stage
4

---

M+: Extending McmoryLLM with Scalable Long-Term Memory
Llama-3.2-38-1286
L'ama-3.1-80 3.1-88-SnapKV
et al., 2024) and consists of 351 tuples in the format (book, question, answer). Each book has an average input length of 192k tokens. The task requires answering ques- tions based on the entire book, and we use the QA-FI score as the evaluation metric.
Llama-3.1-80-166
Libma-311-88-MM55
0.180
0.25
0.175
0.24
QA-11 (Longbook-OM QA)
0.170
H
0.23
LongBook Event QA: We propose this new benchmark to evaluate the model's ability to recall past events and reason chronologically. This dataset is constructed as follows: (I) We use the Named Entity Recognition (NER) tool from SpaCy to identify the ten most frequently mentioned char- acters in each of the first five books from the LongBook- QA dataset. (2) Each book is divided into 4096-token chunks in chronological order, and events experienced by the main characters are extracted using gpt-4o. This results in event lists with 1,016, 221, 644, 348, and 409 events for the five books, respectively. (3) For each event, we con- struct multi-choice question-answering task by prompting gpt-4o to generate five fake events as distractors. The model is provided with the book text, five past events, and asked to identify the ground-truth event from six options. We use accuracy as the evaluation metric.
0.165
0.22
0.21
0.160
0.20
0.155
0.19
0.150
0.18
0.145
Longbook
boog Event- QA
Figure 2. Overall Performance Comparison Longbook Question Answering. Best viewed in colors.
1 training) are short documents under 4k tokens, we need to train on longer documents to enhance the model's long- context modeling abilities. Thus, we extract documents from SlimPajama that range from 4k to 64k tokens and split them into four categories based on their lengths: 4k-8k, 8k-16k, 1 k2k, 32k-64k 64k. The statistics of obtained dataset is shown in Appendix C. For each length range, we randomly sample 200,000 examples, and they are com- bined with snapshot of f fineweb in equal proportions (1:1:1:1:1), with each subset contributing to 20% of the total data. We set this proportion to upsample longer documents, which is important for long context modeling as suggested by Fu et al. (2024). Training runs for one epoch with around one week using the same training tasks in Stage 1.
We compare M+ against the following baselines: (1) Llama- 3.1 1---B-166- The original model with context window fixed as 16k. (2) Llama-3.1 1atiaa...-------appp 1-8B-Sn We processes a 32k token input and dynamically selects 16k key-value caches from the saved 32k caches with the techniques introduced from SnapKV (Li et al., 2024). SnapKV incurs significant memory overhead, as illustrated in Section 4.2. (3) Llama- 3.1-3B-128k: A 3B parameter version of the Llama3 series. This model uses a 128k context window and consumes comparable GPU memory to M+ because of its smaller size.
Training with long-term memory (Stage 3) Building Stage 2, we introduce long-term memory to enhance M+. Note that in Stage 1 and Stage 2, there is only the short- term memory 0 where each layer O: has 12,800 tokens. In stage 3, we adjust the configuration by setting O to 10,240 tokens and retrieving Ko = 2, 560 tokens from the long- term memory, maintaining a total of 12,800 memory tokens as in the previous stages. Now the structure of the memory tokens becomes slightly different, as 2,560 tokens are from the long-term memory, we design Stage 3 to ensure the model 6 understand the tokens from long-term memory - we continuously train from the checkpoint obtained after Stage 2 on a newly constructed dataset sampled from the same long documents extracted from SlimPajama but distinct from the instances used in Stage 2.
(1) 1lmma-3.1-3--18.. A 3B parameter version of the Llama3 series. This model uses a 128k context window and consumes comparable GPU memory to M+ because of its smaller size. (2) Llama-3.1-8B-16k: The original model with context window fixed 16k. (3) Llama-3.1- We use BM25 as the retriever. Specifically, we process the entire book by dividing it into chunks of 4,096 tokens and retrieve four relevant chunks for each question. The model Llama-3.1 is required to answer the question with four retrieved chunks. (4) Llama-3. ----a---------apphte processes a 32k token input and dynamically selects 16k key- value caches from the saved 32k caches with the techniques introduced from SnapKV (Li et al., 2024). SnapKV incurs significant memory overhead, as illustrated in Section 4.2.
We
4. Experiments
We present the primary results in the main paper, while defer- ring extended discussions and supplementary experiments to the appendix. These include: similaritics to attention- based retrieval methods (Appendix 1.1 the structure of long-term memory (Appendix E.2); latency and memory consumption when scaling (Appendix E.3): FLOPs com- parison (Appendix E.4); and the interpretability of memory
4.1. Long Book QA and Event QA
1. EXPERIMENTAL SETTINGS
We evaluate our model on two datasets designed to test long- context understanding and long-term memory capabilities:
LongBook-QAr This dataset is part of oo-Bench (Zhang
5

---

M+: Extending McmoryLLM with Scalable Long-Term Memory
vectors (Appendix E.5).
Table I. GPU Memory Cost Comparison. MemoryI...M-B is the
model obtained after Stage I training, serving as an ablation study.
4.1.2. EXPERIMENTAL RESULTS
Method
GPU Memory Cost (MB)
The results for both benchmarks are shown in Figure 2. From the results, we observe the following: (1) M+ consis- tently outperforms all baselines, demonstrating its efficiency and effectiveness. In LongBook-QAA M+ achieves superior QA-FI scores even while processing the least number of tokens (12,800 tokens in memory and 2,048 tokens in the generation context window). Similarly, in LongBook Event QA, M+ excels at identifying ground-truth events, show- casing its ability to reason over long-term dependencies. (2) Compared to -Jama-3.1-3B--288, the results on dataset Longbook-Event-OK suggest that tasks requiring reasoning capabiliiiiss benefit more from larger models with tailored structures for extended context windows rather than smaller models with longer context capacities. This highlights the importance of balancing model size and memory mecha- nisms under fixed GPU memory budgets. (3) Llama-3.1-80 V underperforms Llama-3.188B-166 on LongBook- QA, indicating that solely relying on attention scores to select key tokens may not consistently yield optimal results. In contrast, M+ leverages a jointly trained retriever to iden- tify and extract memory tokens, resulting in more effective performance on both datasets. (4) M+ outperforms Llama- 3.1 3.1-8B with BM25 retriever, and Llama-3. 1-8B with BM25 does not consistently outperform the original model Llama- 3. 16k, particularly in tasks like Longbook-Eat -venkk-eent-g A which requires a more global understanding of the entire narrative. This highlights the limitations of chunk-level re- trieval in scenarios that demand long-range comprehension. (5) Memory Efficiency: While M+ achieves sfate-o--the-ard results, it does so with highly memory-efficient design. Detailed memory cost comparisons are provided in Section 4.2.
Llama-3.1-8B-Snaal
32574.49
Llama-3.2-3B-128k
30422.70
M+
21177.76
Llama-3.1-8B-16
19239.21
M+ (offload)
17973.34
MemoryLLMMSBD
21176.24
MemoryLLM-SB (offload)
17967.47
out affecting the performance. This leads to "M+ (offload)" which achieves the least GPU memory consumption. We also include the GPU cost of MemoryL.LM-8B, which is the model obtained after Stage described in Section 3.2.4. This shows that the Long- -erm Memory does not incur more GPU costs.
4.3. Knowledge Retention Experiments
1.3.1. EXPERIMENTAL SETTINGS
To evaluate the ability of M+ to recall long-term knowledge, we follow the experimental setup in MemoryLLM (Wang et al., 2024a) on datasets SQuAD and NaturalQA, for- matted as (context, question, answer), where context and question are sentences, and answer is the correct response to the question. Consistent with Wang et al. (2024a), we extract samples with answer lengths of three tokens or fewer for SQuAD and four tokens or fewer for NaturalQA. After filtering out ambiguous examples that gpt-4o-min minni fails answer, we select the first 100 ex- amples from the remaining answerable set to conduct our evaluation. To test the model's long-term retention abil- ity, we insert distracting contexts between cont ext and quest i on. These distracting contexts are sampled from the training set of SQuAD. Both NaturalQA and SQuAD are constructed from Wikipedia and they are within the same domain. Moreover, the contexts in SQuAD training set are of more consistent lengths (around 300-500 tokens for each context), thus sample the distracting contexts from SQuAD training set for both NaturalQA and SQuAD.
4.2. GPU Cost Comparison
In this section, we report the maximum GPU memory al- located during the inference across both datasets for each method mentioned in Section 4.1. The results are shown in Table From the results, we can find that M+ has the low- est GPU memory cost except for Llama-3. 1.-..3..-BB--- 1-8B-166. The reason that M+ uses fewer tokens but costs more GPU is that we have 12,800 tokens in each layer, while 1.ama-3.1--8 16k has only one layer of 16k tokens. Therefore, we propose to offload the memory tokens on CPU, and reload them into GPU when the calculation reaches a certain layer. By "CPU offloading", we specifically mean offloading the memory vectors present in each layer of the model. It is important to note that other models, such Llama-3.1-8B do not have memory vectors, so our CPU offloading can only be applied to MeMoryLLM and M+. With that, we can sacrifice some 1/O time cost but substantially decrease the GPU cost with-
We compare with the following baselines: MemoryLLM- 7B: The proposed model in Wang et al. (2024a), with the backbone Llama2-7B. and trained with C4 dataset. Llama- -1-8B-SnapKV: We fix the cache size to 16384 and dynam- ically adjust the remaining key-value caches in the cache pool according to the newly injected distracting contexts (consistent with the settings from Section 4.1). The maxi- mum prompt length is set to 49,152 (48k), which requires over 70 GB of GPU memory. We use longer prompt length here as we want to explore more knowledge retention abili-
6

---

M+: Extending MemoryLLM with Scalable Long-Term Memory
Table Experimental Results LongBench.
2wikimqa
hotpotqa
qasper
musique
multifieldqa en narrativeqa
Avg
MemoryI.L.M-7B (20k)
27.22
34.03
19.57
13.47
29.56
20.64
24.08
1.ama3.1-80 (8k)
34.87
43.10
29.96
24.96
43.18
24.29
33.39
1lama3.1-8h (16k)
34.11
44.72
30.05
31.96
48.86
25.19
35.81
M+ (8k)
33.12
37.99
29.91
20.68
40.11
24.18
31.00
M+ (16k)
32.71
38.56
30.39
24.58
46.32
24.12
32.78
window. The first 6k tokens (12 chunks) are compressed into 256 + 256 NN N-K +.... 256 (>**)) ) 2755.6
0.7
M+
MemoryLLMM78
0.6
Llama-3.1-88-Sme
tokens (with around 316.4 memory tokens dropped). For 16k tokens, the first 14k tokens compressed 5530 tokens (with around 1638 tokens dropped). As some tokens are dropped, the performance may get affected slightly. Note that this could lead to a longer context window while sacri- ficing some of the performance in relatively shorter context tasks. (2) Limited Cross-Chunk Attention: When process- ing chunks into memory, M+ uses the last K tokens and the hidden states from the new chunk as input (illustrated in Figur el). In contrast, Ilama-3. 1-8B processes each chunk with access to all previous tokens, enabling cross-attention between chunks. While this approach allows Llama-3. 1-8B to maintain full attention across chunks, it comes with sig- nificantly higher computational and memory costs due to the quadratic scaling of transformer computations. In com- parison, M+ achieves linear computational scaling, making for processing extremely
Accuracy
0.5
Barderline (0.09)
0.4
0.3
0.2
0.1
0k
10k
20k
30k
40k
50k
60k
70k
80k
90k
100k
110k
120k
1 13kk
140k
150k
160k
Figure 3. Knowledge Retention Results on SQuAD.
lies of -ama--31-8B-Snappp
4.3.2. EXPERIMENTAL RESULTS
The experimental results on SQuAD are presented in Fig- ure 3. We present the results on NaturalQA in Appendix B.1 (Figure 8) as both figures have similar trends. From these figures, key observations include: (1) M+ significantly outperforms Memoryl..M-77,, demonstrating substan- tial improvement in knowledge retention compared with the last version. (2) M+ surpasses Llama-3.1 equipped with SnapK, indicating that storing knowledge directly in memory is more effective than relying on key-value cache mechanisms. (3) Even though Llama-3. --8B-SnapKV is given the context window 48k, it struggles to recall infor- mation injected more than 30k tokens earlier, highlighting the limitations of key-value cache methods for long-term knowledge retention.
it more S-D--mmmoy-effiiicee long inputs (see Section 4.2). albeit with some trade-off in performance in tasks with relatively shorter contexts.
4.5. Ablation Study
4.5.1. ABLATION STUDY ON LONG-TERM MEMORY
In this section, we study the effectiveness of our long-term memory ensure that the observed performance improve- ments over MemoryLLMe from the integration of LTM rather than solely from the ad- ditional training. Recall from Section 3.2.4 that the first two training stages do not use long-term memory. We compare with three models: (1) MemoryLLM-SB: The model obtained after Stage It shares the same structure as MemoryLLM-P7 (Wang al., 2024a) but upgrades the backbone from Llama-2-1 to Llama-3.. 3.1--B3..BBB and includes changes such as dataset shifts and multi-LoRA settings. (2) MemoryLLM-SB--Logg The model obtained after Stage 2. (3) M+: The final model obtained after Stage 3.
and Llama-3, 1--a-- 16k stem
4.4. Experimental Results on (Relatively) Short Documents
We evaluate the performance of M+ and Llama-3. 1-BBB relatively short documents using the LongBench benchmark, considering input lengths of 8k and 16k tokens. The eval- uation metric is QA-F1, following Bai et al. (2023). The results, presented Table show that M+ could match the performance of Llama-3.1-8 on 4 out of 6 datasets, except for hotpotqa qa and mus musique. This performance differ- ence be attributed two primary factors: (1) Random Dropping Mechanism: M+ employs random dropping mechanism that can lead to partial information loss. For instance, processing an 8k input requires splitting it into 12 chunks (each chunk being 512 tokens), while the last 2k tokens are directly included in the generation context
Long Context Modeling Ability Improves Over Stages We evaluate the three models on held-out subset of Slim- Pajama containing 1000 examples with lengths between 32k and 64k tokens. We compute the validation loss for each model on this subset and report the results in Figure 4. The results demonstrate that long-context modeling abil-

---

M+: Extending McmoryLLM with Scalable Long-Term Memory
1.70
McmoryLLM 8B (Stage 11
0.8
(Stage
0.7
MeimoryLLM-BB-LLo (Stage
1.65
MM-B-LoL (Stage 2) (Stage 3)
MemoryLLM 8B-Attn
0.6
1.60
Accuracy
Borderfine (0.09)
0.5
1.55
0.4
1.50
0.3
1.45
0.2
1.40
1.35
100k
110k
120k
1304
140k
15kk
16ak
1.30
Figure 5, Ablation Study on SQuAD dataset.
10000
20000
30000
40000 50000
Number of Tokens
MemoryL.LM-8B, which be attributed to the inclusion of longer training examples in Stage 2 (4k-64k), whereas E Fineweb-edu (used in Stage 1) contains very few exam- ples longer than 4k.
Figure 4. Validation loss comparison on held-out subset from Slim-Pajama, consisting of 1,000 examples. The three models, MemoryLLMM88. MemoorL.LM-SS---aa and M+, are obtained after Stages 1, 2, and 3, respectively (Section 3.2.4).
ity improves progressively across training stages, with M+ achieving the lowest validation loss, indicating the strongest performance on long-context inputs.
4.5.2. ABLATION STUDY ON RETRIEVER
We conduct ablation study evaluate the performance of our trained retriever compared an attention-based re- trieval method. Using the SQuAD and NaturalQA datasets, we compare M+ with an iicettio-basee retrieval method inspired by H2O (Zhang et al., 2023). In the attention -based method (M+-Attn), the key-value cache of past tokens is stored, and during generation, a fixed number of tokens are retrieved from the cached keys and values based on their attention scores. To match approach, we use the same short-term memory as M+, but the memory tokens in the long-term memory are retrieved according to the attention scores rather than using our retriever. To implement this method, we adapt M+ to store key-value caches instead of hidden states in the long-term memory to avoid any addi- tional computation cost. During generation, for each token, we extract 2,560 keys and values for each head from the long-term memory, along with the 10,240 memory tokens in the current memory pool. The results on SQuAD are shown in Figures 5 and the results on NaturalQA are shown in Ap- pendix B.2 (Figure 9). From the figures we can see that M+ substantially outperforms M+-Attn, showing the advantages of our trained retriever over the attention- based approach in terms of knowledge retention and retrieval efficiency.
Long-term memory Significantly Improves Knowledge Retention We further assess MemoryLLM--88 Lonnoooo and M+ knowledge retention tasks using SQuAD and Natu- raIQA datasets with the same setting as in Section 4.3. In our experiments, we find MemoryL.LM-BB-Looo marginally better than MemoryLLM-8B on knowledge re- tention tasks, thus for simplicity, we omit the results of MemoryL.L here. The results for MemoryL..MV Long and M+ on dataset SQuAD are presented in Figure 5 and the results on NaturalQA are presented in Appendix B.2 (Figure 9). From the results, can observe that (1) Despite having only 12,800 memory tokens, MemoryL..M-88-LLoom retains knowledge for up to 30 tokens in NaturalQA and 50 tokens in SQuAD, demonstrating effective compression of information into memory tokens. (2) Stage 3 signifi-
performs
cantly enhances retention, extending the model's s ability to recall knowledge from 50k over 160 tokens. During inference, 2.560 tokens are retrieved from long-term mem- ory, combined with 10,240 tokens from short-term memory, resulting in 12,800 effective memory tokens. These re- sults underscore the effectiveness of our long-term memory mechanism in improving knowledge retention and handling extremely long contexts.
4.6. Analysis
4.6. MODEL QUALITY WITHIN CONTEXT WINDOW
Long-term memory does not affect the performance on relatively short documents To show whether long-term memory affects the model's performances on relatively short documents, we conduct ablation study with models from three different stages on the dataset LongBench while fix- ing the context window 8k. The results are shown in Table 3. From the table we can see that M+ has similar performance as MemoorLLLM-S--Loo dow. This means adding long-term memory does not af- feet the performance on relatively short documents. Mean- while, MemoryL.LM-88-- -B--LLong is significantly better than
M+ uses 12,800 memory tokens alongside a 2,048-token generation context window. In this section, we evaluate the model's performance within the standard 2, 2,048-token context window to ensure that the addition of memory does not degrade its base capability. We randomly select 1,000 examples from the fineweb-edu dataset (snapshot which does not overlap with the
CC-MA CC-MAIN-2022 1-N-2024-111 training data. For this evaluation, we cap the input sequence length at 2,048 tokens and report perplexity for both M+ and LLaMA-3.1-88 The results show that LLaMA-3.1 achieves a perplexity of 1.9734, while M+ records sim-
on 8k context win-
8

---

M+: Extending McMoryLLM with Scalable Long-Term Memory
Table Ablation study of the effects of different stages on LongBench.
2wikimqa
hotpotqa
qasper
musique multifieldqa.er
narrativeqa
Avg
MemoryLLM-SB (Stage 8k)
32.30
33.39
23.88
12.37
35.91
21.46
26.55
McmoryLLLM-8BBLLLog (Stage 2, 8k)
32.23
37.86
31.62
20.35
42.16
23.49
31.29
M+ (Stage 3, 8k)
33.12
37.99
29.91
20.68
40.11
24.18
31.00
250
Retrieved Tokens
Ground-truth Tokens
Inference Time (seconds)
30
MemoryLLM 8B
Llama-3.1 3.1-8B
200
25
M+
MemoryLLM 8B (Offload)
130
20
M+ (Offload)
100
15
50
10
5
10k
20X
306
40%
50c
66k
70x
55k
50k
100k 110k 120K 1206 140K 150c 100K
Number of Takens Injected
Figure 6. Number of ground-truth tokens in long-term memory and the number of retrieved groud-truth tokens as more tokens are injected into the memory.
16K
32K Sequence Length
64K
128K
Figure Latency Analysis
focus on the following settings: (1) Llama-3.1-8B-
ilar perplexity of 1.9828. These findings indicate that M+ maintains competitive performance on documents shorter than 2,048 tokens, confirming that the base model's quality within the context window remains intact.
To analyze the latency, we use Llama-3. 3.1-8B with a full context window 128k; (2) MemoryL.LMSSB (After Stage 1); (3) M+ (After Stage 3); (4) MemoryLLM-9K (Offload): offload the memory onto CPU and load the corresponding memory tokens into GPU when the computation hits cer- tain layer; (5) M+ (offload): Offloading the memory onto CPU and load them back when necessary. All experiments in this section are conducted on single H100 GPU. The results shown in Figure 7. From the figure, could find that (1) MemoryLLM-SI has slightly higher latency than Llama-3.
4.6.2. RETRIEVAL QUALITY
In implementation, the long-term memory is initially of size 5120, and then it gradually increases to 80k in our knowledge retention experiments (it hits 81,276 when there 160k tokens are injected). To access retrieval quality, we leverage the knowledge retention task with SQuAD dataset, where the first K 256 tokens are critical for answer- ing the questions. These K 256 tokens are denoted ground-truth tokens. We track the number of ground- truth tokens in the long-term memory and how many tokens are retrieved back into the "Extracted LTM" pool in Figure when queried after various numbers of tokens are injected. We present the results in Figure 6, demonstrating the re- trieval quality as more tokens are dropped from the memory pool to the long-term memory. From the figure we can see that around 30% tokens retrieved. For reference, random 3% tokens.
in relatively shorter documents (16k, 32k, 64k) but has lower latency on long documents (128k); (2) M+ has higher latency than MemoryL..M-BB where the latency is mainly introduced by the retrieval process. (3) Offloading the memory onto CPU introduces slightly more latency, while it becomes negligible when the sequence grows longer. In the case of 128k input, the introduced latency for M+ (offload) compared with M+ is 1 second, leading to 3% additional computation time for M+.
5. Conclusion and Future Work
retrieval would retrieve 2, 560/81 /81, 276
In this work, we present M+, an enhanced memory- augmented language model that extends the long-term re- tention abilities of MemoryL..M. By integrating a long- term memory (LTM) mechanism with co-trained re- triever, M+ effectively retrieves and utilizes past informa- tion, significantly extending the knowledge retention abil- ities from MemoryLLMM achieve superior performances in long-context understanding tasks compared with recent baselines given the similar budget of GPU memory. In fu- ture work, we plan to reduce CPU-GPU communication overhead, enabling more efficient generation with M+.
4.6.3. LATENCY ANALYSIS
While M+ introduces additional computation due to the memory token retrieval from the long-term memory, we perform detailed analysis to quantify this latency. Specifi- cally, we analyze latency under the setting of 128k input. For reference, we use the processing time of Llama-3.1 1-8BB performing a forward pass on 131,071 (=128k-1) tokens to generate the final token. To ensure fairness, we inject 131,072 2,048 tokens into the memory and ask M+ to predict the last token using the remaining 2,047 tokens. We

---

M+: Extending McmoryLLM with Scalable Long-Term Memory
Impact Statement
Larimar: Large language models with episodic memory control. In ICMI.. OpenReview.net, 2024.
This work introduces a memory-auummented approach for Large Language Models (LLMs), enabling them to more effectively retain and retrieve long-term information and thereby offering potential benefits in areas such educa- tion, research, and industry. The increased memory capacity could potentially raise concerns regarding AI safety, relia- bility, and fairness. If not carefully managed, these models could propagate biased content over extended text spans or store sensitive information for unintended durations. It therefore crucial to employ robust safeguards, including bias mitigation strategies and ongoing oversight, prevent misuse or the reinforcement of harmful content. Beyond considerations already inherent to LLMs, we do not foresee other significant societal impacts arising from this work.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., Goyal, A., Hartshom, A., Yang, A., Mitra, A., Sravanku- mar, A., Korenev, A., Hinsvark, A., Rao, A., Zhang, A., Rodriguez, A., Gregerson, A., Spataru, A., Rozière, B., Biron, B., Tang, B., Chern, B., Caucheteux, C., Nayak, C., Bi, C., Marra, C., McConnell, C., Keller, C., Touret, C., Wu, C., Wong, C., Ferrer, C. C., Nikolaidis, C.,, Al- lonsius, D., Song, D., Pintz, D., Livshits, D., Esiobu, D., Choudhary, D., Mahajan, D., Garcia- Garcia-Olano, D., Perino, D., Hupkes, D., Lakomkin, E.,, AlBadawy, E., Lobanova, E., Dinan, E., Smith, E. M., Radenovic, F., Zhang, F., Synnaeve, G., Lee, G., Anderson, G. L., Nail, G., Mialon, G., Pang, G., Cucurell, G., Nguyen, H., Korevaar, H., Xu, H., Touvron, H., Zarov, I., Ibarra, A., Kloumann, I. M., Misra, I., Evtimov, I., Copet, J., Lee, J., Geffert, J., Vranes, J., Park, J., Mahadeokar, J., Shah, J., van der Linde, J., Billock, J., Hong, J., Lee, J., Fu, J., Chi, J., Huang, J., Liu, J., Wang, J., Yu, J., Bitton, J., Spisak, J., Park, J., Rocca, J., Johnstun, J., Saxe, J., Jia, J., Alwala, K. V., Upasani, K., Plawiak, K., Li, K., Heafield, K., Stone, K., and et al. The Ilama 3 herd of models. CoRR, abs/2407.21783, 2024.
References
Al Adel, A. and Burtsev, M. S. Memory transformer with hierarchical attention for long document processing. In 2021 International Conference Engineering and Telecom- munication (En&T). pp. 1-7. IEEE, 2021.
Bai. Y. Lv, X., Zhang. J.. Lyu, H., Tang, J., Huang, Z., Du, Z., Liu, X., Zeng, A., Hou, L., et al. Longbench: A bilingual, multitask benchmark for long context under- standing. arXiv preprint arXiv:2308. /4508, 2023.
Fedorenko, E., Piantadosi, S. T., and Gibson, E. A. Lan- guage is primarily a tool for communication rather than thought. Nature, (30(801)):575--)
Belcak, P. and Wattenhofer, R. Tiny transformers excel at sentence compression. arXiv preprint arXiv:2410.23510, 2024.
2024.
Fu, Y., Panda, R., Niu, X., Yue, X., Hajishirzi, H., Kim, Y., and Peng, H. Data engineering for scaling language mod- els to 128k context. Forty-first International Confer- ence on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL https: //openreview....mmmm.e net
Bellard, F. Nnep v2: Lossless data compression with trans- former. Technical report, Technical report, Amarisoft, 2021.
Bulatov, A., Kuratov, Y., and Burtsev, M. S. Recurrent memory transformer. In NeurIPS, 2022.
Ge, T., Hu, J., Wang, I.., Wang, X., Chen, S., and Wei, F. In-context autoencoder for context compression in large language model. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL https: //openreview. net forum?id-uun j 4ZuGJE.
Bulatov, A., Kuratov, Y., Kapushev, Y., and Burtsev, M. S. Scaling transformer to Im tokens and beyond with rmt. arXiv preprint ar 230:4:2304.1 1062, 2023.
Chen, X., Jiang, J.-Y., Chang, W.-C., Hsieh, C.JJ,, Yu, H.-F., and Wang, W. MinPrompt: Graph-based min- imal prompt data augmentation for few-shot question answering. In Ku, L.-W., Martins, A., and Srikumar, V. (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 254-266, Bangkok, Thailand, August 2024. Association for Computational Linguis- tics. doi: 10. 10.18653/w1/2024.aacl long. 16. URL https: //acl flaclantholog ogy. wwt. ac1 ong. 6/.
Gutiérrez, B. J., Shu, Y., Gu, Y., Yasunaga, M., and Su, Y. Hipporag: Neurobiologically inspired long-term memory for large language models. arXiv preprint arXiv:2405. 11.405.1481. 14831, 2024.
Hao, S., Sukhbaatar, S., Su, D., Li, X., Hu, Z., Weston, J., and Tian, Y. Training large language models to reason in a continuous latent space. arXiv preprint arXiv:2412.06769, 2024.
Das, P., Chaudhury, S., Nelson, E., Melnyk, I., Swami- nathan, S., Dai, S., Lozano, A. C., Kollias, G., Chen- thamarakshan, V., Navrátil, J., Dan, S., and Chen,
He, Z., Karlinsky, L., Kim, D., McAuley, J., Krotov, D., and Feris, R. Camelot: Towards large language models
10

---

M+: Extending Mc with Scalable Long-Term Memory
with training-free consolidated associative memory. arXiv preprint arXiv:2402. /3449, 2024.
Rakotonirina, N. C. and Baroni, M. Memoryprompt: A light wrapper to improve context tracking in pre-trained lan- guage models. arXiv preprint arXiv:2402. 2024.
Hu, C.,, Fu, J., Du, C., Luo, S., Zhao, J., and Zhao, H. Chatdb: Augmenting Ilms with databases their sym- bolic memory. arXiv preprint arXiv:2306. 03901, 2023.
Simoulin, A. and Crabbé, B. How many layers and why? an analysis of the model depth in transformers. In Pro- ceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the I Itth International Joint Conference Natural Language Processing: Stu- dent Research Workshop, pp. 221-228, 2021.
Jawahar, G., Sagot, B., and Seddah, D.
What does bert learn about the structure of language? In ACL 2019-57th Annual Meeting of the Association for Computational Linguistics, 2019.
Khandelwal, U., Levy, O., Jurafsky, D., Zettlemoyer, L., and Lewis, M. Generalization through memorization: Nearest neighbor language models. arXiv preprint arXiv:1911.00112.22 2019.
Wang, W., Dong, L.., Cheng, H., Liu, X., Yan, X., Gao, J., and Wei, F. Augmenting language models with long-term memory. arXiv preprint arXiv:2306.007172 2023.
Wang, Y., Gao, Y., Chen, X., Jiang, H., Li. S., Yang. J., Yin, Q., Li, Z.., Li, X., Yin, B., Shang, J., and McAuley, J. J.
Li, Y., Huang, Y., Yang, B., Venkitesh, B., Locatelli, A., Ye, H., Cai, T., Lewis, P., and Chen, D. Snapkv: LLM knows what you are looking for before generation. CoRR, abs/2404. 14469, 2024. doi: 10.48550/ARXIV. 2404.14469. URL https::////... /tcctps://oi... Dexiv.22004, 14469.
MEMORYLLM, towards self-updatable large language models. In ICML. OpenReview.net, 2024a.
48550/
Wang. Y., Han, C., Wu, T., He, X., Zhou, W., Sadeq, N., Chen, X., He, Z., Wang, W., Haffari, G., Ji, H., and McAuley, J. Towards lifespan cognitive systems. CoRR, abs/2409. 2024b.
Modarressi, A., Köksal, A., Imani, A., Fayyaz, M., and Schütze, H. Memllm: Finetuning ilms to use an explicit read-write memory. arXiv preprint arXiv:240..1167 2024.
Wang. Y., Liu, X., Chen. X., O'Brien, S., Wu, J., and McAuley, J. Self-updatable large language models with parameter integration. arXiv preprint arXiv:2410.00487, 2024c.
Packer, C., Fang, V., Patil, S. G., Lin, K., Wooders, S., and Gonzalez, J. E. Memgpt: Towards lims as operating systems. CoRR, abs/2310.08560, 2023.
Wang. Y.. Wu, R., He, Z., Chen, X. and McAuley, J. Large scale knowledge washing. arXiv preprint 2rXiv:2405.16720, 2024d.
Park, S. and Bak, Memoria: Resolving fateful forgetting problem through human-inspired memory architecture, 2024.
Penedo, G., Kydlícek, H., Allal, L. B., Lozhkov, A., Mitchell, M.. Raffel, C., von Werra, L., and Wolf, T. The fineweb datasets: Decanting the web for the finest text data at scale. CoRR, abs/2406. 17557, 2024. doi: URL https: //doi. org/10. 1.0.0..8555 (wt/88550///xii 2406. 17557.
Wu, X., Xiao, L., Sun, Y., Zhang, J., Ma, T., and He, L. survey of comaa--n--thelowee for machine learning. Future Generation Computer Systems, 135:364-381, 2022a.
10.48500AARXI22446..755
Wu, Y., Rabe, M. N., Hutchins, D., and Szegedy, C. Memorizing transformers. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022b. URL //tppreve//peereveeww forum? id-TrjbxaRenf-
Pham, Q. H., Ngo, H., Luu, A. T., and Nguyen, D. Q. Who's who: Large language models meet knowledge conflicts in practice. arXiv preprint arXiv:2410. 15737, 2024.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, J. Explor- ing the limits of transfer learning with unified text- to-text transformer. J. Much. Learn. Res., 21:140:1- 140:67, 2020. URL https: /etls///mmmr.///////////eeeeeerr hammm..rrr/ppee..mmmmmmmom papers/ v21/20- 00m.0mm.7......
Yang, H., Lin, Z., Wang, W., Wu, H., Li, Z., Tang, B., Wei, W., Wang, J., Tang. Z., Song, S., Xi, C., Yu, Y., Chen, K., Xiong, F. Tang. L.,, and E, W. Memory Language modeling with explicit memory, CoRR, abs/2407.01178. 2024.
Rahman, C. M., Sobhani, M. E., Rodela, A. T., and Shatabda, S. An enhanced text compression ap- proach using transformer-baa based language models. CoRR, abs/2412.15250,, 2024.
Yin, Z. Sun, Q. Guo, Q.. Zeng. Z., Cheng, Q., Qiu, X., and Huang, X. Explicit memory learning with expectation maximization. In EMNLP, pp. 16618-16635. Association for Computational Linguisties, 2024.

---

M+: Extending McMoryLLM with Scalable Lonm-Term Memory
Zhang, X., Chen, Y., Hu, S., Xu, Z., Chen, J., Hao, M., Han, X., Thai, Z.., Wang, S., Liu, Z.,, al. oobench: Extending long context evaluation beyond 100k tokens. In Proceed- ings of the 62nd Anmual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp- 15262-15277, 2024.
Zhang. Z., Sheng, Y., Zhou, T., Chen, T., Zheng, L.,, Cai, R., Song, Z., Tian, Y., Ré, C., Barrett, C. W., Wang, Z., and Chen, B. H2O: heavy-hitter oracle for efficient gemerative inference of large language models. In NeurlPS, 2023.
Zhong, W., Guo, L., Gao, Q. and Wang. Y. Memorybank: Enhancing large language models with long-term mem- ory. arXiv preprint arXiv:2305. 130500,00 2023.
Zhou, W., Jiang, Y. E., Cui, P., Wang, T., Xiao, Z., Hou, Y., Cotterell, R., and Sachan, M. Recurrentgpt: Interac- tive generation of (arbitrarily) long text. arXiv preprint arX3v 2023.
12

---

M+: Extending McmoryLLM with Scalable Lerg-Term Memory
A. Justifications of using stags-
Eight A100 GPUs support the following configurations:
Full fine-tuning with an 8k context window using Fully Sharded Data Parallel (FSDP).
6k context window with full attention using deepspeed staaesttae---e -2
32k context window with full attention using accelerate and saving models in this configuration encountered version incompatibility issues and we haven't 't found solutions online.
--offgg---ffoddd. However,
Based on these trails, we do scale up the model with peppapeed-stagee and the cross-attention be of the shape 2048 by 14848.
or FSDP, but choose to use
B. Experiments on datasets NaturalQA
B.1. Knowledge Retention Experiments on NaturalQA
The results of knowledge retention experiments on NaturalQA are shown in Figure 8.
0.8
M+
0.7
MemoryLLM-78
1amma-.1188 SnapKV
Accuracy
0.6
Borderline (0.2)
0.5
0.4
0.3
0.2
0k
10k
20k
30k
40k
50k
60k
70k
80k
90k
100k
110k
120k
130k
140k
150k
160k
Figure 8, Knowledge Retention Results on NaturalQA.
B.2. Ablation Study on NaturalQA
The results of ablation study on NaturalQA are shown in Figure
M+ (Stage 3)
0.7
MemoryLLM-88B (Stage 2) LM-B-AAL
0.6
Accuracy
Borderline (0.2)
0.5
0.4
0.3
0.2
0k
10k
20k
30k
40k
50k
60k
70k
80k
90k
100k
110k
120k
130k
140k
150k
160k
Figure Ablation Study on NaturalQA dataset.
13

---

M+: Extending Mc with Scalable Long-Term Memory
C. Statisticc of the Dataset of Long Documents
We go through the whole dataset SlimPajama-622 of Llama-3. 1-B.. The statistics are shown in Table 4. We show six categories here (4k-8k, 8k- 16,, -2lkkkk,3.2k------ -4k,,,kk64,666 222.6664444 64k- 228k,12kke but we only use the data within the first four categories (4k-8k, 16-33K, examples longer than 64k are mainly from the category Book and lack diversity.
and extract all dataset that have more than 4k tokens using the tokenizer
This is because the
Range
Total
CommonCrawi
GitHub
ArXiv
C4
StackExch.
Wikipedia
Book
4k-S8k
11,189.99 .999
7.759,741 (69.35%)
692.224 (6.19%)
286,537 (2.56%)
1,825,018 (16.31%)
142.457 (1.27%)
481,854 (4,31%)
2.168 (0.02%)
6,706.87
3,273.619 (69. (89)
270.369 (5.74%)
550, 192 (11 69%)
439. (9.33%)
20,284 (0.43%)
146.545
6,535 (0. 14%)
16k-32k
1,607,0044
968,714 (60.28%)
95,445 (5.94%)
423,401 (26.35%)
70,223 (4.37%)
1,510 (0.09%)
34,123 (2.14%)
13,448 (0.845)
443,438
(50.55%)
32,653 (7.36%)
146.582 (33.06%%
3,413 (0.77%)
102
5,940 (1.34%)
30,580 (6.9092)
64k--128k
192,515
72,583 (37.70%)
11,753 (6.10%)
27,942 (14.51 1%)
38 (0.02%)
5 (0.00%)
507 (0.26%)
79,687 (41.39%%
23,721 (24.18%)
4,523 (4.615%
5,167 (5.27%%)
(0.00%)
2 (0.00%)
(0.05%))
64,635 (65.893)
Table 4. Number examples by sequence-length range and source (counts and percentages).
D. Additional Training Details
In our training, we follow MemoryLLM (Wang et al., 2024a) and design three sub-tasks:
Two-Chunk Training: Given document split into two chunks (æ1, æ2), we inject æ1 into the memory and update the transformer ø using the loss on 2-2. Notably, we retain the gradients across both forward passes.
Multi-Chunk Training: For documents with multiple chunks In), we inject T1,..., En-1 into the memory while detaching gradients, then update 6 using the loss on Zn--
Since the memory is continually updated during training, we cache the last chunk on of earlier documents and revisit periodically. When revisiting Tx there are already many chunks injected between "1.-.. 1 and Fn- We denote the number of injected chunks between fn-1 and vn as revisit distance. We carefully tune the probability of deleting and updating the cache after each training step, and we manage to maintain the average revisit distance to be around 60 for Stage 1 & Stage 2, and maintain the average distance to be around 200 for Stage 3.
Revisiting Cached Chunks:
E. Discussions
E.1. Similarities to Attention-Basee Retrieval Methods
In M+, we use a co-trained retriever to retrieve the hidden states. In this process, 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 advantagcous:
Efficiency: Methods such as SnapKV maintain and retrieve key-value pairs per head, which becomes extremely costly when scaled. our setting- with 32 layers and 32 attention heads per layer-this requires 1024 retrievals per query, resulting 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 dramattcaal reducing both computational cost and latency.
Performance: In Figure 6, the curve labeled LLm-oryLLM-B--Ai 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.
Design: Note that our training setup includes both relevant and irrelevant documents (See details 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.
14

---

M+: Extending McmoryLLM with Scalable Long-Term Memory
E.2. The Form of Long-Termi Memory (Hidden States KV)
In our work, we choose the use hidden states as the latent-space memory instead of key-value (KV) caches. This based on the following two considerations:
Compression Efficiency: As detailed in the paper, we compress each 512-token chunk into 256 memory vectors per layer in 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 above, hidden states can be effectively retrieved using our co-trained retriever, requiring only 32 retrievals for each query. In contrast, a KV-ccche 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.
E.3. Latency and Memory Consumption while Scaling
We aim to discuss the scalability of M+ by analysing the latency and memory consumption when scaling up. Theoretically, end-to-end retrieval latency scales linearly with three key variables:
Hidden size of the retriever, denoted by d. In M+, we d = 256, whereas the base model uses d 4096.
(2) Size of long-term memory, denoted by We cap this 150k entries.
33) Number of transformer layers, denoted by L. For LLAMA-3-SHL T= 32.
Hence,
latency & d
Because we hold the long-term memory size fixed when scaling the model, 8 effectively constant:
latency & d d.
Both d and L grow with the model size M, following
M x dL,
which implies a linear relationship between retrieval latency and model size:
latency x M.
Concrete Example Scaling from LLAMA-3-8B (d 4096, L 32) LLAMA-3-70B (d 8192, 80) yields
8192 80
5,
4096 x 32
i.e. a 5x increase in retrieval latency. For comparison, the parameter count rises by scales roughly linearly-ratee rraher than quadratically- model size.
70B
~ 8.75x, showing that latency 8B
The extra memory overhead from our method arises solely from the introduced memory tokens. This overhead also scales linearly with both and L: thus, the move from LLAMAA3-888 to LLAMA-344 3- incurs an analogous ~5x increase in memory usage, mirroring the latency scaling.
Memory Consumption
15

---

M+: Extending McmoryLLM with Scalable Long-Term Memory
E.4. FLOPs Comparison
We report the total FLOPs for generating one token after processing sequence of varying lengths (from 2k 128K), using single H100 GPU. We employ the torch. carch.profiles. library to capture FLOPs during inference. The results are as follows:
Sequence Length
LLaMA-3.1-8B
M+
2048
5.68 10
6.92 10
4096
1.13 1044
1.32 x 10
8192
2.26 x 10
2.55 x 10
16384
4.48 x 10
5.01 x 10
32768
8.88 x 10
9.86 x 10
65536
1.75 x 10
1.94 x 10
131072
OOM
3.78 x 1045
From the results, we observe that M+ LLaMA-3.1-BB exhibit comparable FLOPs across all tested sequence lengths. Notably, while LLaMA-3.1-8B runs of memory (COM) at the 128k setting, M+ remains functional, highlighting its superior scalability for long-context inference.
E.5. Interpretability of Memory Vectors
Our memory vectors can be viewed as hidden states within the transformer layers, with the key difference being that they may store more compressed information due to their persistent role across sequences. As such, the type of information they capture should be similar to the representations observed 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 reported in prior work on transformer interpretability (Jawahar et al., 2019; Simoulin & Crabbé, 2021):
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 dropping vectors from the short-term memory and storing them for extended use. Importantly, long-term memory vectors are structurally identical to short-term ones. This means that, at any point, swapping a vector between long-tern and short-term memory has no immediate effect on model behavior.
In essence, the long-term memory acts as a cache that helps memory vectors persist over time rather than being overwritten quickly.
16


Reviews:
[{"review_id": "ULaoVtYhP3", "reviewer": "Reviewer_WXyo", "summary": "**Main Findings:**\n\nEquipping large language models (LLMs) with latent-space memory has gained significant interest, as it extends the effective context window of existing models. However, preserving and retrieving information from distant past contexts remains challenging. To address this, this paper proposes M+, an enhancement of the existing MemoryLLM model, introducing a latent-space retrieval mechanism designed specifically for long-term memory retention.\n\n**Main Algorithmic/Conceptual Ideas:**\n\nM+ incorporates a co-trained retriever within the MemoryLLM framework, enabling dynamic retrieval of relevant latent space information during text generation. This allows the model to effectively leverage long-term contextual memories.\n\n**Main Results:**\n\nEmpirical evaluations on several long-text benchmarks (including Longbook, SQuAD, and LongBench) demonstrate that M+ outperforms MemoryLLM.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Essential References Not Discussed:**\n\nThank 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.\n\n**Weaknesses**:  \n\n**[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:\n\n**(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.\n\n**(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. \n\n**(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.\n\n**[W2] Representation of Long-Term Memory (Hidden States vs. KV):** \n\nWe 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:\n\n**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.\n\n**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.\n\nWe believe these benefits make hidden states a more efficient and effective choice for long-term memory representation in our system.\n\n[1] Hybrid computing using a neural network with dynamic external memory.   \n[2] Titans: Learning to Memorize at Test Time.   \n[3] Scaling Transformer to 1M tokens and beyond with RMT.", "claims_and_evidence": "Yes, the claims are clear and convincing.", "methods_and_evaluation_criteria": "Yes, the proposed methods make sense for the target problems.", "theoretical_claims": "Not applicable. No theoretical claims.", "experimental_designs_or_analyses": "Yes, I have checked the experimental designs.", "relation_to_broader_scientific_literature": "Introducing a retrieval mechanism into MemoryLLM can effectively address the challenges associated with extremely long-context content.", "essential_references_not_discussed": "[1] is one of the earliest works to discuss the design of extra long-term memory. [2] [3] are closely related and up-to-date works that design memory modules for long-context ability in LLMs.\n\n[1] Hybrid computing using a neural network with dynamic external memory\n\n[2] Titans: Learning to Memorize at Test Time\n\n[3] Scaling Transformer to 1M tokens and beyond with RMT", "other_strengths_and_weaknesses": "**Strengths:**\n\nThe proposed method clearly demonstrates improvements over the original MemoryLLM by introducing an effective retrieval mechanism (M+) for handling extremely long-context content.\n\nThe paper provides detailed descriptions of algorithms and experimental setups, facilitating straightforward reimplementation by the research community.\n\n**Weakness**:\n\n1. The novelty of this paper is incremental. The proposed long-term memory mechanism closely resembles existing retrieval-based designs for long-context processing, such as SnapKV, into MemoryLLM. The primary difference is utilizing KV or latent states for long-term memory, which has not been clearly demonstrated as crucial for improving performance.\n\n2. The paper lacks a rigorous analysis or ablation study comparing KV and latent states as long-term memory.", "other_comments_or_suggestions": "No additional comments.", "questions_for_authors": "1. Could you elaborate further on any additional conceptual or algorithmic insights beyond introducing a retrieval mechanism into MemoryLLM? Specifically, what distinct advantages or innovations does your method offer compared to existing retrieval-based approaches (e.g., SnapKV)?\n\n2. It would be beneficial to explicitly compare the effectiveness of leveraging KV cache versus latent states as long-term memory. Providing experimental results or ablation studies on this comparison could clarify the significance and necessity of your design choice.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Z5TMwTWuRl", "reviewer": "Reviewer_f5Uz", "summary": "Memory model: memory pool (based on MemoryLLM) and a long-term memory with additional temporal information. Every time, the memory pool is updated, and a subset of tokens is dropped to the Long-term memory. For recall from memory, a small subset of the long-term memory vectors is retrieved according to the dot product with the query from input. The search in the long-term memory uses a low dimensional projection of the vectors (keys) and the input query for such recall. These projectors are trained separately as part of the retrieval mechanism for the long-term memory. The M+ model is fine-tuned with 2 sets of weights based on LoRA: for reading and for updating the memory. Also, it was trained in 3 stages. The first stage fine-tunes Llama-3.1-8B following the MemoryLLM setup to incorporate the memory pool. This process is done on shorter documents. The second stage extends this to a balanced set of documents of various lengths. The last stage introduces the long-term memory and adapts the model to it on a new subset of long context documents. The M+ model is then evaluated on the LongBook-QA, a synthetically generated extraction of events called LongBookEvent-QA, SQuAD, NaturalQA, and LongBench. The baselines are Llama3.1-8B-16k, a similar version with SnapKV, and a Llama3.1 3B model with 128k context length. The results show that M+ has higher performance with lower or comparable memory consumption. The authors include ablation study comparing M+ after each fine-tuning stage for validation loss convergence and knowledge retention (SQuAD and NaturalQA).\n\n## update after rebuttal\n\nI appreciate the authors replying to my questions. There is no additional evidence that leads me to update my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank reviewer f5Uz for their recognition of the value of our work. We address the reviewer's concerns below:\n\n**Relation To Broader Scientific Literature:**\n\nTo 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.\n\nAs for the questions:\n\n**[Q1] Evaluation Metric in LongBench:**\n\nThank 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.\n\n**[Q2] Related Work and Comparison with [2]:**\n\nWe 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.\n\nReferences:\n\n[1] LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding.    \n[2] Assessing Episodic Memory in LLMs with Sequence Order Recall Tasks.    \n[3] Titans: Learning to Memorize at Test Time.   \n[4] Memory Layers at Scale.", "claims_and_evidence": "The claims made in this work seem appropriately supported with convincing evidence.", "methods_and_evaluation_criteria": "The methods and the evaluation criteria make sense for the presented application.\nNote that no metric is described for the LongBench results (Table 2). Moreover, given those results, it is unclear what the authors want to convey with that experiment. \nAlso, it is worth mentioning that there are existing datasets that aim to test the quality of memory augmented models like M+ (see [1]). \n\n\n[1] “Assessing Episodic Memory in LLMs with Sequence Order Recall Tasks” by Pink et al., 2024.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Mostly, the design and the analyses look sound. I would encourage the authors to test on datasets that evaluate for long context, instead of selecting documents from a general dataset. The intention behind the ablation study is valuable, however, it is not evident that the long-term memory is the one helping to obtain better results in stage 3. A more correct version of the experiment would be to ablate the tokens retrieved from the memory (i.e., change to padding, zero vectors, or noise). This would diminish any effect of the additional training in Stage 3.", "relation_to_broader_scientific_literature": "The contributions in this work are relevant to this community and valuable. We should note that similar solutions have been proposed and have been shown to work similarly. This work doesn’t compare results beyond its predecessor (MemoryLLM).", "essential_references_not_discussed": "To the best of my knowledge the authors cite many previous work on memory-augmented LLMs. They appropriately introduce the predecessor MemoryLLM.", "other_strengths_and_weaknesses": "* The work is clearly written and nicely presented\n* The experiments evaluate a vast amount of details about the model", "other_comments_or_suggestions": "* Line 156: “tokens with the largest ages” -> “oldest tokens”\n* Figure 3 shows “Llama-3.2-3B…” while the text mentions “Llama-3.1-3B”.", "questions_for_authors": "* What is the metric used for LongBench?\n* What is the performance of M+ on the task of [1]?\n\n[1] “Assessing Episodic Memory in LLMs with Sequence Order Recall Tasks” by Pink et al., 2024.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PfHJNukAm6", "reviewer": "Reviewer_v8Qd", "summary": "This paper presents M+, an enhanced memory-augmented language model that extends long-term memory retention beyond conventional limits. Building on MemoryLLM, M+ integrates long-term memory with a co-trained retriever.  Extensive experiments across long-context understanding, question answering, and knowledge retention benchmarks demonstrate that M+ consistently outperforms prior baselines, offering a robust and efficient solution for processing extremely long documents.\n\n## update after rebuttal\nMost of my concerns are addressed so I have raised my reviewing score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are sincerely grateful for Reviewer v8Qd’s recognition of our work. Below, we address the reviewer’s questions in detail: \n\n**[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. \n\nThat 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: \n\nMemoryLLM-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. \n\nMemoryLLM-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. \n\nAlthough 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:\n\n1. **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.\n\n2. **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+.\n\n3. **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:\n\n    | Model                   | 2wikimqa | hotpotqa | qasper | musique | multifieldqa\\_en | narrativeqa | Avg   |\n    |------------------------|----------|----------|--------|---------|---------------|----|--------------|\n    | MemoryLLM-8B (8k)      | 32.30    | 33.39    | 23.88  | 12.37   | 35.91         | 21.46 | 26.55  |\n    | MemoryLLM-8B-Long (8k) | 32.23    | 37.86    | 31.62  | 20.35   | 42.16         | 23.49 | 31.29  |\n    | M+ (8k)                | 33.12    | 37.99    | 29.91  | 20.68   | 40.11         | 24.18 | 31.00  |\n\nThese 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.\n\nIn 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.", "claims_and_evidence": "Yes. Most of the claims are supported by a range of quantitative experiments and ablation studies. For example, the paper demonstrates through long-book QA, Event QA and knowledge retention benchmarks that M+ outperforms baselines like MemoryLLM and Llama families. Detailed GPU memory cost comparisons and latency analysis further back the claim that M+ achieves extended retention with efficient resource usages.", "methods_and_evaluation_criteria": "Yes. The propose method M+ for integrating a long-term memory module into MemoryLLM match the stated goal of extending context retention beyond 20k tokens, and the chosen evaluation benchmarks (e.g., LongBook-QA, Event QA, and knowledge retention tasks) effectively measure whether the model can recall and use information from far in the past.", "theoretical_claims": "The paper primarily focuses on an empirical framework rather than on detailed proofs of novel theorems.", "experimental_designs_or_analyses": "Yes. The knowledge retention experiments and ablation studies are well suited to testing M+’s ability to remember and retrieve distant information.\nWhile the paper’s approach to evaluating memory retention is well designed, the evaluation dataset and task types are limited. More complex datasets or practical applications are needed to demonstrate the effectiveness of the method [1, 2].\n[1] Maharana, Adyasha, et al. \"Evaluating Very Long-Term Conversational Memory of LLM Agents.\" Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2024.\n[2] Lu, Junru, et al. \"Memochat: Tuning llms to use memos for consistent long-range open-domain conversation.\" arXiv preprint arXiv:2308.08239 (2023).", "relation_to_broader_scientific_literature": "The paper builds on research in memory-augmented language models, particularly latent-space memory approaches (e.g., MemoryLLM) that store and retrieve hidden-state representations instead of raw tokens. The authors go a step further by co-training a retriever with their model, in contrast to attention-based retrieval from key-value caches. Their method also connects with work on long-context modeling by pushing context windows into the 100k+ range. By demonstrating improved knowledge retention and efficient retrieval on challenging long-QA tasks, the paper contributes to broader efforts of making LLMs handle truly extensive contexts while remaining computationally feasible.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "The core ideas of this paper build on existing latent-space memory approaches and chunk-based processing techniques. While the practical enhancements (e.g., CPU offloading to keep GPU usage low and adding a co-trained retriever for selective recall) are valuable, these methods can be viewed as incremental refinements, and I found them limited in originality and novelty, despite the paper’s clear demonstration of utility.", "other_comments_or_suggestions": "1.\tIn the comparison experiments (Sections 4.1 and 4.2), since M+ is largely based on MemoryLLM, it would be beneficial to include direct comparisons showing how much M+ actually improves over MemoryLLM. Observing that margin explicitly could clarify the advantages of the proposed enhancements.\n2.\tIn Table 2, the authors report results on shorter-document tasks, yet M+ still outperforms MemoryLLM by a noticeable margin. Given that M+’s primary benefit appears to be in handling extremely long contexts, it would be helpful to clarify why it achieves such gains even with relatively small datasets.", "questions_for_authors": "No", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TxdDJvKqHm", "reviewer": "Reviewer_zEbf", "summary": "The paper proposes an LLM memory-augmentation method called M+. By building on MemoryLLM, it improves the long-context understanding and information retention of a base LLM. They introduce what they call long-term memory vectors that are extracted by a trained retriever. M+ outperforms the base model and other existing baselines on long-context QA and knowledge retention benchmarks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Claims And Evidence:**\n\nWe 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. \n\n**Supplementary Material:**\nWe will publish the code and model upon acceptance and make sure of the reproducibility. \n\n**Essential References Not Discussed:**\nThank you for bringing these up, we will add these works ([1], [2], [3]) into our paper.\n\n**Comments, Suggestions and Questions:**  \n**Figure 3**: We will revise the scales to ensure consistency across subplots.   \n**Limitation Section:**: We briefly mentioned our plan to reduce CPU-GPU communication overhead in future work, we will add a section to discuss this. \n\n**[Q1]**\n\n(1) **Evaluations on LongBench**: We would like to note that our evaluations on LongBench are included in Section 4.4.\n\n(2) **Experimental Results on NIAH**: We would like to emphasize two key points:\n- **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.\n- **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.\n\n**[Q2]**\n\nWe 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+. \n\n**[Q3]**\n\nOn 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. \n\n**[Q4]**\n\nUsing 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.\n\n**[Q5]**\n \nRetrieval 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$.\n\n**[Q6]**\n\nM+ 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. \n\n**[Q7]**\n\nOur 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.\nAcross layers, we hypothesize that the memory vectors follow a similar pattern to what has been observed in prior work on transformer interpretability [4, 5]:  \n- **Lower layers** tend to encode more **surface-level features**,  \n- **Higher layers** tend to encode more **semantic or abstract information**.  \n\nRegarding 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.  \n\nReferences:\n\n[1] Recurrent Memory Transformer   \n[2] Scaling Transformer to 1M tokens and beyond with RMT   \n[3] MemoryPrompt: A Light Wrapper to Improve Context Tracking in Pre-trained Language Models  \n[4] How Many Layers and Why? An Analysis of the Model Depth in Transformers  \n[5] What does BERT learn about the structure of language?", "claims_and_evidence": "M+ improves the long-context understanding and information retention of LLMs. This has been validated by the experiments in section 4.1 and 4.3. M+ also claims to be an efficient method but this is not clearly supported by the evidence presented in the paper.  It has higher memory consumption and latency than the original base model. Adding CPU offloading to Table 1 is a bit misleading. It could be applied to any model and it is not specific to M+.", "methods_and_evaluation_criteria": "The benchmark datasets used in the paper are appropriate as they evaluate the long-context memory capabilities of models. However, the models were not evaluated on standard long-context benchmarks like the classic Needle-in-a-Haystack test and LongBench.", "theoretical_claims": "None", "experimental_designs_or_analyses": "Concerning analysis, not much is said about the memory vectors. It would greatly help the community to study what kind of information is contained in the memory vectors.", "relation_to_broader_scientific_literature": "The paper studies the long-context extension of LLMs. It also builds on an existing approach called MemoryLLM.", "essential_references_not_discussed": "Reccurrent Memory Transformers and their variants (https://proceedings.neurips.cc/paper_files/paper/2022/hash/47e288629a6996a17ce50b90a056a0e1-Abstract-Conference.html\n, https://arxiv.org/abs/2304.11062) were not discussed in the paper. They also study latent-space memory. \nThe MemoryPrompt (https://aclanthology.org/2024.lrec-main.976/) paper analyzes the content of these memory vectors.", "other_strengths_and_weaknesses": "None", "other_comments_or_suggestions": "Figure 3 is a bit hard to read. The scales on the 2 sides are different but combined in the same figure.\nA Limitations section would help better delineate the actual contributions of the paper.", "questions_for_authors": "1- What is the performance on the following standard long-context benchmarks: Needle-in-a-Haystack test and LongBench?\n\n2- What is the memory consumption of MemoryLLM in your experiments?\n\n3- Is the performance of the base Language Model affected by the introduction of the long-term latent-space memory? (perplexity and other metrics on standard LM benchmarks)\n\n4- What is the performance of a simple RAG baseline?\n\n5- Can you provide an estimate of how the memory consumption and latency would behave for very big models? \n\n6- FLOPs was never reported. What is the FLOPs compared to the base model and other baselines? What are the results if you do FLOP-matched comparisons?\n\n7- From an interpretability perspective, what kind information is contained in these memory vectors? How do they differ across layers? What distinguishes the long-term memory vectors from the short-term ones?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper builds on MemoryLLM by introducing M+, a memory‑augmented transformer that integrates (1) a compressed latent‑space long‑term memory pool and (2) a co‑trained retriever to fetch relevant memory tokens during generation. By storing and retrieving 256 compressed vectors per layer rather than per‑head key‑value pairs, M+ reduces retrieval overhead (32 vs. 1,024 retrievals per query) and extends effective context retention from ~20 K to over 160 K tokens. The authors evaluate on long‑context QA (LongBook‑QA, LongBook‑Event‑QA), standard QA (SQuAD, NaturalQA), and multi‑domain long‑context benchmarks (LongBench), reporting consistent gains over LLaMA‑3.1, MemoryLLM, SnapKV, and RAG baselines. Ablations confirm the benefits of the retrieval mechanism and hidden‑state representation.

Reviewers generally agree that M+ delivers solid, practically relevant improvements in long‑context retention. Concerns about incremental novelty are outweighed by the paper’s scale, efficiency, and thorough empirical validation. The authors’ rebuttal clarified comparisons to KV‑cache baselines (via additional experiments in Figures 3, 5, 6), detailed memory/latency scaling, and provided short‑document performance to contextualize gains. Requests for deeper interpretability and theoretical analysis fall outside this paper’s primarily empirical scope but can be addressed in follow‑up work.
Decision: Accept (poster)