PaperlayerView exact revision on arXiv

Source-linked AI summary

Pretraining with hierarchical memories: separating long-tail and common knowledge

Hadi Pouransari, David Grangier, C Thomas, Michael Kirchhof, Oncel Tuzel

arXiv:2510.02375v3cs.CLcs.AIcs.LG

TL;DR

Language models store world knowledge in parameters, including long-tail facts that may be unnecessary for a given prompt and remain loaded for on-device use. The paper pretrains an anchor model with context-retrieved hierarchical memories for long-tail knowledge, reporting improved knowledge-task results with a small fetched memory and performance comparable to models with more than twice the runtime parameters.

  • Problem

    Long-tail facts can remain permanently loaded in model parameters even when they are unnecessary for on-device tasks, while limited fast memory constrains deployment.

    Source: §1 · p. 2

  • Method

    The approach pairs an anchor model for common knowledge and reasoning with a hierarchically organized memory bank for long-tail knowledge, retrieving relevant parameters by context.

    Source: §1 · p. 2

  • Results

    Knowledge-specific task performance rose from 34.1% to 40.3% for a 160M model and from 40.9% to 45.9% for a 410M model with approximately 10% additional parameters retrieved from memory.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Small models augmented with memory banks match regular transformers with 2× more parameters, and the design offers potential hardware-efficiency and data-privacy benefits.

    Source: §6 · p. 11

    The paper leaves optimal memory scaling laws, architecture search, retriever design, anchor-model design, and improving anchor reasoning for future work.

    Source: §6 · p. 11

Abstract

from arXiv · show

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. During pretraining and inference, we fetch a small, context-dependent memory block and add it to the model. 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. Through extensive experiments, 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.

1 Introduction

The paper separates common knowledge and reasoning in an anchor model from long-tail knowledge in context-retrieved hierarchical memories. It studies this design for performance and deployment efficiency, including memory augmentation of small models.

Motivation and approach

  • The proposed design uses an anchor for common knowledge and reasoning, plus a memory bank dedicated to long-tail knowledge.The paper focuses on capturing long-tail world knowledge with parametric memories, leaving improvements to the anchor’s reasoning capabilities for future work.Source: §1 · p. 2
  • The authors attribute baseline degradation on lower-frequency knowledge to catastrophic forgetting, while topic-specific memory updates reduce susceptibility to forgetting.The proposed memory parameters are activated and updated only on sequences of similar topics.Source: §1 · p. 2

Runtime efficiency

  • For on-device use, context-dependent retrieval fetches only the required knowledge parameters—about 10% of the model—while the larger memory bank can reside in slower storage.The hierarchical memories are intended to align with device memory hierarchies such as RAM, flash, and external disk.Source: §1 · p. 2

Privacy and extensibility

  • The design can associate training-data tokens with memory parameters, enabling targeted memory deletion or editing and access restrictions while keeping the anchor public.The authors also describe adding memories from new data after pretraining, including private user data or organizational databases.Source: §1 · pp. 2–3

Contributions

  • The contributions include hierarchical-memory transformer architectures and clustering-based pretraining, with memory banks scaled to 21B parameters.The paper examines memory type, depth, size, position, memory-to-model ratios, and runtime efficiency relative to hardware storage hierarchies.Source: §1 · p. 3

Results

  • 34.1% →40.3% for a 160M model and 40.9% →45.9% for a 410M model on knowledge-specific tasks with approximately 10% additional parameters retrieved from memory.The paper reports performance comparable to a model with more than 2× the runtime parameters.Source: §1 · p. 3

2 Pretraining with memories

Pretraining combines an anchor model with context-retrieved memory parameters in a next-token prediction objective. A hierarchical clustering tree routes documents and task contexts to memory blocks, while several memory types provide ways to augment transformer layers.

2 Pretraining with memories

  • The pretraining objective predicts each document’s next tokens using anchor parameters augmented with retrieved memories; sparse updates let the anchor learn common knowledge and reasoning while memories learn related long-tail content.Semantically similar inputs are expected to retrieve the same memory subset, whose parameters receive gradients primarily from related documents.Source: §2 · p. 3

2.1 Clustering-based memory retriever

  • Documents are embedded and recursively divided into clusters, and greedy nearest-centroid traversal routes a document or task context through the hierarchy.The described setup uses four levels with a branching factor of 16; document cluster indices are precomputed offline, avoiding training-time retrieval overhead.Source: §2.1 · p. 4
  • Each cluster at each level receives a memory parameter block, and retrieval combines the blocks along the document’s cluster path.The bank size sums blocks across clusters and levels, while the fetched size sums one block per level.Source: §2.1 · p. 4

