Variation-aware Vision Transformer Quantization

ViT Quantization을 빠르게 보려고 하는데, 공개되어있는 github가 있길래 읽어봤다.

 

https://github.com/HuangOwen/VVTQ

GitHub - HuangOwen/VVTQ: Official PyTorch implementation of paper "Variation-aware Vision Transformer Quantization"

 

Accpeted는 안된 것 같다.

 

 

 

---

 

 

논문의 저자들이 하고자하는 이야기는 다음과 같다.

 

"ViT Quantization은 불안정하며, QAT에서 이 불안정함을 완화시켜줄 연구가 많이 되지 않았다."

 

그렇다면 왜 ViT Quantization은 불안정한가?

 

저자들은 단순하게, ViT가 Quantization할때 "variation"가 심하다고 이야기한다.

ViT, 특히 MHSA(multi-head self-attention) module이 quantization에 대해서 매우 민감한데, 위 그림을 보면 layer의 index에 따라서 quantization을 했을때 accuracy가 심하게 하락하는 애들이 있다고 한다. (주로 첫번째와 마지막 layer의 head가 그러하다.)

 

기존 CNN에서 자주 사용하던 방법들을 그대로 사용하기도 어려운게, CNN과 ViT의 weight distribution이 매우 다르기 때문이라고 한다.

 

ViT에서 Quantization을 학습하다보면, 어느 순간 weight oscillation 현상 (quantization 경계에 몰리면서 quantized $W^{q}$가 변한다는 이야기)가 있다고 한다.

 

 

이러한 문제를 바탕으로, 저자들은 module별로 step size를 정해야하며 variation이 심하게 일어나지 않게 최적화를 해야한다고 이야기한다.

 

우선 제안한 방법은 크게 세가지다.

 

1. Multi-crop Knowledge Distillation

걍 image를 random crop해서 full precision teacher에 넣은 output을 저장했다가, student 학습할때 soft label로 사용하겠다는 이야기다.  기존 연구들에서 Kd loss가 reguarlization term으로 동작하면서, 학습중에 variance를 줄여준다는 이야기가 있어서 사용했다고 한다.

2. Module-dependent Quantization

 

ViT에서는 CNN과 다르게, module마다 민감도가 다르다고 이야기했다. 그렇기 떄문에, module-dependent quantization을 사용했다. (MHSA의 query, key, value에 대해서 quantization step size $s$를 독립적으로 학습함)

 

 

3. Oscillation-aware Bin Regularization

 

저자들은 weight의 variance가 oscillation 현상을 만든다고 주장했다.

이러한 현상을 완화하기 위해, weight distribution을 regularize할 수 있는 Oscillation-aware Bin Regularizer(OBR)을 썼다.

모든 모듈 $m$에 대해서 $|w_{m}^{r} - w_{m}^{q}|_{2}$로 real $w$와 qunatized $w$의 차이를 줄여주고, 모든 quantization bin에 대해서 variance를 계산해서 이를 최소화하도록 하는 term이다.

 

 

 

 

 

Method들도 다 이미 있던 내용에 실험도 좀 부족해보인다.

일단 ViT의 weight variance현상과 MHSA에서 layer별로 민감도가 다르다는 분석만 챙겨가면 될 것 같다.