[PaperReview] Training with Quantization Noise for Extreme Model Compression

Abstract

매 iteration forward 때 마다 서로 다른 가중치의 일부분에 양자화 하는 방법이다. 이렇게 할 경우 편향성 없는 gradient를 다른 가중치에 흐르게 할 수 있다.

---

Introduction

실 생활에서 쓰이는 성능 좋은 대부분의 NN architecture는 수 많은 파라미터를 가진다. 또 효율적인 모델 구조를 가진 NN 일지라도 수백 메가바이트가 요구되므로 로봇과 가상 환경에서 사용하기 제한적이다. Quantization은 NN architecture 관점에서 이미 최적화가 잘되어 있는 모델을 압축 할 수 있다는 점이 특징이다.

Scalar quantization 이 흔한 방법이며, 가중치를 FP에서 INT로 대체하는 방법이다. 이럴 경우, gradient가 null을 가진다는 문제점과 이 문제점을 해결하기 위한 방안을 사용할 때 train과 test 때의 불일치를 발생한다.

QAT에서는 이러한 문제(gradient가 null 을 가지는 현상)의 해결 방법으로 STE 소개하였는데 이 방법은 int8 quantization 같은 곳에는 잘 작동하나, 더 낮은 비트 수를 사용 하게 된다면 제대로 작동하지 않을 수 있다. QuantNoise는 가중치의 일부분에만 양자화를 적용하므로 매우 높은 압축에도 안정적인 학습을 할 수 있고, 대부분의 가중치가 편향성 없게 업데이트가 된다.

---

Related Work

• Model compression
  • Weight pruning : 학습 시 또는 fine-tuning 시에 구현되는 방법
  • Weight sharing
  • Lightweight architecture
  • Knowledge distillation
• Quantization weight와 activation을 low-precision으로 학습 시키는 방법.
  

---

Quantizing Neural Networks

Fixed-point Scalar quantization

floating-point 표현을 low bit precision의 fixed-point 표현으로 바꾸는 메소드이다. 적용할 경우 메모리 사용량이 줄어들며, fixed-point arithmetic을 지원하는 하드웨어에서 추론 속도가 빨라진다.

Weight matrix를 (k, l)개수의 block 으로 나눈 뒤, 각 블록의 값을 가능한 $2^n$으로 반올림을 한다. codeword 값으로 대치 된 값에 대해서 quantize와 dequantize를 순차적으로 적용한다.

activation 또한 N-bit fixed-point number로 반올림 하고, 압축율은 X32/N이 된다.

Product Quantization

• Traditional PQ Vector quantization 중 하나로, 블록은 한 종류의 값으로 대치하기 이전에 사전에 가중치들의 그룹으로 그룹화 된 것이다. traditional vector quantization은 matrix W를 p개의 column으로 쪼개고 p개의 vector의 codebook을 학습한다. 반면에, product quantization은 각각의 column을 m개의 subvector로 쪼개고 m x p개의 subvector의 codebook을 학습한다.
  
• Iterative PQ 단일 행렬에 대해서 PQ를 적용하는 것이 아니라, 전체 네트워크에 적용을 할 때, 높은 압축에 유도되는 양자화 표류가 누적된다. 이전 레이어의 변형된 값의 아웃 풋이 다음 레이어에 인풋이 되므로 이렇게 발생되는 표류는 큰 성능 저하로 이어진다. 이러한 문제의 해결 방법으로, iterative PQ 제안 되었고, 맨 아래 계층부터 양자화를 하여 위의 계층을 조정하는 방법이다. 이 때, 압축되지 않은 teacher model이 사용된다.
  각 계층의 codebook은 요소들에 할당된 gradient의 평균으로 학습한다.
  

💡 **Single weight?** 여기서 single weight는 weight 들이 일정 블록 단위로 묶였을 때, 블록 안의 값이 한 종류의 값(codeword)로 되치 되었을 때 상태를 말한다.

Combining fixed-point with product quantization

Code book의 값의 종류(centroids)가 K개라 할 때, 각 centroid는 d의 dimension으로 표현 되고, 8bit로 scalar quantization을 했을 때의 예시이므로 dimension을 구성하는 각 요소는 8bit로 표현되어 codebook의 크기는 $8*Kd$가 된다.

Product quantization을 적용하였을 때, 행의 개수가 p, 열의 개수가 q로 구성된 가중치는 열의 개수대로 q개의 vector quantization으로 바뀌고 각 열을 m개로 나누어 product quantization을 적용한다.

$(p, q) \rightarrow (m, q), \quad (m \leq q)$

크기로 바뀌며 각 블록($b_{mq}$)은 centroid값 K로 대치 되고, K를 표현하기 위해서는 $log_2K$의 비트가 필요하다. ⇒ $log_2K*mp$

input 배치 크기가 1이고, input의 dimension은 n이므로 dimension을 이루는 요소의 비트 수는 8bit로 이루어져 있어 $8*n$ 이다.

---

Method

Weight matrix를 block 단위로 쪼갠 후, 매 forward 때마다 랜덤하게 선택되어진 블록들에 대하여 noise를 가하여 학습한다. backward 시에는 STE를 사용하여 noise가 안된 weight matrix 대신 사용되어 학습 된다. 이렇게 랜덤하게 일부만 선택되어 학습을 하게 되면 QAT와는 다르게 noise가 가해지지 않은 블록이 편향 없는 gradient를 흐르게 한다. 이 때 noise는 dropout과 유사하게 값을 0으로 대체하는 것을 말한다.

또한, QuantNoise 방식은 다른 quantization method와 함께 적용 가능하다는 장점이 있다.
이 논문에서 FIxed-point scalar quantization과 Product quantization 적용하였다.