2.2 Inference with parametric memories

  • The paper considers LoRA memories for transformer projections, learned KV-cache memories, and FFN memories that expand the inner SwiGLU dimension with fetched parameters.The memory configuration sets level-specific block multipliers, generally making coarser levels larger or omitting levels.Source: §2.2 · pp. 4–5

3 Design choices for pretraining with memories

The section compares memory types and hierarchical design choices, then examines the allocation and deployment tradeoffs. It reports gains in Avg-SK accuracy and lower loading latency for hierarchical memories under the tested settings.

Memory type

  • FFN memories significantly outperform LoRA and KV memories across tested memory sizes, so the remaining experiments use FFN memories.The comparison uses a frozen 160M anchor and memories trained from scratch for 275B tokens.Source: §3 · p. 5

Memory depth and size

  • At fixed fetched size, deeper single-level memories yield greater accuracy; at fixed bank size, shallower memories do better because they fetch more parameters.Performance increases with total bank size across all tested memory configurations.Source: §3 · p. 5

Hierarchical memory design

  • Hierarchical configurations independently control total bank size and fetched size; deeper levels receive fewer, more similar-content updates, while shallow levels receive more common-fact updates.The authors describe this gradient pattern as shielding deeper memories from forgetting and favoring a hierarchy from common to specific knowledge.Source: §3 · pp. 5–6

Hierarchical memory size

  • 39.1% to 40.1% Avg-SK resulted when bank size grew from 4.6B to 18.7B with fetched memory size barely changed.The comparison changes the configuration from (256,64,16,0) to (256,64,16,4).Source: §3 · p. 6
  • 44.5% Avg-SK accuracy came from a 160M anchor with ≃240M fetched memory, outperforming a regularly trained 410M model by 3.6 points.The augmented model uses 400M runtime parameters; Figure 4a reports the accuracy result.Source: §3 · p. 6, Figure 4 · p. 7

Memory-to-model ratio

  • In Figure 4b, a 1:10 fetched-memory-to-anchor-model-size ratio appears optimal, but may not generalize across runtime, memory-bank, and training budgets or when anchor and memory parameters are co-trained.The configurations use a 6.3B memory bank; anchor-plus-memory runtime parameters remain fixed at 410M.Source: §3 · pp. 6–7

Deployment

  • Hierarchical memories load in 38ms versus 198ms for a flat bank in the hardware example; partial swaps take 47ms versus 198ms for full reloads.The first comparison places shallow levels on faster components; the second exploits compositionality across related queries.Source: §3 · p. 7, Figure 5 · p. 7

Clustering sensitivity

  • Changing initial clustering causes very small sensitivity in final performance; two-level clustering with p = 2 and k = 256 slightly improves knowledge-specific tasks.The optimal clustering and memory configuration also depends on the hardware hierarchy used for storage and fetching.Source: §3 · pp. 7–8

4 Results

The results evaluate hierarchical parametric memories across training setups, model sizes, retrieval strategies, and deployment conditions. They report gains on knowledge-specific tasks, with co-training, post-hoc augmentation, and combination with RAG each examined.

4 Results

  • +5.1 points on Avg-SK versus A1 and 2 points lower Wiki-En perplexity with fetched memories; anchor+generic memory scores 34.7% on Avg-SK, 4.5 points below fetched memories.The setup uses ≃18M fetched parameters from a ≃4.6B memory bank.Source: §4 · p. 8, Table 1 · p. 9
  • A4 performs worse than A2 at the same training budget, while at 410M the anchor+generic setup beats regular training on Avg-CK, consistent with long-tail knowledge offloading.The authors describe the latter result as preliminary.Source: §4 · p. 8
  • +4.1 points on Avg-SK for the 410M model and +3.6 points for the 1.4B model when fetched memories are compared with generic memories.The paired memory banks are 12.7B and 21.1B parameters, respectively.Source: §4 · p. 8, Table 1 · p. 9
  • Avg-SK rises from 39.2% to 40.3% when the anchor is co-trained rather than frozen; authors attribute the gain to better memory use and additional anchor training.The latter effect should be minor if the anchor is already converged.Source: §4 · p. 8, Figure 6 · p. 9
  • Across three open-weight architectures, post-hoc memories improve specific-knowledge accuracy while common-knowledge results stay the same or slightly decrease; even wrong-cluster retrieval matches or beats the no-memory baseline.The authors suggest the common-knowledge change may relate to the models’ tuned pretraining mixtures and use of DCLM.Source: §4 · p. 9, Table 4 · p. 10
  • Less than 10% additional inference time for a 40-token generation includes routing, fetching, and extra inference FLOPs; adversarially blocking 1/16 of the bank drops atomic-number accuracy from 70% to 20%.The runtime comparison is on an H100 GPU; the blocking result is preliminary.Source: §4 · p. 9, Figure 6 · p. 9, Table 3 · p. 10
  • 45.7% AVG-SK for 410M with RAG-Wiki plus parametric memories surpasses RAG-Wiki alone at 41.6% and parametric memories alone at 44.5%.The paper describes the two approaches as complementary.Source: §4 · p. 10, Table 5 · p. 10

