여러분이 직접 하지 않도록 제가 TurboQuant의 수학적 원리에 31시간을 투자했습니다

TurboQuant는 실제로 어떻게 작동할까요? 과연 그만큼의 가치가 있을까요? Nvidia의 FP4와 같은 현대적인 양자화 (Quantization) 기술과는 무엇이 다를까요?

PolarQuant의 조건화 (Conditioning)
TurboQuant를 이해하려면 먼저 PolarQuant를 이해해야 합니다:

랜덤 사전 조건화 (Random Preconditioning)와 극좌 변환 (Polar Transformation)을 사용하는 새로운 양자화 방법입니다. 우리의 방법은 재귀 알고리즘 (Recursive Algorithm)을 사용하여 KV 임베딩 (KV Embeddings)을 극좌 좌표 (Polar Coordinates)로 변환한 다음, 결과로 나온 각도 (Angles)를 양자화합니다. 긴 문맥 평가 (Long-context evaluation) 결과에 따르면, PolarQuant는 KV 캐시 (KV Cache)를 4.2배 이상 압축합니다.

한꺼번에 받아들이기에는 양이 많습니다. 차근차근 나누어 살펴보겠습니다. 아주 처음부터 시작해 보죠:

KV 캐시 (KV Cache): 문제점

모든 트랜스포머 (Transformer) 기반의 LLM은 동일한 방식으로 어텐션 (Attention)을 계산합니다. 각 토큰에 대해 모델은 세 가지 벡터를 생성합니다: 쿼리 (Query, 내가 무엇을 찾고 있는가?), 키 (Key, 내가 무엇을 포함하고 있는가?), 그리고 값 (Value, 내가 어떤 정보를 담고 있는가?). 어텐션 점수는 softmax(Q·K^T/√d)·V로 계산됩니다. 즉, 쿼리가 모든 키와 내적 (Dot-product)을 수행하여 어떤 토큰이 중요한지 파악한 다음, 값들의 가중 합 (Weighted sum)을 구합니다.

자기회귀 생성 (Autoregressive generation)을 빠르게 만드는 비결은 KV 캐싱 (KV Caching)입니다. 일단 토큰의 키와 값 벡터가 계산되면, 주어진 시퀀스 (Sequence)에 대해 이들은 절대 변하지 않습니다. 따라서 이를 저장해 두었다가 해당 요청의 향후 모든 토큰에 대해 재사용합니다. 문제는 이 캐시가 시퀀스 길이에 따라 선형적으로 증가한다는 점입니다. 새로운 토큰이 추가될 때마다 레이어 (Layer)당, 헤드 (Head)당 하나의 키 벡터와 하나의 값 벡터가 추가됩니다. 32개의 레이어, 8개의 KV 헤드, 헤드 차원(Head dimension) 128, 그리고 128K 문맥(Context)을 가진 Llama-3.1-8B의 경우: 128,000 × 32 × 8 × 128 × 2 (K와 V) × 2 바이트 (FP16) = KV 캐시만 무려 16 GB에 달합니다. 단 한 명의 사용자 세션(Session)을 위해서 말이죠. 동일한 GPU에서 더 많은 동시 세션을 실행하면 KV 캐시는 메모리 병목 현상 (Memory bottleneck)이 됩니다.

양자화 (Quantization): 기존 솔루션

현재 가장 공격적인 양자화 솔루션은 NVFP4 2레이어 양자화입니다. 이는 이전 포스트에서 다루었지만, 간략한 개요를 위해 작동 방식은 다음과 같습니다:

전체 행렬을 스캔합니다. 행렬 내에서 전역 최댓값(global maximum value)을 찾습니다. 그런 다음 16개 요소로 구성된 각 블록의 최댓값을 찾습니다. 이를 통해 지역적(local) 및 전역적(global) 입도(granularity)를 확보할 수 있습니다. 숫자를 지역 최댓값으로 나누고, 스케일(scales)을 전역 최댓값으로 나눕니다. 그런 다음 숫자들을 가장 가까운 4비트 버킷(4-bit bucket)으로 캐스팅(cast)합니다. 이 과정을 두 행렬(가중치(weights)와 활성화 함수(activations)) 모두에 대해 수행합니다. 특수 코어(specialized cores)를 사용하여 두 행렬을 곱하는데, 이 코어는 내부적으로 스케일링 인자(scaling factors)를 가져와 이를 재구성한 뒤, 행렬 곱셈 결과에 스케일링 인자의 역수를 곱하여 다시 전체 정밀도(full-precision)로 끌어올립니다.

버킷팅 문제 (The bucketing problem)

