이 글은 경량화 시리즈의 다섯 번째(#5)입니다. 저비트 양자화 #4가 “학습이 끝난 모델을 사후에 잘 양자화하는(PTQ)” 처방들이었다면, 이번엔 아예 학습 단계에 양자화를 집어넣는 QAT입니다. 코드는 github.com/warpspaceinc/efficient-ml-practice의 qat.ipynb에 있고, 아래 모든 plot과 표는 그 코드를 직접 돌려 측정한 우리 숫자입니다.

PTQ가 무너지는 지점

지금까지의 양자화는 전부 PTQ(Post-Training Quantization), 학습이 끝난 FP32 모델을 사후에 정수 격자로 옮기는 것이었습니다. 큰 모델에는 잘 통합니다. INT8 PTQ로 ResNet-50은 −0.1%, GoogleNet은 손실이 거의 없습니다.

문제는 작은 모델입니다. 같은 INT8 PTQ인데도 수치가 인상적입니다1. 예를 들어:

모델FP32INT8 PTQ (per-tensor)
ResNet-5076.1%−0.1%
MobileNetV170.9%0.1%
MobileNetV271.9%0.1%

MobileNet은 70.9%에서 0.1%로, 즉 완전히 죽습니다. 파라미터를 아껴 만든 컴팩트한 모델일수록 여유(redundancy)가 없어서, 양자화 오차를 흡수할 곳이 없기 때문입니다. 그리고 이 현상은 비트를 낮출수록(4비트 이하) 모델 크기와 무관하게 나타납니다.

우리 MNIST MLP(784→256→128→10)로도 똑같이 재현됩니다. 가중치를 비트폭별로 PTQ 했을 때:

  • 8비트: 97.1% (멀쩡)
  • 4비트: 96.9%
  • 3비트: 94.9%
  • 2비트: 13.6%, 사실상 랜덤(10%)에 가깝게 붕괴

그럼 어떻게 해야 할까요? 답은 간단하면서도 강력합니다. “어차피 양자화된 상태로 추론할 거면, 학습할 때부터 양자화된 상태로 학습하자.”

QAT: 학습 중에 양자화를 흉내 낸다

QAT(Quantization-Aware Training)의 핵심은 fake quantization(모의 양자화)입니다. 학습의 forward pass에, 실제 추론 때 일어날 양자화를 미리 시뮬레이션해서 끼워 넣습니다. 구조는 세 가지 원칙으로 요약됩니다2.

  1. FP32 마스터 가중치 $W$를 계속 유지합니다. 학습되는 실체는 이 연속값입니다.
  2. forward에서만 가중치와 활성값을 양자화했다가 복원(fake-quant)해서, 정수 격자 위의 값으로 연산합니다.
  3. 추론 때는 이 양자화된 가중치만 씁니다.

수식으로 보면, #2의 아핀 매핑을 forward에 그대로 삽입하는 것입니다. 가중치와 출력(활성값) 각각에:

$$Q(W) = S_W\, q_W, \qquad Q(Y) = S_Y\,(q_Y - Z_Y)$$

여기서 $q_W = \mathrm{round}(W/S_W)$처럼 round로 격자에 스냅한 뒤 다시 scale을 곱해 실수로 되돌립니다(그래서 “fake”: 값은 격자 위에 있지만 연산 자체는 FP32로 돕니다). 학습이 이 격자 위에서 진행되니, 모델은 “이 가중치를 2비트로 뭉개도 출력이 좋도록” 스스로를 조정하게 됩니다. 양자화 오차를 사후에 감내하는 게 아니라, 학습 목표에 처음부터 반영하는 것입니다.

그런데 여기서 결정적인 문제가 하나 생깁니다.

STE: round는 미분이 0인데 어떻게 학습하나

fake quantization의 심장은 round 함수입니다. 그리고 round는 계단 함수입니다. 거의 모든 점에서 기울기가 0이고, 격자 경계에서만 무한대입니다.

$$\frac{\partial Q(W)}{\partial W} = 0 \quad (\text{거의 모든 곳에서})$$

역전파의 연쇄법칙을 그대로 따라가면 재앙입니다.

$$g_W = \frac{\partial L}{\partial W} = \frac{\partial L}{\partial Q(W)} \cdot \frac{\partial Q(W)}{\partial W} = \frac{\partial L}{\partial Q(W)} \cdot 0 = 0$$

모든 가중치의 gradient가 0이 됩니다. 학습이 한 발짝도 못 나갑니다. round 하나가 역전파를 통째로 막아버리는 것입니다.

해법이 STE(Straight-Through Estimator)입니다34. 발상은 뻔뻔할 만큼 단순합니다. “forward에서는 round를 쓰되, backward에서는 round가 없었던 셈 치고 gradient를 그냥 통과시킨다.” 즉 양자화 함수를 미분할 때 그것을 항등함수(identity)로 간주합니다.

$$g_W = \frac{\partial L}{\partial W} \;\approx\; \frac{\partial L}{\partial Q(W)}$$

forward는 계단, backward는 기울기 1. 이 “거짓말"이 QAT를 성립시킵니다.

재미있는 건 이 뻔뻔한 트릭의 출처입니다. STE는 딥러닝의 대부 제프리 힌튼(Geoffrey Hinton)이 2012년 Coursera 강의에서 지나가듯 소개한 아이디어이고3, 이듬해 Bengio 등이 형식화했습니다4. 그 힌튼은 튜링상(2018)에 이어 2024년 노벨 물리학상까지 받았고, 최근까지도 AI를 이야기하며 왕성히 활동하고 있습니다. 10여 년 전 강의에서 툭 던진 한 줄이 지금도 저비트 학습의 근간을 떠받치고 있는 셈입니다.

STE — forward는 계단 함수, backward는 항등함수(기울기 1)로 통과

왼쪽이 forward입니다. $Q(w)$는 점선(항등함수 $y=w$) 주변을 계단으로 근사합니다. 오른쪽이 backward입니다. 진짜 미분은 0(빨강)이지만, STE는 표현 범위 안에서 gradient를 1(초록)로 통과시킵니다. 범위 밖(clip된 영역)에서는 0으로 둬서, 격자를 벗어난 값에는 신호를 주지 않습니다.

아래에서 직접 만져보세요. 스칼라 가중치 $w$ 하나짜리 계산 그래프($w \to Q(w) \to \hat y=Q(w)\cdot x \to L=\tfrac12(\hat y-t)^2$)입니다. 가중치 $w$·입력 $x$·정답 $t$를 슬라이더로 바꾸면 loss와 gradient가 실시간으로 바뀝니다. 특히 backward에서 진짜 미분(0)과 STE gradient가 어떻게 갈리는지 보세요.

PyTorch로는 커스텀 autograd.Function 하나면 됩니다. forward에 양자화, backward에 통과(+clip 마스크)를 넣습니다.

class FakeQuantSTE(torch.autograd.Function):
    @staticmethod
    def forward(ctx, w, n_bits):
        qmax = 2 ** (n_bits - 1) - 1
        S = w.detach().abs().max() / qmax + 1e-12       # per-tensor symmetric scale
        q = torch.clamp(torch.round(w / S), -qmax, qmax)  # 격자로 스냅
        ctx.save_for_backward((w.abs() <= S * qmax).to(w.dtype))  # clip 마스크
        return q * S                                     # 복원 (fake-quant)
    @staticmethod
    def backward(ctx, g):
        (mask,) = ctx.saved_tensors
        return g * mask, None          # STE: 범위 안은 통과, 밖은 0

그리고 QAT 학습은 forward에서 이 fq를 가중치에 씌운 채로 평범하게 돌리면 됩니다. 마스터 가중치는 FP32로 남고, STE 덕분에 gradient가 그 마스터로 흘러갑니다.

def qat_forward(x):
    x = F.relu(F.linear(x, fq(fc1.weight, b), fc1.bias))
    x = F.relu(F.linear(x, fq(fc2.weight, b), fc2.bias))
    return F.linear(x, fq(fc3.weight, b), fc3.bias)

직접 측정: PTQ vs QAT

같은 MNIST MLP를, 같은 비트폭에서 PTQ(사후 양자화)와 QAT(FP32 모델에서 fake-quant로 2에폭 fine-tune)로 각각 재봤습니다.

비트폭별 PTQ vs QAT 정확도 — 2비트에서 PTQ는 붕괴, QAT는 회복

bitsPTQQATFP32
897.1%97.9%97.1%
496.9%97.7%97.1%
394.9%97.8%97.1%
213.6%91.9%97.1%

읽어야 할 그림은 명확합니다.

  • 8비트에선 둘 다 무손실. 여유가 있는 구간에선 굳이 QAT가 필요 없습니다.
  • 비트가 낮아질수록 격차가 폭발. 2비트에서 PTQ는 13.6%(랜덤 수준)로 죽는데, QAT는 91.9%까지 살려냅니다. MobileNetV1(INT8 PTQ 0.1%에서 QAT 70%대로)과 똑같은 패턴입니다1.
  • QAT가 baseline을 살짝 넘기도 하는데(3·4비트), fine-tune으로 몇 에폭 더 학습된 효과입니다.

핵심은 이것입니다. PTQ는 “이미 정해진 가중치를 격자에 최대한 잘 맞추는” 문제였습니다. QAT는 “격자 위에서 가장 좋은 가중치를 처음부터 찾는” 문제입니다. 자유도가 다르니 저비트에서 결과가 갈립니다.

직접 돌려보기

위 plot과 표는 아래 노트북을 직접 돌려 나온 것입니다. 런타임 → 모두 실행이면 끝납니다.

  • 📓 qat.ipynb: fake quantization, STE(autograd.Function), PTQ vs QAT 비교

STE 다음: scale·clip도 학습한다, LSQ·PACT

지금까지 우리 fake-quant는 scale $S$를 $|W|_{\max}/q_{\max}$라는 고정 휴리스틱으로 정했습니다. STE는 가중치 gradient만 흘려보낼 뿐, 정작 격자의 간격 $S$ 자체는 학습되지 않았죠. 다음 개선은 그 scale과 clip마저 학습하는 것입니다.

LSQ: step size를 학습한다. LSQ(Learned Step Size Quantization)5는 $S$를 학습 파라미터로 둡니다. $Q(w)=\mathrm{round}(\mathrm{clip}(w/S))\cdot S$에서 $\partial Q/\partial S$를 양자화 상태 전이에 민감하게 계산해 gradient로 최적화합니다. 범위 안에서는 $\partial Q/\partial S = \mathrm{round}(w/S) - w/S$(격자값과 연속값의 차이), 범위 밖에서는 $\pm q_{\max}$입니다. 효과는 저비트에서 큽니다. 우리 MLP를 고정 scale QAT와 비교하면:

bitsQAT (고정 scale)QAT + LSQ (scale 학습)
497.9%97.7%
397.6%97.9%
292.3%96.7%

2비트에서 92%가 97%로. 격자 간격 하나를 데이터가 정하게 했을 뿐인데 마지막 격차가 메워집니다. 여유가 없는 저비트일수록 “scale을 얼마로 둘까"가 정확도를 가릅니다.

PACT: activation의 clip 범위를 학습한다. weight는 학습이 끝나면 min/max가 고정이지만, activation(ReLU 출력)은 범위가 열려 있습니다. #4에서 KL로 자름점을 골랐다면, PACT6는 그 clip 임계값 $\alpha$를 학습 파라미터로 둡니다. ReLU를 $\mathrm{clip}(x, 0, \alpha)$로 바꾸고 $\alpha$를 gradient로 학습해, activation 양자화의 범위를 데이터가 정하게 합니다. LSQ가 weight의 scale을 배우듯, PACT는 activation의 clip을 배우는 셈입니다.

그리고 LLM 시대에 접어들며 이 이야기는 훨씬 더 멀리 갑니다: 삼항 {−1,0,+1}으로 처음부터 학습하는 BitNet, FP4로 프리트레이닝하는 흐름까지. 그건 별도의 LLM 양자화 편에서 따로 다루겠습니다.

정리

  • PTQ는 큰 모델엔 통하지만 작은 모델·저비트에서 무너진다. MobileNet INT8이 0.1%로 죽고, 우리 MLP도 2비트 PTQ에서 13.6%로 붕괴한다.
  • QAT는 학습 forward에 fake quantization을 끼워 넣어, 모델이 양자화된 상태에서 좋도록 스스로를 조정하게 한다. FP32 마스터 가중치는 유지된다.
  • round의 gradient는 0이라 그냥은 학습이 안 되는데, STE가 backward에서 양자화를 항등함수로 간주해 gradient를 통과시킨다.
  • 결과: 2비트에서 PTQ 13.6%에서 QAT 91.9%로. 저비트일수록 QAT의 값어치가 커진다.

이걸로 양자화 편(#2 · #4 · #5)이 마무리됩니다. 데이터타입에서 양자화, 프루닝, 저비트, QAT까지, 모델을 작고 빠르게 만드는 도구들을 개념부터 실측까지 훑었습니다.

References

  1. Krishnamoorthi. Quantizing Deep Convolutional Networks for Efficient Inference: A Whitepaper. arXiv 2018. (MobileNet PTQ·QAT 수치의 원출처) ↩︎ ↩︎

  2. Jacob et al. Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference. CVPR 2018. ↩︎

  3. Hinton et al. Neural Networks for Machine Learning. Coursera Lecture, 2012. (STE의 기원) ↩︎ ↩︎

  4. Bengio et al. Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation. arXiv 2013. ↩︎ ↩︎

  5. Esser et al. Learned Step Size Quantization. ICLR 2020. ↩︎

  6. Choi et al. PACT: Parameterized Clipping Activation for Quantized Neural Networks. arXiv 2018. ↩︎