기술 요약: HeadQ: KV 캐시 양자화를 위한 모델 가시적 왜곡 및 점수 공간 보정

1. 문제 제기

현재 KV 캐시 양자화 방법들은 주로 저장 공간 재구성을 최적화하며, 일반적으로 원본 키/밸류와 양자화된 키/밸류 간의 평균 제곱 오차 (MSE) 를 최소화합니다. 그러나 저자들은 이러한 목표가 어텐션 메커니즘이 실제로 캐시 데이터를 소비하는 방식과 불일치한다고 주장합니다.

• 키 (Keys) 는 로그잇 (logits, 쿼리 - 키 점수) 을 통해 읽히므로, 키의 특정 선형 투영만이 어텐션 분포에 영향을 미칩니다.
• 밸류 (Values) 는 어텐션 가중치에 따른 읽기 ($\sum p_j v_j$) 를 통해 소비됩니다.

본 논문은 지속적인 캐시 오차를 원시 유클리드 저장 좌표가 아닌 모델 가시적 좌표 (model-visible coordinates) 에서 측정해야 한다고 주장합니다. 교란 (perturbation) 은 큰 저장 공간 MSE 를 가지면서도 어텐션에 영감을 미치지 않을 수 있습니다 (소프트맥스-널 방향에 위치하는 경우). 반대로 작은 저장 공간 MSE 를 가지더라도 치명적인 어텐션 손상을 초래할 수 있습니다 (점수 가시적 방향에 영향을 미치는 경우). 이러한 불일치는 관련 없는 구성 요소에 비트를 낭비하면서 행동적으로 중요한 구성 요소는 보호하지 못하는 비최적 양자화 전략으로 이어집니다.

2. 이론적 기반: 모델 가시적 기하학

본 논문은 저장 오차와 행동 오차를 구분하는 기하학적 프레임워크를 정립합니다:

2.1 키 점수 공간 (피셔 기하학)

고정된 쿼리 $q$에 대해, 키 교란 $\delta k_j$는 점수 교란 $\delta s_j = q^\top \delta k_j / \sqrt{d}$를 유발합니다.

• 소프트맥스-널 방향 (Softmax-Null Directions): 모든 키에 상수 벡터 $\mu$를 더하는 토큰 공유 이동은 모든 점수에 상수를 추가합니다. 소프트맥스는 상수 이동에 불변하므로, 이 교란은 소프트맥스-널이며 저장 MSE 가 잠재적으로 크더라도 어텐션 손상을 일으키지 않습니다.
• 점수 가시적 공간 (Score-Visible Space): 어텐션은 쿼리에 상대적인 키 교란의 중심화된 활성 계수 (centered active coefficients) 에만 의존합니다.
• 피셔 정보 (Fisher Information): 소프트맥스 분포의 국소 2 차 기하학은 피셔 정보 행렬로 포착됩니다. 논문은 원본과 교란된 어텐션 분포 간의 KL 발산이 원시 키 MSE 가 아닌 점수 오차의 피셔 가중 분산 (Eq. 6) 에 국소적으로 비례함을 유도했습니다.

2.2 밸류 읽기 기하학

밸류는 가중 합을 통해 소비됩니다. 고정된 어텐션 가중치 $p$에 대해, 읽기 오차는 $\Delta o = \sum p_j \Delta v_j$입니다.

• 논문은 밸류 오차에 적합한 대리 지표로 $A^2$ 가중 토큰 왜곡을 제안합니다. 원시 밸류 MSE 와 달리, 이 지표는 각 토큰의 왜곡을 해당 어텐션 확률의 제곱 ($p_j^2$) 으로 가중치하여, 고도로 어텐션된 토큰의 오차가 출력 오차에 더 크게 기여한다는 사실을 반영합니다.

3. 방법론: HeadQ

HeadQ는 원시 키 벡터가 아닌 "점수 가시적" 잔차를 보존하도록 설계된 키 측 양자화 방법입니다.

알고리즘 개요

1. 보정 (Calibration): 각 레이어와 헤드별로 PCA 를 통해 보정 쿼리 (RoPE 후) 에서 저랭크 기저 $U_h$ (랭크 $r$) 를 학습합니다.
2. 캐시 쓰기 (Cache Write):
   • 표준 저비트 양자화기 $B_b$가 기본 키 $\hat{k}^{base}_j$를 생성합니다.
   • 잔차 $e_j = k_j - \hat{k}^{base}_j$를 계산합니다.
   • 잔차를 학습된 기저에 투영합니다: $z_j = Q_q(U_h^\top e_j)$ (여기서 $Q_q$는 저비트 사이드 양자화기입니다).
   • 캐시는 기본 키와 사이드 코드 $z_j$를 저장합니다.
3. 캐시 읽기 (추론, Cache Read):
   • 현재 쿼리 $q$를 사용하여 사이드 코드를 역양자화합니다.
   • 가법적 로그잇 보정을 적용합니다:
     $$\hat{s}_j = \frac{q^\top \hat{k}^{base}_j}{\sqrt{d}} + \frac{q^\top U_h \hat{z}_j}{\sqrt{d}}$$
   • 이 보정은 전체 키 벡터를 재구성할 필요 없이 점수 오차를 직접 대상으로 합니다.

