Search
📜

[논문 리뷰] KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache (ICML’24)

필자는 현재(26년 4월) 온디바이스(On-device) 환경에서 LLM 추론 효율을 극대화하기 위한 시스템 최적화 연구를 수행하고 있다. 특히 제한된 메모리 자원 내에서 최적의 추론 성능을 확보하는 것을 핵심 목표로 삼고 있다.
최근 LLM의 컨텍스트 길이가 계속해서 확장됨에 따라 KV 캐시가 점유하는 메모리 비중이 시스템의 주요 병목으로 작용하고 있음을 알게 되었다. 이를 해결하기 위해 ‘KV 캐시 양자화(Quantization)’를 세부 연구 방향으로 설정하였다.
이 글은 해당 분야의 랜드마크 논문인 ICML 2024에 올라온 논문 KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache를 정리한 글이며, 필자의 생각이 포함되어 있다.

1. Background

1.1 Quantization 수학적 정의

양자화하려는 값들의 범위 내에서 가장 작은 값과 가장 큰 값을 활용해서 Zero point와 Scaling factor를 구한다.
zXz_X는 zero-point (분포의 offset을 잡는 상수), sXs_X는 scaling factor (실수 범위를 [0, 2B2^B-1] 정수 범위로 mapping하는 scale), ⌊⋅⌉는 반올림 연산자이다.
BB-bit integer quantization-dequantization은 다음으로 정의된다:
Q(X)=⌊X−zXsX⌉,X′=Q(X)⋅sX+zXQ(\mathbf{X}) = \left\lfloor \frac{\mathbf{X} - z_X}{s_X} \right\rceil , \mathbf{X}' = Q(\mathbf{X}) \cdot s_X + z_X
•
zXz_X : minX\text{min} X (zero point)
•
sX=(maxX−minX)/(2B−1)s_X = (\text{max}X - \text{min}X) / (2^B-1) (scaling factor)
양자화 에러(Quantization Error)는 두 원인에서 발생한다:
a.
Rounding(반올림)으로 인한 오차
b.
Scaling factor가 원래 분포의 range를 모든 원소에 공통으로 적용한다는 점. 특히 outlier가 존재하면 max/min outlier에 지배되어 sXs_X가 커지고, outlier가 아닌 원소들의 해상도(정밀도)가 급격히 낮아진다.
[예시로 이해하기]
어떤 토큰의 Key 벡터가 [0.1, 0.2, 0.15, 50.0]이라고 가정하자. (마지막 채널이 outlier임) Per-token으로 2-bit 양자화하면 s≈49.9/3≈16.63s≈49.9/3≈16.63이 되어, 정상 값들인 0.1, 0.2, 0.15가 모두 0으로 양자화된다. 반면 Per-channel 양자화라면 outlier 채널의 큰 scale이 해당 채널에만 적용되고, 정상 채널들은 자기 채널 내 분포에 맞는 작은 scale을 받아 값이 왜곡되지 않는다.

1.2 Per-Token vs Per-Channel Quantization

X∈Rlprompt×dX∈R^{l_{prompt}}×d를 Input tensor라고 하자. 여기서 lpromptl_{prompt}는 프롬프트의 토큰 수, dd는 hidden dimension (채널 수)이다.
Fig 3. Definition of per-token and per-channel quantization
Per-token quantization은 각 행(token) 단위로 (zz, ss)를 계산한다. 즉 zz, ss는 토큰 ii의 원소들 {Xi,1X_{i,1}, Xi,2X_{i,2}, …, Xi,dX_{i,d}}가 동일한 scale을 공유한다.
Per-channel quantization은 각 열(channel) 단위로 (zz, ss)를 계산한다. 채널 jj의 원소들 {X1,jX_{1,j}, X2,jX_{2,j}, …, Xlprompt,jX_{{l_{prompt}}, j}}가 동일한 scale을 공유한다.

1.3 LLM Inference Workflow

