Pretraining with hierarchical memories: separating long-tail and common knowledge
Authors: Hadi Pouransari
David Grangier
C Thomas
Michael Kirchhof
Oncel Tuzel
Venue: ICLR 2026 Poster (2026)
Keywords: Large language models
pretraining
memory
long-tail
knowledge
reasoning
forgetting
Primary area: foundation or frontier models, including LLMs
TL;DR: We pretrain transformers with hierarchical parametric memories that automatically store long-tail world knowledge, fetched by context at inference time to boost small language model performance.

Abstract: The impressive performance gains of modern language models currently rely on scaling parameters: larger models store more world knowledge and reason better. Yet compressing all world knowledge into parameters is unnecessary, as only a fraction is used per prompt, and impractical for edge devices with limited inference-time memory and compute. We address this shortcoming by a memory-augmented architecture and a pretraining strategy aligned with existing hardware paradigms. We introduce small language models that access large hierarchical parametric memory banks encoding world knowledge. Our pretraining learns to store long-tail world knowledge in the memory parameters, while the small language model acts as an anchor capturing common knowledge and general reasoning abilities. Through trillion-token-scale experiments, we show significant gains: a 160M-parameters model augmented with an 18M-parameters memory fetched from a 4.6B memory bank obtains comparable performance to a regular model with more than 2x the parameters. We study the optimal type and size of parametric memories in transformers, scaling them to over 21B parameters. We find that our proposed hierarchical feed-forward memories work robustly across transformer architectures, whether added during pretraining or post-hoc.

Published as a conference paper at ICLR 2026
PRETRAINING WITH HIERARCHICAL MEMORIES: SEPARATING LONG-TAIL AND COMMON KNOWLEDGE
Hadi Pouransari', David Grangier?, C Thomas", Michael Kirchhof & Oncel Tuzel
Work done while at Apple AApple
pouransari 10gm00mmmll. com otuzeleapple..om
ABSTRACT
The impressive performance gains of modern language models currently rely on scaling parameters: larger models store more world knowledge and reason bet- ter. Yet compressing all world knowledge into parameters is unnecessary, only fraction is used per prompt, and impractical for edge devices with limited inference-time memory and compute. We address this shortcoming by memory- augmented architecture and pretraining strategy aligned with existing hardware paradigms. We introduce small language models that access large hierarchical parametric memory banks encoding world knowledge. Our pretraining learns to store long-tail world knowledge in the memory parameters, while the small lan- guage model acts as an anchor capturing common knowledge and general rea- soning abilities. Through rillion-token-ssale experiments, show significant gains: a 160M-parameters model augmented with an 18M-parameters memory fetched from 4.6B memory bank obtains comparable performance to a regular model with more than 2 the parameters. We study the optimal type and size of parametric memories in transformers, scaling them to over 21B parameters. We find that our proposed hierarchical feed-forward memories work robustly across transformer architectures, whether added during pretraining or post-hoc.
INTRODUCTION
Frontier large language models (LLMs) have advanced significantly in recent years, demonstrating strong capabilities across both world knowledge and reasoning tasks. These improvements have largely been driven by scaling the number of parameters and training tokens. However, limited results exist on the role of model parameter count in each specific aspect, and on whether knowledge and reasoning can be disentangled.
100
The atomic number of [Silican]
Acturasy (per 9%)
baseline +memory
Elomont traquency
160M
10'
World Knowledge (pre-training datasett)
10
baseline
ours
10'
0000: 00 anchor model
1.48
00
10'
single model
0000
100
Head (froquont)
Dlmmnts
Long-tal (rano)
Exact Match (36)
Schematic of chor parameters), others are fetched per input document (memory parameters). Middle: Accuracy improvement over baseline when ~ 10% of parameters are allocated as memories for knowledge- intensive task (predicting the atomic numbers of elements), using models with 160M, 410M, and 1.4B parameters, corresponding to rows A2, B2, and C2 in Table 1. Right: Elements sorted by their frequency of appearance the DCL DCL.MBaseline dataset (5 buckets, each with ~ 24 elements). With the proposed pretraining- weetiiiinn-wttt--mmmmee memories,, long-tail data. While the baseline 1.4B model has only 17% accuracy on the least frequent element bucket, augmenting with only 10% memory parameters increases the accuracy to 83%.
Figure 1: Left:
some parameters are always used (an-
observe significant improvements, especially

---

Published as a conference paper at ICLR 2026
LLMs memorize many facts in their parameters during pretraining (Roberts et al., 2020), most of which are long-tail tail knowledge, overly specific, and unnecessary for the intended use. For example, part of the parameters of the Qwen3-2B (Yang et al., 2025) model store the fact that Albert Einstein was born on March 14, 1879, a detail that is not essential for executing on-device personal assis- tant tasks, yet is permanently loaded into RAM and considered in each computation. Ideally, this capacity would be utilized for reasoning and commonsense abilities.
We propose use a base model as an anchor capture common knowledge and reasoning ca- pabilities and augment it with a memory bank that dedicates large set of memory parameters to long-tail tail knowledge. In this work we focus on capturing long-tail world knowledge with parametric memories, leaving improvements to the anchor's 's reasoning capabilities to future work. In Fig. I, we illustrate this through a representative mnowledge-intensing task: predicting the atomic num- ber of elements. Baseline pretraining performance degrades for samples with lower frequency in the pretraining data. This degradation is attributed to catastrophic forgetting (Toneva et al., 2018), which arises from destructive gradient updates caused by dissimilar content (Ghosal et al., 2025) applied to the same set of parameters. In contrast, in the proposed approach, memory parameters are activated and updated only on sequences of similar topics, thereby reducing susceptibility to forget- ting. Besides training dynamics advantages, separating long-tail toigttal knowledge into dedicated memory parameters offers several additional benefits, as discussed below.
Runtime efficiency: For on-device deployment of LLMs, the main bottleneck is the availability of large and fast memory. Methods such mixture-of-experts (Shazeer et al., 2017) (MoEs) improve compute efficiency by activating only a fraction of the feed-forward modules. Still, MoEs require on-demand random access the full set of experts at every layer and for every token, such that all model parameters remain loaded, not just the active ones. This makes on-device deployment of MoEs challenging. In contrast to MoE models, which require refreshing the majority of their parameters (experts) at every token, in our proposed method, depending on the context, only the required knowledge parameters (about 10% of the model) are fetched and attached to the anchor model (see Fig. 2). This allows fast on-device memory to be used primarily for anchor model parameters, while knowledge parameters are stored in a slower but larger storage. Furthermore, we learn knowledge in the form of hierarchical memories, allowing inference to naturally align with existing device memory hierarchies (RAM, flash storage, external disk) and benefit from their compositionality over a session of interaction with the model (see Fig. 5).
Training efficiency: A key bottleneck in large-scale distributed training of standard LLMs is the need to communicate large gradient tensors between compute nodes. In the proposed approach, during pretraining, based on the content of documents in a batch, only a small fraction of memory bank parameters is retrieved and updated together with the anchor model parameters. As result, the gradients are highly sparse, which substantially reduces node-to-node communication. This property makes the proposed method well-suited for co-locating training data and compute in massively distributed setups, similar to 2023a). For instance, pretraining an anchor model with 160 million parameters and 4.6 billion memory parameters requires less than 1.7x the compute budget of training 160M model alone.
(BTM) methods (Li et al., 2022; Gururangan et al.,
Privacy and knowledge editing: Separating knowledge and reasoning abilities enables a direct map- ping between training tokens and specific subsets of parameters (memories). Ownership of training data tokens can therefore be linked to ownership of the corresponding memories. This makes it pos- sible to remove or edit certain data from the model by deleting updating the associated memory parameters, or to restrict access to specific memories while keeping the anchor model public. Al- though knowledge editing and privacy are not the focus of this paper, we present preliminary results on the effect of memory blocking in Section 4. The proposed approach also allows creating new memories from new data after pretraining, as shown in Section 4. Through this mechanism, a strong public reasoning model can be readily augmented with memories built from large private user data or organizational databases, providing a more efficient alternative to very long contextual memories.
Our main contributions are transformer architectures augmented with large hierarchical memories (Fig. 2) and a clustering-based pretraining method (Section 2.1), through which long-tail knowledge is automatically incorporated into the memory bank. We, for the first time, systematically examine how different memory types (Fig. 3), depths (Fig. 3d), sizes (Fig. 4a), positions (Table 12), and memory-to to mode ratios (Fig. 4b) affect performance. We further analyze the runtime efficiency of
2

---

Published as a conference paper at ICLR 2026
was born on March 14, 1879.
Hierarchical memory bank:
W
clusters)
Anchor model block
Fetched memory
W
Level 2 (256 clusters)
R(x; W)
W War
Level (Ak clusters)
Level (65K clusters)
L Blocks
Memory R
Proposed architecture: For a given context z (such as question text), the memory retriever module selects relevant parameters from a large set of memory bank parameters. These memory parameters ing data (with up to 164 65k clusters at level 4) The anchor model, together with the retrieved memories, then responds the question.
Figure 2:
organized hierarchically based on the hierarchical clustering of the pretrain-
the proposed method in relation hardware storage hierarchies (Section 3). Memory learning is scaled to banks containing up to 21B parameters.
In this work, we focus demonstrating the effectiveness of the proposed parametric memories in capturing long-tail world knowledge during pretraining. The results show consistent improve- ments over baseline pretraining (Table 1) as well over raw contextual memories. We further show that the proposed parametric memories can be combined with methods such retrieval- -agmeneed generation (Table 5) for additional gains. Table 4, we also demonstrate robust applicability to arbitrary transformer architectures as post-hoc augmentation. Finally, show that with approx- imately 10% additional parameters retrieved from the memory bank, we observe improvements on knowledge-spede tasks from 34.1% 40.3% for 160M model and from 40.9% + 45.9% for 410M model, confirming the runtime benefits of the proposed approach (achieving the performance of model with more than 2x the runtime parameters).
PRETRAINING WITH MEMORIES
Consider a language model with parameters 0, called anchor parameters, a large bank memory parameters W, and a retriever module R that, given context æ, fetches only the relevant mem- ory parameters R(;; W) from the memory bank. During pretraining, we optimize the next token prediction loss for each document in the dataset D:
[(x) = - IP Po, xeer,
--x)=-E
where (o) I<t) is the model's output probability distribution over the vocabulary for token t in document given its previous tokens, using anchor parameters augmented with the retrieved memory parameters R(I; W).
In terms of parameter counts, we generally assume RRa; W) 0 < W. Therefore, with the training objective in Eq. (1), the memory bank parameters are updated only sparsely. Intuitively, we expect R(; W) to retrieve the same subset of parameters from W for imputs x that are seman- tically similar. As result, memory parameters receive gradients primarily from documents with related content, mitigating forgetting (Ghosal et al., 2025) and enabling them to efficiently memo- rize long-tail world knowledge. In contrast, receives gradients for all r E and expected to primarily capture common knowledge and general reasoning capabilities.
At inference time, the model uses only |0| + RR(x; W) parameters, introducing only a minor overhead compared model that relies solely on 0. Fig. 2 presents the overall architecture. We next describe the construction of each component in this framework.
2.1 CLUSTERING-BASED MEMORY RETRIEVER
Hierarchical clustering: Given a pretraining dataset D, we use an off-the- thel- text embedding model á map each document D to ((x) R We then perform hierarchical clustering using
3

---

Published as a conference paper at ICLR 2026
the document embeddings: first, divide the documents into & clusters using k-means; next, further subdivide each cluster into sub-clusters, and continue this process for p levels. Finally, obtain kf' nested clusters at each level. In this paper, cluster the DCLM-baseline dataset (Li et al., 2024) with 3.2 billion documents into a hierarchy with = levels and dividing factor of k 16 resulting in 16, 162, 163; and 16 clusters at levels 1, 2, 3, and 4, respectively (see Fig. 2). We use Sentence-BERT all 1 MiiiiMin--------- (LL-M6-0- v2 embedding model (Reimers & Gurevych, 2019) with dimension c 384. See more details in Appendix C.
To retrieve a cluster, we map a document x: to an index tuple I(x) (i), ;2,...,;ip ip) by traversing the clustering tree greedily: at each level I, 6(x) is compared to i centroids and the sub-tree corre- sponding to the closest (via L2 distance) centroid (ix) is then visited. This greedy traversal leads to a fast and scalable retrieval within O (pkk) comparisons. For our tree of depth p = 4, we denote the cluster index as [(x) = (i1, 82. i3, 14). We pre-compute the cluster index for each document in the pretraining dataset offline, resulting in no training-time overhead. See details in Appendix A.2. At test time, we use the same traversal to get the cluster index for the task context (e.g.. question text).
We assign a memory parameter block to each cluster in given hierarchi- cal clustering tree, denoted by WL, E RR", where / E (1, 2, 4} denotes the level, in E {l,.... k' is the cluster index at level 1, and st is the size of the memory parameter blocks at level I. The mem- ory bank W consists of all memory parameter blocks Wi, The memory retriever is then:
Hierarchical memories:
R(x; W) i W2,2 W3,ia Wa.ii where I(x) (ii, 12, 13, i4).
The total number of parameters in the memory bank is W| = s1k + s2kh² + szk + sakt, and the number of retrieved parameters (fetched memory size) is [2 W)| = 81 + 82 + 83 84.
2.2 INFERENCE WITH PARAMETRIC MEMORIES
There are multiple ways to add parametric memories R(x; W) to an anchor model 0, i.e., to practically model Po, (R): x<t)). We limit our discussion to decoder-only, transformer- based (Vaswani et al., 2017) language models. For all memory types, we instantiate them such that at the beginning of training they have no effect on the anchor model. See details in Appendix I.
popular approach to augment a model with a (small) set of extra parameters is to patch its linear layers with low-rank adaptation matrices (LoRa, Hu et al., 2022). We consider three types of LoRa memories: LoRa-QK adapts the query and key projection layers, LoRa-VO adapts the value and output projection layers, and LoRa-FFN adapts all three linear layers in the SwiGLU feed-forward network (FFN) (Shazeer, 2020) with rank r matrices.
LoRa-Memories:
The knowledge in the context tokens transformer attends to is ultimately a quence of K-caches. A natural extension of providing context knowledge is thus to learn KV-cache parameters directly. This also be seen as a generalization of prefix tuning (Li & Liang, 2021; Lester et al., 2021), where only input token embeddings are learned. At each transformer layer, data-dependent query tokens cross-attend to fetched KV memories.
KV-Memories:
Previous works have argued that transformer-based language models mainly store their knowledge in the FFN layers (Geva et al., 2020; Dai et al., 2022: Yao et al., 2022). Inspired by this, we introduce FFN memories: we expand the inner dimension of the SwiGLU FFN layers by through concatenation with the fetched memory parameters, which equivalent to fast addition.
FFN-Memories:
The number of parameters in memory block assigned to clusters at level depends on the mem- ory type, the anchor model's architecture (e.g., hidden dimension, depth, etc.), and the memory block size multiplier (rank for LoRa, number of KV memory tokens, or FFN dimension expansion size). Therefore, a hierarchical memory configuration ($1, 53, 84) can be written as co(r1, T2, 73, r4), where is constant determined by the anchor model architecture and memory type, and is the memory block size multiplier for memories at level that we control. For simplicity, we drop the constant cy in the rest of the paper and provide details in Appendix I. In practice, we generally set these multipliers so that coarser levels have larger parameter blocks, i.c., ri 2 r2 2 ra 2 T4, or set rt = O when no memories are assigned to clusters at level
4

