Co-LMLM: Continuous-Query Limited Memory Language Models

📄 arXiv • 🌐 Project page • 🤗 Models & Data

This repository contains the code accompanying the paper "Co-LMLM: Continuous-Query Limited Memory Language Models." It covers the full pipeline: data annotation, pretraining of the three model variants (Co-LMLM and the Standard LM / LMLM-Asker baselines), retrieval-index construction, and the perplexity / generation / downstream evaluation suite.

Co-LMLM is a retrieval-aware pretraining method: factual spans in the training corpus are annotated as <FACT q="..." a="...">span</FACT>, and the model learns to emit a continuous query vector at each <FACT> position that retrieves the span's content from an external index at inference time, rather than memorizing it in-weights. The models in the paper are SmolLM2 135M / 360M trained from scratch; variants trained on FineWeb-Edu are also supported by the same configs. In the code, the Co-LMLM trainer and its package are under colmlm/; the two baselines are the Standard LM (lm_baseline/) and LMLM-Asker (lmlm_asker/).

All runnable scripts use pyrallis for configuration: every script is invoked with --config_path path/to/config.yaml, and the YAML is parsed into a top-level dataclass defined next to the script. The main config class referenced under each section is the authoritative source for which fields a YAML can contain — start there when building your own configs. Example configs are in a configs/ directory next to each component.

Env Setup

Dependencies are managed with uv (Python 3.12). The full dependency set is pinned in pyproject.toml / uv.lock. To create a virtual environment and install everything:

uv sync

All scripts are then invoked via uv run, e.g.

PYTHONPATH=src uv run python -m lmlm.colmlm.train --config_path path/to/config.yaml

PYTHONPATH=src is required for the python -m / direct-script invocations because the source package layout is under src/ (src/annotation, src/lmlm, plus the top-level consts.py / common.py). The example commands below set it explicitly.

Output paths

Global output paths (annotations, training outputs, indices, and W&B logs) are defined in src/consts.py. They are all rooted under PROJECT_DIR, which reads the LMLM_PROJECT_DIR environment variable and defaults to ~/lmlm-project:

export LMLM_PROJECT_DIR=/path/to/your/storage   # default: ~/lmlm-project

API keys

Several stages call hosted LLM APIs and read their keys from environment variables:

GEMINI_API_KEY — Gemini seed annotation and the SimpleQA / FActScore LLM graders (default provider).
OPENAI_API_KEY — the OpenAI-backed alternative for the SimpleQA / FActScore graders.

Quick Start

# Install (see Env Setup) and download the model + Wikipedia index (~113 GB, resumable).
uv sync
hf download lil-lab/CoLMLM-360M-FW --local-dir ./CoLMLM-360M-FW
hf buckets sync hf://buckets/lil-lab/co-lmlm-360m-fw-wiki-index ./co-lmlm-wiki-index

import sys; sys.path.insert(0, "src")   # make the repo's packages importable (or set PYTHONPATH=src)
from lmlm.eval.hf_generate import load_retriever_generator

gen = load_retriever_generator(
    model_path="./CoLMLM-360M-FW",
    index_path="./co-lmlm-wiki-index",
    db_path="./co-lmlm-wiki-index/entries.db",
    max_new_tokens=48,
)

result = gen.generate("The chemical symbol for the element gold is")

print(result.text)
# The chemical symbol for the element gold is<FACT> Au</FACT> Au. It is a<FACT> transition metal</FACT>
# transition metal that is a<FACT> slightly reddish-yellow</FACT> slightly reddish-yellow metal. It is a
# very dense metal and is used in jewelry, coins, and other items. It

print(result.num_retrievals)                      # 3
for e in result.retrieved_entries:
    print(f"{e.text_value!r}  (score {e.score:.3f})")
# ' Au'  (score 0.909)
# ' transition metal'  (score 0.873)
# ' slightly reddish-yellow'  (score 0.844)

