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

1. Introduction

• Decoder-only models entangle long-term memory with reasoning, while short-term context methods do not resolve that entanglement and external retrieval methods incur inference-time processing, storage, or search costs.The proposed modular architecture allows pretrained memories to be scaled independently, swapped across domains at inference, and reused across models while the base model remains frozen.Source: §1 · pp. 2–3
• Memory Decoder [Cao et al., 2026] demonstrated pretrained parametric memory up to 1B parameters on corpora of millions of tokens, leaving larger-scale models and diverse pretraining data unexplored.The paper also identifies open questions about general-purpose memory, scaling with model and data size, and allocating parameters between memory and the base model.Source: §1 · p. 3
• 37.34 average score across 17 benchmarks for frozen Pythia-410M with a 6.9B memory, versus 29.86 for Pythia-410M alone; it surpasses frozen Pythia-12B’s 37.24 with 39% fewer total parameters.These results support the paper’s claim that scaling memory can be more parameter-efficient than scaling the base model alone.Source: §1 · p. 3
• A distributed Faiss pipeline with embedding compression, index sharding, and parallel search addresses indexing and retrieval bottlenecks for 300B-token memory pretraining.Sparse kNN-distribution storage and distributed streaming load only the entries needed by each batch.Source: §1 · p. 3
• 1.7B domain memories improve average scores across BioInst, LawBench, and FinEval by more than 9 points for Qwen3 Base models from 0.6B to 14B.For Qwen3-14B-Base, the improvements are 17.96 points on BioInst, 8.97 on LawBench, and 3.04 on FinEval.Source: §1 · p. 3
• After vocabulary transfer using 20% of the standard memory-training budget, domain-average gains are 4.26 points for OLMo-2-7B and 7.77 for OLMo-3-7B.Further analyses report effectiveness with few-shot prompting, the largest benefits on knowledge-intensive tasks, and gains from larger memories and training budgets.Source: §1 · p. 3

• Memory Decoder [Cao et al., 2026] imitates a retriever by searching corpus-context keys with base-model queries and weighting the next-token values of the K nearest keys.The frozen base model Mθ encodes each context; the retrieved values and normalized distance weights form the kNN distribution.Source: §2 · p. 3
• The memory decoder predicts the retrieved distribution directly from the original context, using retrieval only to construct offline pretraining supervision.The memory and base models share a tokenizer and output vocabulary; retrieved text is not fed to the memory decoder.Source: §2 · p. 3
• The training objective aligns memory predictions with the kNN distribution while retaining corpus language modeling, with β controlling retrieval supervision [Bengio et al., 2003; Van Erven and Harremos, 2014].The KL term teaches retrieval imitation, while the language-modeling term limits deviation from the corpus distribution.Source: §2 · p. 4
• At inference, the base and memory models process the same context in parallel, and α controls their prediction interpolation.The combined prediction incorporates retrieval behavior learned by the memory during pretraining.Source: §2 · p. 4

3.1. Large-Scale 𝑘NN Distribution Construction

• At 207B contextual representations, quadratic indexing and search make standard Faiss infeasible; OPQ compression, centroid sharding, parallel GPU retrieval, and sparse caching address the bottleneck.The pipeline routes queries to shards for local search, then merges candidates into the final K neighbors and caches the retrieved distribution in sparse form.Source: §3 · p. 4, §3.1 · p. 5, Figure 3 · p. 5

3.2. General Memory Pretraining

• General memories with 1.4B, 2.8B, and 6.9B parameters are pretrained on 300B tokens using sparse kNN storage and distributed streaming.On the Pile, rows retain 64.95 token-probability pairs on average, yielding an approximately 250× dense-to-sparse storage ratio; storage scales as O(N + M) rather than O(N|V|).Source: §3.2 · p. 6

3.3. Domain Memory Pretraining

• Domain memories for biology, law, and finance use continued-pretrained datastore models and the same architecture, objective, and sparse training pipeline, with corpora up to 4.4B tokens.Each domain receives a separate memory supervised by kNN distributions constructed from its corpus representations.Source: §3 · p. 4, §3.3 · p. 6

