CoMem Explained — From Paper to Working Code in 10 Minutes

A developer has implemented CoMem, a technique from a recent paper that reduces the memory footprint of long-context LLM inference by caching intermediate layer states instead of full KV caches. The approach exploits the functional hierarchy of transformers, where understanding is complete by the middle layers, allowing upper layers to operate on retrieved relevant states. The implementation, provided in Python, demonstrates significant memory savings and addresses issues like 'Lost in the Middle' in RAG systems.

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 Slice off the query portion return hidden -query hidden.shape 0 : --- Quick memory sanity check --- cache = CoMemCache topk=64 seq len, d model = 128 000, 4096 Simulate caching 128k tokens 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" → CoMem cache 128k tokens : 1024 MB = 1.02 GB just states Paper reports 18.26 GB total including model weights + cache + activations vs 89.36 GB for standard KV cache setup Pre-trained models aren't optimized for receiving retrieved intermediate states. The paper adapts upper layers using LoRA: python 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. """ Freeze everything first for param in model.parameters : param.requires grad = False Apply LoRA config to upper layers lora config = LoraConfig task type=TaskType.CAUSAL LM, r=lora rank, lora alpha=32, target modules= "q proj", "v proj" , Only apply to layers above split point 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" | Metric | CoMem | Baseline | Delta | |---|---|---|---| | RULER score | 97.05 | — | — | | LoCoMo score | 38.27 | 34.59 KV-Direct | +3.68 +10.6% | | Memory @ 128k | 18.26 GB | 89.36 GB | 4.9× less | | Prefill speed @ 128k | 7.83× faster | 1× | — | 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 https://arxiv.org/abs/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. One-liner to estimate whether CoMem is worth it for your model def should use comem current vram gb: float, context tokens: int - str: CoMem approximation: ~d model seq len 2 bytes + model weights vs full KV: n layers d kv seq len 2 2 bytes + model weights 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 → YES — KV cache needs 85.9GB, CoMem needs 1.0GB 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.