IndexMem: Long Context Reasoning Optimization Based on Learnable Index and Latent Memory

IndexMem is a long-context LLM reasoning optimization solution proposed by the arXiv team on May 25, 2026. Its core innovations are: predicting the importance of KV entries via a learnable index to replace heuristic eviction strategies; introducing a lightweight latent memory module to compress evicted tokens, avoiding irreversible information loss. This solution maintains stable Needle-in-a-Haystack retrieval performance even under aggressive eviction strategies, effectively addressing the KV cache memory bottleneck in long-context scenarios of the Transformer architecture.

As LLM capabilities expand, user demand for long-context processing grows (e.g., whole-book analysis, multi-turn dialogues, etc.). The KV cache size of the Transformer attention mechanism grows linearly with sequence length; when processing tens of thousands of tokens, it occupies dozens of GB of video memory, becoming a bottleneck for inference latency and cost. Existing solutions fall into two categories: sparse attention (reduces computation but keeps KV cache unchanged) and KV cache compression (actively discards KV but easily loses information). IndexMem focuses on the KV compression direction and proposes improvements.

Current mainstream KV compression uses heuristic strategies (LRU, attention weight threshold, sliding window), which have two major problems: 1. Lack of precise understanding of token importance; static rules cannot adapt to input-dependent dynamic distributions. 2. Irreversible information loss: key information disappears permanently after eviction, leading to retrieval failures in long-distance dependency scenarios (e.g., Needle-in-a-Haystack).

IndexMem's two core innovations:

1. Learnable Indexer: A lightweight neural network that takes KV and query states as input and outputs importance scores, with advantages of input adaptability, end-to-end optimization, and fine-grained control;
2. Latent Memory Module: Compresses evicted tokens into a compact state, updates online, and provides residual reading to compensate for attention loss, enabling infinite memory under a bounded KV budget.

Technical details:

• Learnable Index Architecture: 2-3 layer MLP, inputting KV statistical features, query similarity, positional encoding, and historical attention patterns, outputting importance scores;
• Latent Memory Compression: Gated recurrent mechanism (hidden_state = gate*hidden_state + (1-gate)*compress(evicted_kv)), residual reading to compensate for loss;
• Training Strategy: Pre-training (learning general patterns from long text data) + fine-tuning (adapting to specific task requirements).

Experimental results:

• RULER Benchmark: Performance improved by 25 percentage points under aggressive eviction, generalizing to Qwen, Mistral, and Llama series models;
• Needle-in-a-Haystack: Can still accurately locate key information under extremely long sequences and aggressive compression, with stable performance;
• LongBench: Compression curves are comprehensively better than baselines, with significant practical application value.

Application scenarios: Long document Q&A, codebase understanding, multi-turn dialogues, real-time stream processing. Deployment considerations:

• Computational overhead: The latent memory module introduces a small amount of additional computation, but the memory bandwidth saved is more worthwhile;
• Model adaptation: The indexer needs to be fine-tuned for the target model, with low migration cost;
• Hyperparameter tuning: Cache size and compression ratio need to be adjusted according to the application and hardware.

Limitations:

• Training relies on long text data; there may be insufficient data for minority languages/domains;
• Latent memory cannot fully compensate for information loss under extreme compression ratios;
• Currently only for text; multi-modal expansion requires additional research. Future directions:
• Hierarchical memory architecture (simulating working memory-long-term memory hierarchy);
• Dynamic cache allocation (adjusting cache size according to input characteristics);
• Cross-session memory (persisting latent memory to enable cross-session context continuation).