Summary (Overview)

• Scaling parametric memory: This paper scales Memory Decoder [Cao et al., 2026] to 6.9B parameters pretrained on 300B tokens, demonstrating that long-term memory can be pretrained and scaled independently of the base language model.
• Parameter efficiency: Pairing a frozen Pythia-410M backbone with a 6.9B general memory achieves 37.34 average score across 17 benchmarks, surpassing Pythia-12B (37.24) with 39% fewer total parameters.
• Distributed Faiss pipeline: To handle 207B tokens, the authors develop a distributed pipeline with embedding compression (OPQ256), index sharding, and parallel GPU search, making $k$NN distribution construction feasible at pretraining scale.
• Plug-and-play domain memories: 1.7B domain memories for biology, law, and finance improve average scores across Qwen3 backbones (0.6B to 14B) by >9 points at every scale, and transfer across vocabularies with only 20% training budget.
• Comprehensive analysis: Memory remains effective with few-shot prompting, benefits knowledge-intensive tasks most, shows strong fidelity to retrieval supervision (86.62% top-token match, Pearson $r = 0.917$), and makes training content more extractable than standard CPT.

Introduction and Theoretical Foundation

Background and Motivation: Standard decoder-only language models entangle long-term memory and reasoning within a single parameter set. This entanglement means:

• Long-term memory cannot be pretrained or scaled independently
• Increasing memory size requires proportional increase in total parameters
• Domain adaptation via continued pretraining risks catastrophic forgetting [Kirkpatrick et al., 2017]
• No standalone memory that can be swapped across domains or reused across models

Theoretical Basis: The human brain organizes memory and reasoning into partially distinct neural systems [Baddeley and Warrington, 1970; Squire, 2009], motivating a modular architecture where long-term memory is pretrained and scaled independently.

Prior Work Limitations:

• Short-term memory approaches (Transformer-XL, StreamingLLM, InfLLM) focus on contextual memory, not long-term knowledge acquisition.
• Non-parametric memory ($k$NN-LM, RAG) requires external retrieval at inference, with substantial storage and latency costs.
• Parametric memory (Memory Decoder, MLP Memory, Titans) has only been studied at small scale (<1B parameters, millions of tokens).

Key Research Questions:

1. Does memory pretraining continue to scale with model and data size?
2. Can a general-purpose memory acquire broad knowledge?
3. How should parameter budget be allocated between base model and memory?

Methodology

Memory Decoder Architecture

The core idea is to train a standalone parametric memory module $M_\psi$ to mimic a non-parametric $k$NN retriever, then interpolate its output with the frozen base model at inference.

$k$NN Distribution Construction:

For each context $c_t = x_{<t}$, the frozen base model $M_\theta$ provides a query $q_t = \phi_\theta(c_t)$ (final-layer hidden state). A datastore is constructed:

$$R_{\text{mem}} = \{(k_t, x_t) \mid k_t = \phi_\theta(c_t), (c_t, x_t) \in P_{\text{mem}}\}$$

The $k$ nearest neighbors $N_K(q_t) \subset R_{\text{mem}}$ under distance $d$ are retrieved, and the $k$NN distribution is:

$$p_{\text{ret}}(y \mid c_t) = \sum_{(k_i, x_i) \in N_K(q_t)} \lambda_i(q_t) \mathbf{1}[x_i = y], \quad y \in \mathcal{V}$$

where:

$$\lambda_i(q_t) = \frac{\exp(-d(q_t, k_i)/\tau)}{\sum_{(k_j, x_j) \in N_K(q_t)} \exp(-d(q_t, k_j)/\tau)}$$

Memory Training Objective:

The memory $M_\psi$ is trained with a combined loss:

$$\mathcal{L}_{\text{mem}} = \beta D_{\text{KL}}\big(p_{\text{ret}}(\cdot \mid c_t) \| p_\psi(\cdot \mid c_t)\big) + (1-\beta)\big[-\log p_\psi(x_t \mid c_t)\big]$$

where $D_{\text{KL}}$ is KL divergence and $\beta \in [0,1]$ controls retrieval supervision strength.

Inference:

At inference, the frozen base model and memory process context in parallel:

$$p_{\text{final}}(y \mid c_t) = (1-\alpha) p_\theta(y \mid c_t) + \alpha p_\psi(y \mid c_t), \quad y \in \mathcal{V}$$

where $\alpha \in [0,1]$ controls memory contribution.

Distributed Faiss Pipeline (for Pretraining Scale)

Constructing $k$NN distributions over 207B tokens is infeasible with standard Faiss. The three-step pipeline (Figure 3):

1. Embedding Compression: OPQ256 maps 4096-dimensional hidden states to 256-dimensional search vectors.

2. Index Sharding: IVF partition with HNSW quantizer assigns vectors to contiguous centroid ranges, creating $S$ independent IndexIVFPQ shards.

