Source-linked AI summary

Pre-training Limited Memory Language Models with Internal and External Knowledge

Linxi Zhao, Sofian Zalouk, Christian K. Belardi, Justin Lovelace, Jin Peng Zhou, Ryan Thomas Noonan, Dongyoung Go, Kilian Q. Weinberger, Yoav Artzi, Jennifer J. Sun

arXiv:2505.15962v3cs.CLcs.AIcs.LG

TL;DR

Language models entangle factual knowledge with linguistic ability in their weights, making specific facts difficult to inspect, verify, or update independently. LMLM externalizes entity-level facts to a database and masks retrieved values during pretraining. A 382M-parameter LMLM matches the factual precision of a LLAMA2-7B, while the approach also supports verifiable updates and unlearning through database operations.

  • Problem

    Factual knowledge and linguistic ability are entangled in model weights, making specific facts difficult to update without affecting language competency.

    Source: §1 · p. 1
  • Method

    LMLM annotates training text with entity-level database lookups, masks retrieved facts from the pretraining loss, and queries the database during inference.

    Source: §1 · p. 2
  • Results

    A 382M-parameter LMLM matches the factual precision of a LLAMA2-7B, and smaller LMLM instances match or exceed much larger off-the-shelf models.

    Source: §1 · p. 2, §6 · p. 9
  • Takeaways & Limitations

    Externalizing facts enables instant, verifiable updates and unlearning through database operations.

    Source: §1 · p. 2

    Experiments are limited to small models and datasets due to compute constraints.

    Source: §6 · p. 9

Abstract

from arXiv · show

Neural language models are black-boxes--both linguistic patterns and factual knowledge are distributed across billions of opaque parameters. This entangled encoding makes it difficult to reliably inspect, verify, or update specific facts. We introduce Limited Memory Language Models (LMLM), a new class of language models that externalizes factual knowledge to external database during pre-training rather than memorizing them. Our pre-training approach strategically masks externally retrieved factual values from the training loss, thereby teaching the model to perform targeted lookups rather than relying on memorization in model weights. Our experiments demonstrate that LMLMs achieve competitive performance compared to significantly larger LLMs on standard benchmarks, while offering the advantages of explicit, editable, and verifiable knowledge bases.

1 INTRODUCTION

The introduction frames factual knowledge and language ability as entangled in model weights, making facts difficult to update independently. LMLM shifts entity-level facts to an external database and reports improved parameter efficiency and knowledge control.

1 INTRODUCTION

  • Entangling factual knowledge with language competency makes facts difficult to update without affecting linguistic ability.Reliable memorization can require repeated exposure, while outdated or inappropriate knowledge can be difficult to unlearn.Source: §1 · p. 1
  • LMLM teaches models to query an external database for entity-level facts rather than memorize them in model weights.Its recipe annotates training text, masks returned facts from the loss, and uses database queries during inference.Source: §1 · p. 2
  • Unlike approaches such as RAG that maximize access to external information, LMLM limits knowledge stored in parameters during pretraining.The introduction says this can free model capacity and enable precise factual control; LMLM remains compatible with RAG, tool calling, and task-specific fine-tuning.Source: §1 · p. 2
  • LMLM relies on factual lookups, so out-of-scope knowledge triggers detectable lookup failures rather than a fallback to internal knowledge.The introduction contrasts this with traditional RAG models without LMLM pretraining, which may fall back on internal knowledge.Source: §1 · p. 2
  • A 382M-parameter LMLM matches the factual precision of a LLAMA2-7B, and LMLMs achieve lower perplexity than comparable models.The comparison is with LLM counterparts of the same size, pretrained on the same corpus without external memory; database operations also enable instant, verifiable updates and unlearning.Source: §1 · p. 2

2 RELATED WORK

The related work distinguishes knowledge stored in model parameters from knowledge retrieved externally, and reviews retrieval, memory, and unlearning approaches. It positions LMLM as restricting factual memorization during pretraining and treating forgetting as database deletion.