5 Related works

Related work includes raw-text and symbolic databases, parametric memories, and mixture-of-experts approaches. The paper contrasts their storage, retrieval, and deployment properties with its hierarchical memories.

5 Related works

  • Raw-text memory approaches include Retro [Borgeaud et al., 2022] and retrieval to an external knowledge base during pretraining [Zhao et al., 2025]; the paper identifies corpus scalability and low compression as limitations.Ahn et al. (2016) use a symbolic knowledge graph.Source: §5 · p. 10
  • KV memories, including Cartridges, underperform FFN memories for large-scale memorization; MemSinks instead aim to discard memorization parameters at inference.The paper relates context-dependent retrieval to task-conditioned parameter selection in instruction-following pruning.Source: §5 · pp. 10–11
  • MoE-related methods include PKM and FlexOlmo; the paper distinguishes its design by context-based offloading, explicit memory control, and post-hoc memory addition.The authors expect some of their insights may carry over to MoE approaches.Source: §5 · p. 11

6 Discussion

The paper proposes hierarchical memory banks to store long-tail world knowledge while small anchor models capture common knowledge and reasoning. It identifies limits and future directions for this approach, including unexplored scaling laws and anchor-model design.

6 Discussion

  • Small language models augmented with memory banks match regular transformers with 2× more parameters.The proposed memories target long-tail world knowledge; the anchor captures common knowledge and reasoning.Source: §6 · p. 11
  • Memory scaling laws, anchor architecture and retrievers, anchor reasoning, multilingual use, and non-text modalities remain future directions.Dense-training scaling laws may not apply because memory parameters update less often and receive gradients from only part of the dataset.Source: §6 · p. 11

A Training details

  • Experiments use BFloat16 parameters, PyTorch FSDP, AdamW, gradient clipping at 1.0, and NVIDIA H100 GPUs.AdamW uses β1 = 0.9 and β2 = 0.95.Source: A Training details · p. 17

A.1 Baseline model training

The appendices document baseline and memory-training setups, ablation budgets, and configurations for co-training, open-weight models, and memory-to-model-size experiments.

A.1 Baseline model training

  • Baseline models train from scratch for 240 ≃1.1T tokens with 8,192-token context, 1M-token batches, and cosine learning-rate schedules.Training uses DCLM-Baseline with dataset decomposition, 10k warmup steps, a 5 × 10−3 maximum learning rate, and 5 × 10−5 cooldown.Source: §A.1 · p. 17

A.2 Memory model training

  • Memory training groups documents by level-4 cluster, packs them into 2,048-token sequences, and masks attention across documents.Frozen-anchor memories train asynchronously in 16 level-1-subtree jobs; memory learning uses a 10−4 maximum and 10−5 cooldown learning rate.Source: §A.2 · p. 17

A.3 Additional details of memory type and size ablations

  • Memory-type experiments use 238 total seen tokens, while memory-level experiments use 236, 238, 240, and 241 for levels 1–4.The experiments in Figure 3 use level-2 memories for all memory types; deeper levels receive larger token budgets because they update less often.Source: §A.3 · p. 18

A.4 Additional details of ablation on memory configuration

  • The configuration ablation uses a frozen 160M anchor and 4.6B or 18.7B memory banks trained for 1.1T or 2.2T tokens, respectively.Source: Table 9 · p. 18, §A.4 · p. 18

