경량화 시리즈 #2 — 뉴럴넷 양자화

이 글은 딥러닝 시대의 데이터타입 편의 후속입니다. 앞 글에서 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 논문¹에서 왔습니다. 비슷한 값의 가중치를 한 덩어리로 묶어 대표값 하나로 퉁치자는 것입니다. 예를 들어 2.09, 1.92, 1.87은 전부 “대략 2.0"이니, 이 값들을 2.00 하나로 대체하고 각 자리에는 “몇 번 대표값을 쓸지” 인덱스만 적어둡니다.

절차는 세 단계입니다.

1. 클러스터링 — 모든 가중치를 $k = 2^N$개 그룹으로 K-Means 클러스터링.
2. 코드북(codebook) — 각 그룹의 중심값(centroid)을 대표값으로 저장. 이 목록이 코드북.
3. 인덱스 — 각 자리엔 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$비트 양자화면:

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.²이 제안했고, TensorFlow Lite의 INT8 양자화가 바로 이 방식입니다.

• $r$ — 원래의 실수(real) 값
• $q$ — 양자화된 정수
• $Z$ — zero point. 실수 $0$에 정확히 대응하는 정수. ReLU 뒤의 0, 패딩처럼 흔한 값을 오차 없이 표현하기 위한 장치.
• $S$ — scale. 정수 한 칸이 실수로 얼마인지 (floating-point).

$S$와 $Z$는 실수 범위 $[r_{\min}, r_{\max}]$의 양 끝을 정수 범위 $[q_{\min}, q_{\max}]$의 양 끝에 맞추는 데서 나옵니다.

코드에서는 이 $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_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) 로 비교

읽어야 할 큰 그림은 이렇습니다.

• 8비트면 사실상 공짜. 세 방식 모두 float32와 구분이 안 될 만큼 정확도가 유지됩니다.
• 낮은 비트로 갈수록 K-Means가 유리. 대표값을 데이터 분포에 맞춰(비균등) 놓기 때문에, 같은 2비트라도 균등 격자인 Linear보다 오차가 작습니다. 아래 MSE 곡선에서 더 뚜렷합니다.

하지만 정확도가 전부는 아닙니다. K-Means가 2비트에서 오차가 더 작아도, 추론 연산은 여전히 float입니다. Linear는 오차를 조금 더 감수하는 대신 연산까지 정수로 만들어 실제 하드웨어에서 빨라집니다. 두 방식은 경쟁이 아니라 목적이 다른 도구입니다.

직접 돌려보기

위 plot과 표는 전부 아래 두 노트북을 Colab에서 실행해 나온 것입니다. 런타임 → 모두 실행이면 끝나고, 자기 모델의 .safetensors를 올려 경로만 바꾸면 그대로 적용됩니다.

• 📓 linear-quantization.ipynb — 아핀 매핑, symmetric/asymmetric
• 📓 kmeans-quantization.ipynb — 코드북/인덱스, 압축비

정리

• 양자화 = 표현 가능한 값의 개수를 줄이는 것. 가중치가 0 근처에 몰려 있으니, 32비트를 다 쓰는 건 낭비입니다.
• K-Means(비균등) 는 대표값을 분포에 맞춰 놓아 저장을 줄이지만 연산은 float. 낮은 비트에서 오차가 작습니다.
• Linear(균등) 는 아핀 매핑 $r=S(q-Z)$로 저장도 연산도 정수화. 하드웨어 가속의 기반입니다.
• 8비트는 거의 공짜, 2비트까지 밀면 방식 선택이 정확도를 가릅니다.

다음 실험 노트에서는 여기서 다루지 못한 활성값(activation) 양자화와 QAT(양자화 인지 학습) 로, 정수 추론 파이프라인을 끝까지 세워보겠습니다.