기술 요약: TurboQuant 에서 영감을 받은 KV 캐시 양자화의 통계적 추론 및 품질 측정

문제 제기

KV(키 - 값) 캐시는 컨텍스트 길이에 비례하여 선형적으로 증가하는 트랜스포머 추론의 주요 메모리 병목 현상입니다. 양자화를 통한 압축은 필수적이지만, 표준 스칼라 양자화 (MSE) 는 종종 어텐션 메커니즘에 필요한 기하학적 특성을 보존하지 못합니다. TurboQuant 는 2 단계 파이프라인 (무작위 회전 + 스칼라 양자화 + 편향 보정을 위한 1 비트 QJL 스케치) 을 제안하지만, 키 (K) 와 값 (V) 캐시 간의 비트 할당 최적화는 여전히 불명확합니다. 구체적으로, 키 (및 값) 에 QJL 잔차 보정을 적용하는 것이 다양한 비트 예산과 데이터 분포 (예: 저랭크, heavy-tailed) 하에서 어텐션 품질을 개선하는지 아니면 저하시키는지는 알려져 있지 않습니다.

평균 제곱 오차 (MSE) 나 심지어 KL 발산과 같은 기존 평가 지표들은 종종 스케일과 방향 오차를 혼동하거나, 단일 잘못 양자화된 지배적 키가 어텐션 분포를 붕괴시키는 치명적인 "argmax-shift" 실패를 포착하지 못합니다.

방법론

이론적 프레임워크

저자들은 고차원 구의 기하학에 기반한 분석을 정립합니다. 초구 $S^{d-1}$ 위에 균일하게 분포된 벡터의 경우, 좌표의 주변 분포는 Beta 분포 $\text{Beta}(\frac{d-1}{2}, \frac{d-1}{2})$를 따릅니다. 차원 $d \to \infty$로 갈 때, 이는 가우시안으로 수렴합니다.

• 젠센 부등식 및 Softmax: 이 논문은 중요한 비선형성을 강조합니다. softmax 함수는 키 내적의 분산을 증폭시킵니다. 양자화 오차가 평균 0 이라 하더라도, 지수 함수 ($e^x$) 의 볼록성으로 인해 편향 (젠센 편향) 이 발생하여 높은 양자화 분산을 가진 키에 대한 어텐션 가중치를 과대평가합니다.
• K-V 비대칭성: 키 텐서는 비선형인 내적을 통해 어텐션 메커니즘에 진입하는 반면, 값 텐서는 선형 가중 합을 통해 진입합니다. 결과적으로, 키에서의 분산 증폭은 치명적이지만, 값에서의 편향 없는 편향 보정은 유익합니다.

양자화 방식

엄격한 동등 비트 예산 제약 ($n$ 비트/벡터) 하에서 세 가지 방식을 비교합니다:

1. KV: K 와 V 모두에 스칼라 MSE 양자화 (Lloyd-Max) 적용. 회전이나 스케치 없음.
2. KQV: K 와 V 모두에 Walsh-Hadamard 변환 (WHT) 적용. K 는 스칼라 MSE ($n$ 비트) 로 양자화. V 는 스칼라 MSE ($n-1$ 비트) + 1 비트 QJL 스케치로 양자화.
3. QKQV: K 와 V 모두에 WHT 적용. K 와 V 모두 ($n-1$ 비트) + 1 비트 QJL 스케치 사용.

평가 지표

스칼라 지표를 넘어설 수 있도록, 이 논문은 키 캐시, 값 캐시, 출력 ($T$) 의 세 지점에 대해 오차를 스케일 (SNR) 과 방향 ($\phi$) 로 분해하는 6 차원 오차 프레임워크를 도입합니다.

• Shannon-$\phi$ 평면: 스케일 오차 대 방향 오차를 플롯합니다.
• 통계적 도구: Mann-Whitney $U$(순위 - 이분 상관), Kolmogorov-Smirnov(CDF 분리), 그리고 에너지 거리 (6 차원 지표).
• 메커니즘적 대리 지표: 참조 양자화 softmax 분포 간의 KL 발산 및 상위 5 개 라우팅 재현율.

주요 기여

1. 키에 대한 QJL 의 해로움에 대한 이론적 증명

이 논문은 키의 스칼라 양자화를 1 비트 QJL 스케치로 대체할 경우, 스칼라 잔차에 비해 내적 분산이 최소 $\pi/2$(약 1.57) 배, 그리고 완전한 $n$ 비트 스칼라 방식과 $(n-1)$ 비트 스칼라 + 1 비트 QJL 방식을 비교할 때 $2\pi$배 증가함을 증명합니다.

• 메커니즘: QJL 은 잔차의 편향 없는 추정치를 제공하지만, 크기 정보를 부호 비트로 대체하여 버립니다. 이 추정치의 분산은 $\frac{\pi}{2} \frac{\|r\|^2}{d}$입니다.
• 결과: 이 증폭된 분산은 softmax 지수 함수에 의해 비선형적으로 증폭되어 (젠센 부등식을 통해) 어텐션 가중치의 체계적인 왜곡을 초래합니다. 따라서 공정한 예산 하에서 키에 QJL 을 적용하는 것은 스칼라 양자화보다 구조적으로 열등합니다.

2. K-V 비대칭성과 최적 방식 (KQV)

실험적 결과는 모든 테스트된 분포와 예산에서 KQV(V 만 QJL 적용) 가 QKQV(K 와 V 모두 QJL 적용) 보다 우수함을 확인합니다.

