PaperlayerView exact revision on arXiv

Source-linked AI summary

Memory Decoder at Scale: A Pretrained, Parametric Long-Term Memory

Rubin Wei, Jiaqi Cao, Jiarui Wang, Junming Zhang, Qipeng Guo, Bowen Zhou, Zhouhan Lin

arXiv:2607.27919v1cs.CL

TL;DR

Decoder-only language models entangle long-term memory and reasoning, limiting independent memory scaling. Memory Decoder at Scale pretrains independently scalable memories and finds that larger memories paired with smaller backbones improve benchmark performance more efficiently, including surpassing Pythia-12B with 39% fewer parameters.

  • Problem

    Decoder-only language models entangle long-term memory and reasoning, preventing memory from being pretrained or scaled independently.

  • Method

    The paper scales parametric memories to 6.9B parameters and 300B tokens using distributed Faiss retrieval and sparse, batch-wise kNN storage.

  • Results

    A 6.9B memory paired with Pythia-410M raises average benchmark performance from 29.86 to 37.34, surpassing Pythia-12B at 37.24 with 39% fewer parameters.

  • Takeaways & Limitations

    Independently scaling pretrained memory offers a more parameter-efficient path to improving language model performance than scaling the base model alone.

  • Takeaways & Limitations

    Memory pretraining still incurs an offline indexing and retrieval burden that grows with corpus scale.

Abstract

from arXiv · show

Decoder-only language models entangle long-term memory and reasoning in a single parameter set, making it difficult to scale memory capacity independently. Memory Decoder introduces a parametric long-term memory module but only studies it at a relatively small scale. In this work, we present Memory Decoder at Scale, scaling memory models up to 6.9B parameters and pretraining them on 300B tokens. At this data scale, the combined cost of indexing and search makes a standard Faiss pipeline infeasible. We address this bottleneck with a distributed pipeline for Faiss indexing and retrieval, together with sparse, batch-wise loading of kNN distributions. Across model scales, we find that allocating more parameters to memory yields a better parameter-performance tradeoff than scaling the base model alone. On 17 benchmarks, pairing a 6.9B general memory with Pythia-410M raises its average score from 29.86 to 37.34, surpassing Pythia-12B (37.24) with 39% fewer total parameters. For Qwen3 Base models ranging from 0.6B to 14B, 1.7B domain memories improve the average score across the three domains by more than 9 points at every scale. Overall, our results demonstrate that independently scaling pretrained memory offers a more parameter efficient path to improving language model performance.

1. Introduction

Memory Decoder at Scale addresses the entanglement of long-term memory and reasoning in decoder-only language models by scaling parametric memory pretraining to 6.9B parameters and 300B tokens. It also introduces a distributed Faiss pipeline to overcome indexing and search costs that make standard processing infeasible at this scale.

  • Motivation: Decoder-only language models entangle long-term memory and reasoning within a single parameter set, unlike the brain’s partially distinct memory and reasoning systems.This entanglement makes memory storage and cognitive computation difficult to dissociate in standard language models.
  • Related work: Existing short-term-memory methods improve contextual access during inference but do not resolve long-term-memory and reasoning entanglement.The cited methods include segment-level recurrence, bounded KV caches, distant-context retrieval, and RoPE rescaling.
  • Related work: RAG and kNN language models provide long-term memory through external retrieval, but incur context-processing, storage, and nearest-neighbor search costs.RAG retrieves passages as additional context, while kNN language models interpolate predictions with nearest-neighbor distributions from an external datastore.
  • Contribution: 6.9B parameters and 300B tokens define the scale of the general-memory pretraining introduced by Memory Decoder at Scale.The work scales parametric memory pretraining to language-model pretraining scale.
  • Contribution: 207B tokens in the deduplicated Pile make standard Faiss processing infeasible because indexing and search costs combine at this data scale.The paper addresses this bottleneck with a distributed Faiss pipeline.

2. Preliminary: Memory Decoder

