🤖 machine learning

SoftWater: Class-Aware Rate Allocation for Softmax Quantization

SoftWater는 빈도가 높고 분산이 낮은 클래스에는 더 세밀한 그리드를 할당하고 희귀한 클래스에는 더 거친 그리드를 할당하도록 KL-발산 기반의 레이트-왜곡 문제를 해결함으로써, 소형 LLM의 저장 오버헤드를 크게 줄이는 동시에 퍼플렉시티 저하를 최소화하는 소프트맥스 양자화를 위한 클래스 인지형 레이트 할당 방식입니다.

원저자: Joao V. Cavalcanti, Ashia C. Wilson

게시일 2026-08-13

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Joao V. Cavalcanti, Ashia C. Wilson

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

거대한 도서관의 책들을 아주 작은 배낭에 담기 위해 크기를 줄이려고 노력한다고 상상해 보세요. 인공지능의 세계에서 이 "책들"은 글을 쓰고, 대화하고, 문제를 해결하도록 훈련된 거대한 언어 모델, 즉 컴퓨터 뇌입니다. 이 모델들을 여러분의 스마트폰이나 작은 노트북에서 실행할 수 있을 만큼 빠르게 만들기 위해, 과학자들은 **양자화(quantization)**라는 기술을 사용합니다. 이것은 고화질 영화를 저해상도의 픽셀화된 버전으로 변환하는 것과 비슷하다고 생각하면 됩니다. 디테일은 일부 손실되지만, 파일 크기는 극적으로 줄어들어 휴대하기 쉬워집니다.

보통 과학자들이 이 AI 모델들을 줄일 때는 모델의 모든 부분을 동일하게 취급합니다. 맥락을 이해하는 부분인 "몸통(body)"과 다음에 올 단어를 결정하는 부분인 "머리(head)"를 똑같이 작은 상자에 구겨 넣습니다. 하지만 여기에는 함정이 있습니다. 현대의 AI에서 "머리"는 놀라울 정도로 무겁습니다. 때로는 전체 모델 메모리의 거의 3분의 1을 차지하기도 합니다! 만약 모델 전체를 아주 작게 압축하면서도 머리는 원래의 무거운 형식을 그대로 유지한다면, 실제로 공간을 별로 아끼지 못하는 셈입니다. 이는 옷은 압축해서 넣었지만, 무거운 부츠는 원래의 커다란 상자째로 캐리어에 넣는 것과 같습니다. 우리가 살펴보고 있는 논문은 바로 이 특정 문제, 즉 AI의 "단어 선택 머리"를 버벅거리거나 로봇처럼 들리지 않게 하면서 줄이는 방법을 다룹니다.

이 논문의 저자인 MIT의 Joao V. Cavalcanti와 Ashia C. Wilson은 머리를 줄이는 기존 방식이 마치 호두를 깨는 데 망치를 사용하는 것과 같다는 점을 깨달았습니다. 그들은 이 문제를 율-왜곡(rate-distortion) 문제로 접근했습니다. 여러분에게 AI의 선택을 설명하는 데 사용할 수 있는 "비트(bits)"(디지털 저장 용량이라는 화폐 단위)라는 한정된 예산이 있다고 상상해 보세요. 기존 방식은 모든 단어에 동일한 양의 주의를 기울이는 것처럼, 이 비트를 균등하게 소비했습니다. 즉, 모든 단어에 똑같은 양의 관심을 주는 식이었죠. 하지만 저자들은 AI가 단어를 균등하게 사용하지 않는다는 점에 주목했습니다. AI는 "the"나 "and" 같은 흔한 단어는 끊임없이 사용하지만, "xylophone(실로폰)"이나 "quintessential(전형적인)" 같은 희귀한 단어는 가끔씩만 사용합니다.

이 논문은 SoftWater라고 불리는 새로운 방법을 소개합니다. SoftWater는 모든 단어에 동일한 저장 공간을 주는 대신, 스마트한 스프링클러 시스템처럼 작동합니다. 자주 등장하고 변동성이 낮은 단어(AI가 항상 사용하고 항상 비슷한 방식으로 사용하는 단어들)에는 비트의 미세하고 고해상도인 안개를 뿌립니다. 반면, 드물고 예측 불가능한 단어들에는 더 거칠고 넓은 분무를 사용합니다. 이러한 "클래스 인식(class-aware)" 접근 방식 덕분에, AI는 동일한 총 공간을 사용하면서도 흔한 단어들은 완벽하게 기억하고, 희귀한 단어들에 대해서는 약간의 모호함을 수용할 수 있게 됩니다.

