Memory Decoder at Scale: A Pretrained, Parametric Long-Term Memory
Rubin Wei, Jiaqi Cao, Jiarui Wang, Junming Zhang, Qipeng Guo, Bowen Zhou, Zhouhan Lin
cs.CL
2026-07-30
Scaling the parametric memory module to 6.9B parameters and 300B tokens lets a frozen Pythia-410M paired with a 6.9B memory beat Pythia-12B across 17 benchmarks, with 39% fewer total parameters.
Standard decoder-only language models entangle long-term memory and reasoning in one set of parameters. Memory cannot be scaled independently; to make a model remember more you grow the whole parameter set. Domain adaptation requires full fine-tuning, which is expensive and risks catastrophic forgetting, and there is no swappable, reusable memory module. Existing fixes have gaps. RAG retrieves external passages as context at inference, and kNN-LM interpolates predictions with a nearest-neighbor distribution from an external datastore; both need retrieval at inference, with extra context or heavy index and search costs. Memory Decoder (prior work) introduced a parametric memory module but only at about 1B parameters and millions of tokens, leaving larger scale untested. This work pushes it to language-model pretraining scale.
The memory module is a transformer decoder trained to imitate the output distribution pret of a kNN retriever rather than a single next token. The per-position loss is beta times the KL divergence of pret and the module distribution plus (1-beta) times a cross-entropy language-modeling term: the KL term pulls the memory toward the retrieval distribution, and the language-modeling term keeps it close to the corpus. Retrieval is used only to build offline supervision, never at inference.
At inference, the frozen backbone and the memory process the same context in parallel, and their next-token distributions are interpolated as pfinal equal to (1-alpha) times the backbone plus alpha times the memory, where alpha controls the memory contribution.
The hard part is scale. To train a 6.9B memory on 300B tokens, the authors must build a kNN distribution for every training context over the deduplicated Pile (207B tokens). At 207 million entries, the combined indexing and search cost makes a standard Faiss pipeline infeasible. Their fix is a distributed Faiss pipeline: embedding compression (OPQ256, 4096 to 256 dimensions), index sharding, and parallel search, plus sparse kNN distribution storage and batch-wise distributed streaming that reads only the entries each batch needs. General memory trains on the deduplicated Pile with a datastore from Pythia-6.9B; domain memories cover biology, law, and finance at 1.7B, built from a Qwen3-4B-Base datastore. Compute is 256 A800 GPUs.
General memory on frozen Pythia backbones across 17 benchmarks raises the average at every scale with an equal-size memory: 1.4B backbone 32.76 to 34.36, 2.8B 33.89 to 35.49, 6.9B 36.30 to 37.79. The matched-budget comparison matters more: a 1.4B backbone with a 1.4B memory (34.36) beats a 2.8B backbone alone (33.89) at the same total parameters and training budget, and 6.9B with 6.9B (37.79) beats a frozen 12B backbone (37.24). Across 51 task-by-scale combinations, memory improves 47 and ties one. Knowledge tasks gain most: TriviaQA on the 2.8B backbone rises 8.30 to 17.11, 2WikiMultiHopQA on 1.4B rises 16.89 to 22.57.
| Config | Total params | AVG (17 tasks) |
| Pythia-410M alone | 410M | 29.86 |
| Pythia-410M + Mem-6.9B | 7.3B | 37.34 |
| Pythia-12B alone | 12B | 37.24 |
The headline result: a 410M backbone with a 6.9B memory reaches 37.34, beating the 12B backbone at 37.24 with 39% fewer total parameters. At matched average score, backbone-plus-memory uses 33%, 32%, and 42% fewer parameters than the 2.8B, 6.9B, and 12B models, and because backbone and memory run in parallel, inference latency can be lower than a single model of the same total size. For domain memory, a 1.7B module lifts Qwen3 from 0.6B to 14B by 9 to 10 average points across biology, law, and finance, beating the strongest of CPT, LoRA, or RAG at every scale by at least 4.05 points; switching domains means swapping the memory while the backbone stays frozen. Transferred across vocabularies to OLMo with 20% of the budget, OLMo-2-7B and OLMo-3-7B gain 4.26 and 7.77 points.
It offers a path to better performance that is distinct from growing the backbone: memory and reasoning can be decoupled, with memory pretrained, scaled, and swapped per domain independently. For domain adaptation, you change domains without touching the backbone or full fine-tuning, sidestepping catastrophic forgetting. The parameter-efficiency story is strong: a small backbone with a large memory saves 30 to 40 percent of parameters at matched score.
The authors note that although inference needs no retrieval, the offline cost of building kNN target distributions remains, a preprocessing burden that grows with corpus scale. The memory weight alpha is fixed at inference, not adaptive per input. Joint or staged training of memory and backbone is unexplored; it might improve coordination but would raise cost and could weaken cross-backbone transfer. A question the paper does not answer head-on: because the memory compresses external knowledge into parameters, it is unclear how it compares with RAG on knowledge freshness and auditability, and the benchmarks emphasize general capability rather than time-sensitive settings.