Fig 4. LLM Inference with KV Cache: Prefiil & Decoding Phase
LLM의 추론은 크게 두 단계로 나뉜다.
Prefill Phase
입력 prompt XX (lprompt×dl_{prompt}\times d)가 주어지면:
XK=XWK,XV=XWV{\boldsymbol{X}}_K = \boldsymbol{X}\boldsymbol{W}_K, \boldsymbol{X}_V = \boldsymbol{X}\boldsymbol{W}_V
WK\boldsymbol{W}_K, WV\boldsymbol{W}_V는 Key, Value projection weight이다. 계산된 XKX_K, XVX_V는 이후 decoding을 위해 메모리에 캐싱된다.
Decoding Phase
현재 입력 토큰 embedding 벡터인 tt (1×d1 \times d)에 대해:
•
[1단계] 새 Key, Value 계산 및 cache append:
◦
tK=tWKt_K = tW_K, tV=tWVt_V = tW_V
◦
XK←Concat⁡(XK,tK)X_K \leftarrow \operatorname{Concat}(X_K, t_K), XV←Concat⁡(XV,tV)X_V \leftarrow \operatorname{Concat}(X_V, t_V)
•
[2단계] Attention output 계산:
◦
tQ=tWQt_Q = tW_Q, A=Softmax⁡(tQ,XK⊺)A = \operatorname{Softmax}(t_Q, X_K^{\intercal})
◦
tO=AXVt_O = AX_V

2. Introduction

2.1 해결하려는 문제: “새로운 병목인 KV cache”

Auto-regressive decoding에서 토큰을 하나 생성할 때마다 이전 모든 토큰의 Key, Value를 재계산하면 O(l2)O(l^2) 연산이 매 step 반복된다. 불필요한 연산을 막기 위해 한 번 계산한 K, V를 메모리에 캐싱해 두는 것이 KV cache이다.
그런데 이 KV cache가 새로운 병목 현상이 되고 있다. Pope et al. (2023)에 따르면, PaLM 540B에서 batch 512, context length 2048 조건으로 KV cache만 3TB(모델 weight의 3배)를 차지한다고 한다.
KV cache의 병목은 메모리와 속도의 두 측면에서 동시에 발생한다.
•
메모리 병목: KV cache가 GPU 메모리 용량 자체를 초과하여 메모리 부족 문제
•
속도 병목: Decoding Phase에서 각 step마다 매번 KV cache를 HBM에서 SRAM으로 로드해야 하며, 이 작업이 끝날 때까지 GPU의 Compute unit은 Idle 상태가 됨
KV cache의 크기 자체를 줄일 수 있는 Quantization이 두 병목을 동시에 해결하는 가장 직접적인 해결책이다.

2.2 기존 접근법 및 한계점

논문에서는 KV cache의 병목 현상을 완화하려는 기존 방식을 크게 세 가지로 분류한다.
a.
Architectural 접근: Multi-Query Attention (Shazeer, 2019), Grouped-Query Attention (Ainslie et al., 2023)이 대표적이다. 절대적인 Key와 Value 개수를 줄여서 KV cache가 차지하는 크기를 감소시킵니다. 하지만 모델을 from scratch 학습하거나 fine-tuning이 필수적이다. 그리고 이미 배포된 모델이 적용할 수 없다는 점이 한계이다.
MHA vs MQA vs GQA
Multi-Head Attention (MHA)
Multi-Query Attention (MQA)
Grouped-Query Attention (GQA)
b.
Token Eviction 접근: H2OH_2O (Zhang et al., 2023), Scissorhands (Liu et al., 2024) 같은 논문이 있다. 중요도가 낮은 토큰을 버리지만(eviction), 본질적으로 정보 손실을 초래한다. 이미 버린 토큰이 나중에 중요해지는 경우에는 복구할 수가 없다는 점이 한계이다.
c.
System-level 접근: FlexGen (Sheng et al., 2023)의 offloading, vLLM (Kwon et al., 2023)의 PagedAttention이 해당된다. 메모리 배치를 최적화할 뿐, KV cache의 byte 수 자체는 줄이지 않는다.

2.3 KV cache Quantization 방법

