Official repository for the paper:
Pre-training Limited Memory Language Models with Internal and External Knowledge

📄 Arxiv  |  🌐 Project Page  |  🤗 Hugging Face Models | 🎤 Talk by Kilian @ Simons Institute |

* LMLM is pronounced “LamLam” 🐑

🧠 What is LMLM?

Neural language models entangle language and knowledge, making it hard to verify, update, or forget facts.

Limited Memory Language Models (LMLMs) address this by combining:

• Internal memory (parameters) for fluency and reasoning
• External memory (database) for accurate, editable knowledge

This enables:

• Improved perplexity during pretraining
• Higher factual accuracy, without sacrificing general language understanding
• Instant editing or forgetting of specific facts via database updates

Table of Contents

• Quick Start
• Huggingface Pretrained Models and Data
• Build Your Own LMLM
• Data Preparation
  • Annotate Pretraining Corpus
  • Build External Knowledge Database
• Pretraining and Inference
• Evaluation
• Implementation Limitations and TODOs
• Repository Structure
• Citation
• Acknowledgment

🚀 Quick Start

We provide a custom Hugging Face model class LlamaForLMLM with a built-in generate_with_lookup method. Below is a minimal example of how to use LMLM for inference. For the full script, see scripts/eval/example_lmlm_inference.sh.

from transformers import AutoTokenizer
from lmlm.modeling_lmlm import LlamaForLMLM
from lmlm.database import DatabaseManager

tokenizer = AutoTokenizer.from_pretrained(model_path)
tokenizer.pad_token = tokenizer.eos_token

db_manager = DatabaseManager()
db_manager.load_database(database_path) if os.path.exists(database_path) else None
model = LlamaForLMLM.from_pretrained_with_db(model_path, db_manager=db_manager, use_special_tokens=True).cuda().eval()

output = model.generate_with_lookup(
    prompt='Tell me a bio of Ko Itakura. Ko Itakura is',
    tokenizer=tokenizer,
    max_new_tokens=256,
    enable_dblookup=True,
    enable_postprocess=False,
)
print(model.post_process(output, tokenizer))

To reproduce our results? Jump to 👉 Evaluation

To build your own LMLM? See 👉 Build Your Own LMLM

🤗 Huggingface Pretrained Models and Data

🛠 Installation

Standard installation (via Conda)

conda create -n lmlm python=3.10 -y
conda activate lmlm
pip install -r requirements.txt

Option 2: Development mode (editable install)

pip install -e .

Optional Dependencies

• TRL – for training
• LightEval – for NLU benchmarks
• FactScore – for factual precision benchmarks
• TOFU – for unlearning benchmarks

🏗 Build Your Own LMLM Pipeline

This section walks through constructing a full LMLM pipeline.

✍️ Annotator Preparation

Seed Annotation via GPT-4

Use GPT-4o to annotate a small seed dataset with [dblookup('Entity', 'Relationship') -> Value] calls and return values.

bash scripts/annotate/annotate_gpt4.sh

Requires OpenAI API key in scripts/account. The prompt can be found in ./prompts/gpt4o-v7.1.md

Train an Annotator Model

Fine-tune an LLaMA-3-style model (e.g., LLaMA3.1-8B-Instruct) on cleaned annotations.

bash scripts/train/ft-llama8b_annotator.sh

The prompt can be found in ./prompts/llama-v6.md

📦 Data Preparation

Once the annotator is trained, use it to label the full pretraining corpus and extract knowledge.

Annotate the Pretraining Corpus

We use the OLMo2 Wikipedia corpus (~3B tokens). Use vllm for faster inference. (Annotation over 64 A6000s takes ~2 days.)

bash scripts/annotate/annotate_annotator.sh

Build the Knowledge Database

Extract structured triplets (entity, relation, value) from the annotations to construct the DB.

bash scripts/train/extract_database.sh

Final DB contains ~54.6M knowledge triplets.

🧪 Pretraining and Inference

Pretrain LMLM

Train from scratch with masked facts and interleaved retrieval. Supports GPT-2 and LLaMA2-style models with special tokens for lookup. All runs use 1024-token context, 8 epochs, and mixed precision.

bash scripts/train/pretrain.sh

Example: LLaMA2-382M with batch size 256 trained for 105k steps (~8 H100-days)

Database Lookup

We use cosine similarity over a sentence transformers (all-MiniLM-L6-v2) embeddings with a threshold of 0.6 for fuzzy matching. See: lmlm/database/topk_retriever.py for implementation details.

Note: Embedding the entire database can be slow. To speed things up, we provide a precomputed FAISS index + mapping cache on Hugging Face Hub. The retriever will:

1. Load from local cache if it exists
2. Otherwise, attempt to download from Hugging (kilian-group/LMLM-database-cache)
3. If both fail, it will build the index from scratch and save it to local cache

Inference with DBLookup

Run LMLM generation with database lookup enabled. The generation examples are in output/eval/examples.

bash scripts/eval/example_lmlm_inference.sh

🔍 Evaluation

bash scripts/eval/eval_nlu_task.sh

bash scripts/eval/eval_factscore.sh

bash scripts/eval/eval_trex.sh

Repository Structure

LMLM/
├── configs/              # Experiment configurations
├── data/                 # Input datasets and database
├── experiment/           # Python scripts for evaluation
├── output/               # Outputs (annotations, model predictions, eval results)
├── scripts/              # Bash scripts for running experiments
├── src/
│   └── lmlm/
│       ├── annotate/         # Annotation pipeline
│       ├── database/         # Database construction and retrieval
│       ├── training/         # Pretraining and fine-tuning logic
│       ├── modeling_lmlm.py  # LMLM custom model class
│       └── constants.py
├── requirements.txt
├── pyproject.toml
├── README.md

Implementation Limitations and TODOs

Citation

If you find this work helpful, please consider citing:

@misc{zhao2025pretraininglimitedmemorylanguage,
      title={Pre-training Limited Memory Language Models with Internal and External Knowledge}, 
      author={Linxi Zhao and Sofian Zalouk and Christian K. Belardi and Justin Lovelace and Jin Peng Zhou and Ryan Thomas Noonan and Dongyoung Go and Kilian Q. Weinberger and Yoav Artzi and Jennifer J. Sun},
      year={2025},
      eprint={2505.15962},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2505.15962}, 
}

Acknowledgments

We gratefully acknowledge the open-source community. This project builds on the contributions and insights of many prior works.

License

This project is licensed under the MIT License. See the LICENSE file for details.