장문 컨텍스트 LLM 추론을 위한 KV 캐시 제거 문제 재정의

Reformulating KV Cache Eviction Problem for Long-Context LLM Inference

대규모 언어 모델(LLM)은 장문 컨텍스트 추론을 지원하지만, 키-값(KV) 캐시의 증가로 인해 상당한 메모리 및 런타임 오버헤드가 발생합니다. 기존 KV 캐시 제거 방법은 주로 로컬 어텐션 가중치에 의존하며, 값 표현, 출력 투영 및 헤드 간 상호 작용의 영향을 고려하지 않습니다. 본 연구에서는 KV 캐시 제거를 기존의 헤드 단위, 가중치 평균화 방식에서 벗어나 출력 정보를 고려한, 레이어 단위의 행렬 곱셈 근사 문제로 재정의합니다. 우리는 어텐션 맵과 투영된 값 상태 간의 승수 상호 작용을 명시적으로 모델링하여 토큰의 기여도를 정확하게 정량화하고 헤드 간 의존성을 고려하는 새로운 제거 전략인 LaProx를 제안합니다. 이 지표를 기반으로, 우리는 전역적으로 비교 가능한 중요도 점수를 토큰에 할당하는 최초의 통합 제거 전략을 제안합니다. 이를 통해 모델 전체 수준의 선택이 가능하며, 로컬 헤드 단위의 결정에 의존하지 않습니다. LongBench 및 Needle-In-A-Haystack과 같은 장문 컨텍스트 벤치마크에서 19개의 데이터 세트에 대한 실험 결과, 우리의 접근 방식은 KV 캐시 용량의 5%만을 사용하면서 모델 성능을 유지하며, 모든 구성에서 기존 연구보다 우수한 성능을 보입니다. 특히, 우리의 방법은 극단적인 압축 시나리오에서 기존의 최첨단 모델에 비해 최대 2배의 정확도 손실 감소를 달성하며, 최소한의 오버헤드를 유지합니다.

Large language models (LLMs) support long-context inference but suffer from substantial memory and runtime overhead due to Key-Value (KV) Cache growth. Existing KV Cache eviction methods primarily rely on local attention weights, neglecting the influence of value representations, output projection, and inter-head interactions. In this work, we reformulate KV Cache eviction from a conventional head-wise, weight-averaging approach into an output-aware, layer-wise matrix multiplication approximation problem. We introduce LaProx, a novel eviction strategy that explicitly models the multiplicative interaction between attention maps and projected value states to accurately quantify token contributions while accounting for inter-head dependencies. Building on this metric, we propose the first unified eviction strategy that assigns globally comparable importance scores to tokens, enabling model-wide selection instead of local, head-wise decisions. Experimental results across 19 datasets on long-context benchmarks LongBench and Needle-In-A-Haystack demonstrate that our approach maintains model performance with only 5\% of the KV cache and consistently outperforms prior works across all configurations. Notably, our method achieves up to 2$\times$ accuracy loss reduction under extreme compression scenarios compared to existing state-of-the-art baselines with minimal overhead.