---

Published as a conference paper at ICLR 2026
Fetched memory to anchor model racio
Fetched memory to anchor model retto
-0-
PFN
1efa-FFN
LoRa-OK
17.0
memory
WilteEn perplesity
Sug.Sk 196)
37
16.5
0
-O- FPN
+ Lofu-FFN
LoRa-OK
15.5
memory
34
10'
10a
10
10t
Fetched memory
Fetched size
(a)
(b)
Fetched memory anchor model rabio
Hemory bank to anchor model ratio
10-1
-O- Levele4
Level:3
memory
Furg Sk (96)
37
36
Level4
- Level1
35
-D- Level3
--- No memory
35
+ Levelzz
30
34
10
10a
10
10"
10t2 Memory bank
10
retched size (M)
(c)
(d)
Figure 3: Effect of memory type on Specific- Knowledge benchmarks (a) and Wikipedia perplexity (b). Effect of memory level on performance function of fetched memory (c) and bank size (d).
DESIGN CHOICES FOR PRETRAINING WITH MEMORIES
We use DCLM-Baseline (Dai et al., 2022) for training. The dataset contains = 3.2 billion documents (z 4.3 trillion tokens). We cluster the dataset tree with levels, having 16, 162; 163,, and 16a clusters at each level. See Appendix A for training details. For evaluation, we consider 13 frequently used benchmarks, including multiple-choice and generative tasks. We divide these tasks into two groups based on the level of specific knowledge required (see Appendices B and G for details):
Kommon-Knowledge (Avg-CK); Lambada-Oper 2019), SQuAD (Rajpurkar et al., 2016), Winograd (Levesque et al., 2012), CoQA (Reddy et al., 2019), and WinoGrande (Sakaguchi et al., 2021).
(Paperno et al., 2016), BoolQ (Clark et al.,
Knowllie (Avg-SK): Hellaswag (Zellers et al., 2019), ArcEasy/Challenge (Clark al., 2018), TriviaQA (Joshi et al., 2017), NaturalQuestions-Opa (Lee et al., 2019; Kwiatkowski et al., 2019), PIQA (Bisk al., 2019), and OpenBookQA (Mihaylov et al., 2018).
For evaluation, we retrieve the memories based only on the question text. As an open-ended gen- eration task, also track the average perplexity the 2022 English Wikipedia (Wiki-En) with ~ 6.5M samples (4B tokens), where memories are retrieved based on the full Wikipedia document.
FFN-Memories outperform LoRa and KY memories: In Fig. 3. we compare different memory types introduced in Section 2.2. We attach the memories to pretrained model with 160M param- eters (row A1 in Table 1) and train them from scratch for 275B tokens with the loss objective in Eq. (1), while the anchor model parameters (0) are frozen. For this set of experiments, we use single-Ievel memory configuration in the form of (0, $2, O, 0). corresponding memories, each with size 82, for different values of 82- We observe that FFN-Memories have a sig- nificant advantage over other forms of memory across memory sizes. Based on this observation, in the rest of the paper we use only FFN-Memories.
total of 162 256
Accuracy improves with deeper and larger memories: Here, we explore the design space of dif- ferent single-level memory configurations (81, $2, 83, 84). where only one of the 's is non-zero. The anchor model is pretrained and frozen (row A1 in Table 1) during memory training (see Ap- pendix A.3 for details). As shown in Fig. 3c, for constant fetched memory size (i.e., 81), deeper memories yield greater accuracy improvements, as they capture more relevant and detailed informa- tion for given query. In Fig. 3d, we show that performance is strictly increasing function of total
5

---

Published as a conference paper at ICLR 2026
memory bank size for all memory configurations, more capacity becomes available capture long-tail knowledge. Shao et al. (2024) recently made a similar observation for regular RAG setups by increasing the size of the datastore from which documents are retrieved. Note that for a fixed total memory bank size, a shallower memory corresponds to a larger fetched memory size. Therefore, in Fig. 3d, a fixed memory bank size, shallower memories achieve higher accuracy.
Hierarchical memories enable an optimal design: Unlike single-level memories used in previous experiments, a general hierarchical memory configuration (31, 82, 53, 84) with possibly multiple non-zero values allows independent control of the total memory bank size (E) 16'st) and the size of fetched memory parameters inference time (Σ) st). For a large total memory bank size with a small number of fetched parameters at inference, we can use larger level-3 and level-4 memories.
Conversely, for smaller total bank size with more fetched parameters at inference, we can increase the sizes of level-1 and level-2 memories.
From a learning dynamics perspective, in regular language modeling, all parameters are updated at every iteration, receiving gradients from wide range of dissimilar documents in the dataset. As result, long-tail information often forgotten (Toneva et al., 2018). When training with hierarchi- cal memories, however, memory bank parameters at level I are activated 16' times less frequently compared to anchor model parameters (which can be considered level-0). Consequently, deeper memory bank parameters receive fewer gradient updates, and those updates come from more simi- lar content, shielding them from forgetting (Ghosal et al., 2025) and enabling them to learn long-tail knowledge. In contrast, shallow memories are updated by many more occurrences of common facts in the dataset, giving commonsense knowledge a greater chance to dominate their gradients. This leads to effective learning of hierarchy of memories, that intuitively range from the most common knowledge at level I to the most specific knowledge at level 4.
Above, we showed that single-level memories benefit from larger fetch size and bank size. We now systematically demonstrate this for a general hicrarchical configuration. We evaluate two groups of hierarchical memories added to our frozen pretrained 160M model (row Al in Table 1): one with a memory bank size of 4.6B and another with 18.7B, both with configurations spanning between 1M and 300M fetched parameters (see Appendix A.4 for details). Results in Fig. 4a confirm that performance increases strictly with fetched memory size (while keeping bank size fixed) and with bank size (while keeping fetched memory size fixed) for the general hierarchical configurations.
The point with highest accuracy in Fig. 4a corresponds to augmenting 1 160M anchor model with ~ 240M fetched memory parameters, achieving 44.5% accuracy on Avg-SK. This is a notable result, showing that an anchor+memory model with a total of 400M runtime parameters outperforms regularly trained 410M model (row B1 in Table 1) by 3.6 points. Building on this observation, we ask the following question: For given runtime parameters budget, how many parameters should be allocated to the anchor model and how many to the fetched memories?
To explore this, we train sequence of anchor models with 260M, 320M, 350M 370M, 385M, and 410M parameters, freeze them, and pair them with fetched memories of sizes 150M, 90M, 60M, 40M, 25M, and O. respectively, such that the total runtime parameter count for anchor+memory re- mains fixed at 410M. All configurations use a 6.3B memory bank (see Appendix A.7). In Fig. 4b, 1:10 ratio of fetched memory anchor model size appears optimal and guides our next experiments. However, this observation may not generalize to different runtime, memory bank, and training bud- gets, or when the anchor model and memory parameters are co-trained.
Models with hierarchical memories have on-device deployment advantages. Assuming a von Neumann architecture with hardware organized in increasing size and decreasing speed (RAM, flash disk, external disk), we can store shallower memory levels (with larger fetch size but smaller bank size) on faster components, while offloading deeper levels to slower storage (Alizadeh et al., 2024). In the example in Fig. 5 (left), with hypothetical hardware setup, a hierarchical memory can be loaded in 38ms, whereas loading a memory of the same fetch size from a flat bank (stored on the external disk due to its excessive total size) takes 198ms more than 5x longer. Additionally, even if both hierarchical and flat memory banks are stored exclusively on the external disk, hierar- chical memories still provide a loading speed advantage due their compositionality. For example, when generating different atomic numbers in Fig. 1, the level-1 and level-2 memories remain mostly unchanged, and only deeper memories need to be swapped. This takes 47ms, compared 198ms when the entire flat memory must be reloaded, as shown in Fig. 5 (right).

---

Published as a conference paper at ICLR 2026
Fetched memory to anchor model racio
Anchor model
410
383370
350
320
260
Hemory bank sim: 4:68 bank memory
Rakens: 65n
Tolens: 6128
Tokeni: 1278
Tolorn:: 4818
Takenz 2068
Tikens: S20B
Tokens:
memery
Wikitn perplesity
Sxg-5 1%)
(256 2236.4416..4
26.0
40-
(256.64.16.0
15.5
38
36
10
10
10
10a
0.0
0.1
0.2
0.3
0.4
0.5
0.6
Fetched memory size (14)
Fetched memory anchor model razia
(a)
(b)
Figure 4: (a) Avg-SK accuracy for different hierarchical memories, demonstrating performance gain with larger bank size and fetched memory size. (b) Wiki-En perplexity for different fetched memory-tooanchhe
model size ratios, with the optimal point at 1:10. The purple curve shows the perplexity of anchor models without memory. The green curves show the perplexity of models with memory, with different shades of green corresponding to the progress of memory training.
Hierarchical
Flat
Hierarchical
HAM
Flat
Compute
RLO.
Memory koad 100 ro 1o1MB
Cache
1:5
Memory Ieavel from B18s8) Dank
Evoe Tovee suB foom (1668)
External Disk
Mamory knading
RAM
151145
Disk
Memory keading ssiency 190 ma
Memory koad 2:47 mm 161M8
Faes SEME from 156GB bark
Externel Disk
hatch (REMS fro B16GBI Dank
Evon 15006
External Disk
1.05
Memory koad 2: 1ESMEl B1aG8 Bark.
Cloud
pum nom (1002)
Figure 5: Deployment advantages of hierarchical memories. Left: Memory loading latency is reduced by using the hardware hierarchy. Right: Latency is reduced by exploiting compositionality over time: larger memories for low-level clusters, once loaded, are less likely to need reloading.
Sensitivity to data clustering. In Table we report sensitivity to the initial data clustering, fixing the 160M anchor model frozen and memory configuration with 4.6M fetch size and 4.7B memory bank size, so clustering is the only difference between experiments. To assess sensitivity to the embedding model, we replace the all all-Min LM-6--22 embedding model (Reimers Gurevych, 2019) with GTE 1 Li et al. (2023), which Grangier et al. (2025) show has better clustering performance; nonetheless, we observe very small sensitivity in the final model performance. We further study the choice of clustering depth p and dividing factor & by considering p 2 and k 256, which yields tasks (Table 2). Note that two-level nested clustering
slight improvements on knowledge-spechee with dividing factor 16 leads suboptimal one-level clustering with dividing factor 256, consistent with this observation. In practice, the optimal choice of clustering and memory configuration also depends the underlying hardware hierarchy used to store and fetch the memory bank.
RESULTS
In this section, building on the findings from Section we provide a comprehensive set of results for different-sized models and compare them with baselines. See Appendix A for training details.
Starting from a 160M anchor model pretrained regularly (row A Table 1), we add memories with the (256, 64, 16, 0) configuration, corresponding a fetched memory size of ~ 18M parameters and a total memory bank size of ~ 4.6B parameters. When memories are learned with a frozen anchor model (row A3), observe +5.1 points improvement on Avg-SK compared to A1 and 2 points reduction in Wiki-En perplexity, demonstrating the effectiveness of memories for tasks requiring specific knowledge, with only ~ 10% additional runtime parameters from fetched memories.
To ensure a fair comparison, we also train a generic memory with the same size as the fetched memories (18M parameters) together with the memory bank parameters. When evaluating with
7

---

Published as a conference paper at ICLR 2026
generic memory, unlike fetched memories that are retrieved based context, simply use the anchor+generic memory parameters. This isolates the effect of merely increasing the number of parameters and training tokens in the anchor model. Anchor+generic memory scores 34.7% on Avg- SK, 4.5 points below fetched memories, showing the benefit of context-based retrieval in isolation.
We next explore co-training the anchor model parameters (0) and the memory parameters (W) together, as in Eq. (1). For a fair comparison, during training we use the generie memory with prob- ability 1/(16- - 1) and the fetched memory with probability 16/ + 1), where 16 is the clustering division factor. This ensures there is no training bias in favor of the memory bank parameters.
Row A2 in Table 1 shows co-training results, with Fig. 6a illustrating training curves for A2 and A3 experiments. A key observation is that when we allow the anchor parameters to be co-trained with the memory bank, we obtain greater improvement compared to the case where the anchor model is frozen (Avg-SK 39.2% - 40.3%). This gain can be attributed to two factors: 1) when co-training, the anchor model learns to utilize the memories more effectively compared to when it is frozen, and 2) the anchor model performance improves simply due to more training. The latter effect should be minor if the anchor model is already converged. We provide additional discussion in Appendix E.
We also explore co-training the anchor model and memory bank together from scratch. Results are shown in row A4 of Table 1. Despite using the same total training budget as row A2, A4 shows lower performance. This result suggests that memories are learned more effectively after the anchor model has been trained to some extent (i.e., the setup of row A2). This is analogous to human memory, which develops only after the brain gains semantic understanding, around age 3 (Shaw, 2016).
Next, we scale the co- training setup of row A2 to anchor models of sizes 410M and 1.4B, corre- sponding rows B2 and C2 in Table 1. These models are paired with memory configurations such that the fetched memory is approximately 10% of the anchor model size, corresponding to memory bank sizes of 12.7B and B1.1B for the 410M and 1.4B anchor models, respectively. We observe similar improvements for these larger models: for Avg-SK, fetched memories outperform generic memories by +4.1 points for the 410M and +3.6 points for the 1.4B model.
Finally, for the 410M model, we train a model regularly (without memories) with the same total training budget as row B2 (2.2T tokens). We observe that Avg-CK performance worse than that of the anchor+generic setup in row B2. These preliminary results suggest that when the anchor model co-trained with a large memory bank (as in row B2), long-tail knowledge is offloaded to the memories, enabling the anchor model to perform better on common-knowlede, tasks.
Inference cost: In Table we show inference time (ms) for the 1.4B anchor model without and with memory (rows C1 and C2 in Table 1) on H100 NVIDIA GPU. We consider two setups: (1) memory bank on GPU High Bandwidth Memory (HBM) and (2) memory bank on CPU DRAM. Routing runs on CPU once per generation and adds only minor overhead. With all overheads included (routing, fetch, and extra inference flops), the total time increases by less than 10% for a 40-token generation.
Blocking part of the memory bank: In Fig. 6b, we show the performance of the 410M model with memories (row B2 in Table 1) on predicting atomic numbers of elements (see Fig. 1) when the best-matching parts of the memory bank are adversarially blocked during retrieval. We observe a significant performance drop, from 70% to 20%, when blocking 1/16 of the bank. This preliminary result highlights the potential of the proposed approach for applications with privacy goals.
Failure analysis of retriever: In Appendix J we present analysis of model performance when wrong memories are fetched. We observe that even when fetched memories are not the nearest neighbors in the clustering, overall performance on par with or better than the no-memory baseline.
Adding memory to other pretrained models: We explore the post-hoc addition of parametric memories to wpen-weight models. We span multiple sizes and architectures, namely Gemma 3 270M (Gemma Team et al., 2025). Qwen 2.5 0.5B (Yang et al., 2024), and Llama 3.2 IB (Meta AI, 2024). We add hierarchical FFN memories of ~ 10% the model size host-hoc to pretrained (and frozen) anchor models and train the memories for 1.IT tokens on DCLM, see Appendix A.6 for additional details. Results are shown Table 4. As above, the specific knowledge accuracy improves with memories. This is consistent across all architectures, showing the generality of adding hierarchical memories across all models. Common knowledge remains the same or slightly decreases; potentially
8

