KV Cache is 1 Bit Per Channel: Efficient Large Language Model Inference with Coupled Quantization

논문 원본: https://arxiv.org/abs/2405.03917

KV Cache is 1 Bit Per Channel: Efficient Large Language Model Inference with Coupled Quantization

 

---

1. Introduction

대규모 언어모델(LLM)은 뛰어난 성능을 보이지만, 추론 시 KV Cache 크기 증가가 GPU 메모리와 대역폭의 주요 병목이 된다.
특히 긴 context 또는 큰 batch로 inference할 때, attention 계산 자체보다 과거 토큰들의 K/V를 메모리에서 읽는 비용이 더 크다.
이 때문에 LLM inference는 종종 메모리 바운드가 된다.

 

기존 KV 압축 방식(채널 단위/토큰 단위 양자화)은 4bit 수준까지는 성능을 유지하지만,
1–2 bit 저비트 영역에서는 모델 품질이 급격히 붕괴하는 한계가 있다.

 

본 논문은 K/V embedding을 분석한 결과,
여러 채널 간에 높은 상관관계와 mutual information이 존재한다는 사실을 발견했다.
즉, 채널을 독립적으로 quantization하는 기존 방식은 구조적으로 비효율적이라는 의미다.

이에 따라 논문은 채널을 묶어서 압축하는 Coupled Quantization(CQ) 기법을 제안한다.
CQ는 기존 방식과 달리 여러 채널을 함께(jointly) quantize하여 정보 손실을 줄이고,
심지어 1bit 수준에서도 모델 품질을 유지할 수 있음을 보여준다.

2. Background

2.1 KV Cache와 Attention 구조

Decoder-only LLM의 attention에서는 현재 토큰을 계산할 때
과거 모든 토큰의 K/V embedding을 읽어야 한다.

context 길이·batch 크기·모델 크기가 증가할수록 KV Cache 메모리는 선형 증가하며,
이는 GPU 메모리를 빠르게 소모한다.

2.2 Von Neumann Bottleneck

Attention의 병목은 FLOPs가 아니라 global memory read다:

• KV Cache는 토큰별로 완전히 별도로 저장됨
• 연산보다 메모리 접근 속도가 훨씬 느림
• 따라서 attention은 메모리 I/O에 의해 지연됨

따라서 KV Cache를 압축하면
(1) 메모리 사용량 감소 + (2) inference latency 개선이라는 두 가지 효과를 얻는다.

2.3 기존 Channel-wise Quantization의 한계

기존 방식은 채널마다 독립적으로 centroids를 학습해 값을 양자화한다.
하지만 이 방식은 채널 간 상관성을 고려하지 않아 극저비트(1–2bit 정도)에서는 심각한 성능 저하가 발생한다.

논문의 분석 결과,

• K/V의 각 채널은 서로 강한 상관관계를 갖고 있고,
• marginal entropy의 합보다 joint entropy 증가율이 훨씬 낮다.

즉, 독립 채널 quantization은 정보이론적으로 비효율적이다.

3. Methodology

3.1 동기

논문은 WikiText-2 기반으로 K/V 채널 분포를 분석한 결과:

• Pearson correlation이 전반적으로 높음
• mutual information이 큼
• joint entropy 증가율 < marginal entropy 합 증가율

이를 통해,
여러 채널을 묶어서 인코딩하는 것이 정보 측면에서 더 효율적임을 보였다.

3.2 Coupled Quantization(CQ)

CQ의 핵심 아이디어는 여러 채널을 하나의 그룹으로 묶고,
그 그룹 전체에 대해 하나의 quantization code를 사용하는 것이다.

1. 왼쪽: Channel-wise Quantization 왼쪽 그림은 전통적인 채널 독립 quantization을 적용한 경우다. channel 0과 channel 1 각각에 별도 centroids를 학습 각 채널이 독립적으로 1bit(2개 centroid)를 가짐 원본 데이터(파란 점)의 대부분이 비선형적이고 기울어진 분포를 보이지만 channel-wise 방식은 이를 “가로/세로 축 방향으로만” 나누므로 구조를 제대로 잡지 못한다 2. 오른쪽: Coupled Quantization 오른쪽 그림은 제안한 CQ를 적용한 경우다. channel 0과 channel 1을 묶어서 2차원 벡터 단위로 quantize 두 채널이 동시에 하나의 centroid로 매핑됨 centroid도 2차원 위치로 존재 (총 4개 centroid = 2bit per 2 channels) 실제 분포(파란 점)의 곡선형 구조를 훨씬 더 잘 따라감 이 때문에 quantized된 값(빨간 X)은 원본 분포 근처에 정확히 위치하며, quantization error는 250.6으로 크게 감소한다.

기본 구조

• 연속된 c개의 채널을 묶어 그룹 작성
• 해당 그룹 전체를 1개의 centroid로 quantize
• 그룹 단위 bit 수 = b bits
• notation: CQ-cC-bB
  • 예: CQ-4c8b → 4채널 그룹, 8bit 코드 → 채널당 2bit와 동일한 저장효율

• 채널 간 상관성을 활용 → 정보 손실 최소화
• 기존 1bit 실패 구간에서도 높은 모델 품질 유지

3.3 Centroid Learning

CQ centroid 학습은 두 방식이 존재한다:

(1) Uniform K-means

기본적인 multi-dimensional k-means로 centroid 생성.

(2) Fisher-guided K-means

LLM의 gradient 정보를 이용해
더 중요한 activation에 더 높은 정밀도(weighting)를 주도록 centroid를 학습.

→ 실험에서 perplexity 감소 효과가 더 큼.

4. Experiments

1) Perplexity 측면

• 4bit: CQ가 INT/NF와 동급 또는 더 좋음
• 2bit: CQ는 성능 유지, 타 방법은 붕괴
• 1bit:
  • 기존 방법: 성능 붕괴 또는 NaN
  • CQ(8채널 coupling):
    • LLaMA-7B → perplexity 8.09
    • FP16 baseline 5.68과 근접한 수준 유지

2) Downstream task 정확도

CQ는 low-bit 상황에서도
KVQuant(비-sparse)보다 항상 안정적이고,
KVQuant(1% sparse)와도 동급 이상 성능.

3) Ablation

• coupling 채널 수 증가 → perplexity 안정적 개선
• Fisher-guided centroids → 추가 성능 향상

4) Overhead

• centroid 학습 시간은 GPU 기반으로 처리 가능
• centroid 저장량은 모델 파라미터 대비 <1%

 

-> 

이 논문은 KV cache 병목을 해결하기 위해
채널 간 상관관계를 활용한 Coupled Quantization을 제안하고,
이를 통해 기존 방식이 불가능하던 1bit 수준의 극단적 압축에서도 모델 품질을 유지함을 입증한다.
실험 결과 또한 다양한 모델·벤치마크에서 CQ가 일관적으로 우수함을 보여준다.