Co-LMLM: Continuous-Query Limited Memory Language Models

Abstract: Limited memory LLMs (LMLMs) externalize factual knowledge during pretraining to a knowledge base (KB), rather than memorizing it in their weights. During generation, the model then fetches knowledge from the KB as needed. This recently introduced paradigm provides multiple advantages, including knowledge control capabilities that remain beyond conventional LLMs. We propose continuous-query LMLM (CO-LMLM), where the KB pairs continuous keys with textual knowledge values, a significant departure from prior reliance on relational KB and queries. CO-LMLM generates flexible vector queries at minimal cost, while still integrating human-readable and attributable retrieved knowledge into its generation. We pair this design with an annotation pipeline that tags free-form factual spans in arbitrary text, removing prior work's restriction to Wikipedia. Across pretraining on Wikipedia and FineWeb-Edu and at multiple model scales, CO-LMLM outperforms prior LMLMs and vanilla LLMs in both perplexity and factual precision. At 360M scale, this includes lower perplexity than models pretrained on 40x more data, and SimpleQA-verified performance that is in line with gpt-4o-mini and higher than Claude Sonnet 4.5.

• The paper introduces continuous-query LMLMs that replace textual queries with vector queries, retrieving facts directly from transformer hidden states.
• The paper demonstrates improved perplexity and factual precision with a 13× inference overhead reduction compared to conventional text-based retrieval.
• The paper details a scalable, automated annotation pipeline that externalizes knowledge for efficient unlearning and application across diverse corpora.

Co-LMLM: Continuous-Query Limited Memory LLMs

Overview and Motivation

Recent trends in LLM research emphasize externalizing factual knowledge to improve controllability, attribution, and knowledge editing. Traditional LLMs entangle factual knowledge within their parameters, leading to challenges in attribution, efficient updating, and unlearning. Limited Memory LLMs (LMLMs) address these issues by extracting factual knowledge from model parameters into editable, human-readable knowledge bases (KBs). The initial relational LMLM (Rel-LMLM) formulation, however, was constrained by the requirement for Wikipedia-style, relationally structured annotations and the reliance on textual query generation for retrieval, limiting scalability and retrieval expressivity.

Co-LMLM proposes a paradigm shift by introducing continuous-query LMLMs. Instead of textual queries over relational KBs, fact retrieval is triggered by vector queries sampled directly from the transformer's hidden state. Factual knowledge is stored in a vector-keyed, span-value KB, allowing knowledge externalization over arbitrary, free-form factual spans from diverse corpora. The approach is extended by a scalable annotation pipeline, enabling knowledge externalization well beyond the Wikipedia corpus.

Figure 1: Knowledge separation across three regimes. LLMs with RAG retrieve documents but retain factual knowledge in parameters (left); prior LMLM uses a relational KB accessed via decoded queries (middle); Co-LMLM retrieves unstructured factual spans from an external, continuous-index KB via internal model states (right).

Architecture and Training Mechanism