---

Published as a conference paper at ICLR 2026
Ancloo model
Seen tokeus
Cotrain anchor
Memory config
Bank size
Fetch
Avg-CK (%)
Avg-SK (
WikiEn Ppix 1
Row
Init.
size
Genenic
Fesched
Genetie
Fetched
Generic
Fetched
A1
Seratch
1.IT
min
10, 0,
0
0
45.9
34.1
17.2
A2
160M
Al
+1.1T
Yes
(256, 0)
4.6B
18M
46.6
47.9
47.4
48.7
35.7
40.3
16.7
14.2
A4
Sentch
2.27
Yos
466
46.7
33.8
39.6
17.8
15.6
R1
Seniich
1.1T
16, 0,
0
0
52.3
40.9
13.9
H2
410M
B1
41.1T
Yes
1512,128, 52,01
12.78
50M
55.5
56.1
41.8
45.9
13,8
12.4
B3
Scratch
2.21
1,,3,0
53.2
13.8
C1
Sontch
w/a
0,
49.7
C2
C1
+1.17
Yes
(168,256,18,0
21:18
153M
64.4
64.5
51.3
54,9
11.0
10.2
Table I: Results for pretraining with memories different scales. See full results in Table 11.
80
40-
Atamic number prediction acc. 1%))
410M enchor-generis frow 82 in teb 21
410M sntchr-fetched frow in Table
80
38
50
Leaerablee memury
40
Learnable anchen generic memory
Frozen anchor fetched memory
16
36-
18
20
30
200
400
600
800
1000
0/16 1/16
2/16
4/16
8/16
takens [B]
Fraction of memory bank blocked
(b)
Figure 6: (a) Performance improvements Avg-SK when co-training memories and anchor model parameters jointly during training. (b) Effect of blocking parts of the memory bank from retrieval.
because these models were pretrained on more tuned data mixtures than the simple open-source data mixture DCLM that we use in these experiments.
Comparing with vanilla approach to parametric memory is retrieval-auge auggeeeee generation (RAG), where relevant texts are retrieved from a datastore and prepended to the context (Lewis et al., 2020; Ram et al., 2023; Izacard al. 2023) to improve performance on knowledge-intene intensivee tasks. To study the two memory mech- anisms, contextual memories as in RAG and parametric memories in the proposed approach, in controlled experiment, we use our Sentence- training data as the datastore, making the memory mechanism the only difference. We then evaluate RAG on the baseline models from rows A1, B1, and C1. See implementation details in Appendix H.
generation: An alternative, yet complementary,
embedding model as the retriever and the DCLM
As shown in Table 5, vanilla retrieval from DCLM performs poorly relative to the baseline models, while adding more than 2x runtime FLOPs and requiring large storage for the raw-document data- store. This is likely due to the low quality of DCLM pretraining dataset) when used as a RAG datastore. To give RAG an advantage, we also retrieve from the higher-quality Wiki-En. RAG-Wiki improves baseline performance on SK (e.g.. from 46.9 49.2 for the 1.4B model) while remain- ing slightly below baseline on CK. By contrast, learned memories (with ~ 10% extra parameters) improve both CK and SK, with lower FLOPs overhead. We note that more sophisticated RAG system is possible for coaructtionfollowingg models after an appropriate post-training stage, and is complementary to the proposed pretraining cearee--emorry approach. The two can be combined for additional gains. As a demonstration of this complementarity. combining RAG-Wiki with 10% parametric memories for the 410M model in Table boosts AVG-SK performance to 45.7%, sur- passing either RAG-Wiki (41.6%) the proposed parametric memories (44.5%) alone.
5 RELATED WORKS
Databases. (Ahn et al., 2016) use symbolic knowledge graph, and Borgeaud et al. (2022) introduce Retro, augmenting language-move predictions with a large raw-text memory bank. More recently, Zhao et al. (2025) propose replacing long-tail knowledge with retrieval base by masking retrieved tokens during pretraining. Limitations of these approaches are scalability to large pretraining corpora and low compression rate because the memory stores raw text.
external knowledge
Parametric memories. Wu et al. (2022) propose memorizing transformers, which use nearest- neighbor lookup to sparsely retrieve cached key-value pairs. For efficiency, Eyuboglu et al. (2025)

---

