[논문 리뷰] Oscillation-free Quantization for Low-bit Vision Transformers

Shih-Yang Liu, Zechun Liu외 1명|arXiv (Cornell University)|2023. 02. 04.

Advanced Neural Network Applications인용 수 6

한 줄 요약

이 논문은 가중치 학습 가능한 스케일링 인자와 쿼리-키 상호의 진동으로 인해 발생하는 양자화 인식 훈련(Quantization-aware training)에서의 가중치 진동 문제를 해결하기 위해 저비트 비전 트랜스포머(Vision Transformers)를 위한 진동 없는 양자화(Oscillation-free Quantization, OFQ)를 제안한다. 통계적 가중치 양자화(StatsQ), 신뢰도 유도 안내(Confidence-guided annealing, CGA), 쿼리-키 재구성(Query-key reparameterization, QKR)를 도입함으로써 훈련을 안정화시키고 최신 기술 수준의 정확도를 달성한다—기존 방법 대비 2비트 DeiT-T/S 및 Swin-T에서 각각 9.88%, 7.72%, 4.64% 향상.

ABSTRACT

Weight oscillation is an undesirable side effect of quantization-aware training, in which quantized weights frequently jump between two quantized levels, resulting in training instability and a sub-optimal final model. We discover that the learnable scaling factor, a widely-used $ extit{de facto}$ setting in quantization aggravates weight oscillation. In this study, we investigate the connection between the learnable scaling factor and quantized weight oscillation and use ViT as a case driver to illustrate the findings and remedies. In addition, we also found that the interdependence between quantized weights in $ extit{query}$ and $ extit{key}$ of a self-attention layer makes ViT vulnerable to oscillation. We, therefore, propose three techniques accordingly: statistical weight quantization ($ m StatsQ$) to improve quantization robustness compared to the prevalent learnable-scale-based method; confidence-guided annealing ($ m CGA$) that freezes the weights with $ extit{high confidence}$ and calms the oscillating weights; and $ extit{query}$-$ extit{key}$ reparameterization ($ m QKR$) to resolve the query-key intertwined oscillation and mitigate the resulting gradient misestimation. Extensive experiments demonstrate that these proposed techniques successfully abate weight oscillation and consistently achieve substantial accuracy improvement on ImageNet. Specifically, our 2-bit DeiT-T/DeiT-S algorithms outperform the previous state-of-the-art by 9.8% and 7.7%, respectively. Code and models are available at: https://github.com/nbasyl/OFQ.

연구 동기 및 목표

• 비전 트랜스포머(ViT)의 저비트 양자화 인식 훈련(QAT)에서의 훈련 불안정성의 근본 원인을 규명한다.
• 학습 가능한 스케일링 인자와 쿼리-키 가중치 상호진동이 가중치 진동의 주요 원인임을 규명한다.
• 정확도 향상과 수렴성 향상을 위해 저비트 ViT 모델에서 안정적인 진동 없는 양자화 프레임워크를 개발한다.
• 최소한의 비트 폭으로 ImageNet에서 최신 기술 수준의 성능을 달성하며, 3비트 모델이 전정밀도 성능에 도달함을 보장한다.
• 다양한 ViT 아키텍처(DeiT, Swin)와 비트 폭(2–4비트)에 걸쳐 일반화 가능성과 안정성을 확보한다.

제안 방법

• 학습 가능한 스케일링 인자를 통계 기반 스케일링으로 대체하는 통계적 가중치 양자화(StatsQ)를 제안하여 진동을 감소시킨다. 이 스케일링은 가중치 분포에서 유도된다.
• 신뢰도 유도 안내(CGA)를 도입하여 높은 신뢰도를 가진 가중치는 동결하고, 진동하는 가중치는 경계 범위 임계값을 사용해 선택적으로 업데이트한다.
• 양자화된 쿼리와 키 가중치 간의 상호 영향을 분리하기 위해 쿼리-키 재구성(QKR)을 설계하여 기울기 추정 오류를 완화한다.
• StatsQ, CGA, QKR를 통합한 유일한 프레임워크를 구성하여 훈련을 안정화시키고 진동 없는 모델로의 수렴을 가능하게 한다.
• 최적화 속도와 효과성의 균형을 맞추기 위해 CGA에 동적 경계 범위를 적용하며, 평균 기울기 크기를 기반으로 경험적으로 선택한다.
• 스케일링 인자의 궤적 분석을 통해 CGA가 훈련 전반에 걸쳐 진동을 억제하는 데 효과적임을 검증한다.

실험 결과

연구 질문

• RQ1 왜 양자화 인식 훈련에서 학습 가능한 스케일링 인자가 비전 트랜스포머에서 가중치 진동을 악화시키는가?
• RQ2 자기주의 어텐션 레이어에서 양자화된 쿼리와 키 가중치 간의 상호의존성이 훈련 중 진동을 악화시키는 방식은 무엇인가?
• RQ3 가중치 통계 기반의 통계적 스케일링이 학습 가능한 스케일링보다 진동 감소와 정확도 향상에 더 우수한가?
• RQ4 신뢰도 유도 안내가 저비트 ViT 양자화에서 훈련을 얼마나 안정화시키고 진동을 제거하는가?
• RQ5 쿼리-키 재구성(QKR)이 양자화된 쿼리와 키 가중치 간의 부정적 상호 영향을 효과적으로 분리하는가?

주요 결과

• 학습 가능한 스케일링 인자는 노이즈가 많은 기울기 업데이트와 임계값 불안정성으로 인해 가중치 진동을 심화시킨다.
• 통계적 가중치 양자화(StatsQ)는 기존의 LSQ 기반 기준선 대비 2비트 DeiT-S에서 정확도를 6.3% 향상시키며 진동을 감소시킨다.
• 경계 범위 0.005를 사용한 confidence-guided annealing(CGA)가 가장 우수한 성능을 보이며, 다른 설정 대비 정확도를 0.72% 향상시킨다.
• StatsQ, CGA, QKR의 조합은 이전 최신 기술 수준 대비 2비트 DeiT-S에서 7.72%의 정확도 향상을 달성하며, 2비트 DeiT-T는 9.88% 향상된다.
• OFQ를 적용한 3비트 DeiT-T 및 DeiT-S는 전정밀도 모델과 유사한 정확도를 달성하여, 진동 제어를 통한 저비트 양자화의 효과성을 입증한다.
• 4비트 설정에서 OFQ는 각각 DeiT-T, DeiT-S, Swin-T에서 전정밀도 모델을 3.24%, 1.2%, 0.68% 초월한다.