M+: Extending MemoryLLM with Scalable Long-Term Memory
M+ extends MemoryLLM with a long-term memory mechanism, enabling over 160k tokens retention, significantly improving long-text understanding.
Key Findings
Methodology
The proposed M+ model builds upon MemoryLLM by integrating a long-term memory module stored on CPU, combined with a jointly trained retriever utilizing a dual-projector architecture. The retriever employs dot-product similarity for efficient matching, trained via contrastive loss to maximize relevance. The long-term memory stores dropped tokens with age annotations, sorted chronologically. During training, a multi-stage curriculum progressively enhances the model’s ability to handle longer contexts, involving short-text, long-text, and long-term memory phases. The model leverages LoRA for parameter efficiency. During inference, the retriever fetches relevant long-term tokens, which are concatenated with short-term memory, enabling the model to process over 160k tokens while maintaining GPU efficiency.
Key Results
- On LongBook-QA and LongBook-Event-QA benchmarks, M+ achieves QA-F1 scores of 85.3% and accuracy of 78.9%, surpassing MemoryLLM’s 70.2% and 62.5%, respectively. It maintains this performance with only a 10% increase in GPU memory cost. The model effectively retrieves and retains information over 160k tokens, demonstrating superior long-term memory capabilities.
- In knowledge retention tasks on SQuAD and NaturalQA, M+ improves accuracy by over 15%, confirming its ability to recall long-term knowledge. The CPU-based long-term storage extends the memory span without significant GPU overhead.
- Ablation studies show that joint training of the retriever and model yields better results than separate or static retrieval methods, validating the design choices.
Significance
This work addresses a fundamental bottleneck in large language models: the inability to retain and utilize information over very long sequences. By combining a scalable long-term memory with an efficient retrieval mechanism, M+ significantly extends the effective context window beyond 20k tokens, enabling applications in long document comprehension, reasoning, and knowledge management. It bridges the gap between model capacity and practical long-term memory, opening new avenues for AI systems in legal, scientific, and educational domains. The approach also offers a blueprint for integrating external memory with end-to-end trainable retrieval, impacting future research and industry deployment.
Technical Contribution
The key technical innovation lies in the integration of a CPU-stored long-term memory with a jointly trained, dual-projector retriever, enabling efficient retrieval over hundreds of thousands of tokens. The model employs a multi-stage curriculum, combining short and long-text training, and leverages LoRA for parameter-efficient adaptation. The retrieval mechanism is optimized for low latency and high accuracy, with a focus on maintaining GPU memory efficiency. This architecture allows handling of over 160k tokens, a significant leap over prior latent-space memory methods, and provides a scalable framework for long-context modeling.
Novelty
This study is the first to combine a long-term CPU-based memory with a jointly trained retriever in a large language model, effectively overcoming the 20k token recall limit of previous latent-space memory approaches. The innovative use of age-based sorting and multi-stage curriculum training enables the model to retain and access information over unprecedented lengths, setting a new standard for long-text understanding in LLMs.
Limitations
- Despite improvements, retrieval efficiency may degrade with extremely long sequences (>200k tokens), requiring further optimization of retrieval algorithms and storage structures.
- CPU-based long-term memory introduces I/O latency, potentially affecting real-time inference in some scenarios.
- Training on very long documents demands substantial computational resources and high-quality long-text datasets, which may limit scalability and generalization.
Future Work
Future research will focus on optimizing retrieval algorithms, possibly incorporating sparse attention or hierarchical retrieval to further scale the context window. Exploring multi-modal long-term memory, integrating visual or structured data, is also promising. Hardware acceleration and distributed architectures could reduce latency and costs, enabling real-time applications in industry. Additionally, extending the framework to multilingual and domain-specific models will broaden its impact.
AI Executive Summary
The rapid development of large language models (LLMs) has revolutionized natural language processing, yet their capacity to remember and reason over very long texts remains limited. Traditional models struggle beyond 20,000 tokens, which hampers their application in domains requiring extensive context understanding, such as legal documents, scientific papers, and historical records. Existing solutions, including token-level memory and latent-space embeddings, offer some relief but often face trade-offs between efficiency, interpretability, and scalability.
In response, this paper introduces M+, a novel architecture that significantly extends the memory capabilities of models like MemoryLLM. Building upon the latent-space memory paradigm, M+ incorporates a long-term memory module stored on CPU, combined with a jointly trained retriever that efficiently fetches relevant information during inference. This retriever employs a dual-projector design, optimized through contrastive learning, to match query embeddings with stored long-term tokens. The long-term memory is managed via an age-based sorting mechanism, ensuring chronological relevance and preventing information decay.
The training process involves multiple stages, starting from short texts and gradually scaling to longer documents, including a dedicated phase for integrating long-term memory. Leveraging LoRA for parameter efficiency, the model is trained on datasets ranging from short articles to multi-thousand-word books. Extensive experiments on benchmarks like LongBook-QA and LongBook-Event-QA demonstrate that M+ outperforms existing models, achieving over 85% QA-F1 and 78% accuracy, while maintaining similar GPU memory costs. The model effectively retains knowledge over 160k tokens, a substantial leap from prior methods.
This advancement opens new horizons for AI applications requiring deep, long-term reasoning and knowledge retention. It paves the way for smarter, more context-aware systems in legal, scientific, and educational fields. Despite these breakthroughs, challenges remain in optimizing retrieval efficiency for extremely long sequences and reducing I/O latency. The authors suggest future directions including algorithmic improvements, hardware acceleration, and multi-modal integration, promising a vibrant research trajectory that could redefine long-text AI capabilities.
Deep Dive
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