혼합 차원 예산 할당으로 효율적인 KV 캐시 압축

초록

본 논문은 토큰 폐기 방식의 한계를 넘어, 각 토큰에 할당할 차원을 세밀하게 조정하는 MixedDimKV와 헤드 중요도 정보를 결합한 MixedDimKV‑H를 제안한다. 후보 압축 비율 집합을 기반으로 손실 점수를 계산하고, 이산 최적화를 이분 탐색으로 풀어 메모리 예산 하에서 전체 손실을 최소화한다. 실험 결과 LongBench에서 전체 KV 캐시의 6.25%만 사용해 기존 방법과 동등한 성능을 달성했으며, 50K 길이의 Needle‑in‑a‑Haystack 테스트에서도 0.26% 캐시 사용량으로 100% 정확도를 유지한다.

상세 분석

MixedDimKV는 기존 토큰 폐기 기법을 “0 차원 또는 전체 차원”이라는 이진 선택으로 보는 관점을 확장한다. 저자는 각 토큰에 대해 0%부터 100%까지 다양한 압축 비율을 사전 정의하고, 토큰‑차원 쌍마다 정보 손실을 추정하는 손실 점수(Loss Score)를 도입한다. 이 손실 점수는 압축 전후의 어텐션 스코어 변화와 값(value) 벡터의 차이를 결합한 상한값으로, 토큰의 어텐션 중요도와 값의 크기를 동시에 고려한다. 이를 통해 단순히 어텐션 가중치만 보는 기존 방법보다 더 정교한 민감도 평가가 가능하다.

차원 할당 문제는 “총 예산 B 이하에서 손실을 최소화”하는 이산 최적화로 공식화된다. 저자는 라그랑지안 듀얼을 이용해 문제를 완화하고, 각 토큰에 대해 λ라는 라그랑지 승수를 도입해 손실 + λ·차원 수를 최소화하는 서브문제로 분해한다. 손실 함수 L_i(d) 가 차원 d에 대해 단조 감소한다는 가정 하에, λ가 커질수록 선택되는 차원이 작아지는 단조성을 보이며, 이 특성을 이용해 이분 탐색으로 최적 λ*를 효율적으로 찾는다.

또한, 헤드‑와이즈 압축을 선택함으로써 프로젝션 행렬 오버헤드를 크게 줄이고, 각 헤드별 중요도에 따라 차원을 비균등하게 할당할 수 있다. 이는 Joint‑Head 압축이 메모리와 연산 측면에서 비효율적인 점을 보완한다. 실험에서는 HeadKV와 비교해 동일한 헤드 중요도 정보를 사용했을 때 MixedDimKV‑H가 일관되게 우수한 성능을 보였으며, KV 캐시를 6.25% 수준으로 축소하면서도 LongBench에서 전체 어텐션과 거의 동등한 정확도를 유지한다. 특히 Needle‑in‑a‑Haystack 테스트에서 50K 컨텍스트 길이에도 0.26% 캐시 사용량으로 100% 정답률을 달성한 점은 장기 컨텍스트 추론에 있어 메모리 효율성을 크게 향상시킬 수 있음을 시사한다.

이러한 설계는 PCA 기반 저차원 투영을 헤드 단위로 적용하고, 손실 점수를 배치 연산으로 효율적으로 계산하도록 구현돼 실제 디코딩 지연을 55% 수준으로 감소시킨다. 전체적으로 MixedDimKV는 토큰‑레벨 차원 할당이라는 새로운 차원의 자유도를 도입함으로써, 메모리 제약이 심한 장기 컨텍스트 상황에서 기존 토큰 폐기 혹은 고정 차원 압축 방식보다 뛰어난 효율‑성능 트레이드오프를 제공한다.