KV cache의 크기를 줄이는 가장 직관적인 방법은 Quantization(양자화)이다. 그러나 Weight(가중치) 양자화(GPTQ, AWQ, SmoothQuant, LLM.int8)와 달리, KV cache quantization에는 심층 연구(in-depth studies)가 없었다.
저자들은 KV cache의 각 Key와 Value 값 분포(distribution) 자체를 분석해서 이 분석 결과를 바탕으로 KIVI 알고리즘을 제안한다.

2.4 KIVI 알고리즘 제안

KIVI는 크게 3가지를 제안한다.
a.
Llama/Falcon/Mistral 모델의 KV cache의 각 Key와 Value 값 분포를 체계적으로 분석하여 Key cache는 channel 단위로 outlier를 가지며, Value cache는 명확한 outlier 패턴이 없음을 밝혔다.
b.
이러한 Key와 Value의 비대칭적인 특징을 고려하여 Key는 per-channel, Value는 per-token 양자화라는 비대칭 설계를 제안한다.
c.
per-channel 양자화는 streaming 특성과 충돌하는 문제를 grouped + residual 분할 구조로 해결한다.
Streaming 특성
Auto-regressive한 디코딩 방식에서 토큰이 하나씩 순차적으로 도착하면서 KV cache가 점진적으로 커지는 데이터 구조를 뜻한다.
이러한 특성 때문에 Per-token 양자화는 각 토큰이 들어올 때마다 양자화하면 되기에 Streaming과 호환이 되지만, Per-channel 양자화는 여러 토큰에 걸쳐 채널별로 Scaling factor가 계산되어야 하므로 Streaming과 호환되지 않는다.
그래서 KIVI에서는 그룹(Group)으로 양자화하는 방식을 채택한다.
결과적으로 KIVI를 통해 peak memory 2.6x 감소, batch size 최대 4x 확장, throughput 2.35x-3.45x 향상을 달성하면서 정확도 하락은 거의 없었다.

3. Method

3.1 왜 Key와 Value는 다른 차원으로 양자화되어야 하는가

