KV Cache Compression Through the Lens of Transform Coding

기술 요약: 변환 코딩(Transform Coding)의 관점에서 본 KV 캐시 압축

문제 정의
대규모 언어 모델(LLM)의 긴 문맥 추론(Long-context inference)에서, 과거 토큰의 표현을 저장하여 생성을 가속화하는 키-값(Key-Value, KV) 캐시는 메모리 병목 현상의 주요 원인이 된다. 문맥 길이가 선형적으로 증가함에 따라, 캐시 크기는 현실적인 메모리 및 지연 시간 제약 조건 하에서의 배포를 불가능하게 만들 정도로 비대해진다. KIVI, KVQuant와 같은 기존의 양자화 방법들은 더 낮은 정밀도의 데이터 타입을 사용하여 KV 캐시를 표현함으로써 메모리를 줄이지만, 이들은 대개 캐시 자체의 원시 재구성 오차(raw reconstruction error)를 최소화하는 데 집중한다. 이러한 방식은 양자화 오차가 어텐션 메커니즘을 통해 어떻게 전파되어 최종 출력에 영향을 미치는지 고려하지 못하는 경우가 많다. 결과적으로, 모델의 예측에 거의 영향이 없는 정보는 보존하면서도, 결정적인 정보는 버리는 비효적인 비트 할당을 초래할 수 있다.

방법론: 어텐션 인식 변환 코딩 (Attention-Aware Transform Coding, AATC)
저자들은 신호 처리 이론과 LLM 추론을 결합한 어텐션 인식 변환 코딩(AATC) 방식을 제안한다. 이 접근 방식은 세 가지 주요 구성 요소에 기반한다:

1. 이론적 기초 (왜곡 분해):
   본 논문은 새로운 **어텐션 인식 왜곡 척도(attention-aware distortion measure)**를 도출한다. 백색 잡음 양자화 모델(white-noise quantization model) 하에서, 저자들은 어텐션 출력의 기대 제곱 오차 왜곡이 키(key)와 값(value)의 가산적 기여분으로 분해됨을 증명한다. 결정적으로, 이 왜곡은 토큰과 채널을 가로질러 요인 분해된다.

• 키 왜곡 ($D_K$): 제곱 어텐션 가중치($a_i^2$), 값 투영의 잔차($\|v_i W_O - o\|^2$), 그리고 제곱 쿼리 성분($q_c^2$)에 의해 가중치가 부여된다.
• 값 왜곡 ($D_V$): 제곱 어텐션 가중치($a_i^2$)와 출력 투영 노름($\|W_{O,c}\|^2$)에 의해 가중치가 부여된다.
  이러한 분해는 모든 채널이나 토큰이 출력 오차에 동일하게 기여하는 것이 아님을 밝혀낸다. 즉, 어떤 것들은 "어텐션 인식"적이며 높은 정밀도를 요구하는 반면, 다른 것들은 공격적으로 압축될 수 있음을 보여준다.

2. 변환 코딩 (백색화/Whitening):
   고전적인 변환 코딩(예: JPEG)에서 영감을 얻어, 이 방법은 먼저 키와 값 특징 간의 상관관계를 제거하기 위해 선형 변환을 적용한다. 저자들은 캘리브레이션 세트를 사용하여, KV 캐시를 차원이 서로 상관관계가 없고 강도(분산)에 따라 정렬된 표현으로 변환하는 백색화 변환(Cholesky 분해 및 가중치 행렬의 SVD 활용)을 계산한다. 이 단계는 신호 에너지를 소수의 계수에 집중시켜, 후속 비트 할당을 더욱 효과적으로 만든다.

3. 비트 할당 (역 워터필링/Reverse Waterfilling):
   데이터가 백색화되면, 이 방법은 변환된 채널 전체에 비트를 할당하기 위해 역 워터필링을 사용한다. 이 최적화 문제는 전역 비트 예산 제약 하에서 어텐션 인식 왜곡을 최소화한다.