A.5 Additional details of co-training memories and anchor model

  • The co-training experiment trains memory and anchor parameters together from scratch with a 10−3 maximum learning rate and 0.1 weight decay.It uses a 10−5 cooldown and 10k warmup steps.Source: §A.5 · p. 18

A.6 Additional details of memory learning on open-weight models

  • Post-hoc memories for frozen Gemma, Qwen, and Llama anchors use configurations sized to roughly 10% of anchor parameters.The memory-learning setup omits cross-attention restrictions and uses 5k warmup steps.Source: §A.6 · p. 18

A.7 Additional details of memory-to-model size experiment

  • For the memory-to-model-size experiment, six anchors are trained for 238 tokens, frozen, then paired memories are trained for 239 tokens.The anchor and memory configurations are designed to keep total runtime parameters and memory-bank size almost fixed; memory training uses a constant 10−4 learning rate.Source: §A.7 · pp. 18–19, Table 10 · p. 19

B Evaluation

Evaluation covers 13 tasks and selects benchmarks meeting baseline-performance and stability criteria. The appendix also reports per-dataset results for the Table 1 models.

B Evaluation

  • 13 tasks are evaluated, retaining only benchmarks where the 410M baseline beats random and the standard-deviation-to-mean ratio is below 0.5.Source: B Evaluation · p. 19
  • The benchmark set includes multiple-choice, generation, language-modeling, and schema tasks, with shot counts and random baselines varying by task.Examples include ArcChallenge, TriviaQA, PIQA, Lambada-OpenAI, and Winograd.Source: B Evaluation · p. 19
  • Memory retrieval uses the question for multiple-choice tasks, context for language modeling, shared choice text for schema tasks, and the full document for Wiki-En perplexity.Source: B Evaluation · p. 20
  • The atomic-number task analyzes accuracy against element-name frequency, counting accepted aliases while noting ambiguities such as “lead.”Source: B Evaluation · p. 20

B.1 Per dataset evaluation results

C Data clustering details

The paper clusters training data hierarchically and studies how tokens per parameter affect memory learning. Deep memories receive topic-specific updates, and specific-knowledge accuracy and Wiki-En perplexity improve as TPP increases.

C Data clustering details

  • The clustering tree assigns each example to the closest child centroid, splitting nodes into 16 clusters at each level to a depth of four.Text is segmented into 2,048-token windows and embedded as 384-dimensional sentence-BERT vectors before clustering.Source: C Data clustering details · p. 20

D Tokens per parameter in memory learning

  • The TPP study freezes a pretrained 160M anchor and trains 4,096 memories totaling 3.5B parameters across TPP values of approximately 1–312.The setup uses memory configuration (0,0,16,0).Source: D Tokens per parameter in memory… · p. 21
  • Deep memory parameters are activated on similar-topic sequences, reducing overwriting by unrelated gradients; level-3 memories update once per 4,096 training sequences.Memory parameters at level l update 16^l times less often than anchor parameters, as stated in the passage.Source: D Tokens per parameter in memory… · p. 21
  • Specific-knowledge accuracy and Wiki-En perplexity steadily improve with increasing TPP, reaching the reported endpoint at TPP = 312 after 1.1 trillion seen tokens.Figure 7 plots common-knowledge and specific-knowledge averages and Wiki-En perplexity against TPP; the study does not scale TPP further for computational reasons.Source: D Tokens per parameter in memory… · p. 21, Figure 7 · p. 21

E Fetched memory vs generic memory (additional results)

Co-training the anchor with memories outperforms freezing it, while the advantage of fetched over generic memories grows during training. Memory placement also matters: early-layer memories are less effective than uniform or deeper placement.

E Fetched memory vs generic memory (additional results)

  • Avg-SK rises from 39.2% with a frozen anchor to 40.3% with co-training, which helps the anchor adapt to memories and adds training tokens.The extra-token benefit should be minor when the anchor is already over-trained and has reached performance saturation.Source: E Fetched memory vs generic memory… · p. 22, Figure 8 · p. 22
  • The performance gap between fetched and generic memories grows during co-training, as the larger memory bank benefits more from extended training.The example uses a 4.6B memory bank, compared with a 160M anchor and an 18M generic memory.Source: E Fetched memory vs generic memory… · p. 22, Figure 8 · p. 22

F Memory location

  • Early-layer memories are less effective than uniform placement or placement in deeper layers.The comparison uses a frozen anchor and the same number of memory parameters across placements.Source: F Memory location · p. 23, Table 12 · p. 23

