M+: Extending MemoryLLM with Scalable Long-Term Memory 1 Introduction 2 Related Work 2.1 Token-level Memory 2.2 Latent-Space Memory 3 Methodology 3.1 Preliminaries 3.2 Equipping MemoryLLM with Long-Term Memory 3.2.1 Memory Structures Update Process Generation Process Multi-LoRA Design 3.2.2 Retriever Design and Training Retriever Design Training the Retriever 3.2.3 Training Details Setting Configurations 3.2.4 Data Curriculum Continual Training of MemoryLLM (Stage 1) Long-Context Modeling with Long Documents (Stage 2) Training with long-term memory (Stage 3) 4 Experiments 4.1 Long Book QA and Event QA 4.1.1 Experimental Settings 4.1.2 Experimental Results 4.2 GPU Cost Comparison 4.3 Knowledge Retention Experiments 4.3.1 Experimental Settings 4.3.2 Experimental Results 4.4 Experimental Results on (Relatively) Short Documents 4.5 Ablation Study 4.5.1 Ablation Study on long-term memory Long Context Modeling Ability Improves Over Stages Long-term memory Significantly Improves Knowledge Retention Long-term memory does not affect the performance on relatively short documents 4.5.2 Ablation Study on Retriever 4.6 Analysis 4.6.1 Model Quality within Context Window 4.6.2 Retrieval Quality 4.6.3 Latency Analysis 5 Conclusion and Future Work A Justifications of using deepspeed-stage-2 B Experiments on datasets NaturalQA B.1 Knowledge Retention Experiments on NaturalQA B.2 Ablation Study on NaturalQA C Statistics of the Dataset of Long Documents D Additional Training Details E Discussions E.1 Similarities to Attention-Based Retrieval Methods E.2 The Form of Long-Term Memory (Hidden States vs. KV) E.3 Latency and Memory Consumption while Scaling Concrete Example Memory Consumption E.4 FLOPs Comparison E.5 Interpretability of Memory Vectors M+: Extending MemoryLLM with Scalable Long-Term Memory Yu Wang Dmitry Krotov Yuanzhe Hu Yifan Gao Wangchunshu Zhou Julian McAuley Dan Gutfreund Rogerio Feris Zexue He 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. We open-source our code at https://github.com/wangyu-ustc/MemoryLLM . mixtural-of-expert, memory, large language model 1 Introduction The integration of memory modules into large language models (LLMs) has gained increasing attention  (Wang et al., 2024b ) . Existing approaches for constructing memory modules can be broadly divided into two main categories: (1) Token-level memory  (Packer et al., 2023 ; Modarressi et al., 2024 ) , where memory is represented as structured text, enabling 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 information. 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 text-based 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 information in text-based memory can be challenging  (Pham et al., 2024 ) . Meanwhile, as noted by Fedorenko et al. ( 2024 ); Hao et al. ( 2024 ) , human reasoning often transcends the token level, leveraging deeper, integrated representations akin to latent spaces. In contrast, Latent-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 compact latent space  (Das et 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 et 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 representations beyond discrete tokens, akin to latent spaces. By encoding knowledge in latent representations, the methods with latent-space memory can more closely mimic the mechanisms of human memory, which store information within neural activations. In this paper, we focus on Latent-Space Memory. MemoryLLM  (Wang et al., 2024a ) , as a representative work in this category, enhances a transformer-based language model by incorporating a large number of memory tokens into each layer, creating a memory pool with 1 billion parameters. This framework introduces a carefully designed update and generate process, achieving superior performance compared to the backbone model Llama-2-7B and other long-context methods. However, MemoryLLM faces limitations in recalling 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 MemoryLLM. Unlike previous 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 retriever and the language model. This allows M+ to retrieve only once per layer for all query heads, significantly improving efficiency. Furthermore, as the long-term memory is stored on the CPU, M+ significantly extends long-term retention capabilities without increasing GPU memory usage. We evaluate M+ across a diverse set of benchmarks, including tasks such as long-book understanding, knowledge retention, and question answering on relatively short documents. Experimental results demonstrate that M+ achieves significant performance improvements in all long benchmarks compared to previous memory-based methods while operating within the same or even smaller inference memory budget. In summary, our contributions are as follows: • We enhance MemoryLLM by incorporating a long-term memory mechanism and introducing a co-trained retriever for efficient and effective memory retrieval. • We design a specialized data curriculum for long-context training, enhancing the long-context modeling ability of M+. • Through extensive experiments on multiple benchmarks, we demonstrate that M+ significantly outperforms the baselines while maintaining a similar or reduced 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 et al. ( 2024 ) where they classify methods into implicit memory and explicit memory. 2.1 Token-level Memory Token-level memory refers to memory structures represented in textual forms, which can include raw context, summaries  (Zhong et al., 2023 ; Zhou et al., 2023 ) , knowledge 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 ) proposes treating context and memory as analogous to traditional memory in operating systems, enabling more flexible and organized memory structures. These approaches typically 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 generated 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 MemLLM (Modarressi et al., 2024 ) fine-tunes the model to generate function calls that initiate searches within a knowledge 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 requires 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 ) . Others store knowledge in the form of triplets, which may be unsuitable for representing complex conversations that are difficult to convert into triplets  (Wang et al., 2024d ) . Figure 1: The left side shows the Update and Generation Process of MemoryLLM (Wang et al., 2024a ) . We process the chunk with ϕ l subscript italic-ϕ 𝑙 \phi_{l} italic_ϕ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT to obtain new K 𝐾 K italic_K tokens during the update process, which is perceived by ϕ italic-ϕ \phi italic_ϕ using cross-attention during the generation process. The right side shows the Update and Generation Process of M+. For layer l 𝑙 l italic_l , during Update, the old memory pool θ l subscript 𝜃 𝑙 \theta_{l} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT is split into two parts: K 𝐾 K italic_K dropped tokens and N − K 𝑁 𝐾 N-K italic_N - italic_K remaining tokens. The dropped tokens are stored in the long-term memory Θ l subscript Θ 𝑙 \Theta_{l} roman_Θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT while the remaining tokens and new K 𝐾 K italic_K tokens are combined to obtain the new memory pool θ l ′ superscript subscript 𝜃 𝑙 ′ \theta_{l}^{\prime} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT . Then during generation, we use our co-trained retriever to retrieve tokens from Θ l subscript Θ 𝑙 \Theta_{l} roman_Θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT , which is fed into the transformer layer ϕ l subscript italic-ϕ 𝑙 \phi_{l} italic_ϕ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT along with the short-term memory θ l subscript 𝜃 𝑙 \theta_{l} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT and the query hidden states. The major difference between MemoryLLM and M+ is the introduction of Long-Term Memory Θ l subscript Θ 𝑙 \Theta_{l} roman_Θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT . 2.2 Latent-Space Memory Latent-space memory stores information in a compressed format, embedding knowledge into soft prompts  (Rakotonirina & Baroni, 2024 ) , hidden states  (Khandelwal et al., 2019 ; Bulatov et al., 2022 , 2023 ; Wang et al., 2024a ) , model parameters  (Wang et al., 2024c ) , or an external latent space  (Das et al., 2024 ) , among other methods. Some approaches use memory slots to encode information  (Al 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 M 3 method  (Yang et al., 2024 ) also stores memory in the hidden-state space, archiving a vast pretraining dataset comprising 1.1 × 10 8 1.1 superscript 10 8 1.1\times 10^{8} 1.1 × 10 start_POSTSUPERSCRIPT 8 end_POSTSUPERSCRIPT text chunks. Distinct 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, demonstrating effectiveness in knowledge-editing tasks. Furthermore, SELF-PARAM  (Wang et al., 2024c ) explores embedding knowledge directly into model parameters without degrading the model’s capabilities or requiring additional parameters. These latent-space memory techniques have shown promising results across various downstream tasks. By saving information in a compressed format and leveraging retrieval during generation, they enable substantial expansions of the context window without incurring excessive GPU memory costs. Despite the advantages and potential of Latent-Space Memory, existing methods within this category typically fall short when dealing with extremely long input  (Das et al., 2024 ; Wang et al., 2024c , a ; He et al., 2024 ) . In contrast, M+ can have much longer retention compared to existing methods. 3 Methodology 3.1 Preliminaries We first introduce the structure of MemoryLLM  (Wang et al., 2024a ) , which serves as the base structure of M+. MemoryLLM comprises two main components: θ 𝜃 \theta italic_θ (the memory pool) and ϕ italic-ϕ \phi italic_ϕ (a transformer-based decoder-only language model). The memory pool θ 𝜃 \theta italic_θ consists of L 𝐿 L italic_L layers: { θ l } l = 1 L superscript subscript subscript 𝜃 𝑙 𝑙 1 𝐿 \{\theta_{l}\}_{l=1}^{L} { italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT } start_POSTSUBSCRIPT italic_l = 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_L end_POSTSUPERSCRIPT where L 𝐿 L italic_L is the number of layers in the transformer ϕ italic-ϕ \phi italic_ϕ . For every layer, θ l subscript 𝜃 𝑙 \theta_{l} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT has N 𝑁 N italic_N memory tokens, where each token is a vector in ℝ d superscript ℝ 𝑑 \mathbb{R}^{d} blackboard_R start_POSTSUPERSCRIPT italic_d end_POSTSUPERSCRIPT , with d 𝑑 d italic_d representing the hidden size of the language model. During the update process, the last K 𝐾 K italic_K tokens from the l 𝑙 l italic_l -th layer’s memory pool, θ l subscript 𝜃 𝑙 \theta_{l} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT , are extracted and combined with the chunk to be injected. The resulting new K 𝐾 K italic_K tokens are then merged back into θ l subscript 𝜃 𝑙 \theta_{l} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT (illustrated in Figure 1 ). Merging is achieved by randomly dropping K 𝐾 K italic_K tokens from θ l subscript 𝜃 𝑙 \theta_{l} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT and appending the new K 𝐾 K italic_K tokens to the end. During generation, the memory pool θ l subscript 𝜃 𝑙 \theta_{l} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT is perceived using cross-attention. 3.2 Equipping MemoryLLM with Long-Term Memory In this section, we explain how we instantiate the long-term memory and how it integrates with the language model ϕ italic-ϕ \phi italic_ϕ and the original memory pool θ 𝜃 \theta italic_θ in MemoryLLM. In this paper, we term the original memory pool θ 𝜃 \theta italic_θ as short-term memory to distinguish it from the new long-term memory. 3.2.1 Memory Structures We denote the long-term memory as Θ Θ \Theta roman_Θ . Similarly, it has L 𝐿 L italic_L layers { Θ l } l = 1 N superscript subscript subscript Θ 𝑙 𝑙 1 𝑁 \{\Theta_{l}\}_{l=1}^{N} { roman_Θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT } start_POSTSUBSCRIPT italic_l = 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_N end_POSTSUPERSCRIPT . Each layer has a long-term memory pool where the size is flexible. We specify a maximum size for the long-term memory. The maximum size of the long-term memory is denoted as M 𝑀 M italic_M and the size of long-term memory is flexible. In practice, we choose M 𝑀 M italic_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 𝐾 K italic_K tokens are dropped from θ 𝜃 \theta italic_θ during updates and are permanently discarded. In M+, the dropped K 𝐾 K italic_K tokens are instead stored in the long-term memory Θ Θ \Theta roman_Θ , ensuring their retention for extended durations (as illustrated in Figure 1 ). We assign each token the variable “age” so that after retrieving tokens from Θ Θ \Theta roman_Θ we can sort these tokens according to age, ensuring that the tokens are chronologically ordered. As for the new K 𝐾 K italic_K tokens, they are obtained with the same process as in MemoryLLM, described in Figure 1 . When the memory tokens in the long-term memory reach the maximum capacity, i.e., M 𝑀 M italic_M tokens, we would drop the tokens with the largest ages. Generation Process During generation, at each layer, we extract K 0 subscript 𝐾 0 K_{0} italic_K start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT tokens from the long-term memory Θ l subscript Θ 𝑙 \Theta_{l} roman_Θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT using a retrieval mechanism described below, sort them by their ages, and concatenate them with the short-term memory θ l subscript 𝜃 𝑙 \theta_{l} italic_θ start_POSTSUBSCRIPT italic_l end_POSTSUBSCRIPT . 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. 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 compresses the information (similar to writing) while the generating 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 2 in Raffel et al. ( 2020 ) ). 3.2.2 Retriever Design and Training Retriever Design The retriever has two projectors: query projector f q subscript 𝑓 𝑞 f_{q} italic_f start_POSTSUBSCRIPT italic_q end_POSTSUBSCRIPT and key projector f k subscript 𝑓 𝑘 f_{k} italic_f start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT , which are all instantiated with a two-layer perceptron. The output dimension of both projectors, denoted as d p ⁢ r ⁢ o ⁢ j subscript 𝑑 𝑝 𝑟 𝑜 𝑗 d_{proj} italic_d start_POSTSUBSCRIPT italic_p italic_r italic_o italic_j end_POSTSUBSCRIPT , is set to be a small number. In our experiments, we set d p ⁢ r ⁢ o ⁢ j subscript 𝑑 𝑝 𝑟 𝑜 𝑗 d_{proj} italic_d start_POSTSUBSCRIPT italic_p italic_r italic_o italic_j end_POSTSUBSCRIPT to be d / 20 𝑑 20 d/20 italic_d / 20 where d 𝑑 d italic_d is the hidden size of the language model ϕ italic-ϕ \phi italic ... [OUTPUT TRUNCATED - 26795 chars omitted out of 76795 total] ... long-term memory to avoid any additional 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 Appendix 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. 4.6 Analysis 4.6.1 Model Quality within Context Window 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,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 CC-MAIN-2024-10 ), which does not overlap with the 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-8B. The results show that LLaMA-3.1-8B achieves a perplexity of 1.9734, while M+ records a similar 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. 4.6.2 Retrieval Quality In our 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 𝐾 256 K=256 italic_K = 256 tokens are critical for answering the questions. These K = 256 𝐾 256 K=256 italic_K = 256 tokens are denoted as 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 1 when queried after various numbers of tokens are injected. We present the results in Figure 6 , demonstrating the retrieval 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 are retrieved. For reference, random retrieval would retrieve 2 , 560 / 81 , 276 = 3 % 2 560 81 276 percent 3 2,560/81,276=3\% 2 , 560 / 81 , 276 = 3 % tokens. 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. 4.6.3 Latency Analysis While M+ introduces additional computation due to the memory token retrieval from the long-term memory, we perform a detailed analysis to quantify this latency. Specifically, we analyze latency under the setting of a 128k input. For reference, we use the processing time of Llama-3.1-8B 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 focus on the following settings: (1) Llama-3.1-8B-128k. To analyze the latency, we use Llama-3.1-8B with a full context window 128k; (2) MemoryLLM-8B (After Stage 1); (3) M+ (After Stage 3); (4) MemoryLLM-8B (Offload): we offload the memory onto CPU and load the corresponding memory tokens into GPU when the computation hits a certain layer; (5) M+ (offload): Offloading the memory onto CPU and load them back when necessary. All experiments in this section are conducted on a single H100 GPU. The results are shown in Figure 7 . From the figure, we could find that (1) MemoryLLM-8B has slightly higher latency than Llama-3.1-8B in relatively shorter documents (16k, 32k, 64k) but has lower latency on long documents (128k); (2) M+ has higher latency than MemoryLLM-8B, 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+. Figure 7: Latency Analysis 5 Conclusion and Future Work In this work, we present M+, an enhanced memory-augmented language model that extends the long-term retention abilities of MemoryLLM. By integrating a long-term memory (LTM) mechanism with a co-trained retriever, M+ effectively retrieves and utilizes past information, significantly extending the knowledge retention abilities from MemoryLLM, achieve superior performances in long-context understanding tasks compared with recent baselines given the similar budget of GPU memory. In future work, we plan to reduce CPU-GPU communication overhead, enabling more efficient generation with M+. Impact Statement This work introduces a memory-augmented 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 as education, research, and industry. The increased memory capacity could potentially raise concerns regarding AI safety, reliability, and fairness. If not carefully managed, these models could propagate biased content over extended text spans or store sensitive information for unintended durations. It is therefore crucial to employ robust safeguards, including bias mitigation strategies and ongoing oversight, to 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. References Al Adel & Burtsev (2021) Al Adel, A. and Burtsev, M. S. Memory transformer with hierarchical attention for long document processing. In 2021 International Conference Engineering and Telecommunication (En&T) , pp.  1–7. IEEE, 2021. Bai et al. (2023) 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 understanding. arXiv preprint arXiv:2308.14508 , 2023. Belcak & Wattenhofer (2024) Belcak, P. and Wattenhofer, R. Tiny transformers excel at sentence compression. arXiv preprint arXiv:2410.23510 , 2024. Bellard (2021) Bellard, F. Nncp v2: Lossless data compression with transformer. Technical report, Technical report, Amarisoft, 2021. Bulatov et al. (2022) Bulatov, A., Kuratov, Y., and Burtsev, M. S. Recurrent memory transformer. In NeurIPS , 2022. Bulatov et al. (2023) Bulatov, A., Kuratov, Y., Kapushev, Y., and Burtsev, M. S. Scaling transformer to 1m tokens and beyond with rmt. arXiv preprint arXiv:2304.11062 , 2023. Chen et al. (2024) Chen, X., Jiang, J.-Y., Chang, W.-C., Hsieh, C.-J., Yu, H.-F., and Wang, W. MinPrompt: Graph-based minimal 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 Linguistics. doi: 10.18653/v1/2024.acl-long.16 . URL https://aclanthology.org/2024.acl-long.16/ . Das et al. (2024) Das, P., Chaudhury, S., Nelson, E., Melnyk, I., Swaminathan, S., Dai, S., Lozano, A. C., Kollias, G., Chenthamarakshan, V., Navrátil, J., Dan, S., and Chen, P. Larimar: Large language models with episodic memory control. In ICML . OpenReview.net, 2024. Dubey et al. (2024) Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., Goyal, A., Hartshorn, A., Yang, A., Mitra, A., Sravankumar, 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., Allonsius, D., Song, D., Pintz, D., Livshits, D., Esiobu, D., Choudhary, D., Mahajan, D., 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, I. 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 llama 3 herd of models. CoRR , abs/2407.21783, 2024. Fedorenko et al. (2024) Fedorenko, E., Piantadosi, S. T., and Gibson, E. A. Language is primarily a tool for communication rather than thought. Nature , 630(8017):575–586, 2024. Fu et al. (2024) Fu, Y., Panda, R., Niu, X., Yue, X., Hajishirzi, H., Kim, Y., and Peng, H. Data engineering for scaling language models to 128k context. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024 . OpenReview.net, 2024. URL https://openreview.net/forum?id=TaAqeo7lUh . Ge et al. (2024) Ge, T., Hu, J., Wang, L., Wang, X., Chen, S., and Wei, F. In-context autoencoder for context compression in a 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=uREj4ZuGJE . Gutiérrez et al. (2024) 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.14831 , 2024. Hao et al. (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. He et al. (2024) He, Z., Karlinsky, L., Kim, D., McAuley, J., Krotov, D., and Feris, R. Camelot: Towards large language models with training-free consolidated associative memory. arXiv preprint arXiv:2402.13449 , 2024. Hu et al. (2023) Hu, C., Fu, J., Du, C., Luo, S., Zhao, J., and Zhao, H. Chatdb: Augmenting llms with databases as their symbolic memory. arXiv preprint arXiv:2306.03901 , 2023. Jawahar et al. (2019) 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 et al. (2019) Khandelwal, U., Levy, O., Jurafsky, D., Zettlemoyer, L., and Lewis, M. Generalization through memorization: Nearest neighbor language models. arXiv preprint arXiv:1911.00172 , 2019. Li et al. (2024) 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://doi.org/10.48550/arXiv.2404.14469 . Modarressi et al. (2024) Modarressi, A., Köksal, A., Imani, A., Fayyaz, M., and Schütze, H. Memllm: Finetuning llms to use an explicit read-write memory. arXiv preprint arXiv:2404.11672 , 2024. Packer et al. (2023) Packer, C., Fang, V., Patil, S. G., Lin, K., Wooders, S., and Gonzalez, J. E. Memgpt: Towards llms as operating systems. CoRR , abs/2310.08560, 2023. Park & Bak (2024) Park, S. and Bak, J. Memoria: Resolving fateful forgetting problem through human-inspired memory architecture, 2024. Penedo et al. (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: 10.48550/ARXIV.2406.17557 . URL https://doi.org/10.48550/arXiv.2406.17557 . Pham et al. (2024) 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 et al. (2020) Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res. , 21:140:1–140:67, 2020. URL https://jmlr.org/papers/v21/20-074.html . Rahman et al. (2024) Rahman, C. M., Sobhani, M. E., Rodela, A. T., and Shatabda, S. An enhanced text compression approach using transformer-based language models. CoRR , abs/2412.15250, 2024. Rakotonirina & Baroni (2024) Rakotonirina, N. C. and Baroni, M. Memoryprompt: A light wrapper to improve context tracking in pre-trained language models. arXiv preprint arXiv:2402.15268 , 2024. Simoulin & Crabbé (2021) Simoulin, A. and Crabbé, B. How many layers and why? an analysis of the model depth in transformers. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing: Student Research Workshop , pp.  221–228, 2021. Wang et al. (2023) 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.07174 , 2023. Wang et al. (2024a) 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. MEMORYLLM: towards self-updatable large language models. In ICML . OpenReview.net, 2024a. Wang et al. (2024b) 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. J. Towards lifespan cognitive systems. CoRR , abs/2409.13265, 2024b. Wang et al. (2024c) 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. Wang et al. (2024d) Wang, Y., Wu, R., He, Z., Chen, X., and McAuley, J. Large scale knowledge washing. arXiv preprint arXiv:2405.16720 , 2024d. Wu et al. (2022a) Wu, X., Xiao, L., Sun, Y., Zhang, J., Ma, T., and He, L. A survey of human-in-the-loop for machine learning. Future Generation Computer Systems , 135:364–381, 2022a. Wu et al. (2022b) 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 https://openreview.net/forum?id=TrjbxzRcnf- . Yang et al. (2024) 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 3 3 {}^{\mbox{3}} start_FLOATSUPERSCRIPT 3 end_FLOATSUPERSCRIPT : Language modeling with explicit memory. CoRR , abs/2407.01178, 2024. Yin et al. (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 Linguistics, 2024. Zhang et al. (2024) Zhang, X., Chen, Y., Hu, S., Xu, Z., Chen, J., Hao, M., Han, X., Thai, Z., Wang, S., Liu, Z., et al. ∞ \infty ∞ bench: Extending long context evaluation beyond 100k tokens. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pp.  15262–15277, 2024. Zhang et al. (2023) 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 generative inference of large language models. In NeurIPS , 2023. Zhong et al. (2023) Zhong, W., Guo, L., Gao, Q., and Wang, Y. Memorybank: Enhancing large language models with long-term memory. arXiv preprint arXiv:2305.10250 , 2023. Zhou et al. (2023) Zhou, W., Jiang, Y. E., Cui, P., Wang, T., Xiao, Z., Hou, Y., Cotterell, R., and Sachan, M. Recurrentgpt: Interactive generation of (arbitrarily) long text. arXiv preprint arXiv:2305.13304 , 2023. Appendix A Justifications of using deepspeed-stage-2 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-stage-2 . • 32k context window with full attention using accelerate and deepspeed-stage-3-offload . However, saving models in this configuration encountered version incompatibility issues and we haven’t found solutions online. Based on these trails, we do not scale up the model with deepspeed-stage-3-offload or FSDP, but choose to use deepspeed-stage-2 and set the cross-attention to be of the shape 2048 by 14848. Appendix 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 . Figure 8: Knowledge Retention Results on NaturalQA. B.2 Ablation Study on NaturalQA The results of ablation study on NaturalQA are shown in Figure 9 . Figure 9: Ablation Study on NaturalQA dataset. Appendix C Statistics of the Dataset of Long Documents We go through the whole dataset SlimPajama-627B and extract all dataset that have more than 4k tokens using the tokenizer of Llama-3.1-8B. The statistics are shown in Table 4 . We show six categories here (4k-8k, 8k-16k,16k-32k,32k-64k,64k-128k,128k+) but we only use the data within the first four categories (4k-8k, 8k-16k,16k-32k,32k-64k). This is because the examples longer than 64k are mainly from the category Book and lack diversity. Range Total CommonCrawl GitHub ArXiv C4 StackExch. Wikipedia Book 4k–8k 11,189,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%) 8k–16k 4,706,687 3,273,619 (69.55%) 270,369 (5.74%) 550,192 (11.69%) 439,143 (9.33%) 20,284 (0.43%) 146,545 (3.11%) 6,535 (0.14%) 16k–32k 1,607,064 968,714 (60.28%) 95,445 (5.94%) 423,401 (26.35%) 70,223 (4.37%) 1,510 (0.09%) 34,323 (2.14%) 13,448 (0.84%) 32k–64k 443,438 224,168 (50.55%) 32,653 (7.36%) 146,582 (33.06%) 3,413 (0.77%) 102 (0.02%) 5,940 (1.34%) 30,580 (6.90%) 64k–128k 192,515 72,583 (37.70%) 11,753 (6.10%) 27,942 (14.51%) 38 (0.02%) 5 (0.00%) 507 (0.26%) 79,687 (41.39%) 128k+ 98,097 23,721 (24.18%) 4,523 (4.61%) 5,167 (5.27%) 0 (0.00%) 2 (0.00%) 49 (0.05%) 64,635 (65.89%) Table 4: Number of examples by sequence-length range and source (counts and percentages). Appendix D Additional Training Details In our training, we follow MemoryLLM  (Wang et al., 2024a ) and design three sub-tasks: • Two-Chunk Training : Given a document split into two chunks ( x 1 , x 2 ) subscript 𝑥 1 subscript 𝑥 2 (x_{1},x_{2}) ( italic_x start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT , italic_x start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT ) , we inject x 1 subscript 𝑥 1 x_{1} italic_x start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT into the memory and update the transformer ϕ italic-ϕ \phi italic_ϕ using the loss on x 2 subscript 𝑥 2 x_{2} italic_x start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT . Notably, we retain the gradients across both forward passes. • Multi-Chunk Training : For documents with multiple chunks ( x 1 , … , x n ) subscript 𝑥 1 … subscript 𝑥 𝑛 (x_{1},\dots,x_{n}) ( italic_x start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT , … , italic_x start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT ) , we inject x 1 , … , x n − 1 subscript 𝑥 1 … subscript 𝑥 𝑛 1 x_{1},\dots,x_{n-1} italic_x start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT , … , italic_x start_POSTSUBSCRIPT italic_n - 1 end_POSTSUBSCRIPT into the memory while detaching gradients, then update ϕ italic-ϕ \phi italic_ϕ using the loss on x n subscript 𝑥 𝑛 x_{n} italic_x start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT . • Revisiting Cached Chunks : Since the memory is continually updated during training, we cache the last chunk x n subscript 𝑥 𝑛 x_{n} italic_x start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT of earlier documents and revisit it periodically. When revisiting x n subscript 𝑥 𝑛 x_{n} italic_x start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT , there are already many chunks injected between x 1 , ⋯ , x n − 1 subscript 𝑥 1 ⋯ subscript 𝑥 𝑛 1 x_{1},\cdots,x_{n-1} italic_x start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT , ⋯ , italic_x start_POSTSUBSCRIPT italic_n - 1 end_POSTSUBSCRIPT and x n subscript 𝑥 𝑛 x_{n} italic_x start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT . We denote the number of injected chunks between x n − 1 subscript 𝑥 𝑛 1 x_{n-1} italic_x start_POSTSUBSCRIPT italic_n - 1 end_POSTSUBSCRIPT and x n subscript 𝑥 𝑛 x_{n} italic_x start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT 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. Appendix E Discussions E.1 Similarities to Attention-Based 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 advantageous: • 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. • 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. • 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. E.2 The Form of Long-Term Memory (Hidden States vs. KV) In our work, we choose the use hidden states as the latent-space memory instead of key-value (KV) caches. This is 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 a lossless manner. In contrast, KV-based methods often require downsampling—e.g., dropping half the keys and values—to control memory size, resulting in unavoidable information loss. • Retrieval Efficiency and Performance : As described above, 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. 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, the end-to-end retrieval latency scales linearly with three key variables: (1) Hidden size of the retriever, denoted by d 𝑑 d italic_d . In M+ , we set d = 256 𝑑 256 d=256 italic_d = 256 , whereas the base model uses d = 4096 𝑑 4096 d=4096 italic_d = 4096 . (2) Size of long-term memory, denoted by s 𝑠 s italic_s . We cap this at 150k entries. (3) Number of transformer layers, denoted by L 𝐿 L italic_L . For LLaMA-3-8B , L = 32 𝐿 32 L=32 italic_L = 32 . Hence, latency ∝ d ⁢ s ⁢ L . proportional-to latency 𝑑 𝑠 𝐿 \text{latency}\;\propto\;d\,s\,L. latency ∝ italic_d italic_s italic_L . Because we hold the long-term memory size s 𝑠 s italic_s fixed when scaling the model, s 𝑠 s italic_s is effectively a constant: latency ∝ d ⁢ L . proportional-to latency 𝑑 𝐿 \text{latency}\;\propto\;d\,L. latency ∝ italic_d italic_L . Both d 𝑑 d italic_d and L 𝐿 L italic_L grow with the model size M 𝑀 M italic_M , following M ∝ d ⁢ L , proportional-to 𝑀 𝑑 𝐿 M\;\propto\;d\,L, italic_M ∝ italic_d italic_L , which implies a linear relationship between retrieval latency and model size: latency ∝ M . proportional-to latency 𝑀 \text{latency}\;\propto\;M. latency ∝ italic_M . Concrete Example Scaling from LLaMA-3-8B ( d = 4096 𝑑 4096 d=4096 italic_d = 4096 , L = 32 𝐿 32 L=32 italic_L = 32 ) to LLaMA-3-70B ( d = 8192 𝑑 8192 d=8192 italic_d = 8192 , L = 80 𝐿 80 L=80 italic_L = 80 ) yields 8192 × 80 4096 × 32 =  5 , 8192 80 4096 32 5 \frac{8192\times 80}{4096\times 32}\;=\;5, divide start_ARG 8192 × 80 end_ARG start_ARG 4096 × 32 end_ARG = 5 , i.e. a ∼ 5 × \sim\!5\times ∼ 5 × increase in retrieval latency. For comparison, the parameter count rises by 70 ⁢ B 8 ⁢ B ≈ 8.75 × , \frac{70\text{B}}{8\text{B}}\approx 8.75\times, divide start_ARG 70 B end_ARG start_ARG 8 B end_ARG ≈ 8.75 × , showing that latency scales roughly linearly—rather than quadratically—with model size. Memory Consumption The extra memory overhead from our method arises solely from the introduced memory tokens. This overhead also scales linearly with both d 𝑑 d italic_d and L 𝐿 L italic_L ; thus, the move from LLaMA-3-8B to LLaMA-3-70B incurs an analogous ∼ 5 × \sim\!5\times ∼ 5 × increase in memory usage, mirroring the latency scaling. E.4 FLOPs Comparison We report the total FLOPs for generating one token after processing a sequence of varying lengths (from 2k to 128k), using a single H100 GPU. We employ the torch.profiler library to capture FLOPs during inference. The results are as follows: Sequence Length LLaMA-3.1-8B M+ 2048 5.68 × 10 13 5.68 superscript 10 13 5.68\times 10^{13} 5.68 × 10 start_POSTSUPERSCRIPT 13 end_POSTSUPERSCRIPT 6.92 × 10 13 6.92 superscript 10 13 6.92\times 10^{13} 6.92 × 10 start_POSTSUPERSCRIPT 13 end_POSTSUPERSCRIPT 4096 1.13 × 10 14 1.13 superscript 10 14 1.13\times 10^{14} 1.13 × 10 start_POSTSUPERSCRIPT 14 end_POSTSUPERSCRIPT 1.32 × 10 14 1.32 superscript 10 14 1.32\times 10^{14} 1.32 × 10 start_POSTSUPERSCRIPT 14 end_POSTSUPERSCRIPT 8192 2.26 × 10 14 2.26 superscript 10 14 2.26\times 10^{14} 2.26 × 10 start_POSTSUPERSCRIPT 14 end_POSTSUPERSCRIPT 2.55 × 10 14 2.55 superscript 10 14 2.55\times 10^{14} 2.55 × 10 start_POSTSUPERSCRIPT 14 end_POSTSUPERSCRIPT 16384 4.48 × 10 14 4.48 superscript 10 14 4.48\times 10^{14} 4.48 × 10 start_POSTSUPERSCRIPT 14 end_POSTSUPERSCRIPT 5.01 × 10 14 5.01 superscript 10 14 5.01\times 10^{14} 5.01 × 10 start_POSTSUPERSCRIPT 14 end_POSTSUPERSCRIPT 32768 8.88 × 10 14 8.88 superscript 10 14 8.88\times 10^{14} 8.88 × 10 start_POSTSUPERSCRIPT 14 end_POSTSUPERSCRIPT 9.86 × 10 14 9.86 superscript 10 14 9.86\times 10^{14} 9.86 × 10 start_POSTSUPERSCRIPT 14 end_POSTSUPERSCRIPT 65536 1.75 × 10 15 1.75 superscript 10 15 1.75\times 10^{15} 1.75 × 10 start_POSTSUPERSCRIPT 15 end_POSTSUPERSCRIPT 1.94 × 10 15 1.94 superscript 10 15 1.94\times 10^{15} 1.94 × 10 start_POSTSUPERSCRIPT 15 end_POSTSUPERSCRIPT 131072 OOM 3.78 × 10 15 3.78 superscript 10 15 3.78\times 10^{15} 3.78 × 10 start_POSTSUPERSCRIPT 15 end_POSTSUPERSCRIPT From the results, we observe that M+ and LLaMA-3.1-8B exhibit comparable FLOPs across all tested sequence lengths. Notably, while LLaMA-3.1-8B runs out of memory ( OOM ) 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-term 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 too quickly. Generated on Fri May 30 14:40:50 2025 by L a T e XML