2 RELATED WORK

  • Table 1 compares selected related work using marks for whether listed capabilities are supported.The caption points to a full version in Table 11.Source: Table 1 · p. 2
  • Parametric knowledge is stored in model weights, while non-parametric knowledge is retrieved externally; parametric knowledge can be under-trained and has issues including hallucination, staleness, and limited adaptability.Prior findings cited in the section report that knowledge seen fewer than a few hundred times may remain under-trained.Source: §2 · pp. 2–3
  • Retrieval-augmented and tool-based models add external information, while LMLM directly restricts memorization during pretraining and supports instant unlearning.The section says these capabilities are difficult to achieve with prior retrieval-augmented methods.Source: §2 · p. 3
  • Prior unlearning methods use gradient-based updates or preference optimization, whereas LMLM externalizes facts so forgetting reduces to deleting database entries.The section describes prior approaches as difficult to scale to larger forget sets and as sometimes degrading utility or erasing related knowledge.Source: §2 · p. 3

3 LIMITED MEMORY LANGUAGE MODELS

LMLM externalizes entity-level factual knowledge to a database, trains on text with lookup calls while masking retrieved values from the loss, and uses database lookups during generation. Its data-preparation pipeline extracts facts and scales annotation through a distilled model.

3 LIMITED MEMORY LANGUAGE MODELS

  • LMLM is designed to offload factual knowledge to an external database rather than store it implicitly in model parameters during pretraining.The framework covers data preparation, lookup-masked pretraining, and inference.Source: §3 · p. 3

3.1 PRETRAIN DATA PREPARATION: SEPARATING KNOWLEDGE

  • The method represents entity-level atomic facts as (entity, relation) → value triplets, which are compact and tractable to extract and verify.The section says these triplets map naturally to a knowledge graph and are suited to external storage.Source: §3.1 · p. 4
  • A scalable ANNOTATOR is distilled from high-quality GPT-4o seed annotations after a filtering stage to externalize facts from raw text.The section describes this pipeline as a response to the difficulty of manually extracting factual triples at pretraining-corpus scale.Source: §3.1 · p. 4
  • Extracted facts form a token-efficient external database, while lookup calls interleaved with the corpus teach when to use internal knowledge or issue a lookup.Source: §3.1 · p. 4

3.2 PRE-TRAINING AND INFERENCE

  • Pretraining excludes retrieved factual values and the ending lookup token from the next-token loss, discouraging memorization; inference triggers database retrieval during autoregressive generation.The method separates learning to generate lookup calls from storing and retrieving factual knowledge in the database.Source: §3.2 · pp. 4–5
  • LMLM reports lower pretraining perplexity, higher factual precision while maintaining NLU performance, and targeted forgetting with preserved utility.Figure 4 summarizes results for LMLM with a LLaMA backbone.Source: Figure 4 · p. 5

4 EXPERIMENTS

The experiments evaluate LMLM on language modeling, unlearning, and factual precision, comparing it with standard pre-training and other baselines. They report lower perplexity, editable factual knowledge, and competitive factual performance, while preserving general language understanding.

4.1 EXPERIMENTAL SETUP

  • Models were pretrained for eight epochs on an approximately 3-billion-token Wikipedia corpus; the database contains 54.6 million knowledge triplets.The models used GPT-2- and LLaMA2-style architectures, and retrieval used fuzzy matching with a rejection threshold of 0.6.Source: §4.1 · pp. 5–6

4.2 LEARNING TO LOOKUP FACTS IS EASIER THAN MEMORIZATION

  • Across model sizes and perplexity variants, LMLM lowers perplexity than STANDARD, averaging a 1.98-point reduction under Dynamic despite imperfect lookup calls.Improvements across variants indicate that LMLM learns both to query and to generate more effectively.Source: §4.2 · p. 6, Figure 5 · p. 6
  • LMLM performs on par with STANDARD across five NLU tasks, indicating that factual offloading preserves general-purpose language understanding.This evaluation is presented as a sanity check on whether separating factual knowledge affects general understanding.Source: §4.2 · p. 6