Published as a conference paper at ICLR 2026
Cluster config
Fmnhed. model
Memory config
AVG-CK (%)
AVG-SK
Num
Model
Bank
Routing
Fetch
Predill and
Total
(%)
tokens
setup
Loc.
time
time
Generation time
time
16
MaiLM
V6
164, 4,
42.8
36.6
20
20
w/ mem.
Baseline
HBM
IVa
mss
n'á
0
507 ms
534 ms
541 ms
507 ms
4 16
small
GTB
4,,42
42.6
36.2
20
w/
Haseline
DRAM
IVa
nia
1017
534
1017 ms
564
== 256
2
MiiiLM V6
184,1)
43.1
37.5
40
40
w/ mem.
w/ DRAM
HBM
7 ms
23
0
1040
1040 ms
1070
1047 ms
Table 2: Clustering ablation.
Table 3: Runtime cost comparison on single H100 GPU.
Pretrained model
Bank
Fetch
Avg-CK (9r
Avg-SK (%)
Atomic Number Acc. (%)
Anchur model
Inference setup
Bank size
Avg-CK 0-shot
Avg-SK
0-sbot
RAG-DCLM
32.6
Gemuna 3 270M
44.2
34.3
1.7
15.0
+ memory
5.98
23 2M
44.8
38.2
49.2
10% Memory
9GB
x1.11
45.3
38.4
Qwen 2.50.55
53.9
40.6
53.4
48.6
Liama
memory
43 43
52.1
58.9
90.4
96.6
410M
10% Monuty
RAG-DCTMM
RAG--Wiki
GB
+ memory
M.JB
1M2.2M
57.6
50.5
96.6
Baseline
46.9
RAG DCLM
70
x2.3
55.8
46.1
Table 4: Learning memory on top of pretrained open-weigh models post-hoc. All trainings use 1.1 trillion tokens from DCLM.
RAG-Wiki
49/2
10% Menory
4B
52.4
Table 5: Comparison with vanilla RAG.
introduce Cartridges: KV-memories (similar to what discussed in Section 2.2) that learn spe- cific long document as a more runtime-efficiee alternative to in-context learning. We find that KV- memories underperform compared FFN-memories, at least for large-scale memorization. Mem- Sinks (Ghosal et al., 2025) uses a type of FFN-memories, where they dedicate a fraction (e.g.. 30%) of FFN neurons per layer memorization. However, their goal is to throw those parameters away inference time for privacy and generalization. Our context-dependen memory retrieval is also con- ceptually related to fnslluction-folloiing pruning Hou al. (2025), which selects the most suitable parameters from a larger model based on the task description.
Mixture of experts. Our approach related MoEs (Shazeer et al., 2017). Jelassi al. (2024) show that increasing the number of experts improves memorization while reasoning saturates when active parameters are fixed, aligning with our anchor-memory decomposition to balance reasoning and knowledge. For privacy, Shi et al. (2025) propose FlexOlmo, combining a publicly trained anchor expert with exchangeable domain experts trained on private data. Product key memory (PKM) (Lample et al., 2019) can be seen as a sparser MoE that combines two selected expert sets; Huang et al. (2024) improve PKM via tensor decomposition, and Berges et al. (2024) augment sub- sets of FFN layers with such memory to boost factual benchmarks. These MoE approaches are similar in kind to ours, and we expect that some of our insights may carry over. However, our mem- ory architecture is vastly different, allowing to offload inactive parameters based on the context, give explicit control over memories during both training and inference, and add memories post-hoc.
6 DISCUSSION
Conclusion. We propose pretraining language models with memories to automatically capture long- tail world knowledge in large hierarchical memory banks. Small language models augmented with memory banks match regular transformers with 2x more parameters. Moreover, we propose and analyze several additional potential benefits of this design, including more efficient hardware imple- mentation and enhanced data privacy.
Limitations and future directions. One unexplored direction is the study of optimal scaling law formulas for learning memories. Compute-optimal scaling laws developed for dense training (Hoff- mann et al., 2022), where training tokens can potentially update all learnable parameters, are not necessarily applicable to pretraining with memories, since memory parameters are updated less frequently and receive gradients only from subset of the dataset. We also mainly focused on archi- tecture design for memories, and leave architecture search and design of anchor models, including dynamic and learned retriever mechanisms, as promising future work. Another future direction is to explore methods to improve the anchor model's 's reasoning capabilities when co-trained with memo- ries. Finally, pretraining with memories can benefit multilingual setups (not considered in this work) or other modalities beyond text, such as vision. We leave this investigation for future work.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
Our implementation is based on the publicly available repository OpenLM (Gururangan et al,, 2023b), which we will release after the review period. In particular, our proposed memories can be added to arbitrary open-weight models and trained on any data of interest. We expect that leasing this code will enable other researchers to further investigate this direction. To ensure full transparency of our setup, especially during the review period, provide detailed implementation descriptions in the supplementary material, including training details in Appendix A, evaluation setup in Appendix B, and clustering details in Appendix C.
ACKNOWLEDGMENTS
We would like to thank Cheng- Yu Hsieh, Samira Abnar, Stephen Pulman, Karen Khatamifard, Chun-Liang Li, and Rick Chang for their feedback.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint wrXiv:2303.08777 2023.
Sungjin Ahn, Heeyoul Choi, Tanel Pärnamaa, and Yoshua Bengio. A neural knowledge language model. ArXiv, 1bs/1608..03318 2016. URL https: /aapi. semanticscholas org/ CorpusID: 2600027.
Keivan Alizadeh, Seyed Iman Mirzadeh, Dmitry Belenko, S Khatamifard, Minsik Cho, Carlo C Del Mundo, Mohammad Rastegari, and Mehrdad Farajtabar. Llm in flash: Efficient large lan- guage model inference with limited memory. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 12562-12584, 2024.
David Arthur and Sergei Vassilvitskii. k-means+4: The advantages of careful seeding. Technical report, Stanford, 2006.
Akari Asai, Zeqiu Wu, Yizhong Wang, Avirup Sil, and Hannanch Hajishirzi. Self-rag: Learning to retrieve, generate, and critique through self-reflection. 2024.
Vincent-Piera Berges, Barlas Oguz, Daniel Haziza. Wen-tau Yih, Luke Zettlemoyer, and Gargi Ghosh. Memory layers at scale. arXiv preprint arXiv:2412.09764, 2024.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. In AAAI Conference Artificial Intelligence, 2019. URL https: //ppi. semant commanticcchoo... Iog/Corppss:: 208200:00099000
Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, al. Gpt-neox-20b: An open-source autore- gressive language model. arXiv preprint 2rX1v:2004.06745, 2022.
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Milli- can, George Bm Van Den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, et al. Improving language models by retrieving from trillions of tokens. International conference on machine learning, pp. 2206-2240. PMLR, 2022.
Christopher Clark, Kenton Lee, Ming-Was Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. BoolQ: Exploring the surprising difficulty of natural yes/no questions. Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Lan- guage Technologies, Volume (Long and Short Papers), pp. 2924-2936, Minneapolis, Min- nesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v 1300. URL https: //ss://co a antho ogy. org/Nl 9

---

Published as a conference paper at ICLR 2026
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint 1003vv880000544 2018.
Damai Dai, Li Dong, Yaru Hao, Zhifang Sui, Baobao Chang, and Furu Wei. Knowledge neurons in pretrained transformers. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings of the 60th Anmual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 8493-8502, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10.1 2022. aww-aclloooooo 581/.
URL https: /aa/aaaaaannntooooooooooao .orgg
Sabri Eyuboglu, Ryan Ehrlich, Simran Arora, Neel Guha, Dylan Zinsley, Emily Liu, Will Tennien, Atri Rudra, James Zou, Azalia Mirhoseini, et al. Cartridges: Lightweight and general-purpose long context representations via self-study. arXiv preprint arXiv:2506. 06266, 2025.
Gemma Team, Aishwarya Kamath, Johan Ferret, Shreya Pathak, Nino Vieillard, Ramona Merhej, Sarah Perrin, Tatiana Matejovicova, Alexandre Ramé, Morgane Rivière, et al. Gemma 3 technical report. arXiv preprint 250.. 2025.
Mor Geva, R. Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forwand layers are key- value memories. ArXiv, abs/2012. 2020. URL https: / /api. semant cacholas lar. org/Corpus D 229 229923720.
Gaurav R Ghosal, Pratyush Maini, and Aditi Raghunathan. Memorization sinks: Isolating memo- rization during Ilm training. arXiv preprint TrXiv:2507.09937, 2025.
David Grangier, Simin Fan, Skyler Seto, and Pierre Ablin. Task Task-adaptive pretrained language mod- cls via clustered-importance sampling. In Proceedings of the International Conference on Learn- ing Representations (ICLR), 2025. URL https: http://...........00/////....... 03735.
Suchin Gururangan, Margaret Li, Mike Lewis, Weijia Shi, Tim Althoff, Noah A. Smith, and Luke Zettlemoyer. Scaling expert language models with unsupervised domain discovery. ArXiv, abs/2303.14177 2023a. URL https://apt http:///appp.aammmmiiccooooooooo arggCorppsID: 257756896.
Suchin Gururangan, Mitchell Wortsman, Samir Yitzhak Gadre, Achal Dave, Maciej Kilian, Wei- jia Shi, Jean Mercat, Georgios Smyrnis, Gabriel Ilharco, Matt Jordan, Reinhard Heckel, Alex Dimakis, Ali Farhadi, Vaishaal Shankar, and Ludwig Schmidt. OpenI.M: minimal but performative language modeling (Im) repository. 2023b. URL https:://ttt whwggthhhh....... comb mlfoundations/ oppnnn////p/e/////eeeaaaa lm/. GitHub repository.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hen- nigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Oriol Vinyals, Jack W. Rae, and Laurent Sifre. Training compute-optimal large language models. In Proceedings of the 36th International Con- ference on Neural Information Processing Systems, NIPS '22, Red Hook, NY, USA, 2022. Curran Associates Inc. ISBN 978171387 1713871088.
Bairu Hou, Qibin Chen, Jianyu Wang, Guoli Yin. Chong Wang, Nan Du, Ruoming Pang, Shiyu Chang, and Tao Lei. Instruction-follovii pruning for large language models. arXiv preprint wwXiv:2501.02088 2025.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. ICLR, 1(2):3, 2022.
Zihao Huang, Qiyang Min, Hongzhi Huang, Defa Zhu, Yutao Zeng, Ran Guo, and Xun Zhou. Ultra- sparse memory network. arXiv preprint arXiv:2411.1236/ 2024.
Gautier Izacard, Patrick Lewis, Maria Lomeli, Lucas Hosseini, Fabio Petroni, Timo Schick, Jane Dwivedi- Yu, Armand Joulin, Sebastian Riedel, and Edouard Grave. Atlas: few-shot learning with retrieval augmented language models. J. Mach. Learn Res., 24(1), January 2023. ISSN 1532-4435.
12

---

Published as a conference paper at ICLR 2026
Herve Jegou, Matthijs Douze, and Cordelia Schmid. Product quantization for nearest neighbor search. IEEE transactions on pattern analysis and machine intelligence, 33(1):117-128, 2010.
Samy Jelassi, Clara Mohri, David Brandfonbrener, Alex Gu, Nikhil Vyas, Nikhil Anand, David Alvarez-Melis, Yuanzhi Li, Sham M Kakade, and Eran Malach. Mixture of parrots: Experts improve memorization more than reasoning. arXiv preprint arXiv:2410. 19034, 2024.
Mandar Joshi, Eunsol Choi, Daniel Weld, and Luke Zettlemoyer. TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension. In Regina Barzilay and Min- Yen Kan (eds.), Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp. 1601-1611, Vancouver, Canada, July 2017. Association for Com- putational Linguistics. doi:
URL https: http///aaaaantotoooooo /
P17-1147/.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453 16633466,245 2019.
Guillaume Lample, Alexandre Sablayrolles, Mare' Aurelio Ranzato, Ludovic Denoyer, and Herve Jegou. Large memory layers with product keys. Curran Associates Inc., Red Hook, NY, USA, 2019.
Kenton Lee, Ming-Wei Chang, and Kristina Toutanova. Latent retrieval for weakly supervised open domain question answering. arXiv preprint arXiv: 2019.
Brian Lester, Rami AI-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. In Marie-Francine Moens, Xuanjing Huang, Lucia Specia, and Scott Wen- tau Yih (eds.), Proceedings of the 2021 Conference on Empirical Methods in Natural Lan- guage Processing, pp. 3045-3059, Online and Punta Cana. Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10. (ttp3//112222.mmmm// mapppmmii...... 243. URL org/ 2021.
Hector J. Levesque, Emest Davis, and Leora Morgenstern. The winograd schema challenge. In Pro- ceedings of the Thirteenth International Conference on Principles of Knowledge Representation and Reasoning, KR 12, pp. 552-561. AAAI Press, 2012. ISBN 9781 57135775560..
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, Sebastian Riedel, and Douwe Kiela. Retrieval-augmeettd generation for knowledge-intensive nip tasks. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713833 38295446.
Jeffrey Li, Alex Fang, Georgios Smyriis, Maor Ivgi, Matt Jordan, Samir Gadre, Hritik Bansal, Etash Guha, Sedrick Keh, Kushal Arora, Saurabh Garg, Rui Xin, Niklas Muennighoff, Reinhard Heckel, Jean Mercat, Mayee Chen, Suchin Gururangan, Mitchell Wortsman, Alon Albalak, Yonatan Bitton, Marianna Nezhurina, Amro Abbas, Cheng- Yu Hsieh, Dhruba Ghosh, Josh Gardner, Macicj Kilian, Hanlin Zhang, Rulin Shao, Sarah Pratt, Sunny Sanyal, Gabriel Ilharco, Giannis Daras, Kalyani Marathe, Aaron Gokaslan, Jieyu Zhang, Khyathi Chandu, Thao Nguyen, Igor Vasiljevic, Sham Kakade, Shuran Song, Sujay Sanghavi, Fartash Faghri, Sewoong Oh, Luke Zettlemoyer, Kyle Lo, Alaaeldin El-Nouby, Hadi Pouransari, Alexander Toshev, Stephanie Wang. Dirk Groeneveld, Luca Soldaini, Pang Wei Koh, Jenia Jitsev, Thomas Kollar, Alexandros G. Dimakis, Yair Carmon, Achal Dave, Ludwig Schmidt, and Vaishaal Shankar. Datacomp-lm: In search of the next generation of training sets for language models. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomezak, and C. Zhang (eds.), Advances Neural Information Processing Systems, volume 37, pp. 14200-14282. Curran Associates, Inc., 2024. / /ttpss:/proceeedingg neur iiw.rriss.... ctpppeeeeeeeeeeleeeeeeeeeeeeeeee Patassassssttettpeaaaaaaatatstss and Benchmarks.
URI.
files/ welll//pppe///ee/ paper/2022 /2024/f122/
Track.pdf.
Margaret Li, Suchin Gururangan, Tim Dettmers, Mike Lewis, Tim Althoff, Noah A. Smith, and Luke Zettlemoyer. Branch-train-merge: Embarrassingly parallel training of expert language
13

---

Published as a conference paper at ICLR 2026
models. ArXiv, abs/2208.03306, 2022. URL CorpusID:25137135
Xiang Lisa Li and Perey Liang. Prefix-tuning: Optimizing continuous prompts for generation. Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Langnage Processing (Volume 1: Long Pa- pers), pp. 4582 4597, 2021. URL httph:://///......mm htt.p:///app.mmmmiccisiii.e 230433941.
Zehan Li, Xin Zhang, Yanzhao Zhang, Dingkun Long, Pengjun Xie, and Meishan Zhang. Towards general text embeddings with multi-stage contrastive learning. arXiv preprint arXiv:2308.03281, 2023.
Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Compulational Linguistics, 12:157-173, 2024.
Xinxi Lyu, Michael Duan, Rulin Shao, Pang Wei Koh, and Sewon Min. Frustratingly simple retrieval improves challenging, ceassnning-nntensive benchmarks. arXiv preprint arXiv:2507.01297 2025.
Kevin Meng. David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances in neural information processing systems. 77777399977777 2022.
Meta AL. Llama 3.2 model card. 2024. URL httpp://wwwt ama com/ com//ocss/
model lavel-casn sarde--card-aadd annidss-nddprrpppp. www.ats/ttaaooooo 2eormats/llamaa221 1lama3_2/.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. In EMNLP, 2018
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring broad discourse context. arXiv preprint
2016.
Hadi Pouransari, Chun-Liang Li, Jen-Hao Rick Chang, Pavan Kumar Anasosalu Vasu, Cem Koc, Vaishaal Shankar, and Oncel Tuzel. Dataset decomposition: faster Ilm training with variable sequence length curriculum. In Proceedings of the 38th International Conference on Neural Information Processing Systems, NIPS '24, Red Hook, NY, USA, 2025. Curran Associates Inc. ISBN 77798331314385
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. In Jian Su, Kevin Duh, and Xavier Carreras (eds.), Pro- ceedings of the 2016 Conference on Empirical Methods Natural Language Processing, pp- 2383 2392, Austin, Texas, November 2016. Association for Computational Linguistics. doi: N0.18653/v1/D166 1264. URL https::/p/// /ccllcccccc http::sss//caaaatoolooooooooo
6 6-1264/
Ori Ram, Yoav Levine, Itay Dalmedigos, Dor Muhlgay, Amnon Shashua, Kevin Leyton- Broown and Yoav Shoham. In-context retrieval-augmennee language models. Transactions of the Association for Computational Linguistices 1:1316-13311 2023. doi: a0.1162/tall URL https:
//acllanthology.coom org/ http://223..aa.... taat-alllaa tall.....clll.......
Siva Reddy, Danqi Chen, and Christopher D. Manning. CoQA: conversational question answering challenge. Transactions of the Association for Computational Linguistics, 7:249-266, 2019. doi:
10.1162/tacl a. 00266.0 URL http::/:////aaaaahhlaoccoo httpahhoo.oog...rr///11111111110
Nils Reimers and Iryna Gurevych. Sentence-bert: Sentence embeddings using siamese bert- networks. In Proceedings of the 2019 Conference on Empiricul Methods in Natural Language Processing. Association for Computational Linguisties, 11 2019. URL https: /ttps://arxiv..
org/abs 1908 10084.
Adam Roberts, Colin Raffel, and Noam Shazeer. How much knowledge can you pack into the parameters of language model? In Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu (eds.). Proceedings of the 2020 Conference on Empirical Methods in Natural Language Process- ing (EMNLP), pp. 5418-5426, Online, November 2020. Association for Computational Linguis- URL ht attps: / wwlaantnnhloooooooee or2/2020.
tics. doi: 10.18653/v 00066////1/22000000 1/2020 w2w00....mlpmmai......
emn i n. 437/.
14

---

Published as a conference paper at ICLR 2026
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: adver- sarial winograd schema challenge at scale. Commun, ACM, 64(9):99-106, August 2021. ISSN 0001-0782. doi: 10.1145/3474381 URL https: //ttpp:://oi...// /ttp://////o.oo//////// 0ttp:///..1111/444333333
Rulin Shao, Jacqueline He, Akari Asai, Weijia Shi, Tim Dettmers, Sewon Min, Luke Zettlemoyer, and Pang Wei W Koh. Scaling retrieval-basee based language models with a trillion-token datastore. Advances in Neural Information Processing Systems, 37:91260-91299, 2024.
Julia Shaw. The memory illusion: Remembering, forgetting, and the science of false memory. Ran- dom House, 2016.
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05200) 2020.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely- gated mixture-of-experr layer. arXiv preprint arXiv: 1at1::1700065333. 006538, 2017.
Weijia Shi, Akshita Bhagia, Kevin Farhat, Niklas Muennighoff, Pete Walsh, Jacob Morrison, Dustin Schwenk, Shayne Longpre, Jake Poznanski, Allyson Ettinger, et al. Flexolmo: Open language models for flexible data use. arXiv preprint (eXXi::507.00022 2025.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
Mariya Toneva, Alessandro Sordoni, Rémi Tachet des Combes, Adam Trischler. Yoshua Bengio, and Geoffrey J. Gordon. An empirical study of example forgetting during deep neural network learning. ArXiv, abs/1812.05159, 2018. URL https: aapi semant cschola .rrg/ CorpusID: 55481903.
V0.9. Llm foundry v0.7.0. https:/ httgggitttt.oooooeoeeeoeooee hot.ii/mmmmmmmooddrrm...mm
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL
neurips cc http/pper__flee/peeeeeeeeeeeeeee wtw.ss/ppppeee/117777 file/ htt.aa/eee222seeeeeeeoooooaaaaaa
påf.
Yuhuai Wu, Markus N Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. arXiv preprint 2rXiv:2203.08111 2022.
An Yang, Baosong Yang. Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, Huan Lin, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang, Jiaxi Yang. Jingren Zhou, Junyang Lin, Kai Dang, Keming Lu, Keqin Bao, Kexin Yang, Le Yu, Mei Li, Mingfeng Xue, Pei Zhang, Qin Zhu, Rui Men, Runji Lin, Tianhao Li, Tingyu Xia, Xingzhang Ren, Xuancheng Ren, Yang Fan, Yang Su, Yichang Zhang, Yu Wan, Yuqiong Liu, Zeyu Cui, Zhenru Zhang, and Zihan Qiu. Qwen2.5 technical report. arXiv preprint arXiv:2412.15115. 2024.
An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint w0X1::2505.09388 2025.
Yunzhi Yao, Shaohan Huang, Li Dong, Furu Wei, Huajun Chen, and Ningyu Zhang. Kformer: Knowledge injection in transformer feed-forward layers. In Wei Lu, Shujian Huang, Yu Hong, and Xiabing Zhou (eds.), Natural Language Processing and Chinese Computing, pp. 131-143, Cham, 2022. Springer International Publishing. ISBN 978-3-031-17120-88
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a ma- chine really finish your sentence? arXiv preprint arXiv: /905.09:90000000 0095000300, 2019.
Linxi Zhao, Sofian Zalouk, Christian K Belardi, Justin Lovelace, Jin Peng Zhou, Kilian Q Wein- berger, Yoav Artzi, and Jennifer J Sun. Pre- training large memory language models with internal and external knowledge. arXiv preprint arXiv:2505. /5962, 2025.
15

---

Published as a conference paper at ICLR 2026
A TRAINING DETAILS
In this section, provide the training details of all experiments. All experiments in the paper conducted using the OpenLM (Gururangan et al., 2023b) repository'. The training data is DCLM- Baseline (Li et al., 2024). Rotary Positional Embedding (RoPE) (Su et al., 2024) is used to encode positions in queries and keys before the attention module, with frequency of 100,000. For all training jobs, anchor and memory parameters stored in BF precision and trained with PyTorch Fully Sharded Data Parallelism (FSDP) using AdamW optimizer (B1 = 0.9, 32 = 0.95) and gradient clipping to 1.0. We use NVIDIA H100 GPUs for all experiments in this paper.
A.1 BASELINE MODEL TRAINING
The baseline models, corresponding to rows A1, B1, and C2 in Table 1, are transformers with the block shown in Fig. 13 and configurations detailed Table 6, Table 7, and Table 8, respectively. Each model is trained from scratch with 240 = 1.1T tokens, using a context length of 8, 8,192 and the dataset decomposition approach (Pouransari et al., 2025) with global batch size of 1M tokens. The learning rate follows cosine schedule with 10k warmup steps, maximum value of 5 x 10 10-3, and a cooldown value of 5 x 10 5 Weight decay set to 0.05.
Model
OPM
Model
OpenLM-410M
Model
OpenLM-
Num layers
35
Num layers
24
Numb layers
24
Hidden dim
512
Hidden dim
1,024
Hidden dimension
2,048
Num heads
12
Num heads
16
Num heads
16
Per head dim
Per head dim
64
Per head dim
128
FEN inner dim
2,048
FFN inner dim
2,816
FFN inner dim
5,632
Head-Embeddiing
tied
Head- Embedding
separate
Head-Emberr Embedding
separate
Num params
163,510,016
Num params
411.66 111.665,60 4088
Num params
1.439,893,504
Table 6: OpenLM- 1600
Table 7: OpenL.M-10MM
Table 8: OpenLM- 1B
A.2 MEMORY MODEL TRAINING
For all memory learning jobs, first perform clustering described in Appendix C. After iden- tifying the corresponding level 4 cluster ID) of each document (a number between and 164), we pack documents from the same cluster into sequences of length 2.048, globally shuffle them, and add the cluster ID as a prefix to each sequence. During training, the cluster ID is simply obtained by separating the first token from the rest of the tokens, which represent the actual data. Note that shallower-level cluster IDs can be inferred from the level-4 cluster ID due to the nested structure of our clustering. Each sequence of length 2,048 can contain subsequences from different documents (within the same cluster), separated by special EOT token. The attention mask is restricted to each document to avoid cross-document attention. The global batch size for all jobs 2M tokens (1 1,0224 sequences of length 2,048 each), except for the 1.4B model (row C2 in Table 1), where use global batch size of 4M to improve GPU utilization.
When the anchor model is frozen, we train memories asynchronously (one job for the memories of each level-1 subtree, resulting in a total of 16 jobs) and merge the checkpoints afterward
In addition, we cosine learning rate schedule with maximum value of 10-4 and cool-down value of 10 When training for 240 tokens, we use 10k warmup steps. and for all other trainings with different numbers of tokens, keep the same warmup ratio. We found that warmup has minimal effect on the performance memory learning. We use a weight decay value 10-3, which found improve the performance of memory learning, consistent with its goal of memorization.
A.3 ADDITIONAL DETAILS OF MEMORY TYPE AND SIZE ABLATIONS
The total number of seen tokens for the memory type and size experiments in Fig. 3 is 238 For all types of memories considered, the memories correspond to level-2 clusters.
hhttps://gi thub. com/m/ Foundat ens/ hons/pppe_1mm
16

---

Published as a conference paper at ICLR 2026
For the experiments in Fig. 3d, the total seen tokens are 2 2388, 240, and 241 for memories at levels 1, 2, 3, and 4, respectively. Since deeper memories are updated less frequently, we increased the total seen token budget for those cases.
A.4 ADDITIONAL DETAILS OF ABLATION ON MEMORY CONFIGURATION
For the experiments in Fig. (corresponding Table 9), we use 240 and 241 total seen tokens, corresponding to memory bank sizes of 4.6B and 18.7B, respectively.
Memory config
Bank size
Fetched size (M)
Common-Knoviedge (%)
Specifie-nowleddee (%)
level
level
level 3
level
0
0
0
45.9
34.1
0
16
4
1.1
46.5
37.8
98
42
18
8.5
46.5
38.6
256
64
18.1
47.4
39.2
4.6B
768
96
12
47.1
47.6
40.2
1792
112
102.7
48.3
41.6
5376
289
49.7
43.2
256
64
18.1
47.2
40.1
328
156
74
30
46.3
40.9
18.7B
1216
164
75.5
47.8
41.6
4096
320
238.4
49.2
44.6
Table 9: Ablation different hierarchical memory configurations corresponding results shown in Fig. 4a. For all experiments, the anchor model the 160M model (row A1 in Table 1), frozen during memory learning. Memories trained for 1.1T and 2.2T tokens when the memory bank size is 4.6B and 18.7B, respectively.
A.5 ADDITIONAL DETAILS OF CO-TRAINING MEMORIES AND ANCHOR MODEL
For the experiments in Table 1, we denote the total number of seen tokens, which is either 2 or 242 For the A4 model in Table 1, where both memory and anchor model parameters are trained together from scratch, cool-down value of 10 with 10k warmup steps and weight decay set to 0.1.
cosine learning rate schedule with maximum value of 10-3 and
A.6 ADDITIONAL DETAILS OF MEMORY LEARNING ON OPEN-WEIGIIT MODELS
For all experiments in Table 4, we use total of 20 tokens. All models are initialized from their public pretrained checkpoints, and these anchors remain frozen. We use a memory config (512, 128, 32, 0) for Gemma and Qwen (which results in higher memory parameters for Qwen because it has higher internal dimension and more transformer blocks) and (768, 256, 32, 0) for Llama. These numbers result in roughly 10% memory to anchor model parameter ratio. We follow the optimizer setup Appendix A.2, except that we do not restrict cross-attention because the baseline models were not all trained this way and reduce warmup steps 5k.
A.7 ADDITIONAL DETAILS OF SEMORY-TO-MODES SIZE EXPERIMENT
Here we provide implementation details for the experiments in Fig. 4b.
We summarize the architecture of six anchor models and their corresponding memories for this experiment in Table 10. The anchor models are first trained for 238 tokens, then frozen, and their memories are trained for another 239 tokens. For the anchor models, schedule with a maximum learning rate 10-3 and a cool-down value of 10-5 Weight decay is set to 0.1. We then freeze the anchor model to train the memories. For this particular experiment, memories are trained with a constant learning rate parameters, even without learning rate annealing.
cosine learning rate
to demonstrate convergence memory
17

---

Published as a conference paper at ICLR 2026
Anchor model
Memory
Runtime params
Name
Num layers
Num params
Config
Fetch size
Bank size
260M
12
257,475,584
(3840, 336, 6,
154,165,248
6,341,787,648
411,640,832
320M
17
321,721,344
(1445, 256, 8,
89,250,816
6,341, 6,341,246,976
410,972,
350M
19
347,419.648
226,, 9, 0)
64,496,640
6,341,000 6,341,099,520
411,916,288
370M
21
373,117,952
(384, 200, 10,
38,320.128
6,341.7876448
411,438,080
385M
22
385,967.104
1264, 94, 0)
25,276.416
6,341.001216
411,243,520
410M
24
411,665,408
0,,00,
411,665,408
Table 10: Anchor models and their corresponding memory configurations for the memory nemory-to-moded -o--moee size experiment in Fig. 4b. The anchor model architecture and memory configuration are designed to keep the total runtime number of parameters and the memory bank size (almost) fixed. The anchor model architecture uses hidden dimension of 1,024 and 16 heads for all rows.
B EVALUATION
We use LLM-Foundry (V0.9) as the evaluation framework. To obtain a stable signal, we only include benchmarks for which the baseline 410M model (B1 in Table 1) performs better than random and the ratio of standard deviation to mean performance is less than 0.5. Below is detailed description of the 13 tasks we evaluate pretrained language models in this work:
Lambada-OpenAl (Paperno et al., 2016) with 5, samples, evaluated 0-shot, is of type language model ing, with random baseline performance equal to 0%.
BoolQ (Clark et al., 2019) with 3,270 samples, evaluated O-shot, is of type multiple choice, with random baseline performance equal to 50%.
SQuAD (Rajpurkar et al., 2016) with 10,570 samples, evaluated 3-shots, is of type language modeling, with random bascline performance equal to 0%.
Winograd (Levesque et al., 2012) with 273 samples, evaluated 3-shots, is of type schema, with a random baseline performance equal to 50%.
WinoGrande (Sakaguchi et al., 2021) with 1,267 samples, evaluated 5-shots, of type schema, with random baseline performance equal to 50%.
CoQA (Reddy et al., 2019) with 7,983 samples, evaluated O-shot, is of type language modeling, with a random baseline performance equal to 0%.
Hellaswag (Zellers et al., 2019) with 10,042 samples, evaluated O-shot, of type multiple choice, with random baseline performance equal to 25%.
ArcEasy (Clark et al., 2018) with 2,376 samples, evaluated 3-shots, is of type multiple choice, with random baseline performance equal to 25%.
ArcChallenge (Clark et al., 2018) with 1,172 samples, evaluated 3-shots, is of type multiple choice, with random baseline performance equal 25%.
TriviaQA (Joshi al., 2017) with 3,000 samples, evaluated 3-shots, is of type generation task with answers, with random baseline performance equal to 0%.
PIQA (Bisk et al., 2019) with 1,838 samples, evaluated 0-shot, is of type multiple choice, with random baseline performance equal to 50%.
OpokBA (Mihaylov et al., 2018) with 500 samples, evaluated 10-shots, is of type multiple choice, with a random baseline performance equal 25%.
NaturalQuestions (Lee et al., 2019; Kwiatkowski al., 2019) with 2,655 samples used by Liu et al. (2024), evaluated O-shot, is of type generation task with answers, with a random baseline performance equal to 0%.
In Appendix G, we present our analysis dividing the above tasks into two groups: common- knowledge (CK) and specific-knowledge (SK).
18

---

Published as a conference paper at ICLR 2026
When evaluating models with fetched memory, retrieve memory based the question for multi- ple choice tasks, the context for language modeling tasks, and the portion of the text that is common across all choices for the schema tasks. To compute perplexity (for the Wiki-En En dataset) full document to retrieve memory.
the
We include task on predicting the atomic number of different elements. The model completes a prompt in the form of "The atomic number of {...} is." where (...) is replaced with the name of an element (from a total of 118). The model's generation is then processed to extract integer numerical values, and response is accepted if it matches the actual atomic number. Random baseline performance for this task is 0%. We use the prompt, including the element name, for memory retrieval.
Elements atomic number:
The elements atomic number evaluation is particularly interesting because each query has a specific keyword: the element name. This allows us to count the frequency of each element's 's name in the dataset (as shown Fig. to analyze model performance as a function of knowledge scarcity in the pretraining dataset. We note, however, that this analysis has some minor caveats. For example, the word "lead" both the name of an element and an English verb. In addition, some elements have multiple names; for instance, "Natrium" is an accepted alias for "Sodium". In our frequency calculations, we count acceptable aliases of each element name.
B.1 PER DATASET EVALUATION RESULTS
In this section, we provide per-dataset results for all models discussed in Table 1 as shown in Ta- ble 11.
Model
AreC.
ArcE
BOOLQ
COQA
HLSW
LOAI
NQ
OBQA
PIQA
SQUAD
TRIVQA
WINGRD
WINGRDE
59.7
21.3
A2-Generi
28.2
56.0
60.0
23.3
46.6
51.4
2.2
35.8
69.0
25.2
12.1
73.3
54,3
A2-Fetched
34,2
63.9
58.3
23.5
53.1
54.1
4.0
36.4
73.2
25.8
17.3
75.1
55.5
62.4
56.2
31.3
55.8
4.4
75.4
32.0
17.2
59.5
B2-generic
02-Fecd
41.1
69.1
64.5
32.6
63.4
64.4
7.3
41.4
76.3
34.1
28.4
81.0
61.5
C1
43.2
70.6
58.7
39.9
68.9
67.6
9.0
41.0
77.1
47.6
35.1
85.77
67.6
C2 Generrii
44.3
72.2
68.7
42.7
70.6
69.4
988
46.2
77.2
49.9
38.9
87.2
68.7
C2-Ferched
Table 11: Per dataset results of all models discussed in Table 1.
C DATA CLUSTERING DETAILS
We cluster the training set with hierarchical clustering (Grangier al., 2025). We build clustering tree: each node in the tree is associated with a cluster centroid. The examples traverse the tree from top to bottom, selecting the node corresponding the closest centroids among the current node's children.
Before training the clustering tree, segment our dataset into non-overlapping 2,048 token win- dows and compute eetence-BERT embedding for every window. We rely on the sentence-BERR MiniLM-L6-v2 model (Reimers & Gurevych, 2019). This process associates each segment of text with 384 dimensional vector.
The training of the tree proceed from root to leaves. Iteratively, a new level is built by applying k-means to subset of the examples belonging to each node. We built a tree of depth up to 4, always splitting nodes in 16 clusters. For k-means, we normalize the Euclidean norm of the vectors prior to clustering. We train the model via Expectation Maximization using k-means++ initialization (Arthur & Vassilvitskii, 2006). At each step, we sample 6,400 new examples. With 20 steps, we visit 128k examples. To ensure cluster distribution close to uniform, monitor the cluster sizes at each assignment steps. a cluster is larger than our balancing limit (0.094 ~ 1.5 x 1/16), we split evenly at random its assignments with the smallest cluster, suggested by Jegou et al. (2010).
19

---

Published as a conference paper at ICLR 2026
.9.
150
With noney merrery
W2 menoo mammrr
415
27.
17.6
37.0
41.6-
16.-
-0-
Wit mammor marory
41.
34.0
Token Poramoter
Tokon Porameter
Tokon Ponsammtrrr
(a)
(b)
(C)
Figure 7: Common- Knowledge (Avg-CK), Kneeilid-KKwweedeeeee (Ayg and perplexity Wiki- En as a function of A1 in Table 1) and is kept frozen during memory learning.
(TPP). The anchor model a pretrained 160M model (row
D TOKENS PER PARAMETER IN MEMORY LEARNING
While scaling laws exist for regular language model pretraining that identify the training bud- get-optimal choice for
(TPP) (Hoffmann et al., 2022), these laws may not hold for memory learning due to its different learning dynamics. Here, we experiment with the effect of TPP when learning memories.
We consider the (0, O, 16, 0) memory configuration on top of a pretrained 160M anchor model (row Al in Table 1). This corresponds 4,096 memories, each with 860,160 parameters, for total of 3.5B parameters. We freeze the anchor model and train the memories with 232 233 240 total tokens, corresponding to TPP values ranging from 1 to 312. Results are shown in Fig. 7.
In regular language modeling, knowledge (as in conditional likelihoods) is picked into parameters when and similar content, repeated often (with many TPP) and receives constructive gradient updates, and forgotten if it is rare and gets destructive gradient updates by dissimilar content (Ghosal al., 2025). This means that long-tai tail information forgotten after seen batch because the following batches send destructive gradient updates. What lasts in joint parameters is common knowledge, which occurs more often and with more aligned gradients. We demonstrate this effect in Fig.
In the proposed memory learning method, however, memory parameters are shielded in that they activated and updated only sequences of similar topic. This both reduces the times where knowledge can be overwritten and the dissimilarity of possible other gradients. On average, memory parameters corresponding to level 1 are updated 16' times less often than anchor parameters. For the setup considered above with memory configuration (0, O, 16, 0), where all memories are at level memory parameters receive one update for every 4,096 sequences in the training set.
Specific knowledge stored in deep memories is shielded from catastrophic forgetting (Toneva et al., 2018), because, unlike regular parameters, it not overwritten frequently by destructive gradients from unrelated content. Instead, deep memory parameters are only activated when there are con- structive gradients of similar content in their clusters, that they memorize this specific knowledge. This behavior is shown in Figs. 7b and 7c, where specific-knolledee wikipedia perplexity steadily improve with increasing TPP. The last point corresponds total of 1.1 trillion seen tokens 312). Due to computational constraints, we did not scale TPP further.
benchmark accuracy and
E FETCHED MEMORY VS GENERIC MEMORY (ADDITIONAL RESULTS)
In Fig. 6a, we showed the performance of the model using fetched and generic memories throughout training, both when anchor parameters are frozen and when they are learnable. A key observation is that when allow the anchor parameters to be co-trained with the memory bank, we obtain greater improvement compared to the case where the anchor model frozen (e.g., Avg-SK improves from 39.2% to 40.3% shown in Table 1). This improvement be attributed two factors: 1) when
20

---

Published as a conference paper at ICLR 2026
co-training, the anchor model learns to adapt to the memories more effectively compared to when it is frozen, and 2) the anchor model is exposed to additional tokens, so its performance improves simply due to more training. The latter effect should be minor if the anchor model is already over- trained (i.e., trained with a high TPP) during pretraining and has reached performance saturation, meaning it longer benefits from additional training.
We also track the gap between the performance the anchor model using generic memories (a fixed set of memory parameters independent of the input context) and fetched memories through- out co-training, as shown in Fig. 8b and Fig. 8d for the Avg-SK and Wiki-En perplexity metrics, respectively. We observe that the performance gap between fetched and generic memories grows over time (despite using a cosine learning rate schedule), indicating that longer co-training benefits fetched memories more than generic ones. This is expected, as the memory bank has significantly more parameters (4.6B in this example) compared to the anchor model (160M) and a single generic memory (18M), and thus benefits more from extended training.
39-
(5)
Avo SK (teachee
-
Leamablc ancher tecchod memory
37
Learnable anchor generic
35
200
400
500
800
1000
0
200
400
600
800
1000
Soon takens [8]
Soen tokens (B)
(a)
(b)
0.00
17.0
-0.25
Wib-En perplexity
16.5
-0.50
Leamable anchor. fetched memory
Leamsble anchor generic memory
-0.75
16.0
Frezan anchor: fetched memory
-1.00
15.5
-1.25
-1.50
15.0
-1.75
200
400
600
800
1000
200
400
600
800
1000
Seen tokens (B)
Seen tokens (B)
(c)
(d)
Figure 8: Additional results for the experiment setup of Fig. 6a. Performance improvements from co-training the anchor model and memories, corresponding to row A2 in Table 1. We also show performance when the anchor model is frozen, corresponding to row A3 in Table 1. (a) Avg-SK using fetched and generic memories, (b) Avg-SK gap between fetched and generic memories, which grows training progresses, (c) Wiki-En perplexity using fetched and generie memories, (d) Wiki- En perplexity gap between fetched and generic memories, which widens as training progresses.
F MEMORY LOCATION
So far, we have considered memory parameters to be uniformly distributed across the layers of the model. Meng et (2022) showed that model knowledge is mainly captured in the middle layers. We use our memory learning setup to further explore this hypothesis by considering different distributions of memory parameters across layers, as shown in Fig. 9.
Starting from baseline memory with configuration (0, 64, 0, 0), trained on top of pretrained and frozen 160M model (row A1 in Table 1), corresponding to the level-2 models in Fig. 3d, we study the effect of distributing memory parameters -on-unffrrlly across the model. We consider three setups-early, mid, and late wheee the same number of memory parameters as the uniform baseline are applied only to the first, middle, last 10 layers of the anchor model (see Fig. 9). As
21

---

Published as a conference paper at ICLR 2026
Memory location
Avg-CK (%)
Avg-SK (%)
Elements atomic number (%)
Wiki-En
Uniorm
Early Middie
Late
Ppix
No Memory
1.7
17:2
Anchor model
Uniform
46.7
36.8
14.4
16.0
Early
2.5
Middle
46.
35.4
1.7
16.5
memories
Late
46.9
36.8
20.3
Different distribution of memory along the layers of anchor model.
Figure 9:
Table 12: Comparing effectiveness of memories with differ- ent distribution over the depth of anchor model, which kept frozen.
discussed in Appendix A.3, for these experiments the anchor model frozen, and memories trained for total of 2 tokens.
Results for each memory placement are shown in Table 12. Consistent with the observations of Lample et al. (2019), using memories in the early layers of the anchor model is less effective than using them uniformly (the default setup) or in the deeper layers.
G WHAT TASKS BENEFIT MORE FROM MEMORIES
In Fig. 1, using the atomic number prediction task, we showed that memories can significantly im- prove performance on tasks requiring long-tail knowledge. To extend this concept to commonly used benchmarks for evaluating pretrained models, we introduce approximate, quantitative mea- sure of the degree of knowledge specificity for each the evaluation benchmarks described in Appendix B. For each dataset, randomly sample 100 entries and prompt GPT-4 (Achiam et al., 2023) to estimate the education level (as a proxy for knowledge specificity) at which a typical person would have acquired the knowledge needed answer each question. The average of these ratings is used as the knowledge specificity score of the task.
Specifically, ask the model to each question based the amount of knowledge required, using the following prompt:
Given the following prompt and answer, what facts should human know in order to answer the question correctly? What phase of life will a typical person know all required facts?
Your response should be in the format of an integer between () and 5 based on the following
scale:
O: Only language understanding, all required information is in the context
1: Commonsense facts learned through sensory experiences childhood
2: Facts learned in elementary school
3: Facts learned in middle school
4: Facts learned in high school
5: More specific facts learned later in life
Respond with only single integer and nothing else.
In Fig. 10, we plot the accuracy difference between fetched memories and generic memories against the knowledge specificity score for each benchmark, using the 1.4B model trained with memories (row C2 in Table 1). The plot shows a clear positive correlation between knowledge specificity and performance improvement from fetched memories.
We further group the datasets into six common-knon -onnonn-nnwwldge (purple) and seven specifie-knowledge (green) tasks using knowledge specificity score threshold of 2.0. On average, fetched memo- ries improve peccific-knowlede (Avg-SK) task performance by 3.6 points, while performance on common -nowlede (Avg-CK) tasks remains comparable (64.4 vs. 64.5). Notably, both CK and SK performance show even greater improvement compared to the baseline model (row C1 Table l): Avg-CK improves from 61.2% to 64.5%, and Avg-SK improves from 49.7% to 54.9%.
Additionally, we show the same analysis for the 160M (row A2 in Table 1) and 410M (row B2 in Table 1) models trained with memories in Fig. 11 and Fig. 12, respectively.
22

---

Published as a conference paper at ICLR 2026
TRIVIAQA
Average gain commonsense Average gain knowledge-speciffc
Average gain (fetched accuracy vs generic accuracy)
PIQA
HELLASWAG
OBQA
BOOLQ
WINDGRAD
COOA AINOORAAND
SQUAD
4
Knowledge specificity
Figure 10: Fetched memories improve performance knowledge-intensive benchmarks. Ac- curacy gain (fetched memory vs. generic memory) for the 1.4B model (row C2 in Table 1) as a function of the knowledge specificity score of each benchmark. Knowledge specificity is determined by GPT-4 ratings of 100 sampled entries per dataset, and error bars reflect the standard error of the mean. The positive correlation highlights the value of fetched memories for knowledge-nntensvve tasks. Note that this plot shows the improvement of fetched memories compared to generic mem- ories; the improvement even greater when comparing fetched memories with the baseline model without memory (row Cl in Table 1).
H RETRIEVAL AUGMENTED GENERATION DETAIL
In this section, we provide further details for the experiments in Section 4. We consider two datasets to retrieve documents from: DCLM-Baseline and English Wikipedia 2022, with 6.5 mil- lion and 3.2 billion documents, respectively. The average document length in DCLM, using our EleutherAI/Se PerAA//ppteoxxxxo -neox (Black et al., 2022) tokenizer, is 1,309 tokens (computed from ran- dom 10% subset of the dataset), while for Wikipedia it is 723 tokens. We set the max sequence length evaluation 2,560.
As reference point, we report FLOPs associated with each approach. For RAG models, we use the average sequence length and compute the additional FLOPs on top of 1024-token context. When using memory, the context length does not increase, but additional memory parameters are fetched and added to the anchor model, increasing runtime FLOPs by 10%.
For these experiments, we compare augmenting the anchor model with: 1) fetched memory from the learned bank of memories, versus raw documents retrieved from the same dataset the memories were trained on. For high-quality RAG performance. many factors matter, including the instruction- following and long-context capability the LLM, the retriever's 's quality, the quality the datastore, and additional techniques such as self-reflection (Asai et al., 2024).
In this study, we mainly aim to compare learned memories (ours) against contextual memories with- out additional confounders. We therefore call our retrieval -urreete-a generations vanilla RAG. We use the same retrieval mechanism for both learned memories and RAG: given a query, we identify the context described in Appendix B. Using the Sentence-BERT (Reimers & Gurevych, 2019)
23

---

Published as a conference paper at ICLR 2026
TRIVIAQA
Average gain commonsense Average gain knowledge-specifg
+
Average gain (accuracy fetched vs accuracy generic)
6
HELLASWAG
Arc-E
AMM-OAI
NO
SQUAD
WINOGRAD
WINOGRANDE
BOOLQ
COOA
Knowledge specificity
Figure 11: Fetched memories improve performance knowledge-intensive benchmarks. Ac- curacy gain (fetched memory vs. generic memory) for the 410M model (row B2 in Table 1) as a function of the knowledge specificity score of each benchmark. Knowledge specificity is determined by GPT-4 ratings of 100 sampled entries per dataset, and error bars reflect the standard error of the mean. The positive correlation highlights the value of fetched memories for Enowledge-intensive tasks. Note that this plot shows the improvement of fetched memories compared to generic mem- ories; the improvement even greater when comparing fetched memories with the baseline model without memory (row B1 in Table 1).
all-MiniLM-L6- embedding model, we first determine the closest level-3 cluster in DCLM (comparing against 4096 centroids). Note that the models with memory in Table 5 also trained with hierarchical memory configurations up to level 3. We then retrieve the nearest-neighhor (NN) document from within that level cluster (on average 750k documents, given 3B/4096). This document is then added to the context. To avoid confounders from few-shot complexity in RAG, all tasks in 0-shot setup. We experiment with the 160M, 410M, and 1.4B models, correspond- ing to rows A1, B1, and C1 in Table 1.
As shown Table 5, a vanilla NN retrieval from DCLM does improve baseline model perfor- mance for either common-knowleee knowledge or specific- -nowleee benchmarks, We argue this is mainly due to the low quality of DCLM (a pretraining dataset) when used as a datastore for RAG. Recently, Lyu et al. (2025) showed that with careful filtering, RAG quality can be improved when using web- scale datastores.
To demonstrate the effect of datastore quality here, we also use English Wikipedia to retrieve higher- quality documents for the given context with the same Centence-BERT model. Results Table show that RAG-Wik Wiki improves baseline performance specific-knowledee benchmarks (c.g., from 46.9 to 49.2 for the 1.4B model). However, for common-knows knowledge benchmarks, RAG-Wiki does not improve
baseline. By contrast, using learned memories (with 10% additional parameters relative to baseline) improves performance on both common-knon -ummon-knowledge and -nowleee benchmarks with lower runtime FLOPs overhead.
Finally, we note that high-quality RAG is complementary to the proposed learned memories and can further enhance performance.
24

---

Published as a conference paper at ICLR 2026
Arc-E
Average gain (accuracy fetched vs accuracy generic)
HELLASWAG
TRIVIAQA
PIOA
Average gain commonsense
LAM-OAI
Average gain wewwlegeespeciffee
NQ
WINOGRAD
CODA
SQUAD
OBQA
BOOLO
4
Knowledge specificity
Figure 12: Fetched memories improve performance knowledge-intensive benchmarks. Ac- curacy gain (fetched memory vs. generic memory) for the 160M model (row A2 in Table 1) as a function of the knowledge specificity score of each benchmark. Knowledge specificity is determined by GPT-4 ratings of 100 sampled entries per dataset, and error bars reflect the standard error of the mean. The positive correlation highlights the value of fetched memories for knowledge-intenivee tasks. Note that this plot shows the improvement of fetched memories compared to generic mem- ories; the improvement even greater when comparing fetched memories with the baseline model without memory (row A1 in Table 1).
I MEMORY AUGMENTED TRANSFORMER ARCHITECTURE
In this section, we provide additional detail on different memory augmented transformer architec- tures that we considered.
The transformer block with SwiGLU FFN has seven linear layers, as shown in Fig. 13. We can augment any subset of these with low-rank memories. To avoid exhaustive search, we group the linear layers into three categories based on their role in the transformer block: 1) query and key projection layers, 2) value and output projection layers, and 3) FFN linear layers.
LoRa-Memories:
Accordingly, we define three types of LoRa memories: LoRa-QK (shown in Fig. 14), LoRa-OV (shown in Fig. 15), and LoRa-FFN (shown Fig. 16). Each LoRa consists of two low-rank matrices, A and B. The target linear layer W is patched additively as W + B A. The size of memories is determined by the rank r of matrices and Note that for a model with hidden dimension d, inner FFN dimension dy, per attention head dimension d,, h heads, and / layers the size of fetched memory with LoRa type is as follows:
LoRa-QK memory size: 2r1(d + hdz))
LoRa-OV memory size: 2r[(d + hdy))
LLRR-FFN memory size: 3rl(d + dy)
For graceful initialization, so that memories initially have no effect, we initialize B with zeros, as suggested in the original work (Hu et al., 2022), and A with uniform Kaiming distribution in all
25