연구진은 이 방법을 10억 개의 파라미터 규모의 작은 모델부터 320억 개의 파라미터를 가진 거대 모델까지 포함된 5가지 서로 다른 AI 모델에 테스트했습니다. 그 결과, SoftWater가 이전의 최고 방법이었던 WaterSIC보다 훨씬 뛰어나다는 것을 발견했습니다. 60번의 테스트 중 59번에서, SoftWater는 더 자연스럽게 들리고 실수를 덜 하는(KL 발산, 즉 AI의 선택이 원래 모델과 얼마나 다른지를 나타내는 지표로 측정됨) 모델을 만들어냈습니다.

여기 마법의 숫자가 있습니다. 10억 개의 파라미터를 가진 모델의 "머리"를 가중치당 단 2비트로 압축했을 때, SoftWater는 모델의 전체 크기를 **45%에서 60%**까지 줄이면서도 AI를 아주 약간 덜 완벽하게 만들 뿐이었습니다(혼란도, 즉 퍼플렉시티(perplexity)의 증가율은 2.9%에서 3.7% 사이였습니다). 이와 대조적으로, 기존 방식은 동일한 공간을 사용했을 때 AI를 훨씬 더 혼란스럽게 만들었습니다. 더욱 놀라운 점은, 4비트 머리를 사용했을 때 모델은 원래의 고품질 버전과 구별이 거의 불가능할 정도였다는 것입니다.

또한 이 논문은 이 방법이 매우 유연하다는 것을 발견했습니다. 만약 AI가 법률 문서를 작성하거나 코딩을 하는 것과 같은 특정 작업을 수행할 것이라는 점을 알고 있다면, 그 특정 주제에 맞춰 스프링클러 시스템을 "교정(calibrate)"할 수 있습니다. 이를 통해 AI가 해당 분야에서 가장 필요한 단어들에 대해 더욱 날카로운 성능을 유지하도록 보장할 수 있습니다. 저자들은 교정 데이터를 실제 사용 사례와 일치시킴으로써 AI의 성능이 더욱 향상된다는 것을 보여주었습니다.

요약하자면, SoftWater는 단순히 AI를 줄이는 것이 아니라, 지능적으로 줄입니다. 이 방식은 모든 단어가 평등하게 만들어진 것이 아님을 깨닫고, 그에 따라 저장 공간을 할당합니다. 이렇게 함으로써, 강력하고 고품질인 AI 모델을 이전에는 감당할 수 없었던 기기에서도 실행할 수 있게 하며, 모델을 처음부터 다시 훈련할 필요도 없이 말입니다. 이는 AI를 더 가볍고, 빠르고, 현실 세계에서 바로 사용할 수 있게 만드는 실용적이고 수학적으로 탄탄한 방법입니다.

기술 요약: SoftWater – Softmax 양자화를 위한 클래스 인지적 비율 할당 (Class-Aware Rate Allocation)

1. 문제 정의

대규모 언어 모델(LLM)을 위한 사후 학습 양자화(PTQ) 파이프라인은 일반적으로 마지막 소프트맥스 출력 레이어(즉, "head")를 고정밀도(예: FP16)로 남겨둡니다. 표준 선형 레이어는 가중 평균 제곱 오차(WMSE)를 최소화하도록 양자화되는 반면, 출력 레이어는 확률 분포를 생성하므로 쿨백-라이블러(KL) 발산(divergence)이 더 적절한 왜곡 지표가 됩니다.

이러한 누락은 어휘 사전(vocabulary)이 큰 현대의 소형 LLM에서 점점 더 큰 비용을 초계합니다. 논문의 표 1에서 보여주듯, 출력 헤드는 모델 전체 파라미터의 15~30%를 차지할 수 있습니다. 결과적으로, 헤드를 FP16으로 유지하는 "2비트" 모델은 명목상의 비트율보다 훨씬 더 많은 비트를 가중치당 실제로 저장하게 됩니다. 이를 해결하려는 기존의 시도들은 헤드 구조 자체를 완전히 교체하거나, 소프트맥스 함수의 특정 기하학적 구조와 비균일한 토큰 분포(Zipfian 분포)를 고려하지 못하는 표준 WMSE 기반 양자화를 적용합니다.