Experimental overview

• The study evaluates four settings: general-memory gains across tasks and scales, transfer across backbone scales, domain memories for biology, law, and finance, and transfer across model families and vocabularies.The base model remains frozen, and memory predictions are combined with the backbone distribution using the interpolation rule in Eq. 2.Source: §4 · p. 6

Datasets

• The general memory is trained on the deduplicated 207B-token Pile [Gao et al., 2020] and evaluated on general, knowledge-intensive, factuality, and hallucination-robustness tasks.The evaluation includes Pythia-suite tasks [Biderman et al., 2023], knowledge-intensive benchmarks, and TruthfulQA [Lin et al., 2022] and HaluEval [Li et al., 2023].Source: §4 · pp. 6–7
• Domain memories are trained and evaluated on matched biology, law, and finance data: Biology-Instructions [He et al., 2024], DISC-Law-SFT [Yue et al., 2023] and LawBench [Fei et al., 2024], and FinTrain [Ke et al., 2025] and FinEval [Ke et al., 2025].The biology and law training datasets are reformatted as unsupervised text; finance training uses the book and finance subsets of FinTrain.Source: §4 · p. 7

Backbones and baselines

• General-memory experiments use frozen Pythia backbones from 410M to 12B parameters and memories of 1.4B, 2.8B, and 6.9B; domain experiments use Qwen3 Base models from 0.6B to 14B.General memories are compared with decoder-only Pythia models at matched total parameter counts and training budgets; domain memories are compared with CPT, parameter-matched LoRA, and top-5-document RAG baselines.Source: §4 · p. 7

Implementation details

• General memories are trained for 300B tokens, while domain memories use a compute-matched budget; vocabulary transfer adapts a Qwen3 memory to OLMo backbones.General memories use a Pile datastore built with Pythia-6.9B, corresponding to approximately 1.5 epochs; domain datastore construction uses one Qwen3-4B-Base continued-pretraining epoch, and vocabulary transfer uses 20% of the standard memory-training budget.Source: §4 · p. 7
• General benchmarks use exact match for NQ-Open, TriviaQA, 2WikiMultiHopQA, Bamboogle, and HotpotQA, and standard accuracy for the others; domain scores average benchmark-defined metrics.FinEval combines exact match, Matthews correlation coefficient (MCC), and ROUGE-1; when validation data is available, the interpolation coefficient is tuned on validation and fixed for test evaluation.Source: §4 · p. 7

5.1. General Memory at Scale

• 34.36 versus 33.89: a 1.4B Pythia backbone with equal-size memory outperforms the 2.8B backbone at the same total parameter count and training budget.Equal-size memory improves AVG at every scale; it improves 47 of 51 task-scale combinations and matches the base in one more. On the 6.9B backbone, AVG rises from 36.30 to 37.79, above the frozen 12B backbone’s 37.24.Source: §5.1 · p. 8, Table 1 · p. 8

5.2. General Memory Transfer Across Backbones

• 29.86→37.34: pairing Pythia-410M with a 6.9B memory surpasses Pythia-12B’s 37.24 AVG with 39% fewer total parameters.At matched AVG, Base + Memory configurations use 33%, 32%, and 42% fewer total parameters than the 2.8B, 6.9B, and 12B base models, respectively; memory and backbone can run in parallel.Source: §5.2 · p. 9

5.3. Domain Memory Across Backbones

• 9.88, 9.64, 10.00, 9.09, and 9.99 points: 1.7B domain memories improve average scores over frozen Qwen3 backbones from 0.6B to 14B.They achieve the highest average across biology, law, and finance at every scale, improve all 15 evaluations, and can be switched by activating the corresponding memory while keeping the backbone frozen.Source: §5.3 · p. 9, Table 2 · p. 10

5.4. Domain Memory Transfer Across Vocabularies