---

Published as a conference paper at ICLR 2026
Projectior
Norm
Projection
MHA
Provection
Norm
WS
FFN W2
Projection
Figure 13: Base architecture of a transformer block with SwiGLU FFN layer.
B
ajecsion
Narm
Drojection
By
VHA
Projection
Norm
FFN
W3
FFN
Projection
Figure 14: A transformer block with LoRa memories on queries and keys projection layers.
Additionally, we found that a scaling factor of = works best for the LoRa memories considered here.
cases.
Aligned with previous observations that knowledge in transformers primarily stored FFN lay- ers (Geva al., 2020; Dai et al., 2022; Yao al., 2022), we also find that LoRa-FFN mostly outper- forms alternative LoRa memories for the same number of learnable parameters, as shown in Fig. 3.
learn additional key and value vectors to augment the input-dependent keys and values. The input-dependent query vectors cross-attend to the learned key and value vectors, and their results are added the output of multi-head attention. Note that do not apply causal masking when attending to the learned keys and values. Additionally, we found that KV memories are slightly more effective when used without positional encoding. To ensure no memory effeec at initialization, initialize the learned value vectors with zeros and the learned key vectors with a truncated normal distribution, consistent with other model parameters.
KV-Memories: With these memories,
The size of KV memories is determined by the number of key-value vectors (r), as shown in Fig. 17, and can calculated as follows:
KV memory size: 2rlhd,
FFN-Memories: for FFN memories, we directly expand the inner dimension of the three linear layers in the SwiGLU FFN as shown in Fig. 18. Similar to other memory types, we initialize memories such that at the beginning of training they have no effect. Therefore, we initialize W1 and W2 FFN memories with truncated normal, and W3 FFN memory with zeros. The size of the FFN memory is determined by their inner dimension and can be calculated as follows:
FFN memory size: 3rld
I.1 MEMORY CALCULATION
We detail the memory size calculation for one memory configuration, (256, 64, 16, 0) with the 160M
anchor model, for additional clarification. As noted Section 2.1 (last paragraph), hierarchical configuration ($1, $2, 83, 84) is given by ($1, $3, s4) Co (21,72, Pa). where co depends the anchor and memory type, and r; is the level-i multiplier we choose. For FFN-Memory, co = 3/d (see Fig. 18). For the 160M anchor, for example, cp 53, 760 (see Table 6).
26