핵심 문제는 소프트맥스 레이어 양자화를 KL 발산 지표 하에서의 레이트-왜곡(rate-distortion) 문제로 정식화하여, 원래의 분포와 양자화된 분포 사이의 발산을 최소화하기 위해 고정된 비트 예산을 가중치 행렬 전체에 어떻게 할당할 것인지 결정하는 것입니다.

2. 방법론: SoftWater

저자들은 KL 발산의 2차 분석을 기반으로 한 클래스 인지적 비율 할당 전략을 도출하는 방법인 SoftWater를 제안합니다.

목적 함수 유도:
왜곡은 원래 분포 $p$ 와 양자화된 분포 $q$ 사이의 기대 KL 발산으로 정의됩니다. 로짓 섭동(logit perturbation) $\delta$ 에 대한 KL 발산의 2차 테일러 전개를 사용하여, 저자들은 다음과 같은 오차 지표를 도출합니다:
$D \approx \frac{1}{2} \text{vec}(\Delta)^\top \mathbb{E}[(\nabla^2 f(0) \otimes XX^\top)] \text{vec}(\Delta)$
여기서 $\nabla^2 f(0) = \text{diag}(p) - pp^\top$ 는 소프트맥스 함수의 헤시안(Hessian)이며, $XX^\top$ 는 입력 피처 공분산(input feature covariance)을 나타냅니다. 이는 클래스( $K$ )와 피처( $n$ )를 결합하는 $Kn \times Kn$ 오차 가중치 행렬을 생성합니다.

다루기 쉬운 근사화:
$Kn \times Kn$ 행렬을 직접 인수분해하는 것은 계산적으로 불가능합니다. 저자들은 두 가지 핵심 근사화를 도입합니다:

디더링(Dither) 논거: 감산적 디더링(subtractive dither)을 가정하면 오차들이 행(row) 간에 상관관계가 없어지므로, 헤시안의 비대각 항(off-diagonal terms)을 무시할 수 있습니다.
분리 가능성(Separability) 가정: 저자들은 소프트맥스 곡률 $\lambda$ 와 입력 공분산 $XX^\top$ 가 서로 상관관계가 없다고 가정합니다. 이를 통해 기댓값을 다음과 같이 인수분해할 수 있습니다: $\mathbb{E}[\lambda \otimes XX^\top] \approx \mathbb{E}[\lambda] \otimes \mathbb{E}[XX^\top]$ .

이 인수분해를 통해 문제는 $K$ 개의 독립적인 $n \times n$ 촐레스키 인수분해(또는 하나의 재조정된 $n \times n$ 인수분해)로 축소되어 계산이 가능해집니다. 그 결과 도출된 클래스 $k$ 와 피처 열 $i$ 에 대한 그리드 간격은 다음과 같습니다:
$\alpha_{(k-1)n+i} \propto \frac{1}{\bar{\lambda}_k^{1/2} |\ell_{ii}|}$
여기서 $\bar{\lambda}_k = \mathbb{E}[p_k(1-p_k)]$ 는 기대 곡률(클래스 빈도 및 분산과 관련됨)이며, $\ell_{ii}$ 는 입력 공분산의 촐레스키 인수의 대각 성분입니다.

주요 알고ographic 특징:

클래스 인지적 할당: 이 방법은 빈도가 높고 분산이 낮은 클래스에는 더 미세한 그리드(더 많은 비트)를 할당하고, 빈도가 낮고 분산이 높은 클래스에는 더 거친 그리드를 할당합니다. 이는 Zipfian 특성을 활용합니다.
스무딩 사전 확률(Smoothing Prior): 캘리브레이션 데이터에 존재하지 않는 클래스(여기서 $\bar{\lambda}_k \to 0$ )에 대한 불안정성을 방지하기 위해, 캘리브레이션 분포를 균등 분포와 혼합하는 스무딩 사전 확률을 적용합니다. 이는 최대 그리드 간격을 제한하여 테일러 전개가 유효하게 유지되도록 합니다.
구현: 알고리즘은 선형 레이어를 위한 WaterSIC 알고리즘과 유사하게 클래스별 재조정(rescaling)을 포함하는 순차적 간섭 제거(Successive Interference Cancellation, SIC)를 사용하여 인코딩합니다. 이는 필요한 통계량( $\Sigma_X$ 및 $\bar{\lambda}$ )을 계산하기 위해 캘리브레이션 데이터에 대한 단 한 번의 순방향 패스(forward pass)만을 필요로 합니다.

