CO-LMLM: Teaching Language Models to Ask Questions in a Language Only Machines Understand

Paper: 2607.07707 Authors: Yair Feldman, Linxi Zhao, Nathan Godey, Dongyoung Go, Yilun Hua, Kilian Q. Weinberger, Jennifer J. Sun, Yoav Artzi Categories: cs.CL, cs.AI, cs.LG

The Gap

Here’s the problem everyone’s been circling: standard LLMs bake all their factual knowledge into their weights. That’s convenient for generation but terrible for control — you can’t easily update facts, audit what the model “knows,” or guarantee factual precision.

Enter Limited Memory Language Models (LMLMs), a recent paradigm that externalizes factual knowledge to a knowledge base (KB) during pretraining. The model fetches facts from the KB on demand rather than memorizing them. Prior work on LMLMs (like the KNN-LM lineage and structured retrieval approaches) relied on relational KBs with structured, text-based queries — essentially asking the KB in human language: “tell me about entity X, relation Y.”

The limitation is sharp: text-based queries are rigid. You have to know the exact schema, the entity names, the relation types. The system can’t express fuzzy, compositional, or latent information needs. Additionally, prior LMLM annotation pipelines were restricted to Wikipedia, limiting the training data diversity.

CO-LMLM’s answer: let the model formulate queries in continuous vector space (the language machines actually think in), while still retrieving human-readable textual values. Plus, build an annotation pipeline that works on arbitrary free-form text, not just Wikipedia.

Problem: Facts baked into LLM weights
         = opaque, hard to update, hard to audit
              |
              v
Prior LMLMs: Externalize facts to relational KB
              BUT queries must be structured text
              = rigid, schema-dependent, Wikipedia-only
              |
              v
Limitation: Can't express latent/fuzzy info needs
              Annotation pipeline limited to one source
              |
              v
CO-LMLM: Continuous vector queries + textual values
              + free-form annotation on any text
              |
              v
Evidence: 360M model beats 14B-param LLMs on perplexity
              SimpleQA parity with gpt-4o-mini
              |
              v
Conclusion: Continuous queries are a strictly
              better interface for knowledge retrieval in LMLMs

The Increment

One sentence: Before this paper, retrieving external knowledge during generation required the model to formulate explicit text queries against structured databases — after this paper, the model can express its information need as a dense vector and still get back readable, attributable text, making knowledge retrieval both cheaper and more expressive.

Core Mechanism

The architecture has three moving parts that work in concert: (1) a query generator, (2) a continuous knowledge base, and (3) a retrieval-augmented generation head.

During pretraining, when the model encounters a factual span (detected by the annotation pipeline), it learns to produce a continuous query vector from its hidden state. This vector is not text — it’s a point in embedding space that captures the latent information need. The KB stores pairs of (continuous key, textual value): each piece of knowledge has a dense vector representation as its “address” and a human-readable text string as its “content.” Retrieval is nearest-neighbor search in continuous space: the query vector finds the closest key, and the corresponding textual value is returned and injected into the generation context.

The annotation pipeline itself is a contribution. Instead of relying on Wikipedia’s hyperlink structure or predefined entity schemas, it uses a tagging model to identify free-form factual spans in arbitrary text. This means the KB can be populated from FineWeb-Edu or any other corpus, not just Wikipedia.

[Pretraining Flow]

  Training text with factual spans
              |
              v
  +---------------------+
  | Annotation Pipeline  |  <-- tags factual spans
  | (works on any text)  |      in free-form text
  +---------------------+
              |
              v
  +---------------------+
  | Query Generator      |  <-- hidden state --> continuous vector q
  | (learned, minimal    |
  |  overhead)           |
  +---------------------+
              |
              | q
              v
  +---------------------+
  | Continuous KB        |
  |  keys: dense vectors |
  |  vals: text strings  |
  +---------------------+
              |
              | nearest neighbor: argmin_i ||q - k_i||
              v
  Retrieved text value (human-readable)
              |
              v
  +---------------------+
  | Generation Head      |  <-- conditions on retrieved text
  | (standard LM)       |      for next-token prediction
  +---------------------+

Structural metaphor: a multilingual research assistant with a photographic filing cabinet.

Imagine you’re a researcher writing a paper. You have a filing cabinet where every document is indexed not by keywords (like a traditional library card catalog) but by a dense “fingerprint” — a summary of what the document is about in a compressed, abstract sense. When you’re writing and hit a gap in your knowledge, you don’t type a search query into Google. Instead, you just feel the shape of what you need — “I need something about population genetics and migration patterns, but specifically the admixture angle” — and your assistant, who speaks your internal language, immediately pulls the right document and reads you the relevant paragraph.

That “feeling the shape” is the continuous query. The fingerprint on each document is the continuous key. The paragraph your assistant reads aloud is the textual value. The assistant’s ability to understand your half-formed intuition (rather than requiring you to articulate exact keywords) is what makes continuous queries more expressive than text queries. And the fact that you still hear the paragraph in plain English (not some compressed vector) is what keeps the system interpretable and attributable.

The annotation pipeline is the librarian who shelves new books: instead of only accepting Wikipedia donations (with their nice hyperlink structure), this librarian can take any book off the street, read it, and tag the factual claims inside. That’s how you go from Wikipedia-only to FineWeb-Edu.