Memory Decoder is a standalone parametric memory trained to imitate a non-parametric retriever using richer next-token distribution supervision. It combines retrieval alignment with corpus language modeling and interpolates memory and base-model predictions at inference.

  • Memory Decoder: Memory Decoder trains a standalone parametric memory to imitate a non-parametric retriever over corpus-context keys and observed next-token values.A frozen base model encodes each preceding context as the retrieval query, which produces a kNN distribution.
  • Memory Decoder: Retrieval provides richer supervision than a single observed target by capturing the diversity of plausible continuations.The memory decoder predicts directly from the original context and does not take retrieved text as input; retrieval only constructs offline pretraining supervision.
  • Memory Decoder: The training objective combines KL-based alignment to the kNN distribution with corpus language modeling.The KL term encourages imitation of retrieval behavior, while the language-modeling term limits deviation from the underlying corpus distribution; β∈[0, 1] controls retrieval supervision.
  • Memory Decoder: At inference, the base model and memory process the same context in parallel, and their output distributions are interpolated.α∈[0, 1] controls the memory contribution, incorporating the retrieval behavior learned during memory pretraining.

3. Memory Decoder at Scale

Memory Decoder at Scale extends parametric memory pretraining to billion-parameter memories and 300B tokens, requiring distributed Faiss construction and sparse, batch-wise retrieval loading. Its pipeline addresses quadratic indexing and search costs while supporting general and domain memories at pretraining scale.

  • Distributed Faiss pipeline: 207B contextual representations make standard Faiss indexing infeasible because constructing N searches over N keys is quadratic and keys are 4096-dimensional.The distributed pipeline addresses this joint indexing and search bottleneck at pretraining scale.
  • Distributed Faiss pipeline: OPQ256 compresses 4096-dimensional keys to 256-dimensional search vectors, while IVF-HNSW sharding routes queries to parallel local IndexIVFPQ searches.Writer workers merge shard candidates into final K-nearest-neighbor sets and cache sparse retrieval distributions.
  • Memory pretraining: 1.4B, 2.8B, and 6.9B general memories are pretrained on 300B tokens using sparse distribution storage and distributed streaming.Workers read only the entries required by each batch.
  • Sparse distribution storage: O(N+ M) sparse storage replaces O(N|V|) dense storage; Pile rows retain 64.95 token-probability pairs on average, yielding an approximately 250× dense-to-sparse storage ratio.Duplicate targets are aggregated and negligible probabilities removed before storing retained identifiers and probabilities in sharded arrays.
  • Domain memories: Domain memories use the same architecture, objective, and sparse training pipeline for biology, law, and finance corpora containing up to 4.4B tokens.Each domain receives a separate memory supervised by distributions constructed from continued-pretrained domain representations.

4. Experimental Setup

The experiments evaluate parametric memory as a plug-and-play long-term memory across general, transfer, and domain settings. They use frozen backbones, varied model families and scales, and large-scale pretraining and retrieval infrastructure.

  • Overview: Four settings test general-memory gains, transfer across backbone scales, domain-memory effectiveness, and transfer across model families and vocabularies.The evaluated domains are biology, law, and finance.
  • Overview: The base model remains frozen, and memory predictions combine with the backbone distribution using the interpolation rule in Eq. 2.
  • Datasets: 207B tokens from the deduplicated Pile train the general memory, which is evaluated on general and knowledge-intensive task categories.The supplied task list includes Pythia-suite benchmarks and knowledge-intensive question-answering tasks.
  • Backbones and Baselines: 410M to 12B parameter frozen Pythia backbones pair with 1.4B, 2.8B, and 6.9B general memories, while Qwen3 Base backbones span 0.6B to 14B parameters.
  • Implementation Details: 256 NVIDIA A800 80GB GPUs support training; general memories use 300B tokens, while domain memories receive one epoch of continued pretraining per domain.General memories use peak learning rates of 3 × 10−4, 2.5 × 10−4, and 2 × 10−4 for 1.4B, 2.8B, and 6.9B memories, respectively.