• 4.26 and 7.77 points: transferred memories raise average scores from 19.57 to 23.83 on OLMo-2-7B and from 18.67 to 26.44 on OLMo-3-7B.Using 20% of the standard memory training budget, memory improves all six evaluations and provides the best overall average, although CPT leads on some individual benchmarks.Source: §5.4 · p. 9, Table 3 · p. 10

6.1. Few-Shot Robustness and Knowledge Task Improvements

• Memory improves AVG by 1.43–1.87 points zero-shot, 1.39–1.52 with three shots, and 1.22–1.62 with five shots across all tested scales.The controlled averages use the same 13 tasks, with Bamboogle retaining its zero-shot protocol; memory also improves all eight zero-shot knowledge tasks at every scale. Averaged over scales, gains on TriviaQA, 2WikiMultiHopQA, and HotpotQA are 4.59, 2.96, and 1.72 points, respectively.Source: §6.1 · p. 11

6.2. Case Study: Memory Strengthens Factual Evidence

• In a HotpotQA context, memory assigns larger probability shares to many entity and attribute words, consistent with its aggregate QA gains.The case study uses frozen Pythia-1.4B and a 1.4B memory; a share above 0.5 indicates greater support from memory than from the base model.Source: §6.2 · pp. 11–12, Figure 5 · p. 12

6.3. Memory Size

• The 6.9B memory yields the largest AVG gain for each Pythia backbone: 4.56, 3.67, and 1.49 points for the 1.4B, 2.8B, and 6.9B backbones.Every tested general-memory size improves AVG over its corresponding base, and gains grow overall with memory capacity. For Qwen3 domain settings, 1.7B memory beats 0.6B memory on BioInst and LawBench for every backbone; gains reach 18.5 points on BioInst and 9.4 on LawBench, with capacity increases adding up to 4.8 and 4.4 points, respectively.Source: §6.3 · p. 12

6.4. Memory Training Budget

• Increasing memory training budget improves aggregate performance in both settings: general-task average rises from 56.67 to 57.33, and knowledge average from 13.99 to 14.17.Using the full rather than 20% domain-memory training budget improves Avg for every OLMo backbone and memory size: gains are 0.34 and 3.51 points for 0.6B and 1.7B memories on OLMo-2, and 0.94 and 1.75 points on OLMo-3.Source: §6.4 · p. 12, Table 4 · p. 13

6.5. Memory Training Objective

• Memory exceeds attached CPT by 10.21 and 8.71 points on BioInst, and by 1.68 and 3.93 points on LawBench, across the two backbones.The matched modules use the same domain training data, FLOPs, capacity, and inference interface; the consistent advantage in all four comparisons attributes the additional gains to the memory objective rather than the interpolation interface.Source: §6.5 · pp. 13–14, Table 5 · p. 14

6.6. Extractable Memorization Evaluation

• Memory raises strict EM@8, 16 from 42.4% to 49.7%, and domain anchor completion from 22.6% to 56.5% versus CPT.Both models use the same 1,024 BioInst training prompts, data, and budget; the probes test exact continuation after an eight-token prefix and completion beginning with a hidden four-word span. The results indicate stronger traceability to training data and more reliable recovery from partial context.Source: Table 6 · p. 14, §6.6 · pp. 14–15

6.7. Memory Fidelity to Retrieval Supervision

• Across 1,024 BioInst samples, memory and kNN modes match in 86.62% of cases, with Pearson r=0.9174, mean KL 0.1820, and mean TV 0.0823.Most samples lie near equality for the kNN mode probability, indicating alignment in probability allocation as well as token choice; residual errors arise mainly for diffuse targets. The learned signal can be interpolated with the frozen base distribution without an online datastore lookup.Source: Table 7 · p. 15, §6.7 · p. 15, Figure 8 · p. 15

Memory-augmented pretraining

• Memory-augmented pretraining methods RETRO [Borgeaud et al., 2022], TRIME [Zhong et al., 2022], and LMLMs [Zhao et al., 2025] retain non-parametric inference memory, unlike our standalone parametric module.The module distills the memory signal learned during pretraining and requires neither external memory construction nor online retrieval at inference time.Source: §7 · p. 16

Short-term memory

