TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate

Summary

이 논문은 고차원 벡터를 낮은 비트폭으로 압축하면서 MSE와 내적 오차를 최소화하는 온라인(데이터-무관) 벡터 양자화 알고리즘 TurboQuant을 제안한다. 핵심 아이디어는 입력 벡터에 랜덤 회전 행렬을 곱해 좌표 분포를 Beta 분포(고차원에서 Gaussian에 수렴)로 균일화하고, 서로 거의 독립인 좌표 각각에 Lloyd-Max 최적 스칼라 양자화를 독립적으로 적용하는 것이다.

MSE 최적화용 TurboQuant_mse에 더해, MSE 최적 양자화가 내적 추정에서 편향을 유발한다는 문제를 해결하기 위해 TurboQuant_prod를 추가로 제안한다. 이는 TurboQuant_mse를 (b−1)비트로 적용한 뒤 잔차(residual)에 QJL(1-bit Quantized Johnson-Lindenstrauss) 변환을 적용하는 2단계 구조로, 불편(unbiased) 내적 추정을 보장한다.

Shannon 하한을 이용한 정보 이론적 증명으로 TurboQuant의 왜곡 상한이 최적 하한의 최대 ~2.7배 이내임을 보인다. 별도의 코드북 학습이 불필요해 인덱싱 시간이 사실상 0에 가까우며, GPU 가속과 완전한 벡터화도 지원한다.

Problem

LLM 추론의 KV 캐시 압축이나 고차원 근사 최근접 이웃(ANN) 탐색에서 벡터 양자화가 핵심 병목이다. 기존 방법들은 두 가지 한계 중 하나를 가진다.

1. 데이터 의존 방법(PQ, OPQ 등): 코드북 학습에 k-means 등 무거운 전처리가 필요해 온라인 또는 스트리밍 환경에 적합하지 않다.
2. 데이터 비의존 방법(기존 온라인 양자화): 이론적 왜곡 보장이 최적 Shannon 하한과 거리가 멀거나, GPU 벡터화를 지원하지 않아 실용적이지 않다.

MSE 최적 양자화가 내적 추정에서 편향을 가진다는 점도 ANN 탐색이나 어텐션 계산처럼 내적 보존이 중요한 응용에서 문제가 된다.

Method

TurboQuant_mse (MSE 최적):

1. 랜덤 회전 행렬 $\mathbf{\Pi }\in {\mathbb{R}}^{d\times d}$ 를 QR 분해로 생성한다.
2. 입력 벡터 $\mathbf{x}$에 회전을 적용: $\mathbf{y}=\mathbf{\Pi }\cdot \mathbf{x}$. 결과 $\mathbf{y}$는 단위 구면 위 균등 분포를 따르며, 각 좌표는 Beta 분포 $f_X(x)=\frac{\Gamma (d/2)}{\sqrt{\pi }\cdot \Gamma ((d-1)/2)}(1-x^2{)}^{(d-3)/2}$를 따른다.
3. 고차원에서 좌표들이 거의 독립이므로, 각 좌표에 독립적으로 Lloyd-Max 최적 스칼라 양자화를 적용한다. 코드북은 미리 계산해 저장한다.
4. 복원 시 역방향 회전 ${\mathbf{\Pi }}^{\top }$을 적용한다.

TurboQuant_prod (내적 최적, 불편 추정):

MSE 최적 양자화는 1비트 시 $\frac{2}{\pi }$ 비율의 편향이 발생한다. 이를 2단계로 해결한다.

1. TurboQuant_mse를 $(b-1)$비트로 적용해 근사 벡터 ${\tilde{\mathbf{x}}}_{\text{mse}}$를 구한다.
2. 잔차 $\mathbf{r}=\mathbf{x}-{\tilde{\mathbf{x}}}_{\text{mse}}$에 QJL 변환 $\mathbf{z}=\text{sign}(\mathbf{S}\cdot \mathbf{r})$을 적용한다. ($\mathbf{S}$: i.i.d. Gaussian 행렬)
3. 최종 복원값: $\tilde{\mathbf{x}}={\tilde{\mathbf{x}}}_{\text{mse}}+\frac{\sqrt{\pi /2}}{d}\cdot \Vert \mathbf{r}{\Vert }_2\cdot {\mathbf{S}}^{\top }\mathbf{z}$

