Q-ViT: Accurate and Fully Quantized Low-bit Vision Transformer(NeurlPS, 2022) paper review

트랜스포머(Transformer)는 원래 자연어 처리 분야에서 큰 성공을 거둔 모델 구조이다. 이 구조는 Self-Attention 메커니즘을 통해 입력 간의 장거리 의존성을 효과적으로 학습할 수 있다는 장점을 가지고 있다. 이러한 강점 덕분에 컴퓨터 비전 분야에도 빠르게 도입되었고, Vision Transformer(ViT)는 이미지 분류, 객체 탐지 등 다양한 시각 과제에서 뛰어난 성능을 보여주었다.

하지만 ViT의 성능은 대규모 데이터와 방대한 파라미터 수에 크게 의존한다. 예를 들어, ViT-Huge 모델은 약 6억 개 이상의 파라미터와 162G에 달하는 FLOPs를 필요로 하며, 메모리 사용량 또한 수 기가바이트에 이른다. 이러한 연산 및 메모리 부담은 일반적인 GPU나 모바일 장치, 엣지 디바이스와 같은 자원 제약 환경에서 ViT를 직접 활용하기 어렵게 만든다.

이 문제를 해결하기 위해 연구자들은 다양한 모델 압축 기법을 시도해왔다. 대표적으로는 경량화된 네트워크 설계, 프루닝(Pruning), 저랭크 분해(low-rank decomposition), 지식 증류(knowledge distillation), 그리고 양자화(Quantization)가 있다. 그중에서도 양자화는 특히 매력적인 접근법이다. 모델의 파라미터와 활성값을 낮은 비트로 표현하면, 연산량과 메모리 사용량을 동시에 줄일 수 있기 때문이다. 또한 AI 전용 칩이나 하드웨어 가속기에서 비트 단위 연산(bit-wise operation)을 활용할 수 있어, 실제 추론 속도 개선에도 직접적인 효과를 기대할 수 있다.

CNN에서는 이미 양자화 기법이 활발히 연구되어 8비트 정밀도로도 원래의 부동 소수점 성능에 근접한 결과를 내는 것이 가능해졌다. 그러나 ViT와 같은 트랜스포머 구조는 상황이 달랐다. ViT의 핵심인 Self-Attention 메커니즘은 Query, Key, Value 간의 미세한 분포 차이를 기반으로 글로벌 관계를 학습한다. 그런데 이를 저비트 정수로 단순히 양자화할 경우, 분포가 크게 왜곡되어 성능이 급격히 저하되는 문제가 발생한다. 특히 4비트 이하 초저비트 환경에서는 그 성능 저하가 더욱 두드러지며, 실제 활용이 제한적이었다.

따라서 연구자들에게 남은 중요한 질문은 다음과 같다.

"ViT를 초저비트 환경에서도 안정적으로 양자화할 수 있는 방법은 무엇인가?"

Q-ViT 연구는 바로 이 질문에 대한 답을 찾고자 출발한다.

문제점

Vision Transformer를 단순히 양자화한다고 해서 바로 효율적이고 정확한 모델을 얻을 수 있는 것은 아니다. 연구자들은 완전 양자화된 ViT(Fully quantized ViT)를 직접 실험해본 결과, 성능이 원본 모델에 비해 크게 떨어지는 현상을 확인했다. 이때 나타나는 성능 저하는 단순히 정밀도가 낮아졌기 때문이 아니라, ViT 곶의 특성과 양자화 과정이 충돌하면서 발생하는 두 가지 근본적인 병목 현상으로 설명할 수 있다.

1. Forward 과정에서의 정보 왜곡 문제

ViT의 핵심은 Self-Attention 메커니즘이다. Query와 Key의 내적을 통해 입력 간의 상관관계를 계산하고, 이로부터 Attention 분포를 얻어 장거리 의존성을 학습한다. 그러나 파라미터와 활성값을 저비트 정수로 양자화하면, Query와 Key 값의 분포가 원래와 크게 달라진다. 다시 말해, Full-Precision 모델에서는 Gaussian 형태의 안정적인 분포가 나타나지만, 양자화된 모델에서는 분산이 크게 달라지고 분포가 일그러진다. 그 결과 Attention 지도 자체가 왜곡되며, 이는 곧 ViT가 가지는 전역적 표현 능력을 심각하게 훼손한다.

