기술 요약: RATEQUANT

문제 정의

대형 언어 모델 (LLM) 은 생성 과정에서 이전에 계산된 모든 키 - 값 (KV) 쌍을 캐싱해야 하므로, 시퀀스 길이, 배치 크기, 모델 깊이에 비례하여 선형적으로 증가하는 메모리 병목 현상이 발생합니다. KV 캐시를 양자화하면 이 메모리 사용량을 줄일 수 있지만, 기존 방법들은 모든 어텐션 헤드에 동일한 비트 너비를 균일하게 할당합니다. 이러한 접근 방식은 헤드의 중요도에 대한 잘 문서화된 이질성을 무시합니다.

레이어 또는 채널 수준에서 혼합 정밀도 할당을 시도한 최근 연구들은 특정 양자화와 연결된 휴리스틱 규칙에 의존하며, 이론적 보장과 전이성을 결여하고 있습니다. 저자들은 더 깊고 이전에 인식되지 않았던 장애물을 확인했습니다: 왜곡 모델 불일치. 서로 다른 양자기는 $D(b) = \alpha \beta^{-b}$로 모델링된 고유한 왜곡 - 비트율 곡선을 따릅니다. 감쇠율 $\beta$는 설계에 따라 크게 달라집니다 (예: TURBOQUANT 의 3.6 에서 QuaRot 의 5.3 까지). 한 양자기의 왜곡 모델에서 도출된 혼합 정밀도 할당 전략을 다른 양자기에 무작정 적용하면 어텐션 헤드의 한계 이득 순서가 반전되어, 종종 균일 양자화보다 성능이 저하되는 결과를 초래합니다.

방법론

RATEQUANT 는 이러한 도전 과제를 해결하기 위해 헤더별 KV 캐시 비트 할당을 왜곡 - 비트율 최적화 문제로 공식화합니다. 이 프레임워크는 네 단계로 구성됩니다:

1. 문제 공식화 및 최적 할당

저자들은 헤더별 양자화 평균 제곱 오차 (MSE) 를 비트 너비의 지수 함수로 모델링합니다: $D(b) = \alpha \beta^{-b}$. 목표는 총 비트 예산 $B$를 제약 조건으로 하여 $N$개의 양자화 그룹 (헤드) 에 걸친 가중치 왜곡의 합을 최소화하는 것입니다:
$$\min_{b} \sum_{i=1}^{N} w_i \cdot D(b_i) \quad \text{s.t.} \quad \sum b_i = B$$
여기서 $w_i$는 헤드 $i$의 민감도를 나타냅니다.

라그랑주 정류성을 사용하여 저자들은 **역수분 (reverse waterfilling)**을 통해 폐쇄형 해를 유도합니다:
$$b^*_i = \bar{b} + \frac{\ln w_i - \ln \bar{w}}{\ln \beta}$$
여기서 $\bar{b}$는 평균 비트 너비이고 $\bar{w}$는 민감도의 기하평균입니다. 이 해법은 민감도가 높은 헤드에 더 많은 비트를 할당합니다. 이론적으로 달성 가능한 왜곡 감소는 헤드 민감도의 산술평균과 기하평균 (AM/GM) 비율에 의해 상한이 결정됩니다.

2. 민감도 추정

가중치 $w_i$를 결정하기 위해 RATEQUANT 는 기반 민감도를 사용합니다. 이는 인과적 언어 모델링 손실에 대한 KV 투영 출력의 기울기 노름 제곱을 계산합니다:
$$w_{l,h} = \mathbb{E}_{x \sim \mathcal{D}} \left[ \frac{1}{T} \sum_{t=1}^{T} \left\| \frac{\partial \mathcal{L}}{\partial K_{l,h,t}} \right\|^2 \right]$$
저자들은 기반 기반 프록시가 오차 증폭을 측정하는 활성화 기반 지표 (예: 채널 크기) 보다 직접적인 손실 영향을 측정하므로 우월함을 입증했습니다. 기반 기반 민감도는 초기 및 후기 레이어에서 높은 민감도를 보이는 U 자형 중요도 패턴을 드러내는 반면, 활성화 기반 지표는 깊이에 따라 단조 증가하는 경향을 보입니다.

3. 정수 할당 및 확장

비트 너비는 정수여야 하므로, 연속 해법은 탐욕적 한계 이득 알고리즘(알고리즘 1) 을 사용하여 정제됩니다. 이 알고리즘은 예산이 소진될 때까지 가중치 왜곡 감소가 가장 큰 헤드에 비트를 반복적으로 할당합니다.

