ReFreeKV: Towards Threshold-Free KV Cache Compression

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.

한 줄 요약

ReFreeKV는 입력에 따라 KV 캐시 예산을 자동 조정하는 Threshold-Free KV 캐시 압축 방법으로, 13개 데이터셋에서 평균 25% 이상의 압축률을 달성한다.

핵심 기여도

입력에 의존하지 않는 KV 캐시 압축 목표를 제안, 기존 방법들의 threshold 의존성을 극복.
DBudgetKV라는 새로운 KV 캐시 압축 알고리즘을 제안, 토큰 순위 결정과 주의도 기반 중단 기준을 사용.
13개 데이터셋에서 평균 25% 이상의 압축률 달성, 일부 데이터셋에서는 85%까지 압축 가능.
Llama3, Qwen2.5, Mistral 모델에서 Full-cache 성능과 유사하거나 더 높은 성능 유지.

핵심 아이디어

기존 KV 캐시 압축 방법들은 대부분 특정 입력에 맞춘 사전 정의된 예산(threshold)을 필요로 하며, 이는 다양한 도메인과 입력 길이에 따라 성능이 크게 변동된다. 예를 들어, H2O는 NarrativeQA에서 50% 예산으로 98%의 Full-cache 성능을 달성하지만, GSM8K에서는 42% 미만으로 급격히 저하된다. 이에 반해, ReFreeKV는 입력에 관계없이 Full-cache 성능을 유지하면서 최대한 많은 KV 캐시를 제거하는 새로운 목표를 제시한다. 이를 위해 DBudgetKV 알고리즘을 제안하며, 이는 토큰 순위 결정과 주의도 기반의 중단 기준을 결합하여 동적으로 예산을 조정한다.

기술적 접근법

**DBudgetKV**: 토큰 순위는 위치 기반으로 결정 (시작과 끝 토큰이 더 중요), 추가 계산 없이 순위를 매김.
**중단 기준**: Frobenius norm을 사용하여 Full attention matrix와 Pruned matrix의 차이를 비교. 차이가 특정 threshold를 넘으면 중단.
**효율성**: PyTorch 내장 연산자로 구현, 메모리 및 시간 절약 효과.
**하이퍼파라미터**: Threshold는 universal하게 설정, 토큰 순위 결정 전략과 attention matrix의 k 값이 실험적으로 조정됨.

주요 결과

13개 데이터셋에서 평균 압축률 25% 이상, 일부 데이터셋에서는 85%까지 압축 가능.
Llama3-8B에서 평균 예산 크기 63.7%로 Full-cache 성능 유지.
수학 문제(GSM8K)에서는 높은 예산을 할당, 독해 태스크에서는 높은 압축률 달성.
Full-cache 대비 시간 감소 효과도 보임, 기존 KV pruning 방법 대비 효율성 향상.

의의 및 한계

ReFreeKV는 입력에 따른 threshold 설정 없이도 Full-cache 성능을 유지하는 KV 캐시 압축 방법으로, 실제 세계의 다양한 입력 환경에서 유용하다. 특히, 수학적 추론, 독해, 코드 생성 등 다양한 태스크에서 일관된 성능을 보임. 그러나, 모든 입력에서 완전히 lossless 압축을 보장하지는 않으며, 일부 복잡한 입력에서는 약간의 성능 저하가 발생할 수 있다. 또한, 현재는 토큰 순위 결정 전략이 위치 기반으로 제한되어 있어, 다른 중요도 지표와의 결합 가능성은 아직 탐구되지 않았다.

실용적 활용

ReFreeKV는 대규모 언어 모델의 추론 과정에서 GPU 메모리 사용을 줄이며, 특히 다양한 도메인의 입력을 처리하는 클라우드 기반 서비스나 모바일 환경에서 유용할 수 있다. 또한, 다양한 모델 아키텍처(Llama3, Qwen2.5 등)에 쉽게 통합 가능하며, 추론 시간 감소 효과로 실시간 응답이 필요한 애플리케이션에도 적용 가능하다.