WUSH-KV가 2비트 KV 캐시 양자화 오차를 줄인다

무엇인가

긴 문맥과 큰 서빙 배치를 쓰는 대규모 언어 모델에서 키-값(KV) 캐시는 메모리 용량과 대역폭 비용을 키우는 주된 원인이다. 자기회귀 생성 중 각 어텐션 레이어는 이전 토큰의 키와 값을 저장하고, 이 캐시는 시퀀스 길이와 배치 크기에 선형으로 늘어나면서 새 토큰마다 반복해서 읽힌다. 저비트 양자화는 이 메모리 사용량과 데이터 이동을 줄이는 손쉬운 수단이지만, 공격적으로 비트를 낮추면 어텐션 점수와 출력 양쪽에 큰 오차가 들어간다. 이 논문은 키 오차는 쿼리-키 내적에, 값 오차는 출력 프로젝션을 거친 어텐션 출력에 영향을 준다는 점에 주목해 키와 값을 따로 다루는 변환을 설계한다.

어떻게 동작하나

제안 방법 WUSH-KV는 가중치-활성화 양자화용으로 제안된 WUSH 변환을 KV 캐시로 옮긴 것이다. WUSH는 행렬 곱의 두 인자 각각의 2차 통계량으로부터 데이터 인지 변환을 닫힌 형태로 구성해 양자화 오차를 줄인다. WUSH-KV는 캘리브레이션 데이터로 키/값 헤드마다 별도의 변환을 만든다. 키 변환은 캐시된 키를 소비하는 쿼리 헤드들을, 값 변환은 캐시된 값을 소비하는 출력 프로젝션 블록들을 반영한다. 값 쪽 변환은 추론 전에 가중치에 접어 넣을 수 있어(값 프로젝션을 W_V·T_V^T로, 출력 프로젝션을 T_V^{-T}·W_O로 교체) 온라인 변환 비용이 전혀 없다. 반면 키 쪽 변환은 헤드별 정규화와 위치 의존적인 RoPE 뒤에 적용되기 때문에 가중치로 접어 넣을 수 없다. 논문은 키 변환과 쿼리 쪽 보상을 완전히 접어 넣으려면 각각이 RoPE 및 학습된 RMS 정규화와 교환 가능해야 하는데, 그 제약을 만족하면 좌표 크기를 재분배하지 못하는 부호 반전 변환만 남는다고 설명한다. RoPE 앞에 무제한 변환을 두는 대안도 좌표 자유도는 유지되지만 보상이 캐시 위치에 의존해 매 디코딩 스텝마다 캐시 전체를 복원·회전해야 하므로 채택하지 않았다.

무엇과 다른가

변환은 양자화 대상 텐서의 그램 행렬과 손실의 헤시안으로부터 만들어진다. 키의 경우 그램 행렬은 K_h·K_h^T, 헤시안은 그룹 내 쿼리들의 합 2·Σ_g Q_h,g·Q_h,g^T이고, 값의 경우 그램 행렬은 V_h·V_h^T, 헤시안은 출력 프로젝션 블록들의 합 2·Σ_g W_O(h,g)·W_O(h,g)^T다. 이 두 행렬에 감쇠 비율 γ를 더한 뒤 콜레스키 분해와 대칭 고유분해, 정규화 Hadamard 행렬을 조합해 T = c·H·Λ^{-1/4}·U^T·L^T 형태의 변환을 닫힌 형태로 얻는다. 변환은 블록 대각이고 블록 폭이 양자화 그룹 크기와 같아 적용 비용이 저렴하며, 두 인자를 양수로 스케일해도 결과가 변하지 않아 정규화를 추적할 필요가 없다. 실험에서는 γ=10^{-2}를 사용한다.

어떻게 쓰나

