ReFreeKV: 임계값 없는 KV 캐시 압축을 향하여

초록

LLM 추론 중 메모리 소비를 줄이기 위해 KV 캐시 가지치기를 위한 여러 방법들이 제안되었다. 이러한 기법들은 많은 데이터셋에서 무손실 메모리 감소를 달성할 수 있지만, 최적의 성능을 얻기 위해 입력/도메인 특화된 KV 캐시 예산 임계값을 사전에 결정해야 한다는 덜 강조된 조건에 종종 의존한다. 그러나 이러한 입력 민감형 설계는 오픈 도메인 입력이 다양한 도메인, 길이 및 난이도에 걸쳐 있으며 임계값 선택에 대한 명확한 경계가 없기 때문에 실제 시나리오에서 상당히 제한될 수 있다. 결과적으로, 이러한 입력 민감형 임계값에 대한 의존성은 임의의 입력에 대해 큰 성능 저하를 초래하는 근본적인 한계가 될 수 있다. 본 연구에서는 임계값 제약을 제거하여 강건한 KV 압축을 달성하는 새로운 목표를 제안하며, 전체 캐시 성능을 유지하면서 예산 할당을 적응적으로 조정하는 "임계값 없는" 방법을 주장한다. 그런 다음 이 목표의 첫 번째 구현체로서 새로운 방법인 ReFreeKV를 제안한다. 다양한 컨텍스트 길이, 작업 유형 및 모델 크기를 포함한 13개 데이터셋에 걸친 광범위한 실험을 통해 그 효율성과 효과성을 입증한다. 코드는 https://github.com/Patrick-Ni/ReFreeKV에서 공개적으로 제공된다.

English

To reduce memory consumption during LLM inference, a handful of methods have been proposed for KV cache pruning. While these techniques can accomplish lossless memory reduction on many datasets, they often hinge on an under-emphasized condition: an input/domain-specific threshold for KV cache budget needs to be pre-determined to achieve the optimal performance. However, such input-sensitive design may be considerably limited in real-world scenarios, as open-domain inputs span diverse domains, lengths and difficulty levels, without clear boundaries for threshold selection. As a result, the dependence of such input-sensitive threshold can be a fundamental limitation that causes large degradation on arbitrary inputs. In this work, we propose a new objective that lifts the threshold constraints for robust KV compression, advocating for "threshold-free" methods that adaptively adjust budget allocation while preserving full-cache performance. We then propose a novel method, ReFreeKV, serving as the first instantiation of this objective. Extensive experiments across 13 datasets with diverse context lengths, task types, and model sizes demonstrate its efficacy and efficiency. Our code is publicly released at https://github.com/Patrick-Ni/ReFreeKV.