Every time the model emits <FACT>, it uses that position's hidden state to query the index and splices the retrieved value in as <FACT>value</FACT>, which the model then copies into its running text — so the facts (Au, transition metal, slightly reddish-yellow, each pulled from the Wikipedia index) come from the store, not the weights. The returned result also carries retrieved_entries (each a SearchResult with .text_value and .score), num_retrievals, failed_retrievals, and timing fields — inspect them to see exactly what was retrieved and where.

To generate over a whole file of prompts (batch, non-interactive), use the vLLM script instead — see §4.2 Generation. For the larger FineWeb-Edu + Wikipedia index and all other options, see Released Artifacts and generation.md.

Codebase Structure

src/
├── consts.py                 # Global output paths (rooted at LMLM_PROJECT_DIR)
├── common.py
├── annotation/               # Annotation pipeline (data side)
│   ├── annotate/             # Annotators + example configs (gemini/mlm/question_hybrid)
│   │   └── large_scale_annotation/   # Slurm job-array annotation at scale
│   ├── training/             # Trainers for the MLM and question generator
│   │   ├── mlm/              # Fact-span detector (token classification)
│   │   └── question_generator/       # Per-span question generator (causal LM)
│   ├── data_selection/       # Subsetting / preparing source corpora
│   └── prompts/              # Versioned Gemini prompts
└── lmlm/                     # Modeling, indexing, and evaluation
    ├── colmlm/               # Co-LMLM trainer
    ├── lmlm_asker/           # LMLM-Asker baseline trainer
    ├── lm_baseline/          # Standard LM baseline trainer
    ├── index/                # Index building (asker / retriever)
    │   └── large_scale_index/           # Sharded, Slurm-based index construction
    └── eval/                 # Perplexity, dynamic replacement, generation, factuality, retrieval
scripts/
└── eval/                     # NLU eval, inference-efficiency, results collection

1. Annotation Pipeline

The pipeline produces a pretraining corpus where factual spans are annotated as <FACT q="..." a="...">span</FACT> (q is the question, a is a paraphrased answer). There are three stages: (1) generate ground-truth annotations with Gemini on a seed set, (2) train smaller annotators on that seed, and (3) run the trained annotators over the full corpus. An overview of all annotator types is in src/annotation/OVERVIEW.md.

Source-corpus subsetting is handled by src/annotation/data_selection/data_selection.py (main config DataSelectionConfig; example configs/initial-5k-set.yaml).

1.1 Gemini seed annotation

Gemini produces the ground-truth annotations for the seed set. The script is annotate_with_gemini.py, and the versioned prompts (system.txt, prompt.txt, optional continue_prompt.txt) are under src/annotation/prompts/ (released version: claude_lmlm_opt_v2/v30.1).

PYTHONPATH=src uv run src/annotation/annotate/annotate_with_gemini.py \
    --config_path src/annotation/annotate/configs/gemini/dolmino-wiki.yaml

Main config class: GeminiAnnotationConfig in annotate_with_gemini.py. It supports three input sources (HuggingFace dataset, prepared directory, local JSONL), a GeminiConfig selecting the model/prompt version, and a sharded on-disk cache so partial runs resume.

1.2 Training the annotators

Three annotator components are trained from the Gemini-generated seed. Each is launched through a run_training.sh <config> wrapper around accelerate launch.

MLM fact-span annotator — a ModernBERT token-classification head predicting the <FACT> / </FACT> span boundaries.

bash src/annotation/training/mlm/run_training.sh \
    src/annotation/training/mlm/configs/modernbert_large.yaml

Main config class: MLMExperimentConfig in src/annotation/training/mlm/config.py. Supports LoRA or full fine-tuning, multi-GPU via accelerate, and seqeval / span-IoU metrics. Example config: modernbert_large.yaml.

Question generator — a causal LM trained to emit the question for a single fact span, given the surrounding annotated context.

bash src/annotation/training/question_generator/run_training.sh \
    src/annotation/training/question_generator/configs/default.yaml

Main config class: QuestionGeneratorExperimentConfig in src/annotation/training/question_generator/config.py. Backed by TRL's SFTTrainer, with optional config-controlled question-leakage filtering, document splitting, and question packing (efficient multi-question-per-document training).

1.3 Running the annotators