G What tasks benefit more from memories

The analysis estimates each benchmark’s knowledge specificity and compares it with the gains from fetched memories. Gains are larger on specific-knowledge tasks, while common-knowledge performance remains comparable.

G What tasks benefit more from memories

  • Knowledge specificity is estimated by GPT-4 ratings of 100 sampled entries per dataset, using average education-level ratings as the task score.The rating scale runs from language understanding or commonsense through more specific facts learned later in life.Source: G What tasks benefit more from memories · p. 23
  • For the 1.4B model, accuracy gains of fetched over generic memories positively correlate with benchmark knowledge-specificity scores, based on GPT-4 ratings of 100 sampled entries per dataset.The same analysis is shown for 160M and 410M models trained with memories.Source: Figure 10 · p. 24, G What tasks benefit more from memories · p. 24
  • Fetched memories improve specific-knowledge task performance by 3.6 points on average, while common-knowledge performance is comparable (64.4 vs. 64.5).Against the baseline, Avg-CK improves from 61.2% to 64.5%, and Avg-SK from 49.7% to 54.9%.Source: G What tasks benefit more from memories · p. 24

H Retrieval augmented generation detail

The section compares learned memories with retrieved documents and describes several memory-augmented transformer designs. It also details how one hierarchical configuration maps to fetched-memory and bank sizes.

H Retrieval augmented generation detail

  • The comparison uses fetched memories or raw retrieved documents from the same dataset, with a shared Sentence-BERT retrieval process and zero-shot tasks.The experiments use DCLM-Baseline and English Wikipedia 2022; retrieval selects a nearest-neighbor document from a level-3 cluster.Source: H Retrieval augmented generation detail · pp. 25–26
  • Learned memories improve common- and specific-knowledge benchmarks with lower runtime FLOPs overhead; vanilla DCLM RAG helps neither, while Wiki RAG improves specific-knowledge performance (46.9 to 49.2 for 1.4B) but not common-knowledge.Memory adds approximately 10% runtime FLOPs, and the paper describes high-quality RAG as complementary to learned memories.Source: H Retrieval augmented generation detail · pp. 25, 27

I Memory augmented transformer architecture

  • LoRa memories add low-rank matrices to selected query/key, value/output, or FFN linear layers, with memory size determined by rank and target layer type.The three variants are LoRa-QK, LoRa-OV, and LoRa-FFN.Source: I Memory augmented transformer… · p. 27
  • LoRa-FFN mostly outperforms alternative LoRa memories for the same number of learnable parameters.The authors relate this observation to prior findings that transformer knowledge is primarily stored in FFN layers (Dai et al., 2022; Yao et al., 2022).Source: I Memory augmented transformer… · p. 27
  • KV memories add learned keys and values for attention, whereas FFN memories expand the inner dimension of the three SwiGLU feed-forward linear layers.KV queries cross-attend to learned keys and values; FFN memories are initialized to have no effect at the start of training.Source: I Memory augmented transformer… · pp. 27–28

I.1 Memory size calculation

  • For the 160M anchor with configuration (256, 64, 16, 0), fetched memory is approximately 18M parameters and the full bank approximately 4.6B.The per-level fetched sizes are 13,762,560, 3,440,640, 860,160, and 0 parameters.Source: §I.1 · p. 28

J Retriever failure analysis

The failure analysis probes memory retrieval errors through corrupted memory IDs and sibling-leaf substitutions. Wrong retrieval reduces gains, but reported performance remains at or above the no-memory baseline.

J Retriever failure analysis

  • For atomic-number generation, the answer remains correct when the level-1 memory is correct, but becomes incorrect when that memory is also wrong.In the example, the correct Fermium answer is 100; with the level-1 memory wrong, the model outputs 92.Source: J Retriever failure analysis · p. 29
  • The stress test uses a frozen 160M anchor and a flat bank of 65,536 non-overlapping leaf memories, replacing each nearest neighbor with one of 15 sibling leaves.Accuracy is averaged over all 16 choices for each query, then over the benchmark.Source: J Retriever failure analysis · p. 31
  • Wrong retrieval reduces the gain, but performance remains at or above the no-memory baseline, indicating graceful degradation in the tested analysis.For the hierarchical case, the authors expect greater robustness because retrievals can share shallower-level memory parameters.Source: J Retriever failure analysis · p. 31, Table 13 · p. 31
Loading 2510.02375v3…