• Short-term memory methods extend usable input context through recurrence, KV caching, context retrieval, or positional rescaling, while our work pretrains standalone long-term memory.Examples include Transformer-XL [Dai et al., 2019], StreamingLLM [Xiao et al., 2024b], InfLLM [Xiao et al., 2024a], and LongRoPE [Ding et al., 2024].Source: §7 · p. 16

Long-term memory

• Long-term memory may be non-parametric or parametric; our work scales pretraining of separate general and domain memories attached to frozen backbones and selected at inference.Non-parametric approaches such as kNN-LM [Khandelwal et al., 2019] and RAG [Lewis et al., 2020] require datastore retrieval, while Memory Decoder [Cao et al., 2026] and MLP Memory [Wei et al., 2025] distill retrieval distributions into neural memories; Titans [Behrouz et al., 2024] writes and reads memory at test time.Source: §7 · p. 16

Conclusion

• Memory models were scaled to 6.9B parameters and pretrained on 300B tokens using distributed Faiss indexing and retrieval over 207B corpus tokens.The pipeline uses embedding compression, index sharding, and parallel search; sparse kNN storage and distributed streaming load only entries required by each batch.Source: §8 · pp. 16–17
• A 6.9B general memory lets Pythia-410M surpass Pythia-12B with 39% fewer total parameters, while 1.7B domain memories improve Qwen3 Base scores by more than 9 points at every tested scale.The Qwen3 result averages scores across biology, law, and finance for models ranging from 0.6B to 14B; further analyses find memory effective with few-shot examples, most beneficial on knowledge-intensive tasks, and increasingly effective with larger memories and training budgets.Source: §8 · p. 17
• The results show that long-term memory can be pretrained and scaled independently rather than remaining entangled with reasoning in one parameter set.Independently scaling pretrained memory offers a more parameter-efficient path to improving language-model performance.Source: §8 · p. 17

9. Limitations

• Memory pretraining still incurs indexing and retrieval overhead that grows with corpus scale; adaptive weighting or joint training could improve coordination but raise costs and weaken cross-backbone transfer.The proposed alternatives are adaptive per-input memory weighting and joint or staged training of memory and backbone; inference itself requires no external retrieval.Source: §9 · p. 17

A. Training Datasets

• The general memory uses the deduplicated Pile, while domain corpora are normalized into plain text; examples are EOS-appended and packed into 4096-token blocks.BioInst becomes task-oriented text, DISC-Law-SFT is deduplicated and formatted as legal instruction text, and FinTrain-unsup is chunked as plain continued-pretraining data.Source: A. Training Datasets · p. 23

B. Implementation Details

• Memory modules train separately from frozen backbones: general Pythia-family memories span 1.4B–6.9B parameters and use 300B Pile tokens, while domain memories initialize from Qwen3-1.7B-Base.General-memory runs use 148,000 steps, a 2M-token global batch, sequence length 2048, AdamW, and cosine learning-rate decay.Source: B. Implementation Details · pp. 23–24

C.1. General Evaluation

• General evaluation uses benchmark-defined prompts without dynamically sampled examples; the reported averages use 13 tasks, omit four specified benchmarks, and retain TruthfulQA’s canonical six-pair prefix.Bamboogle remains zero-shot, and the evaluation uses the LM Evaluation Harness and its accompanying evaluators for knowledge tasks.Source: §C.1 · p. 25

C.2. Domain Evaluation

• Domain evaluations use plain base-model prompts, greedy generation, fixed prompts and evaluators, and top-5 retrieved passages for RAG; BioInst averages 21 tasks after uniform relaxed answer extraction.BioInst uses sample_1k splits and sequence length 16,384; its extractor leaves references and metric definitions unchanged and uses no model judge.Source: §C.2 · pp. 25–26
• For BioInst, the deterministic extractor checks final answer spans and task-specific binary, numeric, structured, and categorical formats, applying fixed fallbacks when no answer is extracted.It handles DNA classification, Protein-Solubility, numeric percentages, RNA-switch fields, and EC-number normalization.Source: §C.2 · p. 26
• LawBench’s primary evaluation uses all 20 tasks without demonstrations, 500 examples per task, and an unweighted task mean; one-demonstration results are supplementary.Seven tasks retain a 1,024-token generation limit to avoid truncating responses before a scorable answer.Source: §C.2 · p. 26
• FinEval evaluates 25 tasks without demonstrations and reports an unweighted mean across 20 exact-match, two MCC, and three ROUGE-1 tasks.It uses greedy decoding with a 32,768-token context; unrecognized MCC predictions count as errors, and numeric answers use an absolute tolerance of 10^-6.Source: §C.2 · p. 26