3. Parallel Search:

   • Routing: HNSW quantizer identifies probed centroids; queries grouped by shard ID
   • Local retrieval: Each GPU searches its assigned shard independently
   • Merge: Writer workers collect outputs, merge candidates, cache sparse $p_{\text{ret}}$

Sparse $k$NN Distribution Storage

Dense distributions scale as $O(N|\mathcal{V}|)$. The sparse representation stores only:

$$S_t = \{y \in \mathcal{V} \mid p_{\text{ret}}(y \mid c_t) > \epsilon\}$$

scaling as $O(N + M)$ where $M$ is total retained token-probability pairs. On the Pile, each row retains 64.95 token-probability pairs on average, yielding ~250× storage reduction for 50,304-token vocabulary.

Training Details

Empirical Validation / Results

General Memory at Scale (Table 1)

Key Result: Knowledge-intensive tasks show largest gains (e.g., TriviaQA: 8.30→17.11 for 2.8B backbone). Across 51 task-scale combinations, memory improves 47 and matches the base in one more.

Memory Transfer Across Backbones (Figure 1)

Critical Finding: Pairing a 410M backbone with a 6.9B memory raises AVG from 29.86 to 37.34, surpassing 12B base model (37.24) with 39% fewer total parameters. At matched AVG, Base+Memory configurations use 33%, 32%, and 42% fewer parameters than 2.8B, 6.9B, and 12B base models respectively.

Domain Memory Results (Table 2)

Memory outperforms the strongest baseline at each scale by at least 4.05 points.

Cross-Vocabulary Transfer (Table 3)

With only 20% of standard training budget, transferred memories achieve best average:

• OLMo-2-7B: 19.57 → 23.83 (+4.26)
• OLMo-3-7B: 18.67 → 26.44 (+7.77)

Memory improves all six evaluations across both backbones.

Key Analyses

Few-shot robustness (Figure 4): Memory remains effective under 3-shot and 5-shot settings, with gains of 1.22-1.87 points across settings.

Memory size scaling (Figure 6): Gains grow consistently with memory capacity. 6.9B memory yields largest AVG gain at every backbone size (+4.56, +3.67, +1.49 points for 1.4B, 2.8B, 6.9B backbones respectively).

Training budget (Table 4): Extending training from 1.5 to 5 epochs improves general tasks (56.67→57.33) and knowledge (13.99→14.17). Full budget consistently outperforms 20% budget in domain transfer.

Attached module control (Table 5): Memory outperforms CPT with matched training on BioInst (23.21 vs 13.00 for 1.7B backbone) and LawBench (30.78 vs 29.10), showing gains come from memory objective, not interpolation interface.

Extractable memorization (Table 6): Memory model achieves 49.7% vs 42.4% (CPT) on verbatim continuation, and 56.5% vs 22.6% on domain anchor completion.

Memory fidelity (Table 7):

Theoretical and Practical Implications

Theoretical Implications:

1. Independent memory scaling: Long-term memory need not be entangled with reasoning in a single parameter set. A small backbone with a large standalone memory is more parameter-efficient than backbone scaling.
2. Knowledge separate from reasoning: The consistent gains on knowledge-intensive tasks (TriviaQA, 2WikiMultiHopQA, HotpotQA) suggest memory supplies factual evidence, complementing the base model's reasoning capabilities.
3. Modularity enables reuse: A single trained memory can augment different backbones without retraining, and domain memories can be swapped at inference while the backbone remains frozen.

Practical Implications:

1. Parameter efficiency: At matched performance, Base+Memory uses 33-42% fewer total parameters than backbone-only scaling, reducing training and storage costs.
2. Inference flexibility: The two components can run in parallel, offering lower inference latency than a single larger backbone.
3. Domain adaptation: Domain memories provide plug-and-play specialization without continued pretraining of the backbone, avoiding catastrophic forgetting.
4. Cross-model transfer: Memories transfer across vocabularies with minimal additional training (20% of standard budget), enabling reuse across model families.

Conclusion

This paper demonstrates that parametric long-term memory can be pretrained and scaled independently of the base language model, offering a more parameter-efficient path to improving language model performance.

Main Takeaways:

• Memory Decoder scales to 6.9B parameters on 300B tokens
• Distributed Faiss pipeline, sparse storage, and distributed streaming enable $k$NN distribution construction at pretraining scale
• Scaling memory is more parameter-efficient than scaling the base model alone
• Both general and domain memories provide substantial, consistent gains across model families and scales
• Memory remains effective with few-shot prompting, improves with capacity and training budget, and shows high fidelity to retrieval supervision

Future Directions:

1. Replace fixed interpolation coefficient $\alpha$ with an adaptive weighting mechanism based on context or model confidence
2. Explore joint or staged training of memory and backbone to improve coordination while preserving modular deployment
3. Address the offline cost of $k$NN target construction (current limitation)

Code and Models: Available at $\texttt{github.com/LUMIA-Group/MemoryDecoder-at-Scale}$