주요 차이점: HeadQ 는 현재 쿼리를 저장하지 않습니다. 대신 보정 학습된 기저에 잔차 코드를 저장하며, 런타임 쿼리가 이를 로그잇 보정으로 디코딩할 수 있게 합니다.

4. 주요 기여

1. 모델 가시적 왜곡의 공식화: 논문은 키 왜곡을 "상수 이동에 대한 모듈로 점수 가시적 오차"로, 밸류 왜곡을 "$A^2$ 대각 대리 지표가 있는 읽기 가시적 오차"로 정의합니다. 이는 원시 MSE 가 행동 성능의 나쁜 대리 지표인 이유를 설명합니다.
2. HeadQ 알고리즘: 쿼리 기저에 저랭크 잔차 사이드 코드를 저장하고 이를 가법적 로그잇 보정으로 적용하는 실용적인 캐시 쓰기/읽기 알고리즘입니다.
3. 반증 가능한 증거 사슬:
   • 지표 우위성: 여섯 개의 모델에서 원시 키 MSE 보다 피셔/점수 공간 오차가 어텐션 KL 발산을 훨씬 더 잘 예측합니다 (스피어만 $\rho \approx 0.62\text{--}0.69$ 향상).
   • 반례: 저장 MSE 가 더 낮은 양자화기가 더 높은 어텐션 KL 및 출력 오차를 가질 수 있는 짝지어진 실험을 통해 저장 MSE 가 충분하지 않은 대리 지표임을 증명합니다.
   • 널 공간 개입: 밀집 어텐션이 무시하는 "소프트맥스-널" 평균 구성 요소를 복원하면 비공변 고정 그리드 양자화기에 대해 피셔 손상이 급격히 증가하여, 저장 좌표 낭비가 행동 성능을 파괴함을 보여줍니다.
   • 배포 가능성: 보정 데이터에서 학습된 쿼리-PCA 기저는 활성 잔차를 성공적으로 근사하는 반면, 동일한 예산의 무작위 기저는 실패합니다.
   • 방향성: HeadQ 보정의 부호를 반전 ("오른쪽 부호" 제어) 하면 치명적인 실패가 발생하여, 이 방법이 올바른 방향으로 특정 점수 객체를 보정함으로써 작동함을 확인합니다.

5. 실험 결과

WikiText-103 에서 여섯 개의 모델 (GPT-2, Pythia, Qwen2, Qwen2.5, TinyLlama, Mistral) 을 사용하여 실험을 수행했습니다.

• 키 전용 실험 (밀집 밸류):
  • 가장 손상된 2 비트 행에서 HeadQ(랭크 4) 는 Qwen2, Pythia, GPT-2 와 같은 모델에 대해 KIVI 유사 기준 대비 초과 퍼플렉시티 (PPL) 의 약 84–94% 를 제거합니다.
  • TinyLlama 는 실질적인 개선을 보이며, Mistral 은 이 프로토콜에서 이미 거의 밀집 성능에 근접합니다.
• 전체 KV 구성 (2 비트):
  • HeadQ(키 측) 를 $A^2$ 가중 밸류 정책과 결합하면 2 비트 정밀도에서 여섯 개의 모든 모델에서 성능이 향상됩니다.
• 이상치 분석:
  • Pythia-160M 은 특정 저엔트로피 "라우트 플립 (route-flip)" 영역 (작은 교란이 톱 -1 토큰을 뒤집는 경우) 에서 피셔와 KL 간의 상관관계가 약했습니다. 이 이상치는 특정 가족 고유의 경계에 국한되었으며 더 큰 Pythia 체크포인트 (410M, 1B) 에서는 약화되었습니다.

6. 중요성 및 주장

논문은 연산자 매칭 (operator matching) 이 KV 캐시 양자화를 위한 핵심 설계 교훈이라고 주장합니다:

• 키는 점수 객체입니다: 압축은 쿼리 가시적 잔차를 보존하고 소프트맥스-널 이동에 해상도를 낭비하지 않아야 합니다.
• 밸류는 읽기 객체입니다: 압축은 $A^2$ 가중 오차를 우선시해야 합니다.

저자들은 이 작업이 기작적이고 진단적이며, 생산 준비가 된 서빙 시스템에 대한 주장이 아님을 명시합니다. 그들은 패킹된 커널, 지연 시간, 또는 처리량을 최적화했다고 주장하지 않습니다. 중요성은 다음과 같습니다:

1. KV 캐시 양자화의 이론적 목표를 저장 재구성에서 모델 가시적 왜곡으로 수정합니다.
2. 올바른 오차 기하학을 대상으로 함으로써 저비트 영역에서 상당한 PPL 회복을 달성하는 배포 가능한 방법 (HeadQ) 을 제공합니다.
3. 표준 저장 지표 (MSE) 가 행동적으로 오해의 소지가 있을 수 있음을 보여주어 점수 공간 대리 지표로의 전환을 필요로 합니다.

논문은 다음 단계가 저장 MSE 를 기본 목표로 취급하는 것이 아니라, 이 연산자 매칭 관점을 완전히 최적화된 전체 KV 코덱으로 확장하는 것이라고 결론짓습니다.