KV 캐시 제거(KV cache eviction)를 통한 긴 문맥 성능 향상

학습된 전역 보정(globally-calibrated) KV 캐시 제거(KV-cache eviction) 정책은 메모리 사용량을 줄일 수 있으며, 역설적으로 긴 문맥 추론(long-context reasoning) 점수를 높일 수 있습니다. 해당 논문은 "올바른 토큰이 제거된다면, 제거(eviction)는 방해 요소(distractors)를 억제하고, 어텐션(attention)을 날카롭게 하며, 생성(generation)을 개선할 수 있다" [1]라고 보여줍니다. 이 연구 이전에는 KV 캐시 프루닝(KV-cache pruning)이 압축 기술로 취급되었습니다. 기존 방식들은 고정된 예산에 맞추기 위해 오래된 항목을 삭제했지만, 추상적 추론(abstractive reasoning) 및 다회차 대화(multi-turn dialogue)에서는 항상 전체 캐시(full-cache) 베이스라인에 미치지 못했습니다. 커뮤니티는 정확도를 희생하여 지연 시간(latency)을 절약하는 트레이드오프(trade-off)를 수용해 왔습니다.

하나의 전역 유지 게이트 네트워크(global retention-gate network)는 모든 토큰에 대한 유틸리티 점수(utility score)를 학습한 다음, 모든 레이어(layer)와 헤드(head)에 걸쳐 단일 공유 투영(shared projection)을 적용합니다. "우리는 모든 유지 게이트(retention gates)의 최종 점수 투영을 결합합니다. 이러한 가중치 공유는 유지 점수를 공통된 척도로 보정합니다" [1]. 이를 통해 어떤 위치의 토큰이라도 동일한 유한한 캐시를 두고 경쟁할 수 있습니다. 긴 문맥 언어, 시각-언어(vision-language), 다회차 대화 벤치마크를 아우르는 실험 세트 결과에 따르면, 학습된 제거 방식은 KV 메모리의 극히 일부만을 사용하면서도 전체 캐시 성능과 일치하는 것으로 보고되었습니다. 제한된 예산 상황에서는 이 방법이 베이스라인을 능가하기도 하며, 이는 선택적 제거(selective eviction)가 단순한 근사치가 아니라 성능 향상 도구임을 확인시켜 줍니다.

저자들이 인정한 한 가지 한계점은 유지 점수(retention scores)가 쿼리 불가지론적(query-agnostic)이라는 것입니다. 즉, 현재의 쿼리에 조건화된 완전한 예측기(predictor)보다는 미래 유틸리티에 대한 기하학적 대리 지표(geometric proxy)에 의존합니다. 이는 다음과 같은 미결 과제를 제시합니다: 쿼리 인식(query-aware) 점수 레이어가 유지된 토큰의 신호 대 잡음비(signal-to-noise ratio)를 더욱 개선할 수 있을 것인가? 고려해야 할 구체적인 변화 중 하나는, 현재 안전을 위해 KV 캐시 제거를 비활성화하고 있는 시스템들이 해당 보호 장치를 전역적으로 결합된 유지 게이트(globally tied retention gates)로 교체하는 것을 검토할 수 있다는 점입니다. 이 접근 방식은 다단계 추론(multi-hop reasoning) 정확도를 희생하지 않으면서—때로는 개선하면서—메모리 사용량을 줄일 수 있기 때문입니다.

References Make Each Token Count: KV Cache Eviction을 통한 긴 문맥(Long-Context) 성능 향상 연구