5. Experimental Results

The experiments show that pretrained general and domain memories consistently improve frozen backbones across model scales and tasks. Memory also provides parameter-efficient scaling and transfers across vocabularies with limited additional training.

  • General memory scaling: AVG rises from 32.76 to 34.36, 33.89 to 35.49, and 36.30 to 37.79 when equal-size memories augment 1.4B, 2.8B, and 6.9B Pythia backbones.The largest gains concentrate on knowledge tasks such as TriviaQA, 2WikiMultiHopQA, and HotpotQA.
  • General memory scaling: 47 of 51 task-scale combinations improve with memory, while one additional combination matches its base model.The improvements extend beyond knowledge tasks, including WinoGrande, GPQA-main, and HaluEval.
  • Parameter efficiency: 33%, 32%, and 42% fewer total parameters achieve matched AVG than 2.8B, 6.9B, and 12B base models, respectively.The experiments use reusable memories trained with the same 300B-token budget across parameter allocations.
  • Domain memory specialization: 9.88, 9.64, 10.00, 9.09, and 9.99 points are the average gains from 1.7B domain memories on 0.6B, 1.7B, 4B, 8B, and 14B Qwen3 backbones.Domain memory achieves the highest average across biology, law, and finance at every evaluated scale.
  • Cross-vocabulary transfer: 4.26 and 7.77 points are the average gains after transferring domain memories to OLMo-2-7B and OLMo-3-7B using 20% of the standard training budget.Average scores rise from 19.57 to 23.83 and from 18.67 to 26.44, respectively, while memory improves all six evaluations.

6. Analysis

The analysis shows that larger and longer-trained memories improve performance, reinforce factual evidence, outperform matched attached CPT modules, and more reliably retain and reproduce training content and retrieval targets. These effects hold across general and domain settings and support independently scaling parametric memory.

  • In-context examples: Memory improves AVG by 1.22–1.87 points across zero-shot, three-shot, and five-shot settings on the same 13 tasks.Gains are 1.43–1.87 points zero-shot, 1.39–1.52 with three shots, and 1.22–1.62 with five shots.
  • Memory capacity: The 6.9B memory produces the largest AVG gains for every tested Pythia backbone: 4.56, 3.67, and 1.49 points for 1.4B, 2.8B, and 6.9B backbones.Every tested memory size improves AVG over its corresponding frozen base, and gains grow overall with memory capacity.
  • Memory capacity: A 1.7B domain memory outperforms a 0.6B memory for every Qwen3 backbone, adding up to 18.5 points on BioInst and 9.4 points on LawBench.Increasing capacity from 0.6B to 1.7B adds up to 4.8 BioInst points and 4.4 LawBench points.
  • Training budget: Full training budgets improve Avg for every OLMo backbone and memory-size combination, while general-memory training from 1.5 to 5 epochs raises general-task average from 56.67 to 57.33.The knowledge average also rises from 13.99 to 14.17.
  • Training mechanism: Memory exceeds matched attached CPT by 10.21 and 8.71 BioInst points and by 1.68 and 3.93 LawBench points for 1.7B and 8B backbones, respectively.The control matches training data, FLOPs, capacity, and inference interface, attributing the advantage to the memory objective.
  • Memorization and retrieval alignment: The memory model raises strict EM@8, 16 from 42.4% to 49.7% and domain anchor completion from 22.6% to 56.5%, indicating stronger training-data traceability.On 1,024 shared BioInst training prompts, it also aligns with retrieval targets: 86.62% mode matches, Pearson r=0.9174, mean KL=0.1820, and mean TV=0.0823.

7. Related Work

