토큰 하나의 활용도 극대화: KV Cache Eviction으로 긴 맥락 성능 향상

긴 맥락에서 KV 캐시가 기하급수적으로 커지며 주의 집중이 분산된다. 본 연구는 학습 가능한 retention 게이트로 캐시 엔트리를 글로벌 예산 아래 동적으로 선택·할당하고, 필요 토큰의 유지가 불필요 토큰의 제거보다 더 큰 효율·정확도 이득을 낳는다는 점을 보여준다. 이로써 풀-캐시 인퍼런스에 근접하거나 이를 넘어서는 성능을 더 적은 KV 메모리로 달성할 수 있다.

출발점: self-attention은 현재 쿼리와 모든 과거 토큰 간의 관계를 계산하므로, 컨텍스트가 길어질수록 distractor 토큰이 증가해 유용한 토큰의 주의 집중이 희석된다. 제안은 두 가지 축으로 작동한다. 첫째, retention gates로 각 캐시 항목의 미래 유용도를 예측하고, 둘째, 하나의 글로벌 KV 예산 아래 모든 계층/헤드의 엔트리를 함께 정렬해 캐시를 동적으로 할당한다. 결과적으로 불필요한 토큰의 제거가 주의 집중을 개선하고, 낮은 메모리 비용으로도 full-cache 성능에 근접하거나 향상될 수 있다. 실험은 텍스트-언어와 멀티모달(MM) 벤치마크에서 일관되게 메모리 사용량을 낮추면서 성능을 유지하거나 향상시켰다.

입력: 토큰 시퀀스와 KV 캐시를 초기화한다. 계산: 각 토큰 i에 대해 retention 게이트가 βℓ,h,i를 예측하고, 이를 rt,i = βℓ,h,i^(t−i) 형태로 정의한다. 패스: 게이트의 최종 Projection는 모든 레이어/헤드에서 공유된다. 학습 목표: Lquality = DKL(p ∥ qθ ) + E(x,y)[−log qθ(y|x)], Lcap = Σℓ, h max(0, Σt i=1 βt−iℓ,h,i − Mglobal). 총손실 L = Lquality + λ Lcap. 추론: G̃ℓ,h,i(t) = ΣT s=t+1 βs−iℓ,h,i / (1 − βT−tℓ,h,i) 값을 이용해 토큰(ℓ,h,i)을 전역 순위화하고 상위 Mglobal 엔트리를 캐시에 보존한다. 구현: 가변 길이 시퀀스의 캐시를 관리하는 paged-attention 구조를 도입하고, retention gate은 xt를 입력으로 사용해 βi를 예측한다. 보정: 최종 점수 Projection를 전적으로 tying해 계층 간 점수의 비교 가능성을 확보한다.

주요 벤치마크에서 DBTrimKV는 TrimKV 및 SnapKV 같은 대비 방법보다 우수하거나 동등한 성능을 달성하고, 메모리 버짓이 낮아질수록 이점이 커진다. 예를 들어 LongBench-V2 벤치마크에서 Full KV 대비 9.20% 평균 향상을 기록했고, 512 예산에서 14% 차이로 full-cache를 상회하는 경우가 관찰된다. 단일 모달 텍스트 태스크에서도 다양한 버짓에서 full-cache 대비 성능 저하를 크게 방지한다. 효율성 측면에서도, 더 긴 컨텍스트/생성 길이에 대해 Vanilla 대비 decodings 시간이 감소하거나 비례적 증가를 억제하는 경향이 있다. ablation에서 weight tying의 제거는 성능 저하를 유발하며, gate 입력으로 xt를 사용하는 구성이 더 나은 성능을 보였다.

아키텍처: L 레이어 × H 헤드의 구조에서 각 토큰에 대해 βℓ,h,t를 예측하는 retention gate를 도입하고, final projection는 모든 레이어/헤드에서 묶여 있다. 학습: Lquality + Lcap 손실로 학습, Lcap은 글로벌 KV 버짓 Mglobal를 넘지 않도록 제약한다. 차별점: per-layer/per-head 고정 예산 대신 하나의 글로벌 예산과 학습 가능한-retention score로 토큰의 미래 유용도를 평가해 순위를 매김. 구현 세부: LLM 가중치는 고정하고 retention gates만 학습시키며, PagedAttention으로 가변 길이 KV 캐시를 실현한다. 이론적 근거: Geometric retention(βi) 기반으로 토큰의 미래 지속성을 모델링하고, Assumptions A.1, A.2를 통해 exponential persistence를 보장하는 경향을 보인다. 데이터/실험: 텍스트 및 비주얼-언어 벤치마크에서 full-cache 대비 메모리 절감과 성능 유지/향상을 확인했다. 한계점으로는 대형 frontier-scale(model>70B)에서의 확장성 및 엔드-투-엔드 학습의 필요성 등이 per-study로 남아 있다.

연구는 retention 게이트의 학습만 수행하고 LLM 자체는 고정된 상태로 실험을 수행한다. 더 큰 모델(70B+)에서의 확장성 및 엔드-투-엔드 학습 시 효과를 확인하지 않았다. lookahead 파라미터의 일반화는 데이터에 의존적이며, 최적의 lookahead를 고정하기보다 케이스에 따라 달라질 수 있다.