Key Concepts

• Limited Memory Language Model (LMLM): Think of it this way. A normal LLM is like a student who memorized the textbook — great for recall, but if the textbook has an error, the student repeats it, and you can’t easily fix just that one fact. An LMLM is like a student who left the textbook on their desk and looks things up as needed. The knowledge lives outside the model. When the textbook gets updated, the student immediately knows the new facts. The model’s weights learn *how to think and retrieve, not what to know.

• Continuous Query vs. Text Query: A text query is like Googling “capital of France population 2024.” It’s explicit, structured, and you need to know roughly what you’re looking for. A continuous query is like the feeling you get when you vaguely remember reading something interesting about European demographics but can’t recall the exact phrasing — you just have a *sense of what you need. The system translates that sense into a vector, does math to find the closest match, and hands you the answer. The key insight: the model’s internal state already “knows” what it needs in a way that’s hard to articulate in words, so let it query in its native language (vectors) instead of forcing it to translate to text first.

• Free-form Factual Span Annotation: Prior systems needed Wikipedia because it has convenient structure — hyperlinks, infoboxes, entity pages — that automatically mark “this is a fact worth storing.” The new annotation pipeline is like having a fact-checker read any document and highlight claims with a yellow marker. It doesn’t need the document to be pre-structured. This matters because most of the world’s knowledge isn’t on Wikipedia: it’s in textbooks, news articles, scientific papers, and forum posts.

Framework Shift

Before (prior LMLMs):

  Text query (structured)          Relational KB
  "entity: Paris, attr: pop" ----> | Paris | pop | 2.1M |
                                   | Tokyo | pop | 14M  |
  Rigid schema, Wikipedia-only     Text in, text out

After (CO-LMLM):

  Continuous query (vector)        Continuous KB
  [0.23, -0.87, ..., 0.41] -----> | key: [0.22,-0.86,...] | val: "Paris pop 2.1M" |
                                   | key: [0.91, 0.12,...] | val: "Tokyo pop 14M"  |
  Flexible, any-source             Vector in, text out

From rigid text-based retrieval to flexible vector-based retrieval, the core shift is: let the model ask questions in the language it actually thinks in (continuous vectors), while keeping the answers in the language humans can read (text).

Expert Assessment

Problem choice: This is a real and well-motivated gap. The tension between “models should externalize knowledge for controllability” and “text-based retrieval is too rigid” is genuine and sits at a productive intersection of retrieval-augmented generation and knowledge editing. The field has been inching toward this — dense retrieval for open-domain QA is mainstream, but applying continuous queries specifically to the LMLM pretraining paradigm is a clean, logical next step. Not manufactured.

Method maturity: The core insight — decouple the query modality (continuous) from the value modality (text) — is elegant and well-motivated. It’s not brute force; it’s a genuine architectural simplification that removes a bottleneck. The annotation pipeline for free-form text is practical and necessary, though I’d want to see more analysis of its precision/recall tradeoffs — false positives in factual span detection could introduce noise. One concern: nearest-neighbor retrieval at scale can be expensive at inference time, and the paper doesn’t deeply discuss the latency implications beyond “minimal cost” for query generation.

Experimental integrity: The baselines are fair — they compare against prior LMLMs, vanilla LLMs, and provide scaling curves. The 360M model beating 14B-parameter models on perplexity is eye-catching but needs context: perplexity on what distribution? If the test set has high factual density, external retrieval has a structural advantage. The SimpleQA results aligning with gpt-4o-mini and exceeding Claude Sonnet 4.5 are impressive if the evaluation protocol is clean, but SimpleQA is a relatively narrow benchmark. I’d want to see broader factual QA suites and a failure mode analysis. The Wikipedia + FineWeb-Edu dual pretraining is a good ablation.

Writing quality: The abstract is well-crafted. However, the paper likely cuts corners in the limitations section — every retrieval-augmented system has failure modes when the KB doesn’t contain the answer, and I’d want to see explicit analysis of KB coverage gaps. The annotation pipeline description could use more detail on edge cases (subjective claims, temporally sensitive facts, implicit knowledge).

Verdict: weak accept — The continuous-query idea is clean and well-motivated with solid empirical results, but the evaluation could be broader and the inference-time analysis is thin.

Takeaways

Three things to steal:

1. Decouple query modality from value modality. You don’t have to retrieve the same type of representation you query with. Let your model think in vectors but read in text. This principle transfers to any retrieval system: the best query format for a neural network isn’t necessarily the best format for humans to audit.

2. The annotation pipeline is the unsung hero. If you’re building any system that needs structured supervision from unstructured data, investing in a good tagging/annotation model that works on free-form text is often more valuable than the downstream architecture. Prior LMLMs were bottlenecked by Wikipedia; the annotation generalization unlocked the whole thing.

3. Externalize what you can’t control. The LMLM paradigm itself — pushing factual knowledge out of weights and into a mutable store — is the deepest takeaway. If you’re building systems where factual accuracy and updatability matter more than raw fluency, this is the architectural direction. The specific mechanism (continuous queries) is one implementation, but the principle of knowledge externalization during pretraining is broadly applicable.