Related work distinguishes memory augmentation during pretraining or inference, short-term retention across long contexts, and non-parametric or parametric long-term memory. Representative methods include RETRO, TRIME, Transformer-XL, StreamingLLM, kNN-LM, and RAG.

  • Memory-augmented pretraining: Memory-augmented pretraining incorporates memory into training, unlike the predominantly inference-time focus of prior language-model memory work.RETRO conditions autoregressive generation on database chunks retrieved through an encoder and cross-attention, while TRIME trains with accessible in-batch examples.
  • Short-term memory: Short-term memory methods improve retention and access over long input contexts.Transformer-XL uses segment-level recurrence and relative positional encoding, while StreamingLLM retains initial-token and recent-token KV states.
  • Long-term memory: Long-term memory is implemented non-parametrically or parametrically, with non-parametric methods storing knowledge externally and retrieving it at inference.kNN-LM retrieves similar training contexts to construct a next-token distribution, and RAG retrieves passages to condition generation.

8. Conclusion

Memory Decoder at Scale demonstrates that pretrained long-term memory can be scaled independently from reasoning, using distributed retrieval infrastructure to support models up to 6.9B parameters and 300B-token pretraining. Across model and domain scales, independently scaling memory improves parameter efficiency and language model performance.

  • Contributions: 6.9B parameters and 300B tokens define the largest memory models and pretraining scale studied in this work.The system uses a distributed Faiss pipeline to support memory pretraining at this scale.
  • Scalable retrieval: 207B corpus tokens require distributed Faiss indexing and retrieval with embedding compression, index sharding, and parallel search.Sparse kNN storage and distributed streaming load only batch-required entries when constructing retrieval supervision.
  • Parameter efficiency: 39% fewer total parameters let Pythia-410M with a 6.9B general memory surpass Pythia-12B.Experiments consistently favor pairing small base models with large memory models over scaling the base model alone.
  • Parameter efficiency: More than 9 points of average improvement across biology, law, and finance arise from 1.7B domain memories at every Qwen3 Base scale from 0.6B to 14B.The result holds across the full range of evaluated Qwen3 Base model sizes.
  • Conclusion: Long-term memory can be pretrained and scaled independently rather than remaining entangled with reasoning in one parameter set.The authors conclude that independently scaling pretrained memory provides a more parameter efficient path to improving language model performance.

9. Limitations … C.2. Domain Evaluation

The paper documents preprocessing and evaluation constraints alongside implementation details for general and domain memories. Memory pretraining uses separately trained modules, while evaluations apply fixed, task-specific protocols across general and domain benchmarks.

  • 9. Limitations: Indexing and retrieval for kNN target distributions add preprocessing overhead that grows with corpus scale despite compressed, sharded construction.The authors identify adaptive interpolation weighting as a direction for future work.
  • A. Training Datasets: Domain corpora are normalized into plain continued-pretraining text before tokenization, with BioInst converted from biology instructions and DISC-Law-SFT deduplicated and formatted as legal text.Training examples append the Qwen3 EOS token and are packed with a block size of 4096 tokens.
  • B. Implementation Details: Memory modules are trained separately from frozen backbones and combined with backbone distributions only at inference, with general memories scaling from 1.4B to 6.9B parameters.General-memory runs use GPT-NeoX configurations from the Pythia family, while domain memories initialize from Qwen3-1.7B-Base.
  • B. Implementation Details: 300B training tokens and 148,000 training steps are used for each general-memory size, with 2M-token global batches and sequence length 2048.The optimizer is AdamW, and learning rates use cosine decay with size-specific values.
  • C.1. General Evaluation: General evaluation uses fixed benchmark prompts without dynamically sampled examples, while the three AVG bars use the same 13 tasks and exclude PopQA, TruthfulQA-MC, HaluEval, and GPQA-main.Bamboogle remains zero-shot, and TruthfulQA retains its canonical prefix of six question-and-answer pairs.
  • C.2. Domain Evaluation: Domain methods use plain base-model prompts, greedy generation, fixed prompts and evaluators, and top-5 retrieved passages for RAG.BioInst and LawBench use OpenCompass, whereas FinEval uses the LM Evaluation Harness.
  • C.2. Domain Evaluation: BioInst reports the unweighted mean over 21 biodata_task_gen tasks after multiplying each task metric by 100 and applying a deterministic relaxed answer extractor uniformly across methods.The extractor handles structured, binary, DNA-classification, protein-solubility, and numeric-answer formats.