모든 양자화(quantization) 방법은 본질적으로 매핑(mapping)입니다. 연속적인 숫자를 가져와 가장 가까운 버킷에 할당하는 것입니다. 4비트의 경우 16개의 버킷이 있고, 2비트의 경우 4개의 버킷이 있습니다. 문제는 매핑 그 자체가 아니라, 버킷을 어디에 배치해야 할지 알아야 한다는 것이며, 이를 위해서는 먼저 데이터를 측정해야 합니다.

각 값의 블록에 대해 스케일 인자(scale factor, 최댓값)와 제로 포인트(zero-point, 오프셋)를 계산하고, 이를 양자화된 값과 함께 16비트 전체 정밀도로 저장한 뒤 나중에 재구성하는 데 사용합니다. 이러한 정규화 상수(normalization constants)는 순수한 오버헤드(overhead)입니다. 하나의 이상치(outlier)가 블록 전체의 정밀도를 망가뜨립니다.

버킷팅 문제를 어떻게 해결할 수 있을까요?

만약 우리가 데이터의 분포를 미리 알 수 있다면 어떨까요? 모든 값이 예측 가능하고 좁은 범위 내에 밀집되어 있다고 보장할 수 있다면 어떨까요?

하지만... 데이터를 효과적으로 매핑할 수 있도록 마법처럼 조작하여 빽빽하게 밀집된 분포로 만들 수 있는 방법은 존재하지 않습니다. 그래서 우리는 2레이어 양자화(2-layered quantization)라는 대가를 치릅니다. 우리는 이러한 평범함을 있는 그대로 받아들입니다.

PolarQuant

Google Research의 저자들은 실제로 데이터를 빽빽하게 밀집된 분포로 조작할 수 있다고 주장합니다. 하지만 어떻게 그럴 수 있을까요? 그들은 다변량 정규 확률 변수(multivariate normal random variables)의 두 가지 특성을 언급합니다:

종 모양 곡선(bell curve)의 항목을 가진 무작위 행렬(random matrix)에 임의의 고정된 벡터를 곱하십시오. 그 결과물은 0을 중심으로 하며, 분산(variance)은 원래 벡터의 길이의 제곱과 동일한 다변량 가우시안(multivariate Gaussian) 분포가 됩니다.

벡터의 모든 좌표가 표준 종 모양 곡선(standard bell curve)에서 추출된다면, 해당 벡터의 길이는 일반화된 감마 분포(generalized gamma distribution)를 따릅니다. 고차원(high dimensions)에서 이 길이는 $\sqrt{d}$ 주변에 매우 조밀하게 집중(concentrate)됩니다.

우리는 증명보다 더 우아한 방식을 취할 것입니다. 바로 코딩입니다.

def fact1_gaussian_norms():
torch.manual_seed(0)
dims = [16, 64, 128, 512]
...

fact 1의 분포 증명

1 tokenizer = AutoTokenizer.from_pretrained("gpt2")
2 model = GPT2LMHeadModel.from_pretrained("gpt2")
3 model.eval()
...

두 번째 사실, 즉 임의의 벡터 $x$에 대하여 $S$가 독립 항등 분포(i.i.d.) 정규 항목을 가진 무작위 행렬일 때, 벡터-행렬 곱(vector matrix product)이 다변량 정규 분포(multivariate normal distribution)를 가진다는 사실을 증명하는 것도 똑같이 간단합니다. 우리는 하나의 토큰을 선택합니다. 이 경우에는 5번 레이어, 0번 헤드, 5번 토큰입니다. 이 토큰은 64차원을 가집니다. 이 경우, 이 토큰의 노름(norm)은 우연히 9.85입니다.

그런 다음 우리는 1만 개의 무작위 행렬을 추출합니다. 이 무작위 행렬 각각을 벡터 $x$와 곱합니다. 그런 다음 각 차원의 분포를 분석합니다.

이는 무작위 사전 조건화(random preconditioning)를 거친 후, 벡터들이 가우시안(Gaussians)처럼 행동한다는 것을 우리에게 증명해 줍니다. 모든 곱셈 결과에 걸쳐 샘플링된 벡터의 모든 특징(feature)은 정규 확률 변수(normal random variable)의 분포를 따릅니다.

위의 내용에 만족하고 그것이 작동하는 것을 확인했으므로, 이제 이에 대한 직관(intuition)을 발전시키기 시작할 수 있습니다.

따라서 우리는 두 가지를 확립했습니다. 첫째, 임의의 벡터를 무작위 행렬과 곱하면 가우시안 출력이 생성됩니다. 둘째, 가우시안 벡터의 길이는 $\sqrt{d}$ 주변에 매우 조밀하게 집중됩니다. 그렇다면 사전 조건화(pre-conditioning)는 실제로 무엇을 하는 것일까요?