• 높은 "어텐션 인식 가중치"(왜곡 공식에서 유도됨)를 가진 채널은 더 많은 비트를 받는다.
• 낮은 가중치를 가진 채널은 더 적은 비트를 받거나 zero 비트를 받는다(실질적인 랭크 감소 수행).
• 할당은 키와 값의 비대칭적 역할을 인정하여 키와 값에 대해 별도로 수행된다.

주요 기여

• 왜곡 척도: 양자화 오차를 토큰 및 채널 의존적 요인으로 분리하는 어텐션 인식 왜곡 척도를 도출함으로써, 기존에 단절되어 있던 토큰 제거(token eviction), 양자화, 저차원 투영(low-rank projection)의 관점을 통합하였다.
• 통합적 관점: 기존의 KV 캐시 압축 방법들(KIVI, PALU, KVTC 등)을 도출된 왜곡 분해 내의 특정 요인들과 연결하는 이론적 프레임워크를 제시하였다.
• AATC 알고리부터: 백색화 변환과 역 워터필링을 결합하여 어텐션 인식 왜곡 지표 하에서 최적의 채널별 비트 할당을 달enc하는 실용적인 알고리즘을 개발하였다.
• 실증적 검증: AATC가 다른 방법들이 성능 저하를 보이는 높은 압축률에서도 근사 무손실(near-lossless) 정확도를 달성함을 입증하는 종합적인 평가를 수행하였다.

실험 결과
이 방법은 Llama-3.1-8B-Instruct 및 Qwen-2.5-7B-Instruct를 대상으로 LongBench, RULER, GSM8K, MMLU-Pro, MATH-500을 포함한 표준 벤치마크에서 평가되었다.

• 압축 효율: AATC는 약 5.8배의 압축(Llama-3.1-8B의 경우 메모리를 ~1.07 GB에서 ~184 MB로 감소)을 달성하면서도, 모든 벤치마크에서 전체 정밀도(FP16) 베이스라인과 통계적으로 구별할 수 없을 정도의 성능을 유지한다.
• 비교: KIVI, KVQuant, PALU와 같이 특정 설정(특히 Qwen의 긴 문맥 검색)에서 성능이 저하되는 베이스라인들과 달리, AATC는 근사 무손실 정확도를 유지한다.
• 어텐션 인식 이점: 절제 연구("AATC var-only")에 따르면, 중간 정도의 압축률에서는 전체 방법과 분산 전용(variance-only) 변형이 Llama에서 유사하게 작동하지만, 공격적인 압축(2-bit 예산) 상황의 Qwen, 특히 긴 문맥 시나리오(RULK-32k)에서는 전체 어텐션 인식 방법이 분산 전용 버전보다 유의미하게 우수한 성능을 보였다.

의의 및 주장
본 논문은 최적화 목표를 원시 캐시 재구성이 아닌 어텐션 인식 왜곡으로 전환함으로써, 높은 압축률에서도 모델 성능을 보존하는 체계적인 비트 할당 전략을 도출할 수 있다고 주장한다. 저자들은 AATC를 LLM 추론의 특정 문제를 해결하기 위해 고전적 신호 처리 도구(변환 코딩 및 역 워터필링)를 활용하는 통합적 접근 방식으로 규정한다. 저자들은 자신들의 방법이 기존의 압축 기술에 대한 "근사 무손실" 대안을 제공하며, 이를 통해 긴 문맥 추론이나 검색의 품질을 희생하지 않고도 메모리가 제한된 환경에 트랜스포머 기반 LLM을 배포할 수 있게 한다고 단언한다.

한계점
저자들은 자신들의 이론적 분석이 토큰 독립성을 가정하고 있으며, 이는 자연어에서 엄격하게 성립하지 않을 수 있으나 실제로는 인코딩 품질에 부정적인 영향을 미치지 않는다는 점을 언급한다. 또한, 현재 구현에는 시스템 수준의 최적화(예: 전용 CUDA 커널)가 부족하며, 평가는 두 개의 모델 제품군(Llama 및 Qwen)으로 제한되어 있다.