D. Additional General Memory Results · E. Additional Domain Memory Results

The additional general-memory results report task-level comparisons under matched total-parameter counts and training budgets. Figure 10 also visualizes base-versus-memory parameter allocation and scaling behavior.

  • D. Additional General Memory Results: Figure 10 reports results for individual tasks under the same total-parameter view used in the main analysis.
  • D. Additional General Memory Results: Figure 10 compares general-memory transfer by task under matched parameter counts and training budgets.
  • D. Additional General Memory Results: Triangles denote base models, while pie markers show the Base and Memory parameter shares.
  • D. Additional General Memory Results: Lines trace scaling for Base and Base + Memory configurations.

E.1. Results with 0.6B Domain Memory … G. Additional Case Study

Additional evaluations show that smaller domain memories improve average scores, 1.7B memory leads Qwen3 on one-shot LawBench, and RAG retrieval depth has no consistent winner. The paper also analyzes interpolation sensitivity and presents factual case studies.

  • E.1. Results with 0.6B Domain Memory: 0.6B domain memory improves the average score for every Qwen3 and OLMo backbone, showing gains are not limited to the larger 1.7B memory.OLMo memory rows use 20% of the memory training budget in Table 2.
  • E.2. LawBench One-Shot Results: +Mem-1.7B achieves the highest score for all five Qwen3 backbones and outperforms RAG on one-shot LawBench.One demonstration is included in each prompt.
  • E.2. LawBench One-Shot Results: On OLMo, CPT achieves the highest score for both backbones, while both memory sizes improve over the frozen base.OLMo memory uses 20% of the memory training budget in Table 2.
  • E.3. Effect of RAG Retrieval Depth: Top-5 retrieval performs better in 11 of 21 backbone-domain comparisons, while top-3 performs better in the remaining 10.Neither retrieval depth consistently dominates across backbones or domains.
  • F. Sensitivity to the Interpolation Coefficient: The final distribution interpolates the frozen backbone and memory as p_final = (1 − α)p_base + αp_mem.This combines the backbone and memory output distributions at inference.
  • F. Sensitivity to the Interpolation Coefficient: When validation data is available, α is tuned on the validation split and fixed for test evaluation, with an additional sweep over 21 coefficients on six representative tasks.The sweep uses α ∈ {0.00, 0.05, ..., 1.00} across general and knowledge-intensive benchmarks.

H. Extractable Memorization Evaluation Details … I.2. Metrics