For the paper we use the question-hybrid annotator: the MLM produces the fact spans, then the question generator fills in a question per span using a shared, prefilled KV cache for efficiency.

CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src uv run \
    src/annotation/annotate/annotate_with_question_hybrid.py \
    --config_path src/annotation/annotate/configs/question_hybrid/example.yaml

Main config class: QuestionHybridAnnotationConfig in annotate_with_question_hybrid.py. Supports five input modes (HF dataset, prepared dir, an already-annotated corpus whose questions should be regenerated, local JSON, local JSONL).

The MLM-only annotator is also available for span detection on its own:

CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src uv run \
    src/annotation/annotate/annotate_with_mlm.py \
    --config_path src/annotation/annotate/configs/mlm/dolmino-wiki.yaml

Main config class: MLMAnnotationConfig in annotate_with_mlm.py.

1.4 Annotation at scale

For very large corpora there is a Slurm-based job-array runner under large_scale_annotation/; the entry point is run_manager.py with top-level config LargeScaleAnnotationConfig.

PYTHONPATH=src uv run -m annotation.annotate.large_scale_annotation.run_manager \
    --config_path src/annotation/annotate/large_scale_annotation/configs/example_config.yaml

2. Pretraining

All three pretraining variants share the same streaming-dataset machinery (see src/lmlm/dataset_base.py and the per-variant dataset.py files) and are launched the same way: uv run python -m <module> for single-GPU, or accelerate launch -m <module> for multi-GPU. Each loads an annotated corpus produced by the annotation pipeline above. The released models are SmolLM2 135M / 360M trained from scratch; FineWeb-Edu variants use the same configs with a different data source.

2.1 Co-LMLM

The retrieval-aware pretraining method introduced in the paper. The trainer masks out fact-span content from the next-token-prediction loss and adds an InfoNCE fact–question contrastive objective so that <FACT> hidden states become useful continuous retrieval queries.

PYTHONPATH=src uv run accelerate launch -m lmlm.colmlm.train \
    --config_path src/lmlm/colmlm/configs/example_config.yaml

Main config class: LMLMConfig in src/lmlm/colmlm/config.py. Sub-configs to look at: ModelConfig, ContrastiveLossConfig, DataConfig, TrainingConfig, OptimizerConfig.

An in-loop retrieval-eval example (retrieval metrics + full-eval perplexity computed during training) is provided in example_inloop_retrieval_eval.yaml.

2.2 Standard LM (baseline)

A plain causal-LM baseline trained on the unannotated text (fact tags are stripped). Loss is broken out per token role (overall / fact-span / rest) so the baseline's perplexity on fact spans can be compared directly to the LMLM variants.

PYTHONPATH=src uv run accelerate launch -m lmlm.lm_baseline.train \
    --config_path src/lmlm/lm_baseline/configs/example_config.yaml

Main config class: LMBaselineConfig in src/lmlm/lm_baseline/config.py.

2.3 LMLM-Asker (baseline)

The "ask, then retrieve" baseline: the model is trained to emit <FACT>question</QUESTION> before each fact span (the question sits between <FACT> and </QUESTION>), so at inference the natural-language question is used as a query against an external sentence-transformer-encoded index. The fact-span content itself is excluded from the loss.

PYTHONPATH=src uv run accelerate launch -m lmlm.lmlm_asker.train \
    --config_path src/lmlm/lmlm_asker/configs/example_config.yaml

Main config class: LMLMAskerConfig in src/lmlm/lmlm_asker/config.py.

3. Index Building

Inference for both LMLM-Asker and Co-LMLM requires a retrieval index over the background corpus (an annotated corpus, typically separate from the evaluation target). The single entry point is src/lmlm/index/build_index.py; exactly one of two sub-configs determines what gets built:

asker — encodes the questions (and stores their fact-span answers) with a sentence-transformers model into an AskerIndex. Used by LMLM-Asker.
retriever — runs a forward pass of the trained Co-LMLM model over the annotated background corpus and indexes the hidden state at each <FACT> token (RetrieverIndex). This is the variant used for Co-LMLM in the paper.

CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src uv run -m lmlm.index.build_index \
    --config_path path/to/your_index_config.yaml

Main config class: IndexBuildConfig in build_index.py (with IndexBuildSettings, AskerBuildConfig, RetrieverBuildConfig).

For corpora too large to embed in a single process, a sharded Slurm-based pipeline is in large_scale_index/ (see its README.md): embedding_workers/ manages distributed embedding extraction, then build_index.py / merge_sharded_index.py build the FAISS index from the shards. The entry point is embedding_workers/run_manager.py with top-level config LargeScaleIndexConfig.

4. Evaluation & Generation

The full evaluation flow is documented in src/lmlm/eval/eval_flow.md and the generation scripts (vLLM and HuggingFace backends) in src/lmlm/eval/generation.md. A summary follows.

4.1 Perplexity & dynamic replacement

For LMLM-Asker and Co-LMLM the recommended path is the unified full_eval.py, which chains the four steps below into one script with auto-derived intermediate paths:

annotate the target and background datasets (using the annotators above);
compute static / normalized perplexity with compute_perplexity.py;
build the background index (lmlm.index.build_index);
run dynamic replacement — vllm_asker_dynamic_replace.py for LMLM-Asker, batched_retriever_dynamic_replace.py for Co-LMLM — and recompute perplexity to obtain the dynamic-PPL number reported in the paper.

CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src uv run src/lmlm/eval/full_eval.py \
    --config_path path/to/your_full_eval_config.yaml

Main config class: FullEvalConfig in full_eval.py. For the Standard LM baseline only step 2 is needed. Each step can also be run on its own — see eval_flow.md for per-step invocations.

4.2 Generation

For batch generation over a file of prompts, use the vLLM scripts — one per model family, and the fastest for throughput. (For interactive, single-prompt use, prefer the load_retriever_generator helper from the Quick Start; once you are running non-interactively vLLM is the better choice.) All scripts read JSONL prompts and write JSONL outputs; templates are in configs/generate/:

Model	vLLM script	Config template
Standard LM	`vllm_generate.py`	`lm_template.yaml`
LMLM-Asker	`vllm_asker_generate.py`	`asker_template.yaml`
Co-LMLM	`vllm_retriever_generate.py`	`retriever_template.yaml`

For LMLM-Asker, when </QUESTION> is emitted the question is queried against an AskerIndex and the retrieved answer is injected (immediately followed by </FACT>) before generation continues. For Co-LMLM, when <FACT> is produced a separate HF model extracts the hidden state at that position; that vector queries a RetrieverIndex (retrieval runs concurrently with vLLM via a thread pool).

# Fill in model_path / index_path / index.db_path / prompts_path / output_path in the template,
# then run the Co-LMLM (retriever) generator:
CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src uv run src/lmlm/eval/vllm_retriever_generate.py \
    --config_path src/lmlm/eval/configs/generate/retriever_template.yaml

The same retrieval loop is also exposed through the HuggingFace transformers backend — interactively via load_retriever_generator(...) (Quick Start) and as a config-driven script hf_generate.py (--model_type asker|retriever). See generation.md for full invocation examples of both backends (including one- vs two-GPU setups for the retriever script).

An eval-only baseline for the original structured-LMLM system (JSON knowledge base / FAISS lookup via explicit <|db_entity|> markup) is provided in vllm_lmlm_baseline_generate.py for comparison; its outputs feed the same downstream scorers below. (Training of structured-LMLM is out of scope for this repository.)

4.3 Downstream factuality

Prompt-prep and scoring scripts sit next to the generation scripts. Prep writes JSONL prompts, a generation script produces continuations, then a scorer grades them.

PopQA / SimpleQA — prepare_popqa_prompts.py (PrepareOpenQAPromptsConfig, --dataset_name popqa|simpleqa) prepares both. PopQA is graded by score_popqa.py (ScorePopQAConfig, substring exact-match); SimpleQA is graded by the canonical LLM grader score_simpleqa.py (ScoreSimpleQAConfig, requires GEMINI_API_KEY or OPENAI_API_KEY).
T-REx — prepare_trex_prompts.py (PrepareTrexPromptsConfig) and score_trex.py (ScoreTrexConfig).
FActScore — prepare_factscore_prompts.py (PrepareFactScorePromptsConfig) prepares prompts; the atomic-fact scorer is in factscore/factscorer.py (FactScorerConfig, requires an LLM API key).

