5일 전

LAVa: 계층별 KV 캐시 제거 기법과 동적 예산 할당

Yiqun Shen, Song Yuan, Zhengze Zhang, Xiaoliang Wang, Daxin Jiang, Nguyen Cam-Tu
LAVa: 계층별 KV 캐시 제거 기법과 동적 예산 할당
초록

KV 캐시는 긴 컨텍스트를 갖는 대규모 언어 모델(LLM)의 추론 속도를 향상시키는 데 일반적으로 사용되지만, 높은 메모리 요구량으로 인해 캐시 압축이 필수적이다. 기존의 압축 방법들은 대부분 히우리스틱 기반이며, 동적 예산 할당 기능이 부족하다. 이러한 한계를 해결하기 위해, 우리는 Transformer 잔여 흐름(Residual Streams) 내 정보 손실을 최소화하는 통합 프레임워크를 제안한다. 이를 기반으로 계층별 주의(Attention) 출력 손실을 분석하고, 다양한 헤드 간 캐시 항목을 비교할 수 있는 새로운 지표를 도출함으로써, 계층별로 동적 헤드 예산을 적용한 압축을 가능하게 한다. 또한, 계층 간 정보를 비교함으로써 동적 계층 예산 할당도 실현하였다. LAVa는 기존 방법과 달리 학습이나 여러 전략의 조합에 의존하지 않고, 캐시 제거 및 동적 예산 할당을 통합적으로 처리하는 최초의 전략이다. LongBench, Needle-In-A-Haystack, Ruler, InfiniteBench 등 다양한 벤치마크를 대상으로 한 실험 결과에서 LAVa의 우수성을 입증하였다. 더불어, 우리의 실험을 통해 새로운 통찰을 발견하였다. 생성 태스크(예: 코드 완성)에서는 동적 계층 예산이 핵심적이며, 추출 태스크(예: 추출형 질의응답)에서는 동적 헤드 예산이 중요한 역할을 한다는 점이다. 완전히 동적인 압축 방식인 LAVa는 다양한 태스크 유형에 걸쳐 일관되게 최고의 성능을 유지한다. 코드는 이 URL에서 공개되어 있다.