이 글은 딥러닝 시대의 데이터타입 편의 후속입니다. 앞 글에서 INT/FP 데이터타입이 비트를 어떻게 숫자로 해석하는지 봤다면, 이 글에서는 그 저비트 타입으로 이미 학습된 모델을 실제로 줄여보고 결과를 측정합니다. 코드는 전부 github.com/warpspaceinc/efficient-ml-practice의 노트북에서 가져왔고, 아래 plot은 그 코드를 Colab(T4)에서 그대로 돌려 뽑은 것입니다.
32비트짜리 숫자가 정말 다 필요할까
지난 글 끝에서 관찰한 게 하나 있었습니다. 신경망의 가중치는 대개 0 근처에 종 모양으로 몰려 있고 꼬리가 깁니다. 실제로 MNIST용 작은 MLP(784→256→128→10)를 2 에폭만 학습시킨 뒤 첫 번째 층 가중치를 히스토그램으로 그려보면 정확히 그렇습니다.
값이 이렇게 좁은 범위에 몰려 있는데, 정말 하나하나를 32비트 실수로 따로 저장하고 32비트로 곱해야 할까요? 양자화(quantization)는 “아니오"라고 답합니다. 이 글에서는 대표적인 두 갈래를 코드로 직접 돌려봅니다.
- K-Means 양자화: 비슷한 값을 대표값으로 묶어 저장을 줄인다 (비균등).
- Linear 양자화: 정수 격자에 매핑해 저장도 연산도 정수로 만든다 (균등).
두 방식 모두 학습이 끝난 모델에 사후 적용하는 PTQ(post-training quantization)이고, 여기서는 가중치만 양자화합니다.
1. K-Means 양자화: 대표값 공유로 저장 줄이기
첫 번째 아이디어는 2016년 Deep Compression 논문1에서 왔습니다. 비슷한 값의 가중치를 한 덩어리로 묶어 대표값 하나로 퉁치자는 것입니다. 예를 들어 2.09, 1.92, 1.87은 전부 “대략 2.0"이니, 이 값들을 2.00 하나로 대체하고 각 자리에는 “몇 번 대표값을 쓸지” 인덱스만 적어둡니다.
절차는 세 단계입니다.
- 클러스터링: 모든 가중치를 $k = 2^N$개 그룹으로 K-Means 클러스터링.
- 코드북(codebook): 각 그룹의 중심값(centroid)을 대표값으로 저장. 이 목록이 코드북.
- 인덱스: 각 자리엔 32비트 실수 대신 “몇 번 클러스터인지” 가리키는 $N$비트 정수만 저장.
코드로는 scikit-learn의 KMeans를 그대로 씁니다. 지역 최적해를 피하려고 min~max를 균등 분할한 값으로 초기화합니다.
from sklearn.cluster import KMeans
import numpy as np
def kmeans_quantize(w, n_bits, fit_sample=20000):
"""가중치를 K-Means로 양자화. 복원값, 코드북, 인덱스를 반환."""
shape = w.shape
flat = w.reshape(-1, 1)
k = 2 ** n_bits
lo, hi = float(flat.min()), float(flat.max())
init = np.linspace(lo, hi, k).reshape(-1, 1) # 선형 초기화
idx = np.random.choice(len(flat), min(fit_sample, len(flat)), replace=False)
km = KMeans(n_clusters=k, init=init, n_init=1, max_iter=50).fit(flat[idx])
codebook = km.cluster_centers_.reshape(-1) # 대표값 목록
indices = km.predict(flat) # 각 가중치 → 클러스터 인덱스
recon = codebook[indices].reshape(shape) # 복원
return recon, codebook, indices
2비트($k=4$)로 돌린 결과를 원본 분포 위에 겹쳐 그리면, 대표값(세로선)이 값이 몰린 0 근처에 촘촘히 놓이는 걸 볼 수 있습니다. 이게 “비균등(non-uniform)“의 의미입니다.
얼마나 줄어드나
압축비는 비트 회계로 계산합니다. 파라미터 $M$개, $N$비트 양자화면:
$$\text{원본} = 32M \text{ bit}, \qquad \text{압축} = \underbrace{N \cdot M}_{\text{인덱스}} + \underbrace{32 \cdot 2^N}_{\text{코드북}} \text{ bit}$$num = w.size; k = 2 ** N_BITS
orig_bits = 32 * num
comp_bits = N_BITS * num + 32 * k # 인덱스 + 코드북
print(f"compression: {orig_bits/comp_bits:.2f}x")
$M \gg 2^N$이면 코드북 항은 무시할 만해지고 압축비는 $32M/NM = \mathbf{32/N}$배로 수렴합니다. 2비트면 약 16배, 4비트면 8배입니다. (첫 번째 층 fc1 기준 실측 압축비: 15.99×)
중요한 한계: 저장만 줄인다
여기서 반드시 짚어야 합니다. K-Means 양자화가 줄이는 건 디스크·메모리에 저장되는 크기뿐입니다. 추론할 때는 인덱스를 코드북으로 디코드해서 원래 실수 가중치를 복원한 뒤 곱해야 합니다.
- 저장: 정수 인덱스 (작다) ✓
- 연산: 복원된 32비트 실수로 여전히 floating-point 연산 ✗
메모리 대역폭은 아꼈지만 곱셈기(MAC)는 여전히 float를 돌립니다. 진짜로 정수 연산까지 빠르게 하려면 다음 방식이 필요합니다.
2. Linear 양자화: 정수 연산만으로 추론하기
Linear(선형) 양자화는 정수와 실수를 아핀(affine) 매핑 한 줄로 잇습니다. 2018년 Jacob et al.2이 제안했고, TensorFlow Lite의 INT8 양자화가 바로 이 방식입니다.
$$r = S \cdot (q - Z)$$- $r$: 원래의 실수(real) 값
- $q$: 양자화된 정수
- $Z$: zero point. 실수 $0$에 정확히 대응하는 정수. ReLU 뒤의 0, 패딩처럼 흔한 값을 오차 없이 표현하기 위한 장치.
- $S$: scale. 정수 한 칸이 실수로 얼마인지 (floating-point).
$S$와 $Z$는 실수 범위 $[r_{\min}, r_{\max}]$의 양 끝을 정수 범위 $[q_{\min}, q_{\max}]$의 양 끝에 맞추는 데서 나옵니다.
$$S = \frac{r_{\max} - r_{\min}}{q_{\max} - q_{\min}}, \qquad Z = \text{round}\left(q_{\min} - \frac{r_{\min}}{S}\right)$$코드에서는 이 $S$, $Z$만 계산해주면 양자화(round·clamp)와 복원(dequant)은 PyTorch 내장 fake_quantize_per_tensor_affine이 처리합니다. 임의 비트 폭은 quant_min/quant_max로 지정합니다.
import torch
def linear_quantize(w, n_bits, symmetric=False):
"""PyTorch 내장 fake_quantize로 선형(아핀) 양자화. 복원값과 S, Z를 반환."""
t = torch.as_tensor(w, dtype=torch.float32)
if symmetric: # Z=0 고정 (연산이 단순)
qmax = 2 ** (n_bits - 1) - 1; qmin = -qmax
S = float(t.abs().max()) / qmax; Z = 0
else: # asymmetric: 범위를 정확히 맞춤
qmin = -2 ** (n_bits - 1); qmax = 2 ** (n_bits - 1) - 1
S = (float(t.max()) - float(t.min())) / (qmax - qmin)
Z = int(round(qmin - float(t.min()) / S))
recon = torch.fake_quantize_per_tensor_affine(t, S, Z, qmin, qmax)
return recon.numpy(), S, Z
2비트로 돌리면 대표값(세로선)이 일정한 간격 $S$로 균등하게 놓입니다. K-Means와 정확히 대비되는 지점입니다. 이쪽은 데이터가 어디 몰렸든 격자가 균일합니다.
왜 “정수 연산까지” 가능한가
Linear 양자화의 진짜 값어치는 저장이 아니라 연산에 있습니다. 행렬곱 $Y = WX$의 각 값을 아핀 매핑으로 바꾸면:
$$q_Y = \frac{S_W S_X}{S_Y}\big(q_W q_X - Z_W q_X - Z_X q_W + Z_W Z_X\big) + Z_Y$$괄호 안은 $q_W q_X$를 포함해 전부 정수 곱·합이고, $Z_W Z_X$처럼 입력과 무관한 항은 미리 계산해둘 수 있습니다. 앞의 스케일 비 $\frac{S_W S_X}{S_Y}$만 실수인데, 이 값은 경험적으로 항상 $(0,1)$ 구간이라 $2^{-n} M_0$ 꼴의 고정소수점 곱 + 비트 시프트로 처리됩니다. 결국 floating-point 유닛이 전혀 필요 없습니다. 1편의 에너지 표에서 본 “정수 연산이 float보다 수십 배 싸다"가 여기서 현실이 됩니다.
가중치는 대개 0 대칭이라 실전에서는 $Z_W = 0$인 대칭(symmetric) 양자화를 자주 씁니다. 그러면 $Z_W$ 관련 항이 통째로 사라져 식이 더 깔끔해집니다.
3. 그래서, 얼마나 잃나: 직접 측정
말로만 하면 감이 안 오니, 세 방식(Linear asymmetric / Linear symmetric / K-Means)을 비트 폭을 바꿔가며 MNIST 테스트 정확도로 비교했습니다. 모델의 세 Linear 층 가중치를 전부 양자화한 뒤(복원값으로 교체) 정확도를 잰 것입니다.
def eval_linear(model, n_bits, symmetric):
m = copy.deepcopy(model)
with torch.no_grad():
for layer in [m.fc1, m.fc2, m.fc3]:
wl = layer.weight.detach().cpu().numpy()
recon, *_ = linear_quantize(wl, n_bits, symmetric=symmetric)
layer.weight.copy_(torch.tensor(recon, dtype=torch.float32, device=device))
return accuracy(m)
bits_list = [2, 3, 4, 8]
acc_asym = [eval_linear(model, b, False) for b in bits_list]
acc_sym = [eval_linear(model, b, True) for b in bits_list]
# K-Means 도 동일 구조의 eval_kmeans(model, b) 로 비교
| bits | Linear (asym) | Linear (sym) | K-Means | float32 |
|---|---|---|---|---|
| 8 | 97.2% | 97.2% | 97.2% | 97.2% |
| 4 | 97.1% | 97.0% | 97.1% | 97.2% |
| 3 | 96.7% | 94.4% | 96.8% | 97.2% |
| 2 | 58.1% | 18.8% | 93.9% | 97.2% |
읽어야 할 큰 그림은 이렇습니다.
- 8비트면 사실상 공짜. 세 방식 모두 float32와 구분이 안 될 만큼 정확도가 유지됩니다.
- 낮은 비트로 갈수록 K-Means가 유리. 대표값을 데이터 분포에 맞춰(비균등) 놓기 때문에, 같은 2비트라도 균등 격자인 Linear보다 오차가 작습니다. 아래 MSE 곡선에서 더 뚜렷합니다.
하지만 정확도가 전부는 아닙니다. K-Means가 2비트에서 오차가 더 작아도, 추론 연산은 여전히 float입니다. Linear는 오차를 조금 더 감수하는 대신 연산까지 정수로 만들어 실제 하드웨어에서 빨라집니다. 두 방식은 경쟁이 아니라 목적이 다른 도구입니다.
| 방식 | 저장 | 연산 | 격자 |
|---|---|---|---|
| 원본 | FP 가중치 | FP 산술 | 해당 없음 |
| K-Means 양자화 | 정수 인덱스 + FP 코드북 | FP 산술 | 비균등 |
| Linear 양자화 | 정수 가중치 | 정수 산술 | 균등 |
직접 돌려보기
위 plot과 표는 전부 아래 두 노트북을 Colab에서 실행해 나온 것입니다. 런타임 → 모두 실행이면 끝나고, 자기 모델의 .safetensors를 올려 경로만 바꾸면 그대로 적용됩니다.
- 📓
linear-quantization.ipynb: 아핀 매핑, symmetric/asymmetric - 📓
kmeans-quantization.ipynb: 코드북/인덱스, 압축비
정리
- 양자화 = 표현 가능한 값의 개수를 줄이는 것. 가중치가 0 근처에 몰려 있으니, 32비트를 다 쓰는 건 낭비입니다.
- K-Means(비균등)는 대표값을 분포에 맞춰 놓아 저장을 줄이지만 연산은 float. 낮은 비트에서 오차가 작습니다.
- Linear(균등)는 아핀 매핑 $r=S(q-Z)$로 저장도 연산도 정수화. 하드웨어 가속의 기반입니다.
- 8비트는 거의 공짜, 2비트까지 밀면 방식 선택이 정확도를 가릅니다.