Co-LMLM: Continuous-Query Limited Memory Language Models
Reviewed by Pith T0 review T1 audit T2 compute T3 formal T4 reserved 2026-07-09 01:28 UTCglm-5.2pith:ICJHRN5Hrecord.jsonopen to challenge →
The pith
360M-param model matches gpt-4o-mini on factual precision
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
Core claim
CO-LMLM replaces structured relational queries with continuous vector queries emitted directly from the model's hidden state, enabling retrieval from an unstructured text knowledge base at the cost of a single forward pass. This design externalizes factual knowledge during pretraining rather than memorizing it in weights, achieving higher factual precision and lower perplexity than both standard LMs and prior LMLM variants with relational queries.
What carries the argument
CO-LMLM: a decoder-only Transformer that jointly trains next-token prediction and a contrastive retrieval objective. At inference, emitting a special token triggers dense retrieval using the hidden state at that position as a query vector, fetching a human-readable text span from an external key-value index built during pretraining.
If this is right
- Small models can achieve factuality competitive with much larger systems by offloading factual knowledge to an external, editable index rather than scaling parameters.
- Training-free unlearning becomes possible: deleting entries from the knowledge base removes facts without gradient updates or utility degradation.
- The annotation pipeline, which uses a frontier model to seed lightweight annotators that tag factual spans in arbitrary text, removes the prior restriction to Wikipedia and extends to general web corpora.
- Continuous queries from hidden states are 13x cheaper per retrieval than generating textual questions, making dense retrieval during generation practical at scale.
- The parametric/non-parametric boundary is learned during pretraining rather than imposed at inference, potentially yielding better integration than RAG.
Load-bearing premise
The evaluation primarily tests Wikipedia-style knowledge using a knowledge base built from Wikipedia, creating a domain match between training data, index contents, and test benchmarks. Whether the factuality gains transfer to diverse, non-encyclopedic domains remains untested.
What would settle it
If disabling KB retrieval causes factuality to remain high, the model has not genuinely externalized knowledge but has memorized it in parameters despite the training objective.
Figures
read the original abstract
Limited memory language models (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.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. This paper proposes CO-LMLM, a continuous-query extension of the Limited Memory Language Model (LMLM) paradigm. Instead of using decoded relational queries to retrieve from a structured KB (as in REL-LMLM), CO-LMLM reads the model's hidden state at a special <FACT> token as a dense retrieval query against a key-value store of free-form factual text spans. The method includes a scalable annotation pipeline (frontier-model seed annotation distilled into lightweight annotators) that removes REL-LMLM's restriction to Wikipedia-style entity-relation tuples. Experiments at 135M and 360M scales show that CO-LMLM outperforms matched STANDARD and REL-LMLM baselines on perplexity and five factuality benchmarks, preserves NLU performance, enables training-free unlearning via KB deletion, and reduces query-formation latency by ~13x compared to textual-query variants. The controlled comparisons (CO-LMLM vs. STANDARD with identical training data, CO-LMLM vs. REL-LMLM, CO-LMLM vs. LMLM-ASKER) are well-designed and form the core of the empirical contribution.
Significance. The paper makes a solid contribution to the LMLM line of work. The continuous-query design is a natural and well-motivated improvement over decoded relational queries, and the annotation pipeline that enables scaling beyond Wikipedia to FineWeb-Edu is a practical advance. The ablations (no-KB, enforced-lookup, RAG comparison, LMLM-ASKER efficiency) are thorough and directly probe the method's mechanisms. The training-free unlearning result and the ~13x query-formation speedup over textual-query variants are concrete, falsifiable claims backed by controlled experiments. Code and data are publicly released, which strengthens reproducibility. The central methodological claim — that continuous hidden-state queries outperform decoded textual queries while retaining the editability benefits of externalized memory — is supported by the controlled experiments.
major comments (2)
- The abstract and introduction prominently state that CO-LMLM-360M-FW achieves SimpleQA performance 'in line with gpt-4o-mini and higher than Claude Sonnet 4.5.' This comparison is open-book (CO-LMLM retrieves from a 240M-entry Wikipedia KB) vs. closed-book (frontier models use parametric knowledge alone). The paper's own Table 2 shows that adding BM25 RAG to STANDARD-360M-FW lifts SimpleQA from 1.3 to 15.3, demonstrating that retrieval access accounts for a large portion of the score. The controlled comparisons that isolate the CO-LMLM method (CO-LMLM vs. STANDARD, CO-LMLM vs. REL-LMLM) are sound, but the headline frontier-model comparison does not isolate the contribution of the continuous-query externalization from the contribution of simply having a Wikipedia KB at inference time. The paper should either (a) add a comparison against STANDARD+RAG with a comparably strong retriever over
- The perplexity comparison to off-the-shelf HF/SMOLLM2 models (Figure 4a, Section 4.2) is acknowledged as having a 'likely domain confounder' because the KB is built from English Wikipedia and the test set is Simple English Wikipedia. The paper states this is 'qualified,' but the abstract still claims 'lower perplexity than models pretrained on 40x more data' without flagging the confounder. This is a load-bearing claim for the paper's positioning. The FineWeb-Edu perplexity experiment (Figure 4b) partially addresses this, but the FW test set is also factually dense educational text. The abstract claim should be qualified, or the perplexity comparison should be restricted to the controlled STANDARD vs. CO-LMLM comparison where training data is matched.
minor comments (7)
- The retrieval similarity threshold of 0.7 (Appendix A.6) is mentioned without justification or sensitivity analysis. Given its role in the fallback mechanism, a brief note on how this value was selected would strengthen the paper.
- The contrastive loss hyperparameters (lambda=0.25, tau=0.07) are listed in Table 6 but not motivated. A brief note on whether these were tuned and how sensitive results are to these choices would be helpful.
- Table 1: The T-REx improvement for CO-LMLM over STANDARD is modest (e.g., +1.1 for 360M, +4.0 for 360M-FW) compared to the large gains on QA benchmarks. This is worth discussing, as T-REx tests knowledge completion rather than question answering.
- Section 4.6, Table 4: The LMLM-ASKER comparison uses a separate sentence encoder (jina-embeddings-v5-text-nano) for query encoding, while CO-LMLM uses its own hidden state. This means LMLM-ASKER has an additional component that CO-LMLM lacks. The efficiency comparison is fair, but the factuality comparison may conflate the query mechanism with the encoder quality. A brief discussion of this would better isolate the comparison.
- Figure 4a: The 'Standard HF Ref.' label is ambiguous — it is not immediately clear this refers to off-the-shelf HF/SMolLM2 models. Renaming to 'HF/SMolLM2 (off-the-shelf)' would improve clarity.
- The paper uses 'SimpleQA' and 'SimpleQA Verified' somewhat interchangeably. Standardizing to 'SimpleQA Verified' (with first-use definition) throughout would improve readability.
- Section 3.3: The choice of Gemini 3.1 Pro for seed annotation is mentioned without discussion of alternatives or sensitivity. Given that the distilled annotators are trained on this seed, a brief note on seed quality validation would strengthen the pipeline description.
Simulated Author's Rebuttal
We thank the referee for the careful reading and constructive feedback. Both major comments raise legitimate concerns about how headline claims in the abstract are framed relative to the strength of the underlying controlled experiments. We address each below.
read point-by-point responses
-
Referee: The frontier-model SimpleQA comparison is open-book (CO-LMLM retrieves from a 240M-entry Wikipedia KB) vs. closed-book (gpt-4o-mini, Claude Sonnet 4.5 use parametric knowledge alone). Table 2 shows BM25 RAG lifts STANDARD-360M-FW from 1.3 to 15.3, so retrieval access accounts for much of the score. The headline comparison does not isolate the continuous-query externalization from simply having a Wikipedia KB at inference time. The paper should either (a) add a comparison against STANDARD+RAG with a comparably strong retriever, or (b) qualify the claim.
Authors: This is a fair point. We agree that the frontier-model comparison as currently stated in the abstract does not adequately distinguish the contribution of our continuous-query externalization method from the more general benefit of having retrieval access at inference time. We note that Table 2 in the manuscript already partially addresses the referee's concern: CO-LMLM-360M-FW achieves 21.7 on SimpleQA, compared to 15.3 for STANDARD-360M-FW + BM25 RAG over the same Wikipedia corpus. This shows that CO-LMLM's advantage is not merely due to retrieval access, since the RAG baseline retrieves from the same knowledge source and still underperforms. However, we acknowledge that BM25 is a weaker retriever than what a frontier model might use, and the comparison to gpt-4o-mini and Claude Sonnet 4.5 in the abstract remains open-book vs. closed-book. We will revise the abstract and introduction to explicitly flag that the frontier-model comparison is open-book vs. closed-book, and to note that Table 2 shows CO-LMLM also outperforms a STANDARD+RAG baseline using the same KB. We believe this contextualization preserves the interest of the result while accurately representing its scope. revision: yes
-
Referee: The perplexity comparison to off-the-shelf HF/SMOLLM2 models has a likely domain confounder (KB from English Wikipedia, test on Simple English Wikipedia). The abstract claims 'lower perplexity than models pretrained on 40x more data' without flagging the confounder. The FineWeb-Edu perplexity experiment partially addresses this, but the FW test set is also factually dense educational text. The abstract claim should be qualified, or the perplexity comparison restricted to the controlled STANDARD vs. CO-LMLM comparison.
Authors: We agree. The manuscript already acknowledges this confounder in Section 4.2 ('this experiment has a likely domain confounder, so we consider this result qualified'), but the abstract does not reflect this qualification. The referee is correct that the abstract's claim — 'lower perplexity than models pretrained on 40x more data' — is load-bearing for the paper's positioning and should not be stated without the caveat. We will revise the abstract to either (a) qualify the claim by noting the domain overlap between the KB and the perplexity test set, or (b) reframe the claim to foreground the controlled STANDARD vs. CO-LMLM comparison (which is matched on training data and does not have this confounder) as the primary perplexity result, with the off-the-shelf comparison presented as contextual rather than headline. We will also add a brief note about the FineWeb-Edu test set being factually dense, as the referee notes. The core controlled result — CO-LMLM consistently outperforms STANDARD at matched training data across all scales — is unaffected and remains the primary claim. revision: yes
Circularity Check
No significant circularity: the central derivation is self-contained and evaluated against external benchmarks.
full rationale
The paper's core methodological contribution—training a decoder-only Transformer to externalize factual knowledge into a continuous-key KB via a joint NTP + contrastive InfoNCE objective—is defined independently of its evaluation results. The contrastive loss (Eq. 1-2) pairs document-side hidden states at <FACT> tokens with question-side hidden states at <FACT-q> tokens; these are training-signal pairs constructed from annotated data, not fitted to the target benchmarks. The KB is built from the training corpus (Wikipedia/FineWeb-Edu), which is standard for retrieval-augmented methods. The factuality benchmarks (TriviaQA, PopQA, SimpleQA, T-REx, FactScore) are external evaluation sets, not used during training. The paper cites Zhao et al. [2026] (REL-LMLM) as prior work by overlapping authors, but this citation provides context and a baseline for comparison, not a load-bearing premise that would make the present results circular. The controlled comparisons (CO-LMLM vs. STANDARD with same training data, CO-LMLM vs. REL-LMLM, CO-LMLM vs. LMLM-ASKER) are sound ablations that isolate the method's contribution. The headline comparison to frontier models (gpt-4o-mini, Claude Sonnet 4.5) on SimpleQA is an open-book vs. closed-book asymmetry, but this is a fairness/confounder concern (correctness risk), not circularity—the result is not forced by construction. The paper transparently acknowledges the domain confounder for perplexity (Section 4.2) and provides ablations (no-KB, enforced-lookup, RAG comparison) that are independent of any self-citation chain. No step in the derivation chain reduces to its inputs by definition.
Axiom & Free-Parameter Ledger
free parameters (4)
- λ (contrastive loss weight) =
0.25
- τ (contrastive temperature) =
0.07
- Similarity threshold for retrieval fallback =
0.7
- PQ quantization parameters =
PQ240 (Wiki), PQ96 (Wiki+FW)
axioms (4)
- domain assumption Factual knowledge can be meaningfully separated from linguistic competence during pretraining by masking fact spans from the NTP loss.
- domain assumption The hidden state at the <FACT> token position contains sufficient information to serve as a retrieval query for the corresponding fact.
- domain assumption Synthetic questions generated by a frontier LLM (Gemini 3.1 Pro) and distilled annotators provide adequate supervision for the contrastive retrieval objective.
- domain assumption Top-1 retrieval is sufficient for fact lookup during both training and inference.
invented entities (2)
-
<FACT> / </FACT> special tokens
independent evidence
-
<FACT-q> query marker token
independent evidence
Reference graph
Works this paper leans on
-
[1]
jina-embeddings-v5-text: Task-Targeted Embedding Distillation
URLhttps://arxiv.org/abs/2602.15547. Akari Asai, Zeqiu Wu, Yizhong Wang, Avirup Sil, and Hannaneh Hajishirzi. Self-RAG: Learning to retrieve, generate, and critique through self-reflection. InInternational Conference on Learning Representations, 2024. URLhttps://openreview.net/forum?id=hSyW5go0v8. Vincent-Pierre Berges, Barlas O ˘guz, Daniel Haziza, Wen-t...
work page internal anchor Pith review Pith/arXiv arXiv 2024
-
[2]
URLhttps://proceedings.mlr.press/v119/chen20j.html. Xin Cheng, Wangding Zeng, Damai Dai, Qinyu Chen, Bingxuan Wang, Zhenda Xie, Kezhao Huang, Xingkai Yu, Zhewen Hao, Yukun Li, Han Zhang, Huishuai Zhang, Dongyan Zhao, and Wenfeng Liang. Conditional memory via scalable lookup: A new axis of sparsity for large language models,
-
[3]
Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models
URLhttps://arxiv.org/abs/2601.07372. Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try ARC, the AI2 reasoning challenge,
work page internal anchor Pith review Pith/arXiv arXiv
-
[4]
Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge
URLhttps://arxiv.org/abs/1803.05457. Damai Dai, Li Dong, Yaru Hao, Zhifang Sui, Baobao Chang, and Furu Wei. Knowledge neurons in pretrained transformers. InAnnual Meeting of the Association for Computational Linguistics,
work page internal anchor Pith review Pith/arXiv arXiv
-
[5]
URLhttps://aclanthology.org/2022.acl-long.581/. Matthijs Douze, Alexandr Guzhva, Chengqi Deng, Jeff Johnson, Gergely Szilvasy, Pierre-Emmanuel Mazaré, Maria Lomeli, Lucas Hosseini, and Hervé Jégou. The Faiss library.IEEE Transactions on Big Data, 2025. URLhttps://arxiv.org/abs/2401.08281. Zhengxiao Du, Aohan Zeng, Yuxiao Dong, and Jie Tang. Understanding ...
work page internal anchor Pith review Pith/arXiv arXiv 2022
-
[6]
URLhttps://openreview.net/forum?id=35DAviqMFo. Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield-Dodds, Robert Lasenby, Dawn Drain, Carol Chen, et al. Toy models of superposition. arXiv preprint arXiv:2209.10652, 2022. Gaurav R. Ghosal, Pratyush Maini, and Aditi Raghunathan. Memorization sinks: Iso...
work page internal anchor Pith review Pith/arXiv arXiv 2022
-
[7]
MemLLM: Finetuning LLMs to Use An Explicit Read-Write Memory
URLhttps://proceedings.mlr.press/v202/kandpal23a.html. Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. InInternational Conference on Learning Representations, 2020. URLhttps://openreview.net/forum?id=HklBjCEKvH. Guillaume Lample, Alexandre Sablayrolles, M...
work page internal anchor Pith review Pith/arXiv arXiv 2020
-
[8]
Measuring short-form factuality in large language models
URLhttps://arxiv.org/abs/2411.04368. Yuhuai Wu, Markus N. Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. InInternational Conference on Learning Representations, 2022. URL https://openreview. net/forum?id=TrjbxzRcnf-. Hongkang Yang, Zehao Lin, Wenjin Wang, Hao Wu, Zhiyu Li, Bo Tang, Wenqiang Wei, Jinbo Wang, Zeyun Tang, Shichao So...
work page internal anchor Pith review Pith/arXiv arXiv 2022
-
[9]
URLhttps://aclanthology.org/P19-1472/. Jintian Zhang, Cheng Peng, Mengshu Sun, Xiang Chen, Lei Liang, Zhiqiang Zhang, Jun Zhou, Huajun Chen, and Ningyu Zhang. OneGen: Efficient one-pass unified generation and retrieval for LLMs. InFindings of the Association for Computational Linguistics: EMNLP 2024, 2024a. URL https://aclanthology.org/2024.findings-emnlp...
-
[10]
Start with the bare factual token (a name, number, date)
-
[11]
Shown alone, is this token ambiguous?
Ask: "Shown alone, is this token ambiguous?" If NO, tag just that token
-
[12]
If YES, include the minimal surrounding words that resolve the ambiguity. - If a bare number or noun is meaningful on its own (e.g., "1990", "Paris", " Einstein"), tag just that token. - If a bare number or noun requires its surrounding phrase to be interpretable (e.g., "25" in "25 medals in total" -- "25" alone is ambiguous), tag the full noun phrase tha...
work page 1990
-
[13]
Cover up the tagged span with your finger. Read the question alone
-
[14]
Ask: could this question be typed into Google by someone who has NEVER read my document? Is it clear what they are asking? (If no: under-specified.)
-
[15]
Ask: does the question itself contain or strongly hint at information that appears LATER in this document, especially information that will become another annotation answer? (If yes: over-specified / leaking.)
-
[16]
What is it the second-highest of?
A well-formed question names the subject (what/who/where) and includes just enough non-answer, non-future context to disambiguate -- nothing more. <example> Input: Mount Kenya is the second-highest peak in Africa. WRONG q="What is it the second-highest of?" (under-specified -- "it" has no referent outside this document) WRONG q="What is the second-highest...
-
[17]
Identify the sentence containing the tagged span
-
[18]
Collect all text from the beginning of the document up to and including that sentence. This is the "available context."
-
[19]
If any word or phrase in your question comes from a later sentence, the question is invalid
Every entity, date, descriptor, and detail mentioned in the question MUST appear in the available context. If any word or phrase in your question comes from a later sentence, the question is invalid
-
[20]
If it does, rephrase to remove the leaking detail
**Cross-contamination check:** Re-read your question and confirm it does not reveal or strongly imply the answer to any OTHER fact that appears later in the document. If it does, rephrase to remove the leaking detail
-
[21]
Where is Cuyamaca Reservoir located?
**Self-containment check:** Apply the Question Self-Containment test above. The question must work for a reader who has never seen the document. If two facts define each other, tag only ONE (the rarer fact). <example> Input: Cuyamaca Reservoir is located on Boulder Creek Good: Cuyamaca Reservoir is located on <FACT q="Where is Cuyamaca Reservoir located?"...
work page 1913
-
[22]
The ‘a‘ value must contain the answer to the question in ‘q‘
-
[23]
It must be snippet form -- a phrase, name, number, date, or short noun phrase. Not a complete sentence
-
[24]
It should convey the fact without looking copy-pasted from the span, but must not inflate into a complete sentence
-
[25]
A single word may become a short phrase; a wordy span may be condensed
Length may be slightly longer or shorter than the span. A single word may become a short phrase; a wordy span may be condensed. Length must stay in the snippet range
-
[26]
Proper nouns, numbers, and dates must be preserved exactly in factual content. Only surface form can change (e.g. "21 May 1925" -> "May 21, 1925"). Additional information may be added but not removed
work page 1925
-
[27]
Lists of items should be shuffled -- do not maintain the original ordering unless the order is inherently meaningful
-
[28]
Proper nouns and venue names in particular should usually be kept as-is
If there is no meaningful way to paraphrase the span, keeping ‘a‘ identical to the span text is acceptable. Proper nouns and venue names in particular should usually be kept as-is
-
[29]
Be diverse -- vary your paraphrasing strategies. Sometimes keep dates/names as-is; sometimes use a surface variant (e.g. "Sep" for "September"); sometimes swap the order of a date’s components. Apply this variety across all fact types. --- ### COMPLETE EXAMPLES <example> Input: Stefan Gierowski Stefan Gierowski (21 May 1925 -- 14 August 2022) was a Polish...
work page 1925
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.