---

Published as a conference paper at ICLR 2026
FPN W1
9w
Norm
Prolettion
MHA
Profcctio
Ao
Norm
FRN
VE
FRN W2
Projection
Figure 15: A transformer block with LoRa memories values and output projection layers.
Projection
Pig
auer
kayer
Projection
Figure 16: A transformer block with LoRa memories on SwiGLU-FFN linear layers.
With (71, r2, F3. 14) = (256, 64, 16, 0):
81 53, 760 x 256 13, 762, 560
$2 53, 760 x 64 3, 440, 640
$3 53, 760 x 16 860, 160
$4 0.
Fetched memory (per : $1 $2 + + = 18, 063, 360 parameters (~18M).. Memory bank
(all blocks): 16s; + 162 $2 16 + 16* 81 = 624, 220, 160 parameters (~4.6B).
J RETRIEVER FAILURE ANALYSIS
In this section, we provide a few examples of using different fetched memories during generation (max length of 30 tokens) for the 1.4B model with memory (the C2 row in Table 1) under the memory configuration (768, , 256, 1 16, 0). In both examples, we see that as we corrupt the correct memory ID (starting from the finest to the coarsest level), the model's response becomes gradually less accurate.
For the atomic number example, the correct fetched memory ID (8, 130, 2080, 0) not only leads to the correct answer but also adds relevant details, such as the name of the scientist who discovered the element. When corrupt the memory ID (by replacing it with index O), the model still gives the correct answer as long as the level-1 I memory remains correct, though the additional generated details become wrong. When the level-1 memory is also wrong, the model outputs incorrect atomic number (92). A similar pattern appears in the example involving scientific names animals.
prompt "The atomic number of Fermium is" # correct memory (8, 130, 2080,
(8, 130, 2080, 0):
1000
Fermium is a synthetic element. It was discovered in 1952 by Albert Ghiorso and his team the University of California"...
27

---

Published as a conference paper at ICLR 2026
Projection
FFN W1
S1:U1
Layer Norm
ection
MHН
Projection
Laver
FFN W/3
Projection
memory
Croz attenton
memory
Figure 17: transformer block with learned KV memories.
SLU
Projection
FFN Marrory
Projection
MHA
Projection
FFN W2 Memory
FEN
Projection
Figure 18: A transformer block with FFN memories.
model-output witth meattt--mmory-id (8, 130, 0, 0): 100.
The atomic weight of Fermium is 262. The melting point of Fermium unknown.
The boiling point of"...
model-outpuse -ithh with me 1000.
Fermium is a radioactive metal.
Fermium is member of the actinide series.
Ferm"...
model-ouee o0 92.
Fermium is radioactive metal.
Fermium is radioactive metal.
Fermium is a"...
Another example:
28

---

Published as a conference paper at ICLR 2026
Prompt
"Panthera
pardus
is the scientific
name
for" # correct memory
(10, 168,2703,
model---uuppui--ttttbb-rro------ 168, 2703, 0):
"the leopard. It is a large cat that is native to Africa and Asia. It is the largest of the four big cats. is also"...
the leopard.
The leopard is large cat, and is the second largest cat in the world.
The leopard is"
podel-outpuse with -emerhmmmmrrr-
the lion.
The lion is a large cat, the largest of the four species of the genus Panthera. Lions are native to Africa"
model-outen with (o the lion.
The lion is a symbol of royalty and power.
The lion is symbol of the sun.
The lion"
For more systematic analysis, probed failure modes by stressing the retriever explained below:
Setup: We use frozen 160M anchor and non-hierarchhcal bank (0, 0, 00,0,0,16) 16), the retriever must pick one leaf memory among 16a = 65, 536 blocks (that are non-overlapping)
Corruption protocol: For each query, replace the nearest-neighbor leaf with each of its 15 sibling leaves under the same level-3 parent (plausible local misrouting) and average query accuracy over all 16 choices (for each query accuracy is either 1). Finally we report the mean of this averaged accuracy over the entire benchmark (here considered Arc-Easy and Are-Challenge with 2376 and 1172 queries, respectively.
Task
Model
Retriever
Acc. (%)
ArcE
Baseline
N/A
55.5
ArcE
w/ mem.
uncorrupted
60.9
ArcE
w/ mem.
corrupted
55,8
AreC
Baseline
N/A
27.1
AreC mem.
uncorrupted corrupted
29.9
AreC w/ mem.
28.2
Table 13: Failure analysis of the memory retrieval process.
As shown Table 13, wrong retrieval reduces the gain but performance remains at or above no-memory baseline, indicating graceful degradation and robustness to retrieval errors. Further, for the hierarchical case, where two different retrievals can share memory parameters at shallower levels, expect more robustness.
29


Reviews:
[{"review_id": "3iSTFVe0De", "reviewer": "Reviewer_hJ7H", "summary": "This paper proposes a memory augmented architecture to improve the efficiency of language models. The system uses a small \"anchor\" model for common knowledge and reasoning, which is augmented by a large \"hierarchical memory bank\" designed to store long tail facts. A clustering based retriever selects a small, context relevant set of memory parameters from this bank to be added to the anchor model during inference. The authors present experiments showing that this method can improve performance on knowledge intensive tasks, particularly for rare facts, and that their 160M parameter model with 18M of fetched memory can perform comparably to a larger, 410M parameter baseline model.", "questions": "The authors would need to address several critical points to make their claims more convincing.\n\n1. Why was a static, offline clustering chosen over a learnable retriever that could co train with the model? can the authors provide any analysis on how sensitive the model's performance is to the quality of the initial clustering or the choice of embedding model?\n\n2. Regarding the privacy and editing claims, can the authors provide direct evidence of knowledge separation? \n\n3. Given that the vanilla RAG baseline is weak, how does this method compare to a more robust, modern RAG implementation using a high quality datastore and standard techniques like multi document retrieval and reranking? Without this, the claims of superiority are difficult to evaluate.\n\n4. When does the memory retriever fail? can the authors show examples of prompts where the retriever fetches the wrong or an irrelevant memory block, and how does the model perform in those scenarios?", "limitations": "", "rating": 2, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper is clearly written and the core concept is easy to understand. The use of figures, especially Figure 1 and Figure 2, helps illustrate the architectural idea and its intended effect on long tail knowledge.\n\nThe authors were thorough in ablating their own method's design choices. The paper includes a systematic study of different memory types (FFN, LoRa, KV), the impact of memory depth, and the effects of bank size versus fetched memory size. The analysis in Figure 1, showing improved accuracy on rare elements, provides a clear demonstration of the mechanism working as intended for that specific task. The paper also shows the method can be applied post hoc to existing open weight models, which suggests some level of generality.", "weaknesses": "The paper's conclusions are undermined by significant weaknesses in its experimental design and unaddressed questions about the core mechanism:\n\n1. Static Retriever: The entire method's effectiveness hinges on a static, offline clustering of the training data using an off the shelf Sentence BERT model. This is a major point of failure. The paper provides no sensitivity analysis on the choice of embedding model or clustering algorithm. If this initial, fixed clustering is suboptimal, the model has no way to adapt, and knowledge separation would presumably fail. The system's performance is completely dependent on this external component.\n\n2. Unsupported Claims on Editing and Privacy: The paper makes strong claims about enabling knowledge editing and privacy. The evidence provided is extremely weak. Figure 6b only shows that blocking memory degrades performance on a specific task. This demonstrates the memory is being used, not that facts can be cleanly removed or edited. It fails to address whether long tail knowledge leaks into the anchor model during co training, which would make any privacy or editing claims invalid. \n\n3. Weak RAG Baseline: The comparison to RAG in Table 3 is not strong. The authors compare their method to a vanilla RAG that retrieves a single document from the low quality DCLM pretraining dataset, which they admit performs poorly. Modern RAG systems use far more sophisticated techniques (multiple documents, reranking, chunking). So, by benchmarking against a weak baseline, the paper's claims of superior efficiency and performance over RAG are not well supported.\n\n4. Limited Novelty: The paper frames itself as a novel solution, but it can also be seen as an incremental combination of existing ideas. The concept of separating parameters for knowledge is central to MoEs, and memory augmented networks which are a long standing field of research. The primary novelty is the hierarchical structure and the clustering based pretraining, which as explained above is a weak and static process.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6hmQ3Fgmz4", "reviewer": "Reviewer_6eti", "summary": "This paper proposes a hierarchical memories pretraining framework that enables small models to achieve large-model-level knowledge coverage. By integrating hierarchical clustering-based memory retrieval, the framework decouples world knowledge from model parameters and separately stores common knowledge and long-tail knowledge. Comprehensive evaluations across 13 benchmarks show that the proposed framework achieves notable improvements in efficient hardware implementation and enhanced data privacy.", "questions": "1. The figure obscures key lines and does not specify which parameters were fixed. Could the authors clarify what factors were controlled? In addition, some key results in the figure are visually obscured.\n2. Is the model performance strongly correlated with the clustering quality of ϕ(x)?\n3. How is the number of levels p or division factor k chosen, and how sensitive is performance to these hyperparameters?\n4. It is unclear whether the retriever R(x;W) is trained jointly with the anchor model, or whether its cluster assignments are frozen after initialization. How are gradients propagated through R(x;W)?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The motivation is clearly presented. The paper designs a hierarchical framework combining an anchor model and a memory bank to separately common knowledge and long-tail knowledge.\n2. The proposed method is carefully designed, and its scalability is validated through empirical results on multiple benchmarks.\n3. The paper is generally well-organized and readable.", "weaknesses": "1. The theoretical contribution is relatively limited. While the paper provides empirical evidence that separating the anchor model and memory bank reduces forgetting and improves training stability, the claim lacks theoretical justification or formal analysis.\n2. Per-dataset results are not fully presented. The paper mainly reports averaged results (e.g., Avg-CK and Avg-SK) across 13 benchmarks, which provides a concise overview but may obscure dataset-specific behaviors. \n3. Lack of empirical evidence for the claimed knowledge separation. The paper presents a clear motivation that hierarchical memories aim to separate common reasoning abilities from long-tail knowledge. However, semantic clustering alone is insufficient to guarantee that higher levels correspond to common knowledge while deeper levels represent long-tail knowledge.\n4. Lack comparison on inference cost. The anchor model+memory bank design may require frequent memory retrieval during each inference task. The authors should provide the average delay cost of their design against baselines.\n5. Experiment hardware. The authors should disclose the hardware usage involved in pre-training and inference.\n6. Typo. In Figure 2, level 4 should have 64k clusters, instead of 65k.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ud3U9y892f", "reviewer": "Reviewer_cCUe", "summary": "This paper proposes a memory-augmented architecture for language models that separates common knowledge (stored in \"anchor\" parameters) from long-tail knowledge (stored in hierarchical parametric memory banks). During pretraining and inference, the model retrieves context-dependent memory blocks from a large hierarchical memory bank organized via clustering of training data. The authors demonstrate that a 160M parameter anchor model augmented with 18M fetched parameters from a 4.6B memory bank achieves performance comparable to a 2 times larger standard model. The paper includes extensive experiments on memory types, scaling, and deployment considerations.", "questions": "Is it possible to get some analysis or visualization of what knowledge is actually stored in memory parameters versus anchor parameters? For example, can you show examples of specific facts that are reliably stored in particular memory blocks?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper presents a well-motivated approach to separating common reasoning abilities from long-tail factual knowledge, with clear empirical evidence showing that hierarchical memories particularly benefit knowledge-intensive tasks (e.g., atomic number prediction improving from 1.7% to 67.8% for the 160M model).\n\n- The experimental evaluation is comprehensive, covering multiple model sizes (160M to 1.4B parameters), memory configurations, and architectural variants, with thorough ablation studies on memory types, depths, and sizes that provide actionable insights for practitioners.", "weaknesses": "- The comparison with retrieval-augmented generation (RAG) seems somewhat unfair, as the authors use \"vanilla RAG\" without standard techniques like reranking or filtering, and the improvement over baseline when using high-quality datastores (Wiki-En) is actually comparable to the memory approach on specific-knowledge tasks.\n\n- There's insufficient discussion of failure modes and limitations - for instance, what happens when the clustering assigns dissimilar documents to the same cluster? The paper would benefit from error analysis showing when and why the approach fails.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YDUWHTwWl9", "reviewer": "Reviewer_Uqsx", "summary": "This paper proposes a new pretraining framework for improved knowledge learning: it separates general and domain-specific knowledge in the pretraining corpus and assigns them to different components of the model. To preserve long-tail knowledge, the authors design a individual hierarchical memory bank that is selectively retrieved and updated when training on similar text corpora, ensuring each submodel in the bank becomes proficient at its assigned tasks. Eperiments show that models ranging in size from 160M to 21B parameters benefit from this memory-bank design while training costs remain moderate.", "questions": "1. In Figure 3a, is the fetched memory size determined by the memory size $s_2$? \n\n2. “For a fair comparison, during training we use the generic memory with probability 1/(16 + 1) and the fetched memory with probability 16/(16 + 1), where 16 is the clustering division factor. This ensures there is no training bias in favor of the memory bank parameters.” Can you clarify more on the design choice here? Why you choose this approach instead of training a standalone generic memory model?\n\n3. I'm curious how you arrived at a 4.6B memory bank size, given the model bank parameters (256, 64, 16, 0). Can you provide a concrete example?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. Although the hierarchical design is not new, applying it to language model pretraining to preserve long-tail information is innovative. \n\n2. The model-bank design choices — including parameter placement, parameter sizes, and retrieval ratios — are well supported by extensive preliminary experiments. \n\n3. The effectiveness of pretraining within this framework has been demonstrated at scale up to 1.4B-parameter models in the main page, indicating strong scaling potential. \n\n4. Plugging the memory bank post hoc into frozen models has also proven effective, demonstrating the design’s generality; moreover, the memory module can be quickly pre-trained on private or corporate data to meet different needs. \n\n5. The memory bank design enables on-device deployment by loading only the anchor model and retrieved parameters while storing the bulk of knowledge parameters in external storage.", "weaknesses": "1. The authors claim that using a base model as an anchor better captures common knowledge and reasoning capabilities, and that augmenting it with a memory bank benefits long-tail knowledge tasks. However, I did not find any evaluation of improvements in the system’s reasoning ability.\n\n2. Although the authors argue that this hierarchical memory bank aligns well with modern computer memory design and therefore offers greater efficiency, they did not provide experiments to support this claim, offering only a brief theoretical analysis at the end of Section 3.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JYA5XcrHWh", "reviewer": "Reviewer_2ftq", "summary": "This paper proposes an efficient approach for learning rare knowledge in LLMs than trying to compress them in LLMs parameter. To that, this paper proposes a hierarchical memory bank augmented architecture and a pretraining strategy which demonstrate to achieve similar performance with a small model(roughly half the size of LLM parameters) + very small amount retrieved memory bank parameters. This approach helps with catastrophic forgetting problem, as memory parameters are activated and updated only on sequences of similar topics, thereby reducing susceptibility to forgetting.  The paper discuss 3 approaches of retrieved memory parameter augmentation with the anchor model parameters. There are other related work such as memory layers where FFNs are replaced by learnable memory layers. This work seems a bit incremental although the hierarchical memory back is different.", "questions": "For the experiments, It would be good to address the following\n1. There are other parametric memory work, such as memory Layers at scale. It would be good to compare the results with such alternative approaches.  \n2. It would be good to share results on how high quality RAG combines with this approach\n\nOther comments:\n1. While there is a section discussing adding memory bank on various open models, It would be good to have a discussion regarding if the “learnt” sparse memory bank is transferable, meaning, can be combined with another model with minimal tuning. \n2. A more rigorous scaling law analysis such as how does the results scale with increasing number of parameters per memory block\n3. How does it compare and combine with MoEs for memorization and reasoning", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Overall, this paper makes a solid contribution by proposing an approach well-suited for on-device inference. The compact anchor model can reside in fast local memory, while the larger set of knowledge parameters is stored in slower, high-capacity memory. By decoupling knowledge from reasoning capabilities, the method establishes a clear link between training tokens and specific parameter subsets (memories). This design allows targeted selective knowledge update without altering the core anchor model, which can remain publicly accessible.", "weaknesses": "This work seems a bit incremental over other memory architecture. There are other parametric memory work, such as memory Layers at scale. It would be good to compare the results with such alternative approaches.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)