IndexMem: 잠재 메모리를 활용한 학습 기반 KV 캐시 제거 기법을 통한 장문 컨텍스트 LLM 추론

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

최근 대규모 언어 모델(LLM)은 점점 더 긴 문맥을 처리해야 하지만, 일반적인 소프트맥스 어텐션 방식은 시퀀스 길이에 따라 선형적으로 증가하는 KV 캐시를 생성하여 장문 컨텍스트 추론의 병목 지점이 됩니다. 실질적인 해결책은 중요도가 낮은 KV 항목을 제거하는 것이지만, 기존 제거 정책들은 대부분 휴리스틱 기반이며 토큰의 복잡하고 입력 의존적인 중요도 분포를 제대로 파악하지 못합니다. 본 연구에서는 학습 가능한 인덱서를 도입하여 KV 중요도를 예측함으로써 중요한 토큰을 더욱 정확하게 유지할 수 있도록 합니다. 또한, 단순히 토큰을 제거하면 해당 정보가 영구적으로 손실되어 장거리에서 회복 불가능한 망각 현상과 성능 저하가 발생합니다. 이를 해결하기 위해, 경량의 잠재 메모리 모듈을 제안하여 제거된 토큰을 압축하여 온라인으로 업데이트되는 상태로 저장하고, KV 제거를 통해 잃어버린 어텐션 기여도를 보완하기 위한 잔차 출력을 제공합니다. 종합적으로, 저희 방법은 제한적인 KV 예산 하에서 정확한 장문 컨텍스트 추론을 가능하게 하며, Qwen, Mistral, Llama 모델에서 RULER (4K/16K) 벤치마크에서 일관된 성능 향상(적극적인 제거 시 최대 25 포인트), Needle-in-a-Haystack 검색 성능의 안정성 향상, 그리고 기존 제거 정책에 비해 우수한 LongBench 점수 및 압축 곡선을 제공합니다.

Large Language Models (LLMs) are increasingly expected to operate over long contexts, yet standard softmax attention incurs a KV cache that grows linearly with sequence length, quickly becoming the bottleneck for long context inference. A practical remedy is to evict less important KV entries; however, existing eviction policies are largely heuristic and struggle to capture the rich, input-dependent distribution of token importance. In this work, we introduce a learnable indexer that predicts KV importance, enabling more accurate retention of critical tokens. Meanwhile, naively evicting tokens permanently discards their information, leading to irreversible forgetting and degraded retrieval over long ranges. To address this, we propose a lightweight latent memory module that compresses evicted tokens into a compact, online-updated state and provides residual readouts to compensate for the attention contributions lost through KV eviction. Collectively, our method enables accurate long-context inference under a bounded KV budget, delivering consistent improvements on RULER (4K/16K) across Qwen, Mistral, and Llama models (up to 25 points under aggressive eviction), markedly more stable Needle-in-a-Haystack retrieval, and superior LongBench scores and compression curves compared to existing eviction policies.