QJL의 불편성 보장에 의해 전체 내적 추정도 불편이 된다.

Results

이론적 보장:

	상한	하한	비율
MSE	$\frac{\sqrt{3}\pi }{2}\cdot \frac{1}{4^b}$	$\frac{1}{4^b}$	~2.7
내적 오차	$\frac{\sqrt{3}{\pi }^2\Vert y{\Vert }_2^2}{d}\cdot \frac{1}{4^b}$	$\frac{\Vert y{\Vert }_2^2}{d}\cdot \frac{1}{4^b}$	~2.7

KV 캐시 압축:

• Needle-in-a-Haystack (Llama-3.1-8B, 4~104k 토큰): 4× 압축에서 전체 정밀도와 동일한 점수(0.997) 달성
• LongBench-E (3.5비트): 전체 정밀도 평균 50.06과 동일한 50.06 달성

ANN 탐색 (DBpedia 1536-dim):

방법	인덱싱 시간	Recall
Product Quantization	239.75초	낮음
RabitQ	2267.59초	중간
TurboQuant	0.0013초	가장 높음

Key Insights

• 랜덤 회전으로 “최악 입력”을 제거하는 것이 핵심이다. 임의의 벡터를 구면 위 균등 분포로 변환하면, 이론 분석이 단순해지고 좌표 간 독립성 가정이 고차원에서 거의 정확히 성립한다.
• MSE 최적화와 내적 불편 추정은 별개의 목표이므로, 용도에 따라 알고리즘을 분리한 설계가 중요하다. 특히 b=1에서 MSE 최적 양자화의 내적 편향이 $2/\pi \approx 0.64$로 크게 나타난다.
• 잔차에 QJL을 적용하는 2단계 구조는 $D_{\text{prod}}\le \frac{\pi }{2d}\Vert y{\Vert }_2^2\cdot D_{\text{mse}}$ 관계를 통해 MSE 성능이 좋을수록 내적 왜곡도 낮아지는 연쇄 효과를 만든다.
• 코드북이 Beta 분포에만 의존하므로 데이터를 보지 않고도 미리 계산할 수 있어, 인덱싱 시간이 사실상 0에 수렴하는 것이 ANN 응용에서 큰 강점이다.
• 엔트로피 인코딩으로 최대 5% 추가 압축이 가능하지만 구현 복잡도 대비 이득이 적어 채택하지 않았다는 판단이 실용적이다.

Limitations

• 분석의 핵심 가정인 “고차원에서 좌표 간 근사 독립”은 차원이 낮아질수록 성립하기 어렵다. GloVe(d=200) 실험에서도 경쟁력 있는 결과를 보이긴 하지만, 이론적 보장이 고차원에서보다 느슨해진다.
• 회전 행렬 $\mathbf{\Pi }$가 $O(d^2)$ 메모리를 요구한다. 대규모 배포 환경에서는 Hadamard 기반 구조적 랜덤 회전(예: SRHT)으로 대체할 수 있으나 본 논문에서는 자세히 다루지 않는다.
• QJL이 사용하는 랜덤 가우시안 행렬 $\mathbf{S}$ 역시 $O(d^2)$ 공간을 차지하며, 이를 효율적으로 대체하는 방법은 향후 과제로 남아 있다.
• 내적 오차 하한 증명이 pigeonhole 논증에 의존하므로 tight하지 않을 수 있으며, 실제 worst-case와의 갭이 이론보다 작을 가능성이 있다.

Related Papers

• QJL: 1-bit Quantized JL Transform for KV Cache Quantization with Zero Overhead
• PolarQuant: Quantizing KV Caches with Polar Transformation
• KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache
• QuIP: 2-Bit Quantization of Large Language Models with Guarantees
• Practical and Asymptotically Optimal Quantization of High-Dimensional Vectors in Euclidean Space for Approximate Nearest Neighbor Search (RabitQ)