D. Additional General Memory Results

• Figure 10 reports individual-task results for general memory under the same total-parameter view as the main analysis.The figure caption specifies matched parameter counts and training budgets.Source: D. Additional General Memory Results · p. 27, Figure 10 · p. 27
• Marker area in Figure 10 denotes total parameters, while triangles represent base models and pie markers show Base and Memory parameter shares.Lines trace Base and Base + Memory scaling.Source: Figure 10 · p. 27

E.1. Results with 0.6B Domain Memory

• Average scores improve for every Qwen3 and OLMo backbone with 0.6B domain memory under the zero-shot domain setup.The OLMo rows use 20% of the memory training budget in Table 2.Source: Table 12 · p. 28, §E.1 · p. 28

E.2. LawBench One-Shot Results

• On Qwen3, +Mem-1.7B achieves the highest one-shot LawBench score for all five backbones and outperforms RAG; on OLMo, CPT leads both backbones.Both memory sizes improve over the frozen OLMo base using 20% of the memory training budget.Source: §E.2 · p. 29, Table 13 · p. 29

E.3. Effect of RAG Retrieval Depth

• Top-5 RAG retrieval performs better in 11 of 21 backbone–domain comparisons, while top-3 performs better in the other 10.Neither retrieval depth consistently dominates across backbones or domains.Source: Table 14 · p. 29, §E.3 · p. 29

F. Sensitivity to the Interpolation Coefficient

• General tasks have interpolation-coefficient median 0.55 (IQR 0.35–0.70), versus 0.80 (IQR 0.55–0.95) for knowledge-intensive tasks.The ranges overlap, indicating a gradual change in emphasis rather than a sharp separation; when validation data is available, α is tuned on its split and fixed for test evaluation.Source: F. Sensitivity to the Interpolation… · p. 30
• Across six tasks, 17 of 18 task–scale curves stay above the within-sweep α=0 endpoint for all 20 nonzero coefficients.WinoGrande, NQ-Open, TriviaQA, PopQA, and HotpotQA do so at all three scales; ARC-Easy’s 1.4B curve remains positive over the central sweep, while both larger configurations remain positive throughout.Source: F. Sensitivity to the Interpolation… · p. 30, Figure 11 · p. 31

G. Additional Case Study

• Figures 12–15 present additional factual-question case studies on four topics.The examples concern the 2000 ICC KnockOut Trophy, Kent County Cricket Club in 2011, No. 2 Squadron RAAF, and the documentary film Gimme Shelter.Source: G. Additional Case Study · p. 32, Figure 12 · p. 32, Figure 13 · p. 32

H. Extractable Memorization Evaluation Details

• The 1.7B memory and Qwen3-1.7B-Base CPT model are compared on raw BioInst examples, isolating the memory training mechanism from exposure to the corpus.Source: H. Extractable Memorization Evaluation… · p. 33

H.1. Metric and Decoding Protocol

• Equality is measured on token IDs with one deterministic query per example; k-eidetic EM restricts the rate to suffixes in at most k training rows, and paired columns count prompts solved by exactly one model.Source: §H.1 · p. 33

H.2. Verbatim Continuation

• The memory model favors exact recovery at all three suffix-frequency thresholds, though the subsets are small; one example has all 16 tokens reproduced by memory while CPT stops short.The FunctionEC evaluation samples 1,024 candidates and prompts each model with the protein, task context, and eight output tokens to generate 16 tokens.Source: Table 17 · p. 33, §H.2 · pp. 33–34

