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

TL;DR

This paper proposes Limited Memory Language Models (LmLm), which offload entity-level factual knowledge to an external database during pre-training and rely on explicit lookups during inference. By masking retrieved values in the pretraining loss, LmLm decouples factual memory from language understanding, enabling instant unlearning and easy knowledge editing through database operations. Experiments show LmLm achieves competitive perplexity and markedly higher factual precision with far fewer parameters than size-matched baselines, and it demonstrates effective forgetting on the TOFU benchmark without harming general capabilities. The approach complements retrieval-based systems and offers a path to scalable, verifiable knowledge grounding in language models, while highlighting limitations such as retrieval noise and current focus on entity-level facts.

Abstract

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.

Figures (12)

• Figure 1: A schematic illustration of LmLm. Unlike RAG (left), which exclusively adds knowledge from external sources, LmLmoffloads knowledge from the LLMs to the external database during pre-training.
• Figure 2: Overview of the LmLm framework. Our framework consists of (Left) Data Preparation, where entity-level facts are automatically annotated and stored in an external database; (Middle) Pretraining, where the model is trained on the annotated text while excluding return values from the loss to discourage memorization; and (Right) Inference, where the model interleaves text generation with database lookups to ground its outputs on retrieved facts.
• Figure 3: Training the Annotator model. We distill high-quality annotations from GPT-4o into a lightweight model that learns to identify and externalize factual knowledge from raw pre-training text, enabling scalable annotation of the full corpus.
• Figure 4: Results overview.(Left)$\textsc{LmLm}$ achieves consistently lower perplexity during pre-training, indicating that offloading factual knowledge improves pre-training efficiency. (Middle)$\textsc{LmLm}$ significantly improves factual precision over its $\textsc{Standard}$ counterparts while maintaining NLU performance. (Right) On the TOFU machine unlearning benchmark, LmLm forgets targeted facts while preserving general model utility. Results shown for LmLm with a LLaMA backbone; * denotes off-the-shelf models.
• Figure 5: Validation perplexity comparison between Standard and LmLm on three variants of perplexity. Lower perplexity indicates better performance.