위 그림은 DeiT-T 모델의 일부 층(Block 0, 3, 6)에서 Query 값의 분포를 비교한 결과이다. 왼쪽은 Full-Precision 모델이고, 오른쪽은 4-bit 완전 양자화된 ViT이다. 파란색 영역은 Query 값의 히스토그램이고, 빨간색 곡선은 Gaussian 분포를 근사한 확률밀도함수(PDF)이다.

위 그림에서 확인할 수 있는 중요한 차이는 다음과 같다.

• 분산($\sigma^2$)의 차이
  • Full-Precision에서는 Query값이 안정적으로 Gaussian 분포를 따르며, 분산 값이 일정한 범위 내에 있다.
  • 그러나 4-bit 양자화를 적용하면 분산이 커지거나, 특정 구간에서 불균형하게 변한다.
• 평균($μ$)의 변화
  • Full-Precision에서는 Query 값의 평균이 0에 가깝게 유지되어 대칭적인 분포를 보인다.
  • 반면 양자화된 경우 평균이 변동하면서, 원래의 분포와 차이를 보인다.

2. Backward 과정에서의 비효율적인 최적화 문제

양자화 모델을 학습시킬 때 자주 활용되는 기법은 지식 증류(Knowledge distillation)이다. 일반적으로 Full-Precision 교사 모델의 출력을 기준으로 양자화된 학생 모델이 그 결과를 모방하도록 지도한다. 그러나 이러한 방식은 출력 레벨에서만 이루어지는 비교적 "거친" 지도 신호에 불과하다. Attention과 같은 세밀한 내부 표현이 제대로 교정되지 않기 때문에, 양자화 모델은 여전히 교사 모델과 큰 분포 차이를 보인다. 결과적으로 최적화가 불안정하고, 성능 향상도 제한적이다.

정리하면, 완전 양자화된 ViT가 실패하는 이유는 구조적 문제와 최적화 문제라는 두 가지 측면에서 비롯된다.

• 구조적 병목 : Attention 모듈의 Query/Key 분포 왜곡 $\to$ 전역 의존성 표현 능력 상실
• 최적화 병목 : Coarse-grained 증류 방식 $\to$ 세밀한 분포 차이를 줄이지 못함

이 두 가지 문제가 해결되지 않는 한, 초저비트 환경에서 ViT가 안정적으로 동작하기는 어렵다.

해결 방안

앞서 살펴본 것처럼, 완전 양자화된 Vision Transformer는 두 가지 큰 병목을 겪는다. 전방향 과정에서는 Attention 모듈의 Query/Key 분포가 심하게 왜곡되고, 역전파 과정에서는 지식 증류가 세밀한 분포 차이를 보정하지 못한다. Q-ViT는 이러한 문제를 해결하기 위해 두 가지 핵심 기법을 도입한다.

1. 정보 보정 모듈(Information Rectification Module, IRM)

IRM은 전방향 과정에서 발생하는 정보 왜곡 문제를 해결한다.

기본 아이디어는 양자화된 Query/Key 분포를 원래의 Full-Precision 분포에 최대한 가깝게 유지하는 것이다.

이를 위해 정보 이론을 활용하여 정보 엔트로피를 최대화하는 방식으로 설계된다.

Query와 Key 값이 Gaussian 분포를 따른다고 가정했을 때, 단순 양자화는 분포의 평균과 분산을 무너뜨린다. IRM은 추가적인 학습 가능한 파라미터(γ, β)를 도입해 분포를 재정렬하고, 그 결과 엔트로피가 극대화된다.

이 과정은 곧 양자화된 Attention이 다시 원래와 비슷한 표현력을 갖추도록 보정하는 효과를 낸다.

즉, IRM은 전방향에서 정보가 손실되지 않도록 Query/Key 분포를 교정하는 장치라고 할 수 있다.

2. 분포 기반 증류(Distribution Guided Distillation, DGD)

DGD는 역전파 과정에서 발생하는 부정확한 최적화 문제를 해결한다.

기존 지식 증류는 모델의 최종 출력(logits)을 맞추는 데 초점이 맞춰져 있어, Attention과 같은 내부 표현의 왜곡을 줄이지 못했다.

Q-ViT는 이를 보완하기 위해 Query/Key 수준에서 Teacher-Student 분포의 유사성을 직접 맞추는 방식을 도입한다.

