Research map / Architectures and efficiency / Attention and transformers
M+: Extending MemoryLLM with Scalable Long-Term Memory
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Extends MemoryLLM with a CPU-offloaded long-term hidden-state memory and a co-trained query/key retriever, retaining knowledge from under 20k to over 160k tokens at similar GPU memory.
Problem
Latent-space memory models like MemoryLLM lose information injected beyond roughly 20k tokens, limiting long-term knowledge retention and long-context understanding.
Approach
Dropped memory tokens are moved into a per-layer long-term memory (stored on CPU) instead of being discarded, and a small co-trained retriever (query/key projectors) selects a fixed number of them per layer at generation time, so retrieval happens once per layer rather than per head as in KV-cache methods.
Key results
- Knowledge retention extended from under 20k tokens (MemoryLLM) to over 160k tokens on SQuAD, per the paper's stage-3 analysis.
- On LongBench at 16k, M+ averages 32.78 QA-F1 versus 35.81 for Llama3.1-8B (16k) and 24.08 for MemoryLLM-7B (20k).
- M+ (offload) uses 17973 MB peak GPU memory versus 19239 MB for Llama-3.1-8B-16k and 32574 MB for Llama-3.1-8B-SnapKV.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add long-document training in Stage 2 (4k–64k examples) (ablation) | MemoryLLM-8B · LongBench (8k) | QA-F1 Avg: 26.55 -> 31.29 | Helped |
| add long-term memory in Stage 3 (ablation) | MemoryLLM-8B-Long · LongBench (8k) | QA-F1 Avg: 31.29 -> 31.00 | No clear effect |
| add long-term memory in Stage 3 (ablation) | MemoryLLM-8B-Long · SQuAD and NaturalQA knowledge retention | knowledge retention | Helped |
| use the co-trained retriever instead of attention-score retrieval (ablation) | M+ · SQuAD and NaturalQA | knowledge retention | Helped |
| offload memory vectors to CPU (ablation) | M+ · Inference across both datasets | GPU memory cost (MB): 21177.76 -> 17973.34 | Helped |
| add long-term memory and co-trained retrieval | M+ · LongBook-QA and LongBook Event QA | QA-F1 / accuracy | Helped |
| use SnapKV attention-score selection instead of the original 16k context | Llama-3.1-8B · LongBook-QA | QA-F1 | Hurt |
| use M+ memory compression and retrieval instead of full-context processing | Llama-3.1-8B · LongBench (8k), Avg | QA-F1: 33.39 -> 31.00 | Hurt |
| use M+ memory compression and retrieval instead of full-context processing | Llama-3.1-8B · LongBench (16k), Avg | QA-F1: 35.81 -> 32.78 | Hurt |
| add memory augmentation | Llama-3.1-8B · FineWeb-edu documents capped at 2,048 tokens | perplexity: 1.9734 -> 1.9828 | No clear effect |
Surprising findings
- Llama-3.1-8B-SnapKV underperforms Llama-3.1-8B-16k on LongBook-QA, suggesting attention-score selection of key tokens is not reliably optimal.
- M+ underperforms the Llama-3.1-8B baseline on short-document LongBench tasks hotpotqa and musique, attributed to random dropping and limited cross-chunk attention.
Limitations
- Limited cross-chunk attention and random dropping cause performance drops on relatively short-context tasks (hotpotqa, musique).
- CPU-GPU communication adds latency during generation, which the authors plan to reduce.
- Long-context experiments are limited to a 16k generation/memory context due to GPU constraints; 128k scaling not trained.
Experimental setup
Datasets: fineweb-edu SlimPajama SQuAD NaturalQuestions LongBench InfiniteBench LongBook-QA
Benchmarks: InfiniteBench LongBook-QA LongBook Event QA SQuAD knowledge retention NaturalQA knowledge retention LongBench
Metrics: QA-F1 accuracy validation loss perplexity GPU memory latency FLOPs
Baselines: MemoryLLM-7B Llama-3.1-8B-16k Llama-3.1-8B-SnapKV Llama-3.2-3B-128k Llama-3.1-8B-BM25 M+-Attn
Builds on: MemoryLLM SnapKV H2O T5
Closest papers by meaning
- MEMORYLLM: Towards Self-Updatable Large Language Models (ICML 2024 · similarity 0.92)
- A Little Goes a Long Way: Efficient Long Context Training and Inference with Partial Contexts (ICLR 2025 · similarity 0.89)
- LightMem: Lightweight and Efficient Memory-Augmented Generation (ICLR 2026 · similarity 0.89)
- Peripheral Memory for LLMs: Integration of Sequential Memory Banks with Adaptive Querying (ICML 2025 · similarity 0.88)
- RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval (NeurIPS 2025 · similarity 0.88)
- ChunkKV: Semantic-Preserving KV Cache Compression for Efficient Long-Context LLM Inference (NeurIPS 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.