양자화는 토큰 단위로 이뤄지며, 키나 값 헤드의 d개 채널이 하나의 그룹으로 스케일을 공유한다. 이론 분석은 QuEST 양자화기의 투영 단계를 사용한다. 이는 그룹의 RMS로 스텝 Δ를 정하고 2^b개 레벨을 ±α_b×RMS 사이에 균등 배치하며 범위를 벗어난 값은 클리핑한다. 논문의 정리 2는 가우시안 꼬리와 클리핑 정렬 조건, 덧셈적 반올림 잡음 모델 아래에서 감쇠 없는 WUSH 변환이 민감도 균형(sensitivity-balanced)을 만족하는 모든 가역 변환 가운데 비트폭이 커질수록 (1+O(α_b^{-2})) 배 이내로 근최적임을 보인다. 또한 캐시 앞쪽의 고정 싱크 윈도와 최근 토큰의 롤링 윈도를 완전 정밀도로 유지하고, 나머지 구간을 일정 단위로 묶어 양자화하는 캐시 관리 방식을 함께 쓴다.

전제와 한계

재구성 오차 실험은 Qwen3-8B의 어텐션 모듈 36개 전체를 대상으로, FineWeb-Edu 128개 시퀀스로 캘리브레이션하고 서로 겹치지 않는 길이 1024 시퀀스 32개로 평가했다. 모든 변환에 동일한 QuEST 양자화기를 써서 변환 자체의 효과만 분리했다. 키와 값을 모두 2비트로 양자화했을 때 출력 프로젝션 이후 잔차 연결 이전의 모듈 출력 오차 기하평균은 WUSH 0.208, 어텐션 인지 헤시안을 쓴 WUSH-A 0.195였고, 정규화 Hadamard는 0.325, OSCAR는 0.309였다. WUSH-A의 개선이 크지 않아 주 방법으로는 더 단순한 WUSH 헤시안을 사용한다.

WikiText-2 퍼플렉시티 평가에서는 길이 32768의 FineWeb-Edu 128개 시퀀스로 캘리브레이션하고(단일 NVIDIA L40S에서 약 12분, 최대 39GiB), 2048 길이 시퀀스를 16토큰 청크로 처리하며 S_sink=16, S_keep=128, S_flush=16 윈도를 적용했다. WUSH는 모든 비트폭에서 양자화 변환 중 가장 낮은 퍼플렉시티를 냈고, 2비트에서 10.51로 OSCAR의 13.74보다 크게 낮았다. Hadamard의 2비트 퍼플렉시티 급등은 첫 어텐션 모듈 키 벡터의 스파이크가 변환 후 인접 재구성 레벨 중간값 근처에 고르게 퍼지면서 조밀한 오차를 만들기 때문이라고 분석한다.

다운스트림 평가는 SGLang에 WUSH-KV를 통합하고 OSCAR식 백분위 클리핑 아핀 양자화기(κ_K=0.96, κ_V=0.92, 32B는 0.96)를 사용해 Qwen3-4B-Thinking-2507, Qwen3-8B, Qwen3-32B를 AIME 2025, MATH-500, GPQA Diamond, LiveCodeBench v6에서 시드 3개로 확률적 샘플링해 측정했다. S_sink=64, S_keep=256, S_flush=8을 썼다. 8B에서는 네 벤치마크 모두 WUSH-KV가 OSCAR보다 높았고, OSCAR가 크게 무너지는 구간(8B와 32B의 LiveCodeBench v6)에서는 격차가 더 벌어졌다. BF16 대비 열화 폭도 모델 크기 전반에서 비교적 일관됐다. 생성 토큰 수는 BF16의 최대 약 1.5배로, INT2 전환의 명목상 8배 저장 절감을 지우지 않는 수준이다. 긴 문맥 RULER NIAH와 MRCR에서도 가장 긴 길이에서 OSCAR보다 높은 정확도를 유지했다.

저자들이 밝힌 한계는 명확하다. 조밀한 키 쪽 변환이 RoPE 뒤에 적용되어야 해서 온라인으로 남고, 이는 추론 시 추가 연산을 만든다(값 쪽은 가중치에 접혀 비용이 없다). 또한 WUSH-KV는 OSCAR의 튜닝된 클리핑 비율을 재조정 없이 그대로 재사용했기 때문에 WUSH 변환된 K/V 분포에는 최적이 아닐 수 있고, 여기에 개선 여지가 남아 있다고 본다. 이론적 근최적성 보장은 감쇠 없는 이상적 변환에 대한 것이며 실험은 감쇠를 사용한다는 점도 구분해 둔다. 향후 과제로 더 저렴한 구조적 변환, 더 강한 어텐션 인지 민감도 측정, 시스템 수준 최적화와 처리량 평가를 제시한다.