• KL 발산: 실질적으로 지배적인 4 비트 예산 ($n=4$) 에서 KQV 는 QKQV 보다 KL 발산이 2.25 배 낮습니다 ($p < 10^{-23}$).
• 라우팅: KQV 의 낮은 KL 은 높은 상위 5 개 라우팅 재현율과 낮은 출력 방향 오차와 직접적으로 상관관계가 있습니다.
• 결론: 1 비트 QJL 스케치는 편향 보정은 가능하지만 분산 패널티가 없는 값 캐시에만 예약되어야 하며, 치명적인 분산을 도입하는 키 캐시에는 절대 적용되어서는 안 됩니다.

3. 저랭크 병리

이 논문은 스칼라 지표가 놓치는 저랭크 키 분포에 대한 치명적인 실패 모드를 식별합니다.

• 주변 분포 vs 결합 분포: WHT 는 주변 좌표를 코드북과 일치하는 Beta 분포로 만들지만, 결합 저랭크 구조를 파괴하지는 않습니다.
• 치명적인 KL: 고차원 ($d=1024$) 에서 저랭크 키는 주변 히스토그램이 동일함에도 불구하고 heavy-tailed 키보다 KL 발산이 수배 더 나쁩니다. softmax 는 부분공간 내의 지배적 키에 대한 작은 오차를 증폭시켜 "argmax-shift" 실패를 초래합니다.
• 시사점: 표준 회전 기반 양자화는 저랭크 어텐션 행렬에 불충분하며, 결합 구조를 해결해야 합니다.

4. 예산 의존적 교차점

기하학적 K 재구성 품질의 교차점이라는 반직관적인 발견이 있습니다:

• 낮은 예산 ($n \in \{2, 3\}$) 에서 QKQV(K 에 QJL 적용) 는 때때로 KQV 보다 더 나은 기하학적 K 재구성 (낮은 방향 오차) 을 달성합니다.
• 그러나 이 기하학적 개선은 분산 증폭이 softmax 를 여전히 해치기 때문에 더 나은 어텐션 품질 (KL 또는 라우팅) 로 이어지지 않습니다.
• $n=4$에서 KQV 는 모든 측정 항목에서 승리합니다. 이 교차점은 랭크와 꼬리 가중치에 무관하며, 스칼라 정밀도와 JL 잔차 보정 간의 근본적인 왜곡 - 왜곡 트레이드오프를 시사합니다.

결과

• KQV 의 우세: $n=4$에서 KQV 는 모든 분포 (가우시안, heavy-tail, 저랭크, 집중) 에서 KL 발산, 기하학적 K 오차, 6 차원 거리 측면에서 QKQV 와 KV 를 능가합니다.
• 키에 대한 QJL 의 해로움: 키의 스칼라 비트를 QJL 로 대체하는 것은 분산 증폭으로 인해 KL 발산을 일관되게 증가시키며, 최종 출력에 도움이 되지 않는 특정 저예산 기하학적 재구성 시나리오를 제외하고는 유익하지 않습니다.
• Fat-Tail 견고성: WHT 는 주변 집중 (heavy tails) 에 대한 무조건적인 보호를 제공합니다. KQV 는 꼬리 가중치가 증가 ($\nu \to 1.05$) 함에도 안정적인 방향 오차 ($\epsilon_{dir}^K \approx 0.027$) 를 유지하는 반면, KV 는 치명적으로 실패합니다.
• 저랭크 실패: 저랭크 분포는 주변 정규화에도 불구하고 결합 부분공간 구조가 보존됨에 따라 어텐션 품질에서 가장 심각한 저하를 초래하며, heavy-tail 문제보다 훨씬 더 큽니다.

중요성 및 주장

이 논문은 TurboQuant 의 원래 키에 대한 QJL 적용 (대칭적으로 확장된 경우) 이 해로울 이유에 대한 엄격한 통계적 및 기하학적 설명을 제공한다고 주장합니다.

1. 지표 정제: 평균 KL 은 낙관적인 지표이며, 최악의 경우 KL 과 라우팅 재현율이 KV 캐시 양자화의 신뢰성에 대한 진정한 척도라고 주장합니다.
2. 설계 원칙: softmax 의 비선형성으로 인해 K 와 V 캐시는 근본적으로 다른 양자화 전략이 필요함을 확립합니다. QJL 은 선형 출력 (V) 에 적합한 편향 보정 도구이지만, 비선형 출력 (K) 에는 분산 주입 도구로 해를 끼칩니다.
3. 주변 보장의 한계: WHT 를 통해 주변 분포가 코드북과 일치하도록 하는 것은 필요하지만 불충분함을 보여줍니다. 결합 구조 (특히 저랭크) 가 궁극적인 실패 모드를 결정합니다.
4. 열린 문제: 기하학적 재구성의 예산 의존적 교차점을 열린 왜곡 - 왜곡 문제로 식별하며, QKQV 는 매우 낮은 비트 ($n=2,3$) 에서 기하학적으로 승리하지만 분산의 젠센 증폭으로 인해 기능적으로는 패배함을 지적합니다.

이 연구는 견고한 KV 캐시 양자화를 위해서는 저랭크 구조를 명시적으로 처리 (예: SVD 또는 부분공간 인식 코드북을 통해) 해야 하며, KQV 방식 (K 에 WHT + 스칼라, V 에 WHT + 스칼라 + QJL) 이 현재 비트 예산에 대한 최적 구성임을 결론지었습니다.