3.1.1 예비 실험: 양자화 구성의 비교
Table 1. The results of simulated KV cache group-wise quantization with various configurations
저자들은 CoQA(Conversational Question Answering, 질문과 답변이 꼬리에 꼬리를 무는 대화형 데이터셋)와 TruthfulQA(오개념, 거짓 정보에 대해 얼마나 진실성 있게 답하는지 테스트)를 사용하여 16bit, 4bit(Key와 Value 모두 per-token), 2bit(Key와 Value 각각 per-channel, per-token인 모든 경우의 수) 모든 양자화 설정을 시뮬레이션(Fake quantization)을 돌려서 비교한다.
이 실험을 통해 관찰한 3가지는 다음과 같다:
OB 1. Key와 Value 값 둘 다 per-token으로 양자화하면, INT4에서는 정확도가 유지된다. 하지만, INT2에서는 정확도가 확 떨어진다. ⇒ 즉, INT2까지 양자화하려면 Key와 Value 모두 per-token으로 하면 안 된다.
OB 2. Value는 per-channel로 양자화하면, Key가 어떻게 양자화되든 상관없이 정확도가 상당히 떨어진다. ⇒ 즉, Value는 per-channel로 양자화해서는 안 된다. (per-channel로 하면 안 되는 이유는 3.1.4에서 설명)
OB 3. INT2로 양자화할 때, 가장 정확한 접근법은 Key는 per-channel로, Value는 per-token으로 양자화하는 것이다. ⇒ 즉, Key는 per-channel로, Value는 per-token으로 양자화하는 것이 가장 정확도가 높다.
3.1.2 Key와 Value 값의 분포도
Fig 5. Magnitude of key and value cache for Llama-2-13B and Falcon-7B
Llama-2-13B와 Falcon-7B의 여러 Layer 및 Head에서 Key와 Value의 값 분포를 3D로 보면 다음 2가지를 관찰할 수 있다:
•
Key cache: 일부 고정된 채널에서 값이 10-12정도로 튀는 반면, 나머지 채널은 1 미만이다. ⇒ Key 값에는 특정 채널이 매우 큰 값으로 나타난다.
•
Value cache: 분포가 비교적 고르며 특별한 outlier 패턴이 없다. 값의 절대 크기도 작고 상대적으로 균일하다.
3.1.3 Key 값 분석
Key의 outlier가 “채널에 고정”되어 있다는 특징이 중요하다. Per-channel 양자화는 하나의 scaling factor를 하나의 채널로 제한한다. 따라서 outlier 채널의 큰 ss가 해당 채널에서 사용되고, 정상 채널은 자신의 range에 맞는 정확도를 보존한다.
반대로 per-token 양자화를 쓰면 한 토큰 벡터 안에 outlier 채널의 큰 값과 정상 채널의 작은 값이 섞여 있어서 max에 큰 영향을 받은 ss가 모든 채널에 공통 적용되므로 정상 채널은 정확도를 잃는다.
Table 2. The reconstruction error by Key cache
위의 Table에서 Reconstruction error(양자화한 값을 역양자화했을 때, 양자화한 값과의 차이)는 per-channel이 3배(13.67 → 4.55) 낮고, Attention score의 Reconstruction error는 5배(47.00 → 9.60) 더 작다.
Frobenius Norm (프로베니우스)
∣∣A∣∣F=∑i=1m∑j=1n∣aij∣2||A||_F = \sqrt{\sum_{i=1}^{m} \sum_{j=1}^{n} |a_{ij}|^2}
행렬 A의 모든 원소를 각각 제곱해서 더한 뒤, 마지막에 루트를 씌우는 방식이다.
•
Step 1: 행렬 안에 있는 모든 숫자(aija_{ij})를 꺼낸다.
•
Step 2: 모든 숫자를 제곱(2^2)한다. (음수를 양수로 만들고, 큰 값에 가중치를 줌)
•
Step 3: 제곱한 값들을 모두 더한다. (Σ\Sigma)
•
Step 4: 합계에 루트 (\sqrt{})를 씌운다.
IF) L1L_1 Norm을 사용했다면, 제곱을 하지 않기 때문에 큰 값(Outlier)에 특별한 가중치를 두지 않고, 모든 오차를 동등한 비중으로 합산한다. 즉, “Channel별 양자화와 Token별 양자화의 차이”가 수치상으로 덜 극적으로 보일 수 있다.
3.1.4 Value 값 분석
Table 3. The reconstruction error by Value cache
Fig 5에 따르면 Value는 channel별 Outlier가 없으므로, 단순 reconstruction error만 보면 per-token과 per-channel이 비슷할 것이고 실제로 Table 3을 보면 reconstruction error 자체는 크게 변하지 않는다. (per-token 4.57 vs per-channel 3.73)
그런데 Table 1의 OB2(Value는 per-channel로 양자화 X)를 보면, Value를 per-channel로 양자화하면 Key가 per-channel이든 per-token이든 정확도가 크게 떨어진다는 것을 알 수 있다. 그 이유는 Value가 사용되는 연산의 구조를 보면 알 수 있다. Query와 Key를 내적해서 Attention score(AA)를 구한 뒤, 그 AA에 VV를 내적함으로써 Output을 만들어낸다.
그렇기에 Value 값 자체의 Reconstruction error가 아니라, Ouput의 error를 보는 것이 더 정확하며 실제로 Δ\Delta의 값을 비교해 보니 per-token이 Output error에서 per-channel보다 약 15배(49.89 → 3.55) 더 작다는 것을 알 수 있다.
Δ=∥AXV−AXV′AXV∥F\Delta = \left\| \frac{AX_V - AX'_V}{AX_V} \right\|_F : Attention output의 error
그 이유는 Attention sparsity 때문이다. Output은 Attention score와 Value 벡터의 가중합이다. 여기서 Attention score는 Table 2에 따르면(attention sparsity 84.3%) 매우 희소하다. 대부분의 score는 0이거나 매우 작은 값이고, 소수의 “중요한 토큰”만 유의미한 가중치를 갖는다.
즉, 이 소수의 “중요한 토큰”이 Output에 크게 기여하기에 Value를 per-channel로 양자화하냐 per-token으로 양자화하냐 에 따라 Error가 큰 차이를 보이는 것이다:
•
Per-channel 양자화: 모든 토큰이 동일한 채널별 Scaling factor(ss)와 Zero point(zz)를 공유한다. 한 채널의 ss는 모든 토큰의 해당 채널 값들로 결정되기에 결과적으로 “중요하지 않은 토큰”으로 만들어진 ss가 “중요한 토큰”에 영향을 주게 되어 성능 저하에 원인이 된다.
•
Per-token 양자화: 각 토큰마다 독립적인 ss, zz를 갖기에 “중요하지 않은 토큰”의 양자화 오차가 아무리 커도 다른 “중요한 토큰”에는 영향을 주지 않기에 성능 저하가 크게 일어나지 않는다.
Attention Sparsity (어텐션 희소성)
•
정의: 모델이 다음 단어를 예측할 때, 이전의 모든 토큰을 중요하게 보는 것이 아니라 아주 적은 수의 특정 토큰에만 점수(가중치)를 집중시키는 현상
•
예시: 문장이 아무리 길어도, 다음 단어를 예측하는 데 결정적인 힌트가 되는 단어는 대략 5~10개 내외임. 나머지 1,990개 단어의 Attention score는 Softmax를 거치면서 거의 0이 됨
•
계산 방법:
이 논문에는 정확히 어떤 함수나 수식으로 이를 구했는지는 생략되어 있음. ‘대부분의 어텐션 가중치가 매우 낮아 무시 가능’이라는 사실을 강조하기 위해 통계적 근거로 사용됨
⇒ 보통 어텐션 스코어 행렬에서 매우 작은 값들을 ‘0’으로 간주하여 그 비율을 계산한다.

3.2 KIVI algorithm

3.2.1 해결해야 할 구현의 어려움
Key는 per-channel, Value는 per-token으로 설계 방향은 정해졌다.
그런데 Decoding Phase에서 Auto-regressive하게 토큰이 순차적으로 도착하는 Streaming 구조에서는 per-channel 양자화를 구현하기 어렵다. per-token은 토큰이 하나씩 들어올 때마다 양자화를 하면 되지만, per-channel은 하나의 토큰으로는 양자화하지 못하기 때문이다.
•
Per-token (Value): 새 토큰 tvt_v가 도착하면 즉시 quantize하고 Token 축으로 concat. Streaming에 자연스럽게 더해짐
•
Per-channel (Key): Scaling factor가 여러 토큰에 걸쳐 채널별로 계산되어야 하므로, 단일 토큰만으로는 양자화가 불가능. Channel 축 따라 group을 형성하려면 여러 토큰이 먼저 도착해야 함
3.2.2 Grouped + Residual 분할
여러 토큰이 있어야 Per-channel 양자화가 되는데, Streaming 구조에서 이를 어떻게 해결할 수 있을까? Key 값을 Grouped & Residual 두 부분으로 나눈다.
GG : 그룹에 들어가는 토큰 수, RR : 잔여 토큰의 최대 길이(수), ll : 현재 key 또는 value 캐시의 토큰 수(길이), rr = l%Rl\%R : 현재 ll에 대해 남는 잔여 토큰 수 (RR은 GG의 배수)
•
XKg=XK[:l−r]{X}_{K_g} = {X}_K[:l - r] (Grouped part) : GG 토큰 단위 group이 여러 개 있고, group별로 per-channel 양자화되어 저장
•
XKr=XK[l−r:]{X}_{K_r} = {X}_K[l - r:] (Residual part) : 아직 GG 단위 group을 채우지 못한 가장 최근 토큰들은 Full precision 유지
여기서 rr은 l%Rl\%R이며, RR은 rr이 도달할 수 있는 최댓값이다.
Value 값도 XVg=XV[:l−R]X_{V_g}=X_V[:l-R], XVr=XV[l−R:]X_{V_r}=X_V[l-R:]로 각각 Grouped, Residual로 나눈다. 단, Value는 per-token이기에 residual을 둘 필요는 없지만, 가장 최근 토큰들을 full precision으로 유지하는 Sliding window 효과를 위해 RR만큼 residual을 두도록 알고리즘이 설계되었다.
3.2.3 Prefill & Decoding Workflow with 예시
Group size GG = 4, Residual length RR = 8, Hidden dim dd = 8, Input prompt 길이 lpromptl_{prompt} = 10
KIVI 알고리즘을 보다 이해하기 쉽도록 예시와 함께 설명하자면:
[Prefill Phase]
Fig 6. The overview of KIVI algorithm
Algorithm 1. Prefill Phase
•
[Step A] Projection
XK=XWK∈R10×8\mathbf{X}_K = \mathbf{X} \mathbf{W}_K \in \mathbb{R}^{10 \times 8}, XV=XWV∈R10×8\mathbf{X}_V = \mathbf{X} \mathbf{W}_V \in \mathbb{R}^{10 \times 8}
◦
10개 토큰 x 8차원 Key와 Value 텐서를 만든다.
•
[Step B] Value의 grouped/residual 분할
XVg=XV[:lprompt−R]=XV[:2]X_{V_g}=X_V[:l_{prompt}−R]=X_V[:2] XVr=XV[lprompt−R:]=XV[2:]X_{V_r}=X_V[l_{prompt}−R:]=X_V[2:]
◦
뒤 8개 토큰(최신)은 residual로, 앞 2개 토큰(과거)은 grouped로 배치한다.
◦
Value는 “최신 R개를 항상 full precision으로 보호”하는 설계이다.
Value 분할 (prompt 10 tokens): Token idx: [0 1 | 2 3 4 5 6 7 8 9] └──────┘ └─────────────────────┘ X_Vg X_Vr (2 tokens) (8 tokens = R) ↓ ↓ quantize full precision per-token sliding window
Plain Text
복사
•
[Step C] XVgX_{V_g} per-token 양자화
Q(XVg)←GroupQuant(XVg,dim=token,numGroup=d//G)Q(X_{V_g}) \leftarrow \text {GroupQuant}(X_{V_g}, \text {dim=token}, \text {numGroup}=d//G)
◦
한 토큰의 8차원 벡터를 2개 그룹으로 쪼갬
◦
각 그룹 (4차원씩) 안에서 (zz, ss)를 계산함
◦
즉 per-token이지만, 토큰 한 개 안에서도 hidden dim을 G 크기의 sub-group으로 나눠 그룹별 (zz, ss)를 사용함
•
[Step D] Key의 grouped/residual 분할 및 양자화 (KeyQuant 호출)
Algorithm 1. KeyQuqant
r=l%Rr = l \% R XKg=XK[:l−r],XKr=XK[l−r:]X_{K_g} = X_K[: l - r], X_{K_r} = X_K[l - r :] Q(XKg)←GroupQuant(XKg,dim=channel,numGroup=l//G)Q(X_{K_g}) ← GroupQuant(X_{K_g}, dim=channel, numGroup=l//G)
◦
l=10l = 10, R=8R = 8이므로: r=2r = 2
◦
XKg=XK[:8]X_{K_g} = X_K[:8] (앞 8토큰), XKr=XK[8:]X_{K_r}=X_K[8:] (뒤 2토큰)
◦
Key는 per-channel 양자화이기에 여러 토큰에 걸쳐 같은 채널의 값들을 모아야 (zz, ss)를 계산할 수 있음
[Decoding Phase]
Fig 6. The overview of KIVI algorithm
Algorithm 1. Decoding Phase
•
XKrX_{K_r} 크기 = RR : Key residual이 RR에 도달하면, 전체가 per-channel 양자화되어 grouped 쪽에 이동하고 residual은 비워진다.
•
XVrX_{V_r} 크기 > RR : Value residual은 매 step마다 1개씩 초과되며, 이때 가장 오래된 1개가 per-token 양자화되어 grouped로 이동한다.

3.2 System-Level 구현

2비트로 압축된 데이터를 다시 원래의 정밀도로 되돌리는 Dequantization 과정은 추가적인 연산 시간이 필요하다. KIVI는 이를 해결하기 위해 Fusion 기술을 사용하여 역양자화와 행렬 곱셈을 별도의 단계로 나누지 않고 하나의 연산 단위(Q_MatMul)로 합쳐서 실행함으로써 데이터 로딩 시간을 줄이고 속도를 높였다.
KIVI는 가중치 양자화 방식과 동시에 사용할 수 있어, 모델 자체의 크기와 KV Cache의 크기를 모두 줄이는 ‘이중 압축’이 가능하다.

4. Experiments

4.1 실험 세팅

•
Models: Llama/Llama-2-7B, 13B (MHA), Falcon-7B (MQA), Mistral-7B (MHA)
•
Hyperparams: G=32G = 32, R=128R = 128 (최적 값들)
•
Tasks:
◦
LM-Eval (Normal context)
LM-Eval : 다양한 생성 및 질의응답 성능을 평가하기 위한 프레임워크
▪
CoQA (Exact match accuracy), TruthfulQA (BLEU), GSM8K (Exact match accuracy)
◦
LongBench (Long context)
LongBench : Long context 이해능력을 측정하기 위해 고안된 다중태스크 벤치마크
▪
Single-Documents QA tasks: Qasper (F1 score)
▪
Summarization tasks: QMSum (ROUGE score), MultiNews (ROUGE score)
▪
Few-shot Learning tasks: TREC (classification score), TriviaQA (F1 score), SAM-Sum (ROUGE score)
▪
Code Completion task: LCC (similarity score), RepoBench-P (similarity score)
◦
Needle-in-a-Haystack (NIAH) (Retrieval)

4.2 Accuracy Results

4.2.1 LM-Eval (Normal context)
Table 4. Performance comparison between 16bit, 4-bit per-token quantization, four fake 2bit KV cache quantization similar to those in Table 1, KIVI-2 (2bit) / KIVI-4 (4bit) across various models
Llama/Mistral 계열(MHA) 모델은 2-bit에서도 2% 이내로 하락했다. Falcon-7B는 MQA로 이미 KV cache가 head 수만큼 압축된 상태라 2-bit 양자화의 정확도가 크게 떨어지는 것을 알 수 있다. 이를 통해 MHA vs MQA에 따라 양자화 bit 선택이 달라져야 한다는 것을 알 수 있다.
4.2.2 LongBench (Long context)
Table 5. Performance evaluation of KIVI on various models across a range of benchmarks in LongBench
긴 context에서 KV cache 압축의 이점이 크게 보인다. 8개 task 평균 기준 (16bit → KIVI-2):
•
Llama-2-7B: 44.52 → 44.27, 0.56% 감소
•
Llama-2-13B: 44.85 → 44.69, 0.36% 감소
•
Llama-2-7B-Chat: 45.95 → 45.67, 0.61% 감소
•
Llama-2-13B-Chat: 45.96 → 45.52, 0.96% 감소
•
Falcon-7B: 8.71 → 7.95, 8.73% 감소
•
Mistral-7B: 46.58 → 45.85, 1.57% 감소
4.2.3 NIAH (Needle-in-a-Haystack)
Fig 7. Needle-in-a-Haystack results on Llama-3-8B-Instruct and Mistral-7B-Instruct-v0.2
Needle-in-a-Haystack은 ‘건조더미에서 바늘 찾기’라는 뜻으로, 방대한 양의 텍스트 사이에 아주 작은 특정 정보를 숨겨두고 모델이 이를 정확히 찾을 수 있는지(Retrieval)를 측정하는 실험이다.
Llama-3-8B-Instruct와 Mistral-7B-Instruct-v0.2에서 KIVI-2와 KIVI-4는 Baseline과 거의 동일한 retrieval 패턴을 보인다. Baseline과 KIVI 둘 다 찾아내지 못한 부분이 같다는 말은 KIVI로 인한 성능 저하가 원인이 아니라 Baseline 모델 자체의 고유한 한계점이기 때문이다.

4.3 Efficiency Result

Fig 8-1. Memory usage comparison between 2bit KIVI and 16bit baseline
Fig 8-2. Throughput comparison between 2bit KIVI and 16bit baseline
단일 NVIDIA A100 80GB를 사용해서 Llama-2-7B 모델로 메모리 사용량과 처리량에 대한 비교 실험을 진행했다. 같은 최대 메모리 조건에서 Batch 크기는 최대 4배 이상 확장이 가능했고 처리량은 2.35배에서 3.47배 증가했다.

4.4 Ablation study

Table 6. Ablation study of KIVI by changing group size GG and residual length RR
[Group size 영향]
R은 128로 고정하고, Llama-2-13B 모델로 GSM8K를 활용했다. G = 32, 64는 유사한데, G = 128에서 급락했다. Group 크기가 커질수록 한 그룹 내에 매우 큰 값(Outlier)과 작은 값이 섞일 확률이 높아져 양자화 오차가 심해진다. 따라서 적절한 그룹 크기를 선택하는 것이 정확도 유지에 필수적이다.
[Residual length 영향]
R = 32에서 R = 128까지 변화시켰을 때 성능이 크게 변하지 않는다. 적절한 길이의 Residual을 유지하는 것은 GSM8K와 같은 복잡한 추론 작업에서 매우 중요하므로, 완전히 양자화하는 것보다 최근 토큰들을 Full Precision으로 두는 슬라이딩 윈도우 전략을 사용한다.

5. 논문에 대한 생각

KIVI 논문을 읽으며 느낀 점들을 Pros, Cons, 그리고 연구 방향으로 정리할 수 있을 것 같다.

5.1 Pros

KIVI는 KV cache를 2-bit까지 압축하면서도 정확도를 거의 유지할 수 있음을 보여주었다. 그리고 “Key와 Value는 근본적으로 다른 통계적 특성을 가지므로, 서로 다른 차원(per-channel vs per-token)으로 비대칭 양자화해야 한다”는 인사이트를 체계적인 분석을 통해 증명했다. 이 인사이트는 이후 KVQuant(NeurIPS), KVTC(NVIDIA), TurboQuant(Google) 등 많은 후속 연구에 채택되어 KV cache 양자화 분야의 핵심 원칙 중 하나로 자리 잡았다.

5.2 Cons

Table 4에서 Llama-2-7B의 TruthfulQA에서 16-bit baseline이 30.76(높을수록 좋음)인데 KIVI-2는 33.95를 기록한다. Falcon-7B의 TruthfulQA에서도 23.20 → 24.98로 개선된다. 양자화를 했는데, 오히려 성능이 좋아진다…? 저자들이 이 현상에 대해 명시적인 분석이나 해석을 제공하지 않은 점이 아쉽다.
가장 아쉬운 부분은 양자화의 시스템 오버헤드에 대한 분석이 부재하다는 점이다. KIVI는 양자화하는 token 또는 channel 내에서 Subgroup을 여러 개 만들어서 최대한 촘촘하게 양자화하여 성능을 방어했다. 역양자화할 때, 양자화된 각 Subgroup마다 Scaling factor와 Zero point가 필요한데, 이를 저장하는 비용을 무시할 수 없다고 생각한다. 또한 매 attention 계산 시 역양자화가 실시간으로 발생하는데, 이 오버헤드가 전체 latency에 미치는 영향이 정량적으로 분석되지 않았다.

5.3 연구 방향

KIVI를 비판적으로 읽는 과정에서, 다음과 같은 부분들을 고려해야 한다는 생각을 했다.
시스템 오버헤드를 우선순위에 두는 설계
실제 온-디바이스 환경에서는 Scaling factor와 zero-point 저장 비용, 역양자화의 latency 등이 정확도나 처리량만큼이나 중요한 지표가 된다. 앞으로 내가 하려는 연구에서는 시스템 레벨의 오버헤드를 처음부터 고려하면서 아키텍처를 설계할 계획이다.
Attention Sparsity 특성을 활용한 동적 정밀도 적용
KIVI 논문을 읽으면서 크게 배운 개념 중 하나는 Attention Sparsity이다. 이는 일종의 “AI 버전의 80/20 법칙”이다. 모든 토큰이 동등하게 중요한 것이 아니라, 소수의 중요 토큰이 attention의 대부분을 차지한다는 것이다.
이 인사이트를 양자화에 적용하여 토큰의 중요도에 따라 정밀도를 동적으로 적용하면 어떨까 하는 생각이 들었다.
“50대의 추교현이 20대의 추교현에게 감사할 수 있도록 하루하루 최선을 다해 살고 있습니다.”
The End.