CoMem Explained — From Paper to Working Code in 10 Minutes

You've hit the wall: your long-context LLM pipeline eats 89 GB of VRAM for 128k tokens, your RAG system loses the thread of a long document, and every compression approach you try trades accuracy for memory. There's a new paper that reframes the whole problem — and the fix is surprisingly elegant.

Standard KV caching stores key-value pairs for every layer across the entire context:

Memory = 2 × n_layers × seq_len × n_heads × head_dim × dtype_bytes

For a 32-layer model, 128k tokens, FP16: that's ~85–90 GB. One GPU. Gone.

RAG alternatives chunk and re-embed context — but retrieval at the embedding level discards the structural relationships the model built during the original read. The result: disjointed responses, hallucinations about inter-paragraph relationships, the "Lost in the Middle" failure mode.

CoMem asks: if we've already encoded a document through the lower layers, why re-encode it?

The paper's central empirical claim: transformers have a functional hierarchy.

This means the "understanding" phase is mostly complete by the middle of the network. The upper layers are doing task-specific translation, not comprehension.

Write (offline, once per document):

Run context tokens through layers 1 to L* (split layer). Cache the residual stream state h[L*]

for each token. That's it — no upper layers needed.

Retrieve (at query time):

Run the query through layers 1 to L* to get its intermediate state. Cosine-similarity search over the cache returns top-k most relevant context states.

Recompute (at query time):

Feed [retrieved_k_states + query_states] through layers L*+1 to L. The upper layers now compute over semantically relevant context — without ever re-running the lower layers.

Standard inference: O(n × L) per query (full context, all layers)
CoMem inference:    O(n × L*) once + O(k × (L - L*)) per query
                    where k << n
python
import torch
import torch.nn.functional as F

class CoMemCache:
    """
    Caches intermediate layer states for CoMem inference.

    Memory footprint: O(n * d_model) — no layer dimension.
    Compare: KV cache is O(n * n_layers * d_kv).
    """

    def __init__(self, topk: int = 64):
        self.topk = topk
        self._keys = []   # normalized (d,) tensors for similarity search
        self._vals = []   # raw (d,) tensors to feed into upper layers

    def write(self, split_layer_output: torch.Tensor) -> None:
        """
        Store all token states from split layer output.
        Call once per context document.

        split_layer_output: (seq_len, d_model)
        """
        for i in range(split_layer_output.shape[0]):
            v = split_layer_output[i].detach().cpu()
            self._vals.append(v)
            self._keys.append(F.normalize(v, dim=-1))

    def retrieve(self, query_split_state: torch.Tensor) -> torch.Tensor:
        """
        Get top-k context states relevant to this query.
        Fixed cost regardless of how many tokens are cached.

        query_split_state: (d_model,)
        returns: (k, d_model)
        """
        if len(self._keys) == 0:
            return query_split_state.new_zeros(0, query_split_state.shape[0])

        keys = torch.stack(self._keys).to(query_split_state.device)
        q = F.normalize(query_split_state, dim=-1)

        scores = keys @ q                                    # (n,)
        k = min(self.topk, len(self._vals))
        _, idx = torch.topk(scores, k=k)

        vals = torch.stack(self._vals).to(query_split_state.device)
        return vals[idx]                                     # (k, d)

    def memory_mb(self, d_model: int = 4096) -> float:
        """Approximate memory usage in MB (fp16)."""
        return len(self._vals) * d_model * 2 / 1e6

def run_upper_layers(
    layers: list,            # model.layers[split+1:]
    query_hidden: torch.Tensor,      # (q_len, d)
    retrieved: torch.Tensor,         # (k, d)
) -> torch.Tensor:
    """
    Prepend retrieved context states and run upper layers on combined sequence.
    Returns output for query tokens only.
    """
    if retrieved.shape[0] > 0:
        hidden = torch.cat([retrieved, query_hidden], dim=0)  # (k+q, d)
    else:
        hidden = query_hidden

    for layer in layers:
        hidden = layer(hidden)

    return hidden[-query_hidden.shape[0]:]

cache = CoMemCache(topk=64)
seq_len, d_model = 128_000, 4096

fake_states = torch.randn(min(100, seq_len), d_model)  # sample
cache.write(fake_states)

approx_mb = seq_len * d_model * 2 / 1e6  # full 128k
print(f"CoMem cache (128k tokens): {approx_mb:.0f} MB = {approx_mb/1000:.2f} GB")

Pre-trained models aren't optimized for receiving retrieved intermediate states. The paper adapts upper layers using LoRA:

from peft import get_peft_model, LoraConfig, TaskType
import torch.nn as nn

def setup_comem_lora(model, split_layer: int, lora_rank: int = 16):
    """
    Apply LoRA to upper layers only for CoMem adaptation.
    Lower layers stay frozen — we only need to adapt how
    upper layers process retrieved intermediate states.
    """
    for param in model.parameters():
        param.requires_grad = False

    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=lora_rank,
        lora_alpha=32,
        target_modules=["q_proj", "v_proj"],
        layers_to_transform=list(range(split_layer + 1, model.config.num_hidden_layers)),
    )

    model = get_peft_model(model, lora_config)

    print(f"Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
    return model

def self_distillation_loss(
    student_logits: torch.Tensor,  # (seq, vocab) — CoMem output
    teacher_logits: torch.Tensor,  # (seq, vocab) — full-context output
    temperature: float = 1.0,
) -> torch.Tensor:
    """
    KL divergence between student (CoMem) and teacher (full KV) distributions.
    No ground-truth labels needed — the frozen full model supervises itself.
    """
    student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)
    teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)
    return F.kl_div(student_log_probs, teacher_probs, reduction="batchmean")

Hardware: NVIDIA H20. Model: Qwen3-8B with continued training.

The RULER score is notable — 97.05 is near-ceiling performance. Standard RAG systems typically score in the 70–80% range on RULER because retrieved chunks lose inter-document context.

Only tested on Qwen3-8B. The layer specialization assumption needs verification on Llama-3, Mistral, MoE architectures (Mixtral, DeepSeek). Different models may have different functional split points.

L selection is manual.* The paper uses L* ≈ L/2 but doesn't provide a principled method for finding the optimal split. You'll need to experiment.

Cosine similarity retrieval has blind spots. Multi-hop reasoning often requires context that's semantically distant from the query but logically necessary. Pure cosine similarity misses these cases. A hybrid dense+sparse retrieval would be stronger.

LoRA training required. Off-the-shelf models aren't adapted for receiving retrieved intermediate states — you need the self-distillation fine-tuning step. The paper doesn't report how much performance degrades without it.

Generation quality not evaluated. RULER and LoCoMo are retrieval-heavy. Open-ended generation quality, hallucination rate, and coherence in long documents aren't assessed.

The paper is at arxiv:2607.28263. While the official code isn't out yet, you can start experimenting with the concept:

Profile your model's layer specialization: Measure representational similarity (CKA) across layers on your target model — find where semantics stabilize.

Prototype the cache: The CoMemCache

class above is a drop-in starting point. Hook it into your model's forward pass at the split layer.

Test without LoRA first: Grab the upper-layer outputs without adaptation. Measure quality degradation — this tells you whether LoRA training is strictly necessary for your use case.

Benchmark your target task: Before fine-tuning, verify that RULER/LoCoMo improvements generalize to your domain.

def should_use_comem(current_vram_gb: float, context_tokens: int) -> str:
    kv_overhead_gb = 32 * 128 * context_tokens * 2 * 2 / 1e9  # 32L, 128 d_kv
    comem_overhead_gb = 4096 * context_tokens * 2 / 1e9        # just d_model

    if current_vram_gb + kv_overhead_gb > 80:  # A100/H100/H20 limit
        return f"YES — KV cache needs {kv_overhead_gb:.1f}GB, CoMem needs {comem_overhead_gb:.1f}GB"
    return f"Optional — KV cache needs {kv_overhead_gb:.1f}GB (within 80GB budget)"

print(should_use_comem(current_vram_gb=20, context_tokens=128_000))

What's your experience with long-context inference? Drop a comment — especially if you've tried alternative approaches like StreamingLLM or KV compression in production.