H.3. Domain Anchor Verbatim Completion

• At target-span frequencies of at most 50, the memory recovered 194 of 312 targets (62.2%), versus 98 (31.4%) for CPT; the models exclusively recovered 111 and 15 prompts, respectively.The advantage persisted after restricting by target-span frequency in the candidate pool.Source: Table 18 · p. 34, §H.3 · p. 34

H.4. Additional Qualitative Examples

• Four additional domain-anchor completions and four additional FunctionEC continuations are presented as examples solved only by the memory model.Source: Figure 16 · p. 35, Figure 17 · p. 35

I. Additional Details on Memory Fidelity to Retrieval Supervision

• The appendix expands comparison of retrieval targets with the trained 1.7B BioInst domain memory, showing how learned kNN supervision contributes to final predictions through interpolation with a frozen base model.Source: §I · p. 36

I.1. Sampling Protocol

• The fidelity evaluation uniformly samples 1,024 context–next-token pairs without replacement from the BioInst training datastore and compares stored kNN targets P_i with memory distributions Q_i.It uses the same trained 1.7B BioInst domain memory as the main results.Source: §I.1 · p. 36

I.2. Metrics

• KL emphasizes underweighting likely kNN tokens, total variation measures the probability mass to reassign (bounded from zero to one), and Pearson correlation tests confidence consistency on the target token across samples.Source: §I.2 · p. 36

I.3. Target Distribution Structure and Memory Fidelity

• Single-token targets achieve 99.50% top token agreement and mean TV 0.0050, while multiple-support targets reach 68.47% agreement, mean KL 0.4095, and mean TV 0.1913.TV exceeds 0.5 for 63 of 1,024 multiple-support samples (6.15%), showing that these targets are more difficult to reproduce.Source: §I.3 · p. 36, Table 19 · p. 36
• 10,000-resample bootstrap intervals are narrow across all four metrics, indicating stable aggregate estimates.Table 20 reports 95% confidence intervals over evaluation samples.Source: §I.3 · p. 37, Table 20 · p. 37

I.4. How Memory Contributes to Model Predictions

• At inference, the memory processes the frozen base model’s context and its distribution is linearly interpolated with the base distribution at each next-token prediction, without an online datastore lookup.The distribution comparison attributes the BioInst retrieval knowledge to the trained 1.7B memory’s parameters.Source: §I.4 · p. 37

I.5. Additional Examples

• The DNA example has shared mode enh with KL 0.0020 and total variation 0.0249; in the RNA example, does ranks second in both distributions, with KL 0.0019 and total variation 0.0235.RNA is the shared mode in the RNA example; these are qualitative examples, while quantitative conclusions use the complete random sample.Source: §I.5 · p. 37, Figure 18 · p. 38

J. Per-Task Domain Results

• Per-task domain tables compare Base, CPT, LoRA, RAG, and two memory sizes; evaluations use zero-shot prompting, and RAG retrieves five top-ranked passages.The six OLMo tables use 20% of the memory training budget in Table 2 for both memory columns.Source: J. Per-Task Domain Results · p. 39

J.1. BioInst

• BioInst per-task results include Qwen3-0.6B-Base and OLMo-3-7B backbones.The supplied BioInst tables also cover other Qwen3 and OLMo backbones; the captions do not provide their result values.Source: Table 21 · p. 39, Table 27 · p. 42

J.2. LawBench

• LawBench per-task results include Qwen3-0.6B-Base and OLMo-3-7B backbones.The supplied LawBench tables also cover other Qwen3 and OLMo backbones; the captions do not provide their result values.Source: Table 28 · p. 43, Table 34 · p. 46

J.3. FinEval

• FinEval reports the primary harness metric scaled by 100 for each task, with AVG defined as the macro average over 25 task scores.EM is normalized exact match, MCC is Matthews correlation coefficient over extracted labels, and ROUGE-1 is the unigram F-measure for generated text; all subtasks use generate_until.Source: §J.3 · p. 47