구체적으로, Query와 Key를 기반으로 패치 단위 유사도 행렬(similarity matrix)을 계산하고, 이를 Teacher 모델과 비교하여 차이를 줄인다.

이렇게 하면 단순히 "정답을 맞추는" 수준을 넘어서, Attention 모듈 내부의 분포 자체를 정밀하게 보정할 수 있다.

즉, DGD는 Teacher의 Attention 분포를 학생 모델이 그대로 학습하도록 유도하는 정교한 증류 기법이다.

IRM과 DGD는 서로 보완적인 역할을 수행한다.

IRM은 전방향에서 분포를 교정하여 정보 표현력을 살려주고,
DGD는 역전파에서 Teacher의 지식을 효과적으로 전달하여 최적화 정확도를 높인다.

실험적으로도 두 기법을 함께 적용했을 때 가장 큰 성능 향상이 나타났으며, 2~4bit 초저비트 환경에서도 Full-Precision 모델에 근접하거나 오히려 능가하는 성능을 달성할 수 있었다.

Q-ViT의 전체 구조

먼저 왼쪽 부분은 IRM(Information Rectification Module)을 나타낸다.

입력 특징은 Query(q), Key(k), Value(v)로 변환된 후 각각 양자화된다. 문제는 단순 양자화가 Query와 Key의 분포를 크게 왜곡시켜 Attention 점수 A 계산에 오류를 일으킨다는 점이다.

IRM은 이 과정을 바로잡기 위해 정규화 및 보정 파라미터(γ, β)를 추가한다.

• 수식적으로는 Query/Key를 평균과 분산에 기반해 재정렬하여 정보 엔트로피를 최대화한다.
• 이렇게 하면 양자화 후에도 분포가 Full-Precision 모델과 유사한 Gaussian 형태를 유지할 수 있다.

그 결과 Attention Score $Q_a(A)$가 왜곡되지 않고, 모델의 전역적 표현력이 유지된다.

우측은 DGD 과정을 보여준다.

기존 증류 방식은 최종 출력 로짓만을 비교하는 coarse-grained 접근이었는데, 이는 Attention 내부의 세밀한 분포를 보정하기에 부족하다.

DGD는 Teacher 모델과 Student 모델(양자화 모델) 사이의 Query/Key 분포 자체를 비교한다. 구체적으로는 Query와 Key를 기반으로 유사도 행렬(Similarity matrices) $G_q, G_k$를 생성한다. Teacher 모델에서 얻은 $G_{q;T}, G_{k;T}$와 Student 모델의 $G_q, G_k$를 직접 비교하여 최적화 손실을 계산한다.

이렇게 하면 Attention 분포가 Teacher와 Student 사이에서 세밀하게 일치하도록 학습할 수 있다.

즉 전체적인 Q-ViT 구조를 살펴보면,
Patch Embedding $\to$ Transformer Blocks(MHSA + MLP) $\to$ Classifier 라는 ViT 기본 구조 위에 IRM과 DGD가 추가된 형태이다.

IRM은 Forward 과정에서 Attention 입력을 보정하고,
DGD는 Backward 과정에서 Teacher 모델의 분포 정보를 세밀하게 전달한다.

두 기법을 결합하면, 초저비트(2~4bit) 환경에서도 Full-Precision 모델과 유사하거나 더 나은 성능을 낼 수 있다.'

실험 및 결과

아래 그림은 DeiT-Small 모델로 성능 비교를 하는 그래프이다. 녹색선은 Full-Precision, 파란색은 Fully-Quantized ViT, 주황선은 Q-ViT이다. x축은 Layer Index(Transformer 층 번호)를 의미하고, y축은 Attention Distance(각 층의 Attention 분포 간 L2-norm거리)를 의미한다.

ViT는 층이 깊어질수록 다양한 수준의 Attention 패턴을 학습한다.

• 초기 층: 국소(local)와 전역(global) 정보를 동시에 활용
• 후기 층: 주로 전역적 의존성을 강조

따라서 Full-Precision ViT에서는 층마다 서로 다른 Attention 분포 간 거리가 뚜렷하게 나타난다. 이것이 ViT가 복잡한 시각 표현을 학습할 수 있는 근거이다.