RATEQUANT 를 양자기 무관하게 만드는 두 가지 중요한 확장이 있습니다:

• 왜곡 보정: 보편적인 $\beta$를 가정하는 대신, RATEQUANT 는 소규모 보정 세트 (16 개 시퀀스) 를 사용하여 양자기별 왜곡 모델 ($\alpha, \beta$) 을 적합시킵니다. 이는 "왜곡 모델 불일치" 실패 모드를 방지합니다.
• K/V 분리: 키와 값이 종종 다른 양자화 방식을 사용한다는 점 (예: KIVI 에서 키는 채널별 대칭, 값은 토큰별 비대칭) 을 인식하여, 프레임워크는 이를 서로 다른 왜곡 곡선을 가진 $2N$개의 독립 구성 요소로 취급합니다. 이를 통해 알고리즘은 비대칭 비트 예산을 발견할 수 있습니다 (예: 값보다 키에 더 많은 비트 할당).

4. 파이프라인

이 과정은 일회성 오프라인 비용을 수반합니다:

1. 민감도 보정: 소규모 데이터셋에서 순전파 및 역전파 수행 (8B 모델의 경우 약 1.6 초).
2. 왜곡 모델링: 5 개의 비트 너비에서 $\alpha$와 $\beta$ 적합 (<0.1 초).
3. 할당: 탐욕 알고리즘 실행 (<0.01 초).
4. 추론: 정적 비트 너비 적용으로 런타임 오버헤드 제로.

주요 결과

실험은 WikiText-2 퍼플렉시티 (PPL) 와 하류 작업 (ARC-Challenge, HellaSwag 등) 을 사용하여 Qwen3 모델 (4B, 8B, 32B) 에서 수행되었습니다.

• 불일치 해결: 불일치된 왜곡 모델을 적용하는 것 (예: KIVI 에 TURBOQUANT 의 $\beta$ 사용) 은 2.5 평균 비트에서 KIVI 의 PPL 을 49.3 에서 87.0 으로 저하시켰습니다. 보정된 RATEQUANT 는 이를 해결하여 PPL 을 14.9로 감소시켰습니다 (70% 감소).
• 성능 향상:
  • KIVI (2.5 비트): RATEQUANT 는 PPL 을 49.3 에서 14.9 로 개선했습니다.
  • QuaRot (2.5 비트): PPL 을 6.6 포인트 개선했습니다.
  • TURBOQUANT (4.0 비트): Qwen3-8B 에 대해 양자화로 인한 저하의 66% 를 회복했습니다.
• K/V 분리: 2.5 비트 KIVI 의 경우, RATEQUANT 가 발견한 최적 할당은 키에 2.85 비트, 값에 2.15 비트로, 결합 할당보다 훨씬 우수한 성능을 보였습니다.
• 효율성: 이 방법은 런타임 오버헤드를 추가하지 않습니다. 보정은 단일 GPU 에서 2 초 미만 소요됩니다.
• 하류 정확도: 4.0 비트에서 RATEQUANT 는 하류 작업에서 64.2% 의 평균 정확도를 달성하여 FP16 기준 (64.4%) 과 거의 일치하며 정확도 격차의 89.8% 를 회복했습니다.

의의 및 주장

이 논문은 RATEQUANT 가 헤드 단위 세분성에서 KV 캐시 양자화에 왜곡 - 비트율 이론을 적용하고 키와 값 간의 구조적 차이를 해결한 최초의 프레임워크라고 주장합니다.

핵심 발견은 왜곡 모델 불일치가 혼합 정밀도 할당에서 치명적인 실패 모드라는 점입니다. 특정 기본 양자기에 맞춰 왜곡 곡선 ($\beta$) 을 보정하지 않으면 혼합 정밀도 전략이 해로울 수 있습니다. 다음을 결합함으로써:

1. 기반 기반 민감도(활성화 기반 프록시보다 우월),
2. 실증적 왜곡 보정(양자기에 맞춤),
3. K/V 분리(비대칭 오차 특성 처리),

RATEQUANT 는 혼합 정밀도 KV 양자화를 범용적이고 양자기 무관한 할당 레이어로 변환합니다. 저자들은 헤드 민감도의 AM/GM 비율이 혼합 정밀도 할당이 유의미한 이득을 가져올 시기를 예측하는 실용적 지표로 작용한다고 주장합니다. 이 연구는 왜곡 - 비트율 관점이 KV 캐시를 넘어 다른 이질적 신경망 구성 요소의 자원 할당에도 일반화될 수 있음을 시사합니다.