3. 주요 기여

문제 정식화: 본 논문은 소프트맥스 양자화를 KL 발산 하에서의 레이트-왜곡 문제로 프레임화하고, 클래스 통계와 피처 공분산을 결합하는 유도된 오차 지표를 도출했습니다.
SoftWater 알고리즘: 분리 가능성 근사를 활용하여 복잡도를 $O(K^3n^3)$ 에서 $O(n^3)$ (및 선형 패스)으로 줄이는 실용적인 알고리즘을 제안합니다. 이는 클래스 빈도 및 교차 컨텍스트 가변성에 따라 격자 그리드를 적응시키는 클래스 측 스케일링 인자를 도입합니다.
경험적 검증: Llama 및 Qwen 제품군의 5개 모델(1B ~ 32B 파라미터)을 통해 방법론을 검증했습니다.

4. 실험 결과

저자들은 SoftWater를 WMSE를 최적화하는 WaterSIC 양자화기 및 풀 프리시전(full-precision) 베이스라인과 비교 평가했습니다.

헤드 전용 양자화(Head-Only Quantization): 헤드만 양자화하고 바디는 FP16으로 유지할 때, SoftWater는 다양한 비트율과 모델에 대해 60개 중 59개의 테스트 포인트에서 WaterSIC보다 우수한 성능을 보였습니다. 2비트에서 SoftWater는 WikiText2에 대해 WaterSIC 대비 KL 발산을 6.5배에서 8.3배까지 감소시켰습니다.
전체 양자화 모델: 바디가 양자화된(GuidedQuant 사용) Llama-3.2-1B-Instruct 모델에서, 2비트 SoftWater 헤드는 총 저장 바이트의 **45~~60%**를 제거하면서도 퍼플렉시티(perplexity) 증가를 **2.9~~3.7%**로 제한했습니다. 반면, WaterSIC 베이스라인은 유사한 압축률에서 9.7~17.6%의 퍼플렉시티 증가를 보였습니다.
도메인 적응: 논문은 클래스 측 통계( $\bar{\lambda}$ )가 캘리브레이션 데이터로부터 도출되기 때문에, 캘리브레이션 도메인을 배포 도메인과 일치시키면 해당 특정 도메인에서 가장 낮은 KL 발산을 얻을 수 있다고 언급합니다.
분리 가능성: 실험을 통해 분리 가능성 가정이 잘 작동함을 확인했으며, 프록시(proxy)가 실제 왜곡을 최대 10% 미만으로 과소평가함을 입증했습니다.

5. 의의 및 주장

본 논문은 SoftWater가 헤드가 모델 크기의 상당 부분을 차지하는 모델에서 헤드 양자화를 실용적으로 만든다고 주장합니다. 헤드를 WMSE가 아닌 KL 하의 레이트-왜곡 문제로 취급함으로써, 이 방법은 4비트에서 거의 손실 없는 성능을, 2비트에서 최소한의 퍼플렉시티 저하와 함께 상당한 압축을 달로합니다.

저자들은 자신들의 이득이 오로지 레이트 할당 전략 덕분임을 강조합니다. 이는 WaterSIC 파이프라인에 존재하는 추가적인 정교화 작업들(예: 활성화 드리프트 수정, 잔차 스트림 수정)을 사용하지 않았음에도 불구하고 달성된 결과이며, 이러한 작업들은 레이트 할당과는 독립적인(orthogonal) 요소입니다. 이 방법은 기존 PTQ 파이프라인의 헤드 양자화 단계에 바로 적용 가능한 드롭인 교체(drop-in replacement) 방식으로 제시되며, 구조적 변경이나 파인튜닝 없이도 모델 크기를 획기적으로 줄일 수 있습니다.