하지만, 완전 양자화된 Baseline 모델(파란 선)은 이 거리가 거의 구분되지 않는(flat한 패턴)모습을 보인다. 이는 곧 Attention 메커니즘이 층별로 구별된 의미 있는 표현을 학습하지 못한다는 의미이다. 즉, 정보 열화(Information Degradation) 현상이 발생한 것이다.

Q-ViT는 Full-Precision 모델과 매우 유사한 패턴을 보인다.
층이 깊어질수록 Attention Distance가 다르게 유지되며, 각 층이 의미 있는 표현 차이를 학습하고 있음을 확인할 수 있다.

이는 Q-ViT의 IRM과 DGD가 결합되어 Attention 분포를 보정하고, Teacher 모델의 분포를 효과적으로 전달했기 때문이다.

우측은 Q-ViT로 IRM과 DGD를 적용한 4bit 모델에서도 Query/Key 분포가 Full-Precision과 거의 동일하게 복원됨을 확인할 수 있다.

즉, Q-ViT는 IRM을 통해 Forward 과정에서 분포를 재정렬하고, DGD를 통해 Teacher 부포를 세밀하게 전달받으면서, 양자화 모델에서도 Full-Precision 수준의 분포 특성을 회복할 수 있었다.

위 테이블을 보면, Baseline 대비 IRM 단독, +1.7%(2bit), DGD 단독 +2.3%, IRM+DGD 결합, +3.8%의 성능 향상을 이루었다.

즉, 두 기법은 독립적으로도 유효하지만, 함께 사용했을 때 상승 효과(synergy)가 극대화된다.

위 테이블을 보면, Q-ViT는 초저비트(2~4bit)에서도 Full-Precision 성능에 근접하거나 오히려 능가함을 확인할 수 있다. 특히 Swin-Transformer에서는 FP를 초월하는 결과를 달성하였다.

결론

본 논문은 Vision Transformer를 초저비트 환경에서 완전 양자화할 때 발생하는 성능 저하 문제를 분석하고, 이를 해결하기 위한 새로운 접근법을 제시하였다. 기존의 단순 양자화 방식은 Attention 모듈에서 Query와 Key 분포가 심각하게 왜곡되며, 지식 증류 또한 세밀한 분포 차이를 보정하지 못해 최적화가 불안정하다는 한계를 보였다.

이러한 문제를 해결하기 위해 Q-ViT는 Information Rectification Module (IRM)과 Distribution Guided Distillation (DGD)라는 두 가지 핵심 기법을 도입하였다. IRM은 Forward 과정에서 분포 엔트로피를 최대화하여 양자화된 Query/Key 값이 Full-Precision과 유사한 분포를 유지하도록 교정하였다. 동시에 DGD는 Teacher 모델과 Student 모델의 Attention 유사도 행렬을 직접 비교하여, Backward 과정에서도 정교한 분포 정렬이 이루어지도록 최적화를 개선하였다.

실험 결과, Q-ViT는 ImageNet에서 다양한 Transformer 구조(DeiT, Swin)에 대해 2~4bit 초저비트 양자화에서도 Full-Precision 성능에 근접하거나 이를 초과하는 성능을 달성하였다. 특히 Swin-T와 Swin-S 모델에서는 4bit 양자화 환경에서 Full-Precision 대비 각각 +1.3%, +1.2%의 정확도 향상을 기록하며, 초저비트 ViT의 실현 가능성을 입증하였다.

결론적으로, Q-ViT는 “정확하고 완전한 초저비트 Vision Transformer”를 위한 새로운 패러다임을 제시하였다. 이 연구는 Transformer 기반 모델의 극한 압축을 가능하게 하고, 엣지 디바이스나 모바일 환경과 같은 자원 제약 플랫폼에서도 강력한 성능을 발휘할 수 있는 기반을 마련한다. 향후 연구는 1-bit ViT와 같은 극단적인 양자화 환경, 그리고 다양한 비전 과제(예: 객체 탐지, 분할)로의 확장을 통해 더 큰 응용 가능성을 열어갈 것이다.

References

Li, Y., Xu, S., Zhang, B., Cao, X., Gao, P., & Guo, G. (2022). Q-ViT: Accurate and fully quantized low-bit vision transformer. Advances in Neural Information Processing Systems (NeurIPS 2022), 35, 21511–21524.