음, 그것은 단지 위의 두 가지 사실을 결합한 것뿐입니다. 무작위 행렬 $S$를 생성하고, 이를 우리의 벡터에 곱하여 이론적으로 0을 중심으로 하는 다변량 가우시안 (multivariate Gaussian) 분포를 따르는 출력을 생성합니다. 이는 새로운 결과 벡터의 모든 좌표가 마치 표준 종형 곡선 (standard bell curve)에서 샘플링된 것처럼 작동함을 의미합니다.

첫 번째 사실에 따라, 우리 벡터의 길이(또는 표준 종형 곡선에서 샘플링된 임의의 벡터의 길이)는 일반화된 감마 분포 (generalized gamma distribution)를 따르며, 이는 고차원에서 이 길이가 $\sqrt{d}$ 주변에 매우 조밀하게 집중됨을 보장합니다.

이를 저장된 KV 캐시(KV cache), 즉 저장된 $K$ 행렬과 $V$ 행렬에 적용하면, KV 임베딩 (embedding)의 각 좌표는 마치 동일한 종형 곡선에서 추출된 것처럼 행동할 것입니다. 각각은 정규 분포 (normal distribution)에서 추출된 확률 변수 (random variable)가 됩니다.

이것이 왜 중요한가

이제 알고리즘의 핵심으로 들어갑니다.

높은 수준에서 살펴보면, 우리의 접근 방식은 $d$차원 벡터 $x$의 좌표 쌍을 그룹화하고 각 쌍을 2D 극좌표 (polar coordinates)로 변환하는 것으로 시작합니다. 이를 통해 $d/2$개의 반지름 (radius)과 각도 (angle) 쌍이 생성됩니다. 다음으로, $d/2$개의 반지름을 모아서 이들에 극좌표 변환을 적용합니다. 이 절차는 $\log(d)$번 재귀적으로 반복되며, 최종 출력은 단 하나의 최종 반지름과 $1, 2, 4, \dots, d/2$차원의 각도 벡터 모음으로 구성됩니다.

즉, 우리는 데카르트 좌표 (cartesian)에서 극좌표 (polar)로 변환합니다. 좌표가 두 개일 때는 이것이 쉽습니다...

하지만 $d_{embed}$개의 좌표는 어떻게 변환할까요? 정말 간단합니다. 단순히 좌표를 2개씩 그룹으로 짝을 지으면 되고, 각 $(x_1, x_2)$ 쌍은 $(r, \Theta)$ 쌍이 됩니다. 그러면 $d_{embed}/2$개의 반지름 리스트와 $d_{embed}/2$개의 $\Theta$ 리스트를 갖게 됩니다. $\Theta$들은 따로 보관해 둡니다. 모든 반지름을 가져옵니다. 이것들이 우리의 새로운 좌표 쌍 $(r_1, r_2)$가 됩니다. 각 쌍은 다시 $(R, \Theta)$로 변환됩니다. 그러면 $d_{embed}/4$개의 반지름 리스트와 $d_{embed}/4$개의 $\Theta$ 리스트를 갖게 됩니다. $\Theta$들은 따로 보관합니다. 모든 반지름을 가져옵니다. 이것들이 우리의 새로운 좌표 쌍이 됩니다... 이래서 for 루프가 만들어진 것 같네요. 감이 오실 겁니다. 우리는 마지막 쌍, 즉 마지막 $R$과 $\Theta$가 나올 때까지 이 과정을 계속 반복합니다.

만약 제가 여러분을 놓치지 않았다면, 우리의 값들의 $x_1$과 $x_2$ (또는 어떤 쌍이든)가 무엇이 될지 다시 한번 생각해 봅시다. 전처리 (pre-conditioning)를 거친 후, 이들은 모두 동일한 분산 (variance), 동일한 평균 (mean)을 가지며, 어떤 $\sigma$에 대한 동일한 $N(0, \sigma^2)$ 분포로부터 추출됩니다.

우리는 어떤 연산을 하고 있는 걸까요?

$x_2/x_1$은 무엇일까요?

동일한 분포에서 샘플링된 두 숫자를 나누면, 그 비율은 통계적으로 1에 수렴하게 됩니다. 그리고 $\arctan(1) = \pi/4 = 45^\circ$입니다.

논문에서 이를 위해 도출된 보조정리 (lemma)는 다음과 같습니다:

이것이 첫 번째 단계 (level 1)입니다. 하지만 이들은 여전히 변할 수 있습니다. 각도 (angles)는 $0$에서 $2\pi$ 사이의 어디로든 변할 수 있지만 (좌표 $x_2$ 또는 $x_1$ 중 어느 것이든 음수일 수 있기 때문), 우리는 이 각도들이 $\pi/4$를 중심으로 집중되어 (centered) 있다는 점에는 동의합니다.

이제 2단계 (level 2)로 넘어갑니다. 1단계는 우리에게 반지름 (radius)과 각도 (angle) 쌍을 제공했습니다. 우리는 모든 반지름을 가져와 다음 단계로 넘깁니다.

이제 흥미로운 부분이 나옵니다. 우리는 $\arctan(r_2/r_1)$을 수행합니다. 각 반지름은 이전 단계의 서로 다른 두 쌍으로부터 나온 $\sqrt{x_1^2 + x_2^2}$입니다. 이는 두 반지름 중 어느 하나라도 음수일 가능성이 $0$임을 의미합니다. 우리는 오직 $0$과 $\pi/2$ 사이의 값만을 갖게 될 것임을 보장받습니다. 따라서 수학적 계산을 수행하고, 각 반지름 쌍에 대해 새로운 반지름과 각도를 생성합니다. 1단계에서는 가공되지 않은 좌표 (raw coordinates)가 음수일 수 있었기에 각도가 $[0, 2\pi)$ 전체에 걸쳐 있었습니다. 이제 반지름은 항상 양수이므로, 우리가 생성하는 각도는 $[0, \pi/2]$ 안에 고정됩니다. 각도는 이전 단계만큼 멀리 벗어날 수 없습니다. 또한, 여기서 생성된 각도는 이전 단계보다 $\pi/4$ 주변에 훨씬 더 조밀하게 집중되어 있습니다. 왜 그럴까요?

재귀 (recursing)를 계속할수록, 각 반지름은 더 긴 부분 벡터 (sub-vector)의 노름 (norm)이 됩니다. 2단계에서 각 반지름은 2개의 좌표를 요약합니다 ($r_1 = \sqrt{x_1^2 + x_2^2}$, $r_2 = \sqrt{x_3^2 + x_4^2}$). 만약 $r_3 = \sqrt{r_1^2 + r_2^2}$라면, 이는 $\sqrt{x_1^2 + x_2^2 + x_3^2 + x_4^2}$가 됩니다. 이것이 왜 중요할까요?

사실 1 (Fact 1)은 우리에게 말해주었습니다: 부분 벡터가 길어질수록, 그 노름은 기대값 (expected value) 주변에 더 조밀하게 집중된다는 것을 말입니다. 우리는 이를 그래픽으로 확인할 수 있습니다:

레벨 3에서는 각 반지름(radius)이 4개를 요약합니다. 레벨 4에서는 각 반지름이 8개를 요약합니다. 그리고 사실 1(Fact 1)은 우리에게 다음과 같이 말해주었습니다: 하위 벡터(sub-vector)가 길어질수록, 그 노름(norm)은 대값(expected value) 주변에 더 조밀하게 집중됩니다. 2D 노름은 여전히 크게 변할 수 있습니다. 하지만 8D 노름은 거의 변하지 않습니다. 따라서 두 8D 노름을 나누면 그 비율은 1에 매우 가깝고, 아크탄젠트(arctan) 값은 45°에 매우 가까워집니다.

레벨 3은 4D 하위 벡터의 노름인 반지름들을 쌍으로 묶습니다 (d=4 곡선, 더 좁음, 각도가 더 조밀함). 레벨 4는 8D 하위 벡터의 반지름들을 쌍으로 묶습니다 (d=8, 훨씬 더 조밀함).

레벨 7은 128D 벡터이며, 당신은 64D 하위 벡터의 노름들을 쌍으로 묶게 됩니다. 기본적으로 45도에서 일정합니다. 이들은 모두 동일한 값으로 매핑되므로, 1비트로 양자화(quantized)할 수 있거나 기본적으로 하나의 버킷(bucket)만 필요합니다.

이것이 보조정리 2(Lemma 2)가 포착하는 내용입니다. 레벨 ℓ에서의 분포는 다음과 같습니다:

가우시안 벡터(Gaussian vector)의 극각(polar angles)의 결합 분포(joint distribution)는 각 레벨별 독립적인 분포로 분해됩니다. 여기서 레벨 1은 균등 분포(uniform)를 따르며, 레벨 2 이상은 $\sin^{2^{(\ell-1)}-1}$을 따르는데, 이는 레벨이 올라갈수록 $\pi/4$에서 더 강력하게 정점(peak)을 형성합니다.

그리고 이것을 말해주는 무시무시한 수학 공식이 여기 있습니다:

알고리즘 (The algorithm)