PYTHONPATH=src uv run src/lmlm/eval/prepare_popqa_prompts.py \
    --dataset_name simpleqa \
    --output_path output/eval/simpleqa/simpleqa_prompts.jsonl

4.4 NLU evaluation

Multiple-choice / language-understanding benchmarks via lighteval (with optional masking of LMLM factual special tokens). Driver: scripts/eval/eval_nlu_task.sh, which invokes scripts/eval/eval_nlu_masked.py. Task definitions are in src/lmlm/eval/lighteval_tasks.py.

CUDA_VISIBLE_DEVICES=0 CKPT=/path/to/checkpoint SETTING=smollm2-setting \
    bash scripts/eval/eval_nlu_task.sh

4.5 RAG comparison baseline

A retrieval-augmented-generation baseline that prepends retrieved passages to the prompt: prepare_rag_prompts.py (PrepareRagPromptsConfig) builds the prompts and retrieve_rag_passages.py (RetrieveRagPassagesConfig) attaches the retrieved passages; the resulting prompts feed the Standard LM generation script.

4.6 Inference efficiency

Throughput / overhead measurements (decode rate, dynamic-PPL overhead, encoder rate) are under scripts/eval/inference_efficiency/; see its README.md for the run_timed.sh / run_dynppl.sh drivers and the analysis scripts.

4.7 Retrieval-precision eval

Measures how often the correct entry is retrieved for a given query: eval_retrieval_precision.py (EvalRetrievalPrecisionConfig).

CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src uv run src/lmlm/eval/eval_retrieval_precision.py \
    --config_path src/lmlm/eval/configs/eval/retrieval_precision.yaml

Released Artifacts

Released under the Co-LMLM collection on the Hugging Face Hub (lil-lab).

Co-LMLM-360M (FineWeb-Edu)

Model — lil-lab/CoLMLM-360M-FW: the 360M SmolLM2-based Co-LMLM retriever trained on FineWeb-Edu (HF model repo: weights + tokenizer). Loaded with the standard HF loader; the fact special tokens (<FACT> etc.) are in the tokenizer.

Two retrieval indices are released for this model, both as Hugging Face Storage Buckets and both context-keyed on the model's <FACT> hidden states. Each bucket holds the FAISS index (faiss.index), the faiss-id → entry-id mapping, manifest.json + index_config.json, and the fact-span value store (entries.db):

Index	Bucket	Background corpus	Entries	FAISS factory	Size
Wikipedia	`co-lmlm-360m-fw-wiki-index`	full Wikipedia	236M	`OPQ240,IVF65536,PQ240`	~113 GB
FineWeb-Edu + Wikipedia	`co-lmlm-360m-fw-fineweb-wiki-index`	FineWeb-Edu (100BT) + full Wikipedia	2.2B	`OPQ96,IVF524288_HNSW32,PQ96` (~2.5× more compressed)	~1.07 TB

Download the model and whichever index you need:

# Model (weights + tokenizer)
hf download lil-lab/CoLMLM-360M-FW --local-dir ./CoLMLM-360M-FW

# Option A — Wikipedia index (~113 GB; resumable)
hf buckets sync hf://buckets/lil-lab/co-lmlm-360m-fw-wiki-index ./co-lmlm-wiki-index

# Option B — FineWeb-Edu + Wikipedia index (~1.07 TB; resumable)
hf buckets sync hf://buckets/lil-lab/co-lmlm-360m-fw-fineweb-wiki-index ./co-lmlm-fineweb-wiki-index

Then point a Co-LMLM generation/eval config (or the Quick Start loader) at the model plus the chosen index:

Wikipedia index — model_path: ./CoLMLM-360M-FW, index_path: ./co-lmlm-wiki-index, index.db_path: ./co-lmlm-wiki-index/entries.db.
FineWeb-Edu + Wikipedia index — model_path: ./CoLMLM-360M-FW, index_path: ./co-lmlm-fineweb-wiki-index, index.db_path: ./co-lmlm-fineweb-wiki-index/fineweb_with_fullwiki_entries.db. This index ships its faiss-id → entry-id mapping as a SQLite database (faiss_id_to_entry_id.db) rather than a .txt file, so also set index.use_sqlite_id_mapping: true.

The FineWeb-Edu + Wikipedia FAISS file alone is ~228 GB. If it does not fit in RAM, memory-map it instead of loading it fully by exporting LMLM_FAISS_MMAP=1 — pages are then paged in on demand from local disk (lower RAM, slightly slower search):
```
LMLM_FAISS_MMAP=1 CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src uv run \
    src/lmlm/eval/hf_generate.py --model_type retriever \
    --model_path ./CoLMLM-360M-FW \
    --index_path ./co-lmlm-fineweb-wiki-index \
    --index.db_path ./co-lmlm-fineweb-wiki-index/fineweb_with_fullwiki_entries.db \
    --index.use_sqlite_id_mapping true \
    --prompts_path src/lmlm/eval/configs/generate/prompts_example.jsonl --output_path out.jsonl
```

See retriever_template.yaml and generation.md.

Other checkpoints & data

The remaining checkpoints (135M / 360M Co-LMLM, Standard LM, and LMLM-Asker variants) and the annotated pretraining corpora will be added to the collection as they are released.

Citation

If you use this code or the released artifacts, please cite:

@misc{feldman2026colmlmcontinuousquerylimitedmemory,
      title={Co-LMLM: Continuous-Query Limited Memory Language Models}, 
      author={Yair Feldman and Linxi Zhao and Nathan Godey and Dongyoung Go and Yilun Hua and Kilian Q. Weinberger and Jennifer J. Sun and Yoav Artzi},
      year={2026},
      eprint={2607.07707},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2607.07707}, 
}

Navigation Menu

Uh oh!

Repository files navigation

Co-LMLM: Continuous-Query Limited Memory Language Models

Contents

Env Setup

Output paths

API keys

Quick Start

Codebase Structure

1. Annotation Pipeline

1.1 Gemini seed annotation

1.2 Training the annotators

1.3 Running the annotators

1.4 Annotation at scale

2. Pretraining

2.1 Co-LMLM

2.2 Standard LM (baseline)

2.3 LMLM-Asker (baseline)

3. Index Building

4. Evaluation & Generation

4.1 Perplexity & dynamic replacement

4.2 Generation

4.3 Downstream factuality

4.4 NLU evaluation

4.5 RAG comparison baseline

4.6 Inference efficiency

4.7 Retrieval-precision eval

Released Artifacts

Co-LMLM-360M (FineWeb-Edu)

Other checkpoints & data

Citation

About

Releases

Packages

Contributors

Languages

Name		Name	Last commit message	Last commit date
scripts/eval		scripts/eval
src		src
.gitignore		.gitignore
README.md		README.md
pyproject.toml		pyproject.toml
uv.lock		uv.lock

Uh oh!

Latest commit

History

Folders and files

Repository files navigation

Co-LMLM: Continuous-Query Limited Memory Language Models

Contents

Env Setup

Output paths

API keys

Quick Start

Codebase Structure

1. Annotation Pipeline

1.1 Gemini seed annotation

1.2 Training the annotators

1.3 Running the annotators

1.4 Annotation at scale

2. Pretraining

2.1 Co-LMLM

2.2 Standard LM (baseline)

2.3 LMLM-Asker (baseline)

3. Index Building

4. Evaluation & Generation

4.1 Perplexity & dynamic replacement

4.2 Generation

4.3 Downstream factuality

4.4 NLU evaluation

4.5 RAG comparison baseline

4.6 Inference efficiency

4.7 Retrieval-precision eval

Released Artifacts

Co-LMLM-360M (FineWeb-Edu)

Other checkpoints & data

Citation

About

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages