KV 캐시 압축을 위한 새로운 증류 방식 KVSculpt

본 논문은 대규모 언어 모델(LLM)의 장기 컨텍스트 추론 시 메모리 병목 현상을 일으키는 KV 캐시를 압축하는 새로운 방법인 KVSculpt을 제안한다. 기존 연구는 KV 쌍의 크기를 줄이는 양자화·저‑랭크 기법과, 시퀀스 길이를 줄이는 선택( eviction)·병합( merging) 방식으로 크게 두 축으로 나뉘었다. 선택 방식은 원본 캐시에서 중요한 k개의 위치만 남기고 나머지를 버리는 것이며, 병합 방식은 유사한 벡터들을 가중 평균해 수를 줄이는 것이었다. 두 방법 모두 원본 키 위치에 얽매여 있어, RoPE(Position‑Encoding) 후의 키가 이미 위치 정보를 포함한다는 사실을 활용하지 못한다. KVSculpt은 이러한 제약을 완전히 해제하고, 압축 구역을 k개의 완전 자유로운 KV 쌍(K_c, V_c)으로 재구성한다. 목표는 압축된 캐시가 원본 캐시와 동일한 어텐션 출력을 생성하도록 하는 것이다. 이를 위해 두 가지 손실을 정의한다. 첫 번째는 압축된 어텐션 결과와 원본 어텐션 결과 사이의 L2‑MSE이며, 두 번째는 어텐션 스코어의 로그‑합‑지수(LSE) 차이이다. 두 손실을 동일 가중치(λ=1)로 합산해 전체 손실을 구성한다. 키 최적화는 L‑BFGS를 사용한다. L‑BFGS는 1차 그라디언트와 제한된 히스토리를 이용해 2차 정보를 근사함으로써, 소프트맥스가 만든 비볼록 지형에서도 효율적으로 최소점을 탐색한다. 값 최적화는 현재 키에 의해 결정된 어텐션 가중치를 고정하고, 최소제곱(릿지 회귀)으로 V_c를 직접 해석적으로 구한다. 이 과정을 5번의 키 업데이트마다 한 번씩 반복한다. 초기 키는 어텐션 점수 누적값이 높은 원본 위치에서 추출해, 좋은 초기 베이스를 제공한다. 또한, 레이어·헤드별 압축 난이도가 크게 다르다는 사실을 관찰하고, 파일럿 압축 단계에서 균등 예산을 적용해 각 컴포넌트의 MSE를 측정한다. 이 파일럿 MSE를 제곱근 스케일링하여 예산을 재분배한다. 즉, 어려운 레이어·헤드에 더 많은 k를 할당하고, 쉬운 부분은 적게 할당한다. 이 적응형 예산 할당은 추가적인 추론 비용 없이 수행되며, 전체 KL 발산을 약 1.3배 추가 감소시킨다. 실험 설정은 Qwen2.5‑1.5B‑Instruct 모델(28 레이어, GQA, 12 쿼리 헤드, 2 KV 헤드, 차원 128)과 PG19 테스트 셋을 사용한다. 컨텍스트 길이 N=2048, 디코딩 토큰 128개를 대상으로 압축 비율 r∈{0.1,0.2,0.3,0.5,0.7}에서 KL 발산을 측정한다. Baseline으로는 Random, Attention‑Score, Select+Fit, Joint Optimization을 비교한다. 결과는 다음과 같다. - r=0.3,0.5,0.7에서 KVSculpt은 Select+Fit 대비 KL을 3.5‑4.1배 낮추었다. - 특히 r=0.3(압축률 70%)에서는 KL이 4.1배 감소해, 가장 공격적인 압축 상황에서도 원본 출력과 거의 일치한다. - Joint Optimization은 키를 원본 위치에 제한하기 때문에 KVSculpt에 비해 미미한 개선만 보였으며, 이는 “키를 자유롭게 이동시키는” 것이 핵심임을 입증한다. - L‑BFGS를 Adam으로 교체하면 MSE가 17‑95배 악화되고, 최종 KL도 8‑15배 증가한다. 이는 소프트맥스 지형이 급격히 변하는 특성 때문이며, 2차 정보 활용이 필수적임을 보여준다. - 파일럿 기반 적응형 예산 할당은 레이어별 MSE가 100배, 헤드별 MSE가 467배 차이 나는 것을 확인하고, 이를 반영한 재분배가 추가적인 KL 감소를 가져왔다. 결론적으로 KVSculpt은 (1) KV 캐시를 연속적인 임베딩 공간에서 자유롭게 재구성한다는 개념적 전환, (2) L‑BFGS와 최소제곱을 이용한 효율적 교대 최적화, (3) 입력‑특정 파일럿 신호를 활용한 레이어·헤드 수준의 예산 재분배라는 세 가지 핵심 기법을 통해 기존 선택·병합 기반 압축 방법을 크게 능가한다. 향후 연구는 이 방식을 양자화·저‑랭크와 결합하거나, 더 큰 모델·다양한 토큰 길이에 적용해 메모리·연산 효율을 극대화하는 방향으로 진행될 수 있다.