The appendix evaluates extractable memorization through deterministic BioInst traceability probes and compares a 1.7B memory model with a matched Qwen3-1.7B-Base CPT model. It also measures how closely the trained memory reproduces kNN retrieval supervision using distributional metrics.

  • H. Extractable Memorization Evaluation Details: Two BioInst traceability probes compare the 1.7B memory model with a Qwen3-1.7B-Base CPT model trained on the same corpus.This comparison isolates the memory training mechanism from mere exposure to BioInst data.
  • H.1. Metric and Decoding Protocol: Exact recovery uses token-ID equality, one deterministic greedy query per example, and k-eidetic EM for suffixes appearing in at most k BioInst training rows.Paired CPT-only and memory-only counts include prompts solved by exactly one model.
  • H.2. Verbatim Continuation: 1,024 FunctionEC candidates were sampled after requiring unique 16-token continuations, and the memory model favored the CPT model at all three suffix-frequency thresholds.The frequency-restricted subsets are treated as supporting rather than standalone evidence.
  • H.2. Verbatim Continuation: The memory model reproduced all 16 tokens in a FunctionEC example whose gold suffix occurred in one training row, while the CPT model produced a shorter continuation.The example corresponds to training sample 52,021.
  • H.3. Domain Anchor Verbatim Completion: 62.2% of low-frequency EMP targets were recovered by the memory model versus 31.4% by the CPT model, with 111 versus 15 prompts recovered exclusively.These figures apply to spans observed at most 50 times.
  • H.4. Additional Qualitative Examples: Additional qualitative examples show domain-anchor completions solved only by the memory model, including outputs beginning with exact DNA or epigenetic continuations.The appendix provides two examples and identifies four additional domain-anchor completions.
  • I. Additional Details on Memory Fidelity to Retrieval Supervision: Aggregate metrics and examples assess whether the trained 1.7B BioInst domain memory learns kNN supervision and contributes it through interpolation with a frozen base model.The comparison directly relates the retrieval target to the trained memory.
  • I.1. Sampling Protocol: 1,024 uniformly sampled BioInst datastore examples compare each stored kNN target distribution P_i with the trained memory’s distribution Q_i using KL, total variation, and Pearson correlation.KL emphasizes underassigned likely kNN tokens; total variation measures probability mass reassignment, while Pearson correlation tests confidence consistency on the same target token.

I.3. Target Distribution Structure and Memory Fidelity … J.3. FinEval

The memory reproduces concentrated and multi-token target distributions with high agreement, contributes its learned retrieval signal directly through interpolation, and shows close qualitative matches on DNA and RNA examples. Per-task evaluations organize zero-shot comparisons across BioInst, LawBench, and FinEval backbones, with FinEval reporting macro-averaged scores across 25 tasks.

  • I.3. Target Distribution Structure and Memory Fidelity: 99.50% top token agreement and mean TV 0.0050 characterize single-token targets, whereas multiple-support targets reach 68.47% agreement, mean KL 0.4095, and mean TV 0.1913.Total variation exceeds 0.5 for 63 of 1,024 samples, or 6.15%, indicating that aggregate metrics combine near-exact concentrated targets with more ambiguous cases.
  • I.3. Target Distribution Structure and Memory Fidelity: 95% bootstrap confidence intervals based on 10,000 resamples are narrow across all four metrics, indicating stable aggregate estimates.The intervals are reported for evaluation samples in Table 20.
  • I.4. How Memory Contributes to Model Predictions: The trained 1.7B BioInst memory absorbs datastore knowledge into its parameters, as shown by high top-token agreement, strong probability correlation, and small distribution distances.The kNN targets encode supported continuations and probability allocation among plausible tokens.
  • I.4. How Memory Contributes to Model Predictions: At inference, the memory processes the base model’s context and linearly interpolates its distribution with the frozen base distribution for each next-token prediction, eliminating online datastore lookup.This directly connects the training objective to the retrieval contribution in final predictions.
  • I.5. Additional Examples: KL 0.0020 and total variation 0.0249 show that the memory closely matches the kNN target in the DNA example, where enh is the shared mode.The displayed probability allocations remain close across candidate continuations.
  • I.5. Additional Examples: KL 0.0019 and total variation 0.0235 show close reproduction in the RNA example, where does is ranked second by both distributions and RNA is their shared mode.These figures are qualitative illustrations; quantitative conclusions use the complete random sample.
  • J. Per-Task Domain Results: Tables 21 through 41 report zero-shot per-task domain results, comparing Base, CPT, LoRA, RAG, and 0.6B and 1.7B memory sizes across BioInst, LawBench, and FinEval backbones.The tables are grouped by domain and include Qwen3 and OLMo backbones; the six OLMo tables use 20% of the memory training budget in Table 2.
  • J.3. FinEval: FinEval uses generate_until and reports task-specific primary lm-evaluation-harness metrics multiplied by 100, with AVG as the macro average over 25 task scores.EM, MCC, and ROUGE-1 denote exact match, Matthews correlation coefficient, and unigram F-measure, respectively.
Loading 2607.27919v1…