Co-LMLM employs a single decoder-only transformer, which integrates language modeling and dense retrieval capabilities. The external memory is a key-value store: dense vector keys mapped to human-readable factual text spans. During generation, emitting a special <[FACT](https://www.emergentmind.com/topics/fibril-analysis-for-cellulose-technology-fact)> token at any position in the output triggers a dense retrieval, leveraging the last-layer model hidden state as a query. The top-1 retrieved span is inserted into the generation context and followed by a closing </FACT> marker before decoding resumes. This approach minimizes token overhead per retrieval, since no textual query is generated or decoded.

Figure 2: The method pipeline: fact spans and paired questions are annotated for the pre-training corpus without altering the base text; the model is trained for next-token prediction over non-fact tokens and with a bidirectional InfoNCE loss on query-question pair alignment; at inference, emitting <FACT> triggers dense retrieval using the hidden state.

The training objective combines standard next-token prediction (NTP) and a symmetric, bidirectional InfoNCE contrastive loss. Only non-fact tokens (outside <FACT>...</FACT>) are optimized by NTP, ensuring the model does not memorize externalized spans in its parameters. The contrastive objective aligns the hidden state at each <FACT> location with the state at a corresponding <FACT-q> in a paired, context-dependent synthetic question, simulating retrieval intent.

KB construction is efficient: the model is applied over the corpus, storing L2-normalized hidden states at <FACT> positions as keys, and the annotated spans as values. This process enables rapid re-indexing for new, large-scale corpora.

Scalable, Automated Fact Annotation Pipeline

To scale beyond Wikipedia and support free-form factual retrieval, Co-LMLM uses an automated annotation pipeline: first, a seed set of documents is annotated by prompting a frontier LLM to identify fact spans and generate corresponding question-answer pairs. Two lightweight models are fine-tuned on this seed: a fact span annotator (BIO-tagger) and a question generator. Only these models are used to annotate the full pre-training corpus, enabling economical large-scale fact span/question annotation and generalization to arbitrary text domains.

Figure 3: Automated annotation pipeline. A frontier LLM annotates a small seed set; two models (fact span BIO-tagger and question generator) are fine-tuned and then used for scalable annotation of the full corpus.

Key annotation strategies ensure coverage, minimal span boundaries, question self-containment, and the avoidance of future context leakage. This results in a large corpus annotated with context-targeted fact retrieval events, providing high-quality supervision for both language modeling and retrieval learning.

Experimental Results

Perplexity and Knowledge Scalability

Co-LMLM demonstrates superior perplexity over standard baselines and Rel-LMLM across model scales, even when the baseline is trained on up to 40× more data. With dynamic retrieval and realistic (non-oracle) knowledge bases, Co-LMLM-360M achieves lower perplexity than HF/SmolLM2-360M trained on 4T tokens.

Figure 4: Perplexity evaluation across model sizes and index scales, demonstrating strong performance gains over baselines—even with far less training data.

Scaling the KB from 240M spans (Wikipedia) to 2.2B (Wikipedia + FineWeb-Edu) yields further perplexity reduction on out-of-Wikipedia datasets, underlining robust knowledge externalization and scalable retrieval generalization.

Factuality and NLU Performance

On five factuality benchmarks—TriviaQA, PopQA, SimpleQA, T-REx, and FactScore—Co-LMLM consistently produces large increases in factual precision relative to baseline models, notably improving long-tail and verified factual QA. The improvement over Rel-LMLM is attributed to the non-relational, contextual vector queries and broader knowledge scope.

Figure 5: Factual precision vs. NLU accuracy, showing large factuality improvements with negligible NLU degradation over comparable baselines.

Crucially, Co-LMLM achieves this while fully retaining NLU performance (CommonsenseQA, ARC Easy, etc.), as evidenced by extensive few-shot probe results. The tradeoff between factuality and general language understanding is extremely favorable.

Unlearning and Knowledge Control

Factual knowledge is truly externalized: disabling the KB during inference sharply degrades factual precision, demonstrating minimal factual memorization by parametric weights. Directly removing facts from the external KB suffices for unlearning, yielding effective selective forgetting on the TOFU benchmark—without utility tradeoff observed in retraining-based unlearning (e.g., NPO).

Figure 6: Direct KB manipulation in Co-LMLM enables targeted unlearning/forgetting with no retention loss, in contrast with parameter-based unlearning approaches.

The approach provides immediate, efficient means for both knowledge injection and suppression at deployment without retraining.

Retrieval and Query Mechanism Efficiency

Replacing decoded natural language queries with continuous vector queries yields a roughly 13× reduction in inference-time per-retrieval overhead, sharply decreasing computation and latency. The vector queries enable flexible, context-dependent retrieval mechanisms otherwise limited by linguistic expressivity and generation step overhead of text-based queries.

Practical and Theoretical Implications

Co-LMLM's architecture explicitly disentangles factuality from parametric weights, opening up transparent, source-attributable, and rapidly updatable knowledge integration within generative models. This paradigm has direct applicability for domains requiring verifiable information, knowledge updating, compliance (e.g., redaction/unlearning), and efficient adaptation.

From a learning-theoretic perspective, Co-LMLM establishes a platform to systematically study knowledge scaling, external memory attention, and retrieval learning. The pipeline—data annotation, dual-objective optimization, scalable KBs—generalizes to arbitrary corpora, offering a pathway to LLMs with hybrid or fully externalized memory architectures.

Moreover, retrieval-augmented generation (RAG) remains complementary: document retrieval can be layered atop Co-LMLM for broad context, while the factual KB ensures controlled, atomic information injection.

Limitations and Future Directions

Despite strong empirical results, the present study is limited to sub-billion-parameter models and factuality probes closely approximated by Wikipedia and related corpora. Future research should extend Co-LMLM's joint training and annotation pipeline to multi-billion parameter regimes, more diverse knowledge domains, and continuous/incremental KB updates (e.g., for continual or online learning scenarios). Dynamic memory integration, efficient KB updates post fine-tuning, and non-English/unstructured text domains are promising directions.

Conclusion

Co-LMLM advances the field of LLM design by integrating continuous vector queries for efficient, broad-domain, and transparent knowledge externalization. It outperforms prior LMLMs and standard LLMs both in perplexity and factuality, retaining NLU capability, enabling efficient inference, and supporting fine-grained unlearning via simple database operations. Its scalable annotation pipeline and architectural flexibility position it as a foundation for future research on large-scale, controllable, and trustworthy generative models.