Co-LMLM: Continuous-Query Limited Memory Language Models

Overview§

CO-LMLM extends Limited Memory Language Models (LMLMs) by introducing a knowledge base (KB) with continuous keys and textual values. This eliminates the need for a relational schema, allowing flexible vector queries derived from the model's hidden states. The retrieved text is integrated into generation via cross-attention, improving factual accuracy while maintaining readability.

Methodology§

Continuous-Query Knowledge Base§

During pretraining, the model learns to project its hidden state $\mathbf{h}_t$ into a query vector $\mathbf{q}_t = \mathbf{W}_q \mathbf{h}_t$. The KB consists of $N$ key-value pairs: keys $\mathbf{k}_i \in \mathbb{R}^d$ and values $v_i$ (text spans). Retrieval computes similarity:

$$ s_i = \frac{\mathbf{q}_t^\top \mathbf{k}_i}{\|\mathbf{q}_t\| \|\mathbf{k}_i\|} $$

The top-$k$ values are retrieved and their embeddings are concatenated or cross-attended with the language model.

Training with Span Annotation§

An annotation pipeline tags factual spans in arbitrary text (e.g., Wikipedia, FineWeb-Edu). Each span is paired with a continuous key learned jointly with the model. The training objective combines next-token prediction $\mathcal{L}_{LM}$ and a contrastive retrieval loss $\mathcal{L}_{ret}$:

$$ \mathcal{L} = \mathcal{L}_{LM} + \lambda \mathcal{L}_{ret}, \quad \mathcal{L}_{ret} = -\log \frac{\exp(s_+) }{\sum_{i} \exp(s_i)} $$

where $s_+$ is the similarity with the correct key.

Inference with Retrieval§

At each step, the model produces a query from the current hidden state, retrieves top-$k$ text values, and attends to them as additional context. This provides explicit, attributable knowledge without storing facts in weights.

Implementation Details§

The KB is stored as a matrix $\mathbf{K} \in \mathbb{R}^{N \times d}$ and value embeddings are precomputed from a frozen text encoder. During forward pass, retrieval is performed via efficient approximate nearest neighbor search (e.g., FAISS).

Code Snippet (PyTorch)§

import torch
import torch.nn as nn
import torch.nn.functional as F

class COLMLM(nn.Module):
    def __init__(self, llm, kb_keys, kb_vals, d_model, top_k=5):
        super().__init__()
        self.llm = llm
        self.query_proj = nn.Linear(d_model, d_model)
        self.kb_keys = nn.Parameter(kb_keys)  # (N, d)
        self.kb_val_emb = nn.Embedding.from_pretrained(kb_vals, freeze=True)  # (N, d_val)
        self.top_k = top_k

    def forward(self, input_ids):
        hidden = self.llm(input_ids)  # (B, L, d)
        query = self.query_proj(hidden[:, -1, :])  # (B, d)
        # Compute cosine similarity
        q_norm = F.normalize(query, dim=-1)
        k_norm = F.normalize(self.kb_keys, dim=-1)
        scores = torch.mm(q_norm, k_norm.t())  # (B, N)
        topk = torch.topk(scores, self.top_k, dim=-1).indices  # (B, k)
        # Retrieve value embeddings
        val_emb = self.kb_val_emb(topk)  # (B, k, d_val)
        # Cross-attend: simplified as mean pooling and addition
        context = val_emb.mean(dim=1)  # (B, d_val)
        hidden_aug = hidden[:, -1, :] + context  # fuse
        return hidden_aug  # used for next token prediction

Results§

CO-LMLM achieves lower perplexity than prior LMLMs and vanilla LLMs at multiple scales. At 360M parameters, it outperforms models trained on 40x more data and achieves SimpleQA factual accuracy comparable to gpt-4o-mini and exceeding Claude Sonnet 4.5. This demonstrates the effectiveness of continuous-query retrieval for knowledge-intensive generation.