Technical Summary: Co-LMLM (Continuous-Query Limited Memory Language Models)

Problem Statement

Large Language Models (LLMs) typically internalize factual knowledge within their parameters. This approach presents significant limitations: factual recall scales with model size and exposure frequency, making long-tail knowledge difficult to capture; factual associations are entangled with linguistic representations, making individual facts difficult to attribute, edit, or remove; and this entanglement contributes to hallucinations and stale knowledge.

Recent work introduced Limited Memory Language Models (LMLMs), which externalize factual knowledge into a human-readable Knowledge Base (KB) during pre-training rather than memorizing it. However, the existing implementation, referred to as REL-LMLM (Relational LMLM), relies on a relational KB where facts are stored as structured subject-relation-object tuples. This design imposes several constraints:

1. Data Restriction: It restricts externalizable data to items expressible as relational tuples, limiting the scope primarily to Wikipedia-style structured data.
2. Query Expressivity: Retrieval depends on generating explicit, decoded textual queries (natural language subject-relation questions), which introduces inference overhead and limits the retrieval mechanism to simple queries.
3. Pre-determined Pathways: The training process synthesizes exact queries during data pre-processing, holding retrieval pathways constant and restricting the expressivity of the learned retrieval mechanism.

Methodology

The authors propose Co-LMLM (Continuous-Query Limited Memory Language Models), a paradigm that addresses the limitations of relational LMLMs while retaining the benefits of externalized, editable memory.

Core Architecture

Co-LMLM functions as a decoder-only Transformer that acts as both a language model and a dense retriever.

• Continuous Queries: Instead of decoding a textual query, the model emits a special <FACT> token. The hidden state at this token position is extracted and used directly as a dense vector query.
• Unstructured KB: The Knowledge Base stores pairs of vector keys (the continuous queries) and string values (verbatim factual text spans). This allows for the storage of free-form factual spans without the constraints of subject-relation-object tuples.
• Inference: During generation, when the model emits <FACT>, it triggers a retrieval step. The top-1 fact snippet is spliced into the context immediately after <FACT>, followed by a closing </FACT> token, and decoding resumes. This adds only one autoregressive step per retrieval.

Training Objective

The model is trained jointly on two objectives:

1. Next-Token Prediction (NTP): The model predicts tokens in the sequence, but crucially, it is not optimized to memorize the factual spans enclosed by <FACT> and </FACT>. The loss is calculated only on non-fact tokens and the opening <FACT> token (which triggers retrieval).
2. Contrastive Retrieval Loss: A bidirectional InfoNCE loss aligns the query representation at the <FACT> token in the document with the query representation at a <FACT-q> token in a paired natural language question. This teaches the model to emit a hidden state that effectively retrieves the correct fact span when queried.

Data Annotation Pipeline

To scale beyond Wikipedia, the authors developed an automated annotation pipeline to process arbitrary text (e.g., FineWeb-Edu):

1. Seed Annotation: A frontier model (Gemini 3.1 Pro) annotates a small seed set of documents, tagging factual spans with <FACT q="..." a="...">span</FACT>.
2. Distillation: Two lightweight annotators are trained on this seed set:
   • A Fact Span Annotator (ModernBERT encoder) that performs BIO tagging to identify factual spans in raw text.
   • A Question Generator (Qwen2.5 decoder) that generates a standalone question and answer pair for each span.
3. Corpus Processing: These annotators process the full pre-training corpus, converting raw documents into the annotated format required for Co-LMLM training. This pipeline allows for the extraction of facts from general web text, not just structured encyclopedic entries.

Key Contributions

1. Continuous-Query Mechanism: Replacing decoded textual queries with continuous vector queries emitted from hidden states. This reduces inference overhead (no query decoding) and allows the retrieval mechanism to utilize the full generation context rather than a short, restricted decoded query.
2. Free-Form Knowledge Externalization: Moving from relational tuples to unstructured, free-form factual spans. This enables the externalization of a broader set of knowledge without role or relational constraints, facilitating scaling to general corpora like FineWeb-Edu.
3. Scalable Annotation Pipeline: A method to annotate hundreds of billions of tokens with factual spans and paired questions, removing the restriction to Wikipedia and enabling training on diverse web data.
4. Controllability and Unlearning: Demonstrating that externalizing knowledge allows for "unlearning" via simple database operations (deleting entries) without retraining the model or sacrificing general utility, a capability often lost in parametric knowledge editing.

Experimental Results

The authors evaluated Co-LMLM against standard LLMs (STANDARD) and the relational baseline (REL-LMLM) at 135M and 360M parameter scales, trained on Wikipedia and FineWeb-Edu.

• Perplexity: Co-LMLM consistently achieves lower perplexity than standard models of the same size. At the 360M scale, Co-LMLM trained on Wikipedia achieves lower perplexity than off-the-shelf HF/SMOLLM2-360M models trained on 40× more data.
• Factual Precision: Co-LMLM substantially improves factual precision across benchmarks (TriviaQA, PopQA, SimpleQA, T-REx, FactScore).
  • At 360M scale, Co-LMLM achieves a SimpleQA score of 21.2 (21.7 with FineWeb-Edu), which the authors note is in line with gpt-4o-mini and higher than Claude Sonnet 4.5 according to public leaderboards.
  • It outperforms REL-LMLM, demonstrating that continuous queries and free-form spans are more robust than relational tuples and textual queries.
• NLU Performance: The model retains downstream Natural Language Understanding (NLU) performance comparable to standard models of the same size, indicating that offloading facts does not degrade general language capabilities.
• Unlearning: In TOFU unlearning experiments, Co-LMLM successfully forgets designated facts via direct KB deletion, maintaining high utility on retained knowledge, whereas training-based unlearning methods (like NPO) degrade utility due to parameter entanglement.
• Efficiency: Co-LMLM reduces the cost of query formation to a single hidden-state forward pass (~2.2ms), compared to ~28ms for methods requiring the decoding of a textual question and subsequent encoding.

Significance and Claims

The paper claims that Co-LMLM provides a scalable method to train LLMs with verifiable, controllable, and transparent knowledge usage. By externalizing knowledge into a schema-free index and using continuous queries, the approach overcomes the expressivity and scaling limitations of previous relational LMLMs.

The authors emphasize that this design establishes a foundation for studying knowledge scaling laws across broader domains and memory sizes. They acknowledge current limitations, noting that while the pipeline extends to general corpora, systematic evaluation across diverse knowledge domains remains future work. The paper concludes that Co-LMLM offers a viable path toward LLMs that are not only factually accurate but also editable and attributable, addressing critical trustworthiness issues in current models.