IndexMem: Learned KV-Cache Eviction with Latent Memory for Long-Context LLM Inference

IndexMem targets the KV-cache memory wall in long-context LLM inference by replacing heuristic eviction with a learnable indexer and adding a latent memory that preserves evicted information instead of deleting it permanently. The indexer is a lightweight MQA-style module trained by streaming KL distillation to match the backbone’s attention-derived token-importance distribution, using QK-Norm features and max-over-queries aggregation; during decoding it can reuse cached keys and even compress before prefill ends. The memory module stores evicted tokens in a fixed-size fast-weight state updated online with decay, then returns a gated residual readout, so the final output is attention over retained KV plus compensation from latent memory rather than brittle “memory-as-tokens” softmax insertion. On Qwen3-8B, Mistral-7B-v0.3, and Llama-3.1-8B-Instruct, evaluated with NVIDIA’s KVPRESS on RULER, Needle-in-a-Haystack, and LongBench, IndexMem is consistently the strongest eviction method: it nearly matches full-cache accuracy under 10–25% compression and holds up much better at 50–90% compression, with gains up to 25 points on RULER. The memory residual reduces catastrophic retrieval failures, improves LongBench average scores to 55.26 at 50% compression versus 42.35 for PyramidKV and 42.10 for TOVA, and adds only 0.52M parameters beyond the 19.92M-parameter indexer.