4.3 MACHINE UNLEARNING: LMLM SUPPORTS INSTANT FORGETTING

  • On TOFU’s 5% Forget Set, LMLM achieves effective forgetting without utility loss and preserves knowledge outside the set through database operations, without model updates.Figure 6 compares LMLM with NPO on forgetting quality, utility, and retention.Source: Figure 6 · p. 7, §4.3 · p. 7

4.4 EXTERNALIZING KNOWLEDGE IMPROVES FACTUAL PRECISION

  • On LLAMA2-382M, LMLM gains +17.9% FactScore, +6.1% T-REx EM, and +28.1% PopQA Acc over STANDARD, its knowledge-dense counterpart, approaching PYTHIA-1B and LLAMA2-7B despite far fewer parameters.Source: §4.4 · p. 8
  • Against controlled RAG at the current scale, LMLM achieves higher FactScore and PopQA accuracy, while RAG performs better on T-REx.The authors suggest T-REx’s verbatim Wikipedia queries make retrieval from that corpus especially effective.Source: §4.4 · p. 8, Table 3 · p. 8

5 FURTHER ANALYSIS

Further analyses examine whether masked loss reduces factual memorization and how the amount of externalized knowledge relates to performance. The results support offloading triplet-representable knowledge, while broader and more abstract knowledge remains an open direction.

5 FURTHER ANALYSIS

  • LMLM maintains high loss on factual return-value tokens while standard SFT rapidly lowers it; disabling retrieval substantially reduces FactScore and T-REx performance.These findings are presented as preliminary evidence that LMLM reduces factual memorization in parameters.Source: §5 · p. 8, Table 4 · p. 8, Figure 7 · p. 9
  • The offloading criterion ranks facts by the loss difference between STANDARD and LMLM, prioritizing facts that STANDARD models struggle to learn.The authors characterize such facts as typically long-tail or highly specific; common or easily inferred facts may offer limited upside to offload and may harm general capabilities.Source: §5 · pp. 8–9
  • For knowledge representable as triplets, greater offloading improves language modeling and factual precision while preserving natural language understanding; broader abstract knowledge remains untested.The authors identify extending the analysis to broader and more abstract forms of knowledge as future work.Source: §5 · p. 9, Figure 8 · p. 9

6 DISCUSSION

The discussion reports that smaller LMLMs match or exceed much larger off-the-shelf models, while identifying limits in knowledge scope, factuality, and evaluation.

Toward Efficient Scaling via Knowledge Offloading

  • Even smaller LMLMs match or exceed much larger off-the-shelf models, though experiments remain limited to modest scales by compute constraints.The authors say knowledge offloading may maintain strong factual accuracy with fewer parameters.Source: §6 · p. 9

Beyond Entity-Level Knowledge Offloading

  • Extending knowledge offloading beyond entity-level facts remains challenging because effective formats are unclear and simple QA annotations risk hallucinated facts.The paper focuses on entity-level knowledge triplets.Source: §6 · p. 9

Limitations

  • Some factual knowledge remains unremoved, and more comprehensive probing benchmarks and distinct knowledge-versus-reasoning benchmarks remain underexplored.Source: §6 · p. 9
  • LMLM cannot guarantee perfect factuality, adds lookup-token costs, covers only entity-level facts, and has been tested only on small models and datasets.Database noise and fuzzy-matching errors can introduce inaccuracies; the authors say these issues are traceable and verifiable.Source: §6 · p. 9

7 CONCLUSION

The conclusion presents LMLMs as an alternative that externalizes facts to a database, with potential for verifiable updates and integration with other LLM approaches. The appendices describe the annotation pipeline, model training, and formalized training and evaluation objectives.

7 CONCLUSION

  • LMLMs externalize facts to a database, offering a more memory-efficient way to learn lookups and the potential for verifiable updates and integration with other LLM approaches.The conclusion also describes LMLM as an alternative way to store facts during pre-training.Source: §7 · p. 10

A.1 KNOWLEDGE EXTRACTION DETAILS

  • The annotation pipeline cleans GPT-4o-generated factual triples with a CORRECTOR, then trains an ANNOTATOR to identify grounded lookup calls and apply them across the pre-training corpus.The CORRECTOR filters calls that are not inferable from prior context, overly specific or inconsistent, or syntactically malformed.Source: §A.1 · pp. 16–17
  • GPT-4o annotations can depend on future context, leaving some facts unrecoverable from preceding text alone under left-to-right generation.The paper attributes this bimodal entity- and relation-token loss pattern as likely due to GPT-4o accessing future context.Source: §A.1 · p. 17

A.2 MODEL ARCHITECTURE AND TRAINING DETAILS

  • LMLMs are pretrained from scratch with GPT-2- or LLaMA2-style decoder-only architectures and four added lookup tokens, using 1,024-token contexts for eight epochs.The models use their original tokenizer and vocabulary, extended with the four special tokens.Source: §A.2 · p. 17

A.3 FORMALIZATION OF TRAINING AND EVALUATION OBJECTIVES

  • Training omits loss on retrieved factual values and the closing lookup marker, while normalized perplexity excludes retrieved values but normalizes by original-text length.Negative log-likelihood matches the training-loss computation.Source: §A.3 · p. 18

A.4 DATABASE AND RETRIEVAL SETTING

The database contains annotated entity-relation-value triplets, and retrieval uses fuzzy matching with a rejection threshold or an alternative constrained-generation approach.

A.4 DATABASE AND RETRIEVAL SETTING

  • The database contains 54.6M knowledge triplets; fuzzy matching uses cosine similarity over sentence embeddings and returns unknown below 0.6, while prefix-tree constrained generation keeps lookup calls covered by structured knowledge.Figure 9 contrasts unconstrained and constrained query generation.Source: §A.4 · p. 19, Figure 9 · p. 19

B.1 EVALUATION BENCHMARKS

The appendix specifies evaluation setups for language modeling and factual benchmarks, a TOFU unlearning comparison, and a study of selective knowledge offloading.

B.1 EVALUATION BENCHMARKS

  • Perplexity is evaluated on held-out Wikipedia passages, while FactScore measures supported atomic facts in generated biographies using the benchmark’s first 100 queries.Dynamic perplexity uses model-generated lookup arguments but fixes lookup positions to the reference.Source: §B.1 · p. 19
  • T-REx uses 11,615 left-to-right-compatible examples with Exact Match and Precision@1, while long-tail PopQA evaluates 1,399 rare-entity queries rewritten as knowledge-completion prompts.The PopQA subset covers entities with fewer than 100 monthly Wikipedia page views.Source: §B.1 · pp. 19–20

B.2 MACHINE UNLEARNING SETTING

  • The TOFU setup tests forget 5% using LLAMA-3.2-1B-INSTRUCT against NPO; LMLM removes Forget Set entries from its database, assuming the synthetic knowledge is represented there.The TOFU data are annotated to construct the database, and LMLM receives warm-up training before fine-tuning on the annotated training set.Source: §B.2 · pp. 20–21
  • TOFU judges unlearning by matching the retain model while preserving utility; likelihood-based metrics are not directly comparable across methods because raw and annotated references differ.Lower Forget Set ROUGE alone does not establish successful unlearning, since it can reflect general performance degradation.Source: §B.2 · p. 21

B.3 EXPERIMENTAL SETTING OF FURTHER ANALYSIS

  • Selective offloading ranks facts by the loss gap between LMLM and STANDARD on five following tokens, then progressively externalizes the largest-gap triplets across offloading ratios from 0% to 100%.A larger gap suggests STANDARD struggles to memorize a fact; a smaller gap may indicate easy memorization, contextual clues, or noisy or less useful lookups.Source: §B.3 · p. 21

B.4 EXPERIMENTAL SETTINGS BY FIGURE

The appendix describes evaluation and training configurations associated with Figures 4, 6, and 7, alongside the hyperparameter settings table.

B.4 EXPERIMENTAL SETTINGS BY FIGURE

  • Figure 4 evaluates held-out Wikipedia perplexity during pretraining, reports FactScore and NLU results in tables, and uses LLAMA-3.2-1B-INSTRUCT for its right panel.The held-out validation set contains 100 passages.Source: §B.4 · p. 21
  • Figure 6 reports NPO unlearning results averaged across five random seeds and shows LMLM only before and after unlearning, with retain-model performance as the ideal reference.The figure caption describes the five-seed NPO setup and retain-model marker.Source: §B.4 · p. 21
  • Figure 7 tracks loss on retrieved-value tokens in the TOFU synthetic training set to compare training objectives, evaluating models every 125 steps.Source: §B.4 · p. 21
  • Table 8 lists training hyperparameters for different experimental settings, including the 4k-pair Full set and 3.8k-pair Retain set.Source: Table 8 · p. 22

C DETAILED RESULTS

The detailed results compare LMLM with memory-based and retrieval-augmented approaches, test the role of database access, and examine language understanding, entity coverage, and query decoding.

C DETAILED RESULTS

  • LMLM scores higher than RAG on FactScore and PopQA, while RAG performs better on T-REx, whose queries may overlap verbatim with Wikipedia.LMLM uses structured entity-level lookups, whereas RAG retrieves unstructured text for the model to interpret; the approaches are described as complementary.Source: C DETAILED RESULTS · p. 23
  • Disabling database access reduces factual precision in FACTSCORE and T-REx, often below the STANDARD baseline, consistent with LMLM relying on retrieval rather than memorization.The results connect database edits with direct control over what the model knows and forgets.Source: C DETAILED RESULTS · p. 23, Table 9 · p. 23
  • Training objectives show no significant difference on Torg, entity, and relationship tokens, according to the additional loss curves.Figure 10 reports losses across those token types and return-value tokens.Source: C DETAILED RESULTS · p. 23, Figure 10 · p. 24
  • LMLM achieves consistently higher precision (Exact Match) on T-REx than prior memory-augmented and tool-augmented models, using fewer parameters and no large external tool APIs.Table 10 presents the comparison against prior memory-based work.Source: Table 10 · p. 24, C DETAILED RESULTS · p. 24

D.1 DOES LMLM AFFECT LANGUAGE UNDERSTANDING?

  • On selected high-signal NLU benchmarks, separating factual knowledge during pretraining does not compromise overall performance; Wikipedia-only pretraining limits some off-the-shelf comparisons.The selected benchmarks were intended to assess general language understanding in smaller models.Source: Table 12 · p. 25, §D.1 · p. 25

D.3 ABLATION: UNCONSTRAINED VS. CONSTRAINED QUERY GENERATION

  • Fuzzy matching outperforms prefix-tree decoding in this setting; trie constraints ensure valid queries and reduce hallucinations but can constrain diversity with a relatively small database.Fuzzy matching is used in the reported experiments.Source: Table 13 · p. 26, §D.3 · p. 26

D.2 ENTITY FREQUENCY ANALYSIS IN THE KNOWLEDGE DATABASE

  • Matched database entities span common and long-tail frequencies, with a substantial portion less frequent in the overall training corpus.Entity matching used fuzzy string matching against Wikipedia entries.Source: §D.2 · pp. 25–26, Figure 11 · p. 26

D.4 QUALITATIVE RESULTS

Qualitative examples contrast LMLM outputs with standard and off-the-shelf models, while triplet examples relate loss-difference quantiles to how difficult facts are to memorize.

D.4 QUALITATIVE RESULTS

  • LMLM produces concise, factually grounded responses, while standard and off-the-shelf models often include hallucinated content.The examples are presented as supporting the finding that LMLM achieves factual precision through external knowledge.Source: §D.4 · p. 26, Table 15 · p. 28
  • Higher per-triplet loss-difference quantiles correspond to facts harder for STANDARD to memorize; lower quantiles indicate easier or contextually obvious facts.Table 14 gives qualitative examples of triplets selected for externalization across quantile buckets.Source: Table 14 · p. 27, Figure 12 · p. 27
Loading 2505.15962v3…