Joint Post-Training Quantization of Vision Transformers with Learned Prompt-Guided Data Generation

이 논문은 학습된 프롬프트로 Stable Diffusion Turbo 를 활용한 데이터 생성 기법을 통해 라벨 없는 데이터로 비전 트랜스포머의 전체 레이어를 동시에 최적화하는 종단 간 양자화 프레임워크를 제안하여, 극저비트 설정에서도 최첨단 정확도를 달성하고 에지 배포 가능성을 입증했습니다.

핵심

이 논문은 **"거대한 인공지능 (비전 트랜스포머) 을 작고 가벼운 스마트폰이나 엣지 기기에 넣기 위해, 실수 없이 압축하는 새로운 방법"**을 소개합니다.

기존의 방식은 마치 **"거대한 도서관의 모든 책을 한 권씩 따로따로 요약해서 책장에 꽂는 것"**처럼, 각 페이지 (레이어) 를 따로따로 다듬는 방식이었습니다. 하지만 이 방법은 책 전체의 흐름을 잃어버리기 쉽고, 특히 최신 AI 모델처럼 책과 책이 서로 긴밀하게 연결된 구조에서는 잘 작동하지 않았습니다.

이 논문이 제안하는 방법은 다음과 같은 세 가지 핵심 아이디어로 요약할 수 있습니다.

---

1. "한 번에 통째로 다듬기" (Joint Optimization)

비유: 오케스트라의 지휘자

기존 방법들은 악기 하나하나 (각 레이어) 를 따로 튜닝했습니다. 하지만 비전 트랜스포머는 악기들이 서로 소리를 맞춰야 하는 오케스트라와 같습니다. 한 악기의 소리가 조금만 달라져도 전체 곡의 분위기가 망가질 수 있죠.

이 연구는 **지휘자 (최적화 알고리즘)**가 모든 악기를 한 번에 보고, 서로의 소리가 잘 어울리도록 동시에 조율합니다.

• 장점: 레이어 간의 연결고리를 끊지 않고, 전체적인 흐름을 유지하면서 압축합니다.
• 결과: 기존에 불가능하다고 여겨졌던 '3 비트'나 '1.5 비트' 같은 극도로 낮은 비트에서도 높은 정확도를 유지합니다. 마치 거대한 오케스트라를 작은 라디오 스피커에 담으면서도 음악의 감동을 잃지 않는 것과 같습니다.

2. "실제 사진 없이, AI 가 상상한 그림으로 학습하기" (Data-Free Calibration)

비유: 요리사의 레시피 테스트

AI 를 압축할 때는 보통 "이게 뭐야?"라고 물어보는 수만 장의 실제 사진 (데이터) 이 필요합니다. 하지만 이 사진들을 구하기 어렵거나, 개인정보 문제가 있을 수 있습니다.

이 논문은 **"실제 사진이 없어도 AI 가 스스로 상상해서 그림을 그리게 한다"**는 아이디어를 썼습니다.

• 문제점: 기존에는 AI 에게 "개 사진 그려줘"라고만 했더니, 모든 개가 똑같은 자세로 똑같은 배경에 그려지는 경우가 많았습니다. (다양성 부족)
• 해결책: 연구진은 AI 에게 **"다양한 개를 그려줘"**라고 가르쳤습니다.
  • "산책하는 개", "비 오는 날의 개", "화려한 배경의 개"처럼 **다양한 시나리오 (프롬프트)**를 자동으로 학습시킵니다.
  • 이렇게 AI 가 그린 수많은 다양한 상상 그림을 이용해 압축 과정을 테스트합니다.
• 결과: 실제 사진으로 테스트한 것과 거의 똑같은 성능을 내며, 개인정보나 데이터 수집 없이도 AI 를 최적화할 수 있게 되었습니다.

3. "무거운 짐을 가볍게 나르는 기술" (Channel-Wise Rescaling)

비유: 배낭 정리의 달인

AI 의 내부 데이터는 어떤 부분은 아주 크고 (무거운 짐), 어떤 부분은 아주 작습니다. 이를 무조건 같은 크기의 상자에 담으려다 보면 중요한 정보가 깨지거나 버려집니다.

이 연구는 각 레이어마다 짐의 크기를 미리 조절하는 기술을 썼습니다.

• 무거운 짐은 상자를 크게, 가벼운 짐은 상자를 작게 맞춰서 균형 있게 압축합니다.
• 이렇게 하면 중요한 정보 (큰 값) 는 잘 보존되고, 작은 값들도 사라지지 않게 되어, 압축된 AI 가 원래 AI 와 거의 똑같은 일을 할 수 있게 됩니다.

---

🌟 한 줄 요약

이 논문은 **"거대한 AI 모델을 실제 데이터 없이도, AI 가 상상한 다양한 그림으로 훈련시켜, 스마트폰처럼 작은 기기에서도 고화질로 작동하도록 통째로 최적화하는 새로운 방법"**을 제시했습니다.

이는 앞으로 우리가 스마트폰이나 시계 같은 작은 기기에서도 무거운 AI 모델을 빠르고 정확하게 사용할 수 있는 길을 열어줍니다.

기술 요약

1. 문제 정의 (Problem)

비전 트랜스포머 (Vision Transformers, ViT) 는 이미지 분류 등 다양한 시각 인식 작업에서 뛰어난 성능을 보이지만, 높은 계산량과 메모리 요구사항으로 인해 리소스가 제한된 엣지 장치나 실시간 애플리케이션에 배포하기 어렵습니다. 이를 해결하기 위해 모델 양자화 (Quantization) 가 필수적이지만, 기존 방식에는 다음과 같은 한계가 존재합니다.

• 양자화 인식 학습 (QAT) 의 비효율성: QAT 는 라벨이 있는 데이터와 긴 파인튜닝 주기를 요구하여 시간과 에너지 소모가 큽니다.
• 기존 PTQ (Post-Training Quantization) 의 한계:
  • ViT 는 레이어 간 및 블록 간 의존성 (Inter-block dependencies) 이 강하고, 활성화 값의 분포가 비균일하며 (Non-uniform), 어텐션 메커니즘이 복잡합니다.
  • 기존 CNN 기반 PTQ 방법 (블록 단위 재구성 등) 은 ViT 에 적용 시 전역 상관관계를 무시하여 성능이 급격히 저하됩니다.
  • 특히 초저비트 (Extreme low-bit, 예: W1.58A8) 설정에서 ViT 를 성공적으로 양자화한 선행 연구가 부재했습니다.
• 데이터 의존성: 대부분의 PTQ 방법은 양자화를 위해 실제 ImageNet 과 같은 라벨이 있는 칼리브레이션 데이터가 필요합니다. 데이터 프라이버시 문제나 데이터 부재 상황에서는 적용이 어렵습니다.

2. 제안 방법론 (Methodology)

이 논문은 라벨 데이터 없이 전체 레이어와 블록 간 의존성을 동시에 최적화하는 엔드 - 투 - 엔드 (End-to-End) 공동 PTQ 프레임워크와 학습된 프롬프트 기반의 데이터 프리 (Data-Free) 칼리브레이션 전략을 제안합니다.

가. 엔드 - 투 - 엔드 공동 양자화 최적화 (Joint PTQ Optimization)

• 전체 네트워크 최적화: 블록 단위 재구성이 아닌, 모든 트랜스포머 블록과 레이어를 하나의 통합 목적 함수 하에 공동으로 최적화합니다. 이를 통해 블록 간 오차 보상 (Cross-block compensation) 을 가능하게 하여 저비트 설정에서도 안정성을 확보합니다.
• 채널 단위 재스케일링 (Channel-Wise Rescaling): ViT 의 활성화 값은 채널 간 편차가 크고 아웃라이어 (Outlier) 가 존재합니다. 이를 해결하기 위해 SmoothQuant 와 RepQ-ViT 에서 영감을 받아, 입력 활성화와 가중치에 학습 가능한 스케일 ($\alpha$) 및 시프트 ($\beta$) 벡터를 도입합니다. 이는 활성화의 동적 범위를 평탄화하여 양자화 손실을 줄입니다.
• 손실 함수: 라벨 없이도 학습이 가능하도록, 정밀도 모델 (Full-precision) 과 양자화 모델 간의 **중간 특징 재구성 손실 (MSE)**과 **최종 로그이트 KL 발산 손실 (Distillation Loss)**을 결합합니다. 또한 가중치 미세 조정 (Weight refinement) 을 통해 초기 정밀도 모델과의 일치를 유지합니다.
• 효율성: 단일 GPU 에서 ViT-small 기준 약 1 시간, Swin-Base 기준 약 2.5 시간 내에 수렴합니다.

나. 학습된 프롬프트 가이드 데이터 생성 (Learned Prompt-Guided Data Generation)

실제 데이터 없이 양자화를 수행하기 위해 Stable Diffusion Turbo 를 활용한 데이터 프리 칼리브레이션 전략을 도입합니다.

• 다중 모드 프롬프트 학습 (Multi-Prompt Optimization): 단순히 "a photo of "와 같은 수동 템플릿을 사용하는 대신, 각 클래스당 **여러 개의 학습된 프롬프트 임베딩 (Prompt Embeddings)**을 학습합니다.
• 최적화 목표:
  1. 분류 손실: 생성된 이미지가 사전 학습된 ViT 분류기에 의해 올바른 클래스로 인식되도록 합니다.
  2. 다양성 정규화: 프롬프트 임베딩의 직교성 (Orthogonality) 과 생성된 이미지, 특징, 어텐션 맵의 분산 (Variance) 을 최대화하여 객체의 배치, 배경, 스타일 등을 다양하게 생성합니다.
• 효과: 학습된 프롬프트는 실제 ImageNet 데이터의 활성화 분포를 효과적으로 근사하며, 단순 텍스트 프롬프트보다 훨씬 풍부한 시각적 다양성을 제공합니다.

3. 주요 기여 (Key Contributions)

1. 라벨 없는 엔드 - 투 - 엔드 PTQ 프레임워크: ViT 의 모든 블록과 레이어를 공동으로 최적화하여, 라벨 데이터 없이도 W4A4 및 W3A3 에서 SOTA 성능을 달성했습니다.
2. 생성 기반 데이터 프리 칼리브레이션: Stable Diffusion Turbo 와 학습된 다중 프롬프트를 결합하여, 실제 데이터 없이도 실제 데이터와 동등한 성능을 내는 칼리브레이션 데이터를 생성하는 전략을 제시했습니다.
3. 초저비트 양자화 달성: ViT, DeiT, Swin-T 모델에서 **W1.58A8 (3-비트 가중치)**과 같은 극단적인 저비트 설정에서도 강력한 정확도를 유지하는 최초의 PTQ 결과를 달성했습니다.

4. 실험 결과 (Results)

• 성능 (ImageNet Top-1 Accuracy):
  • W4A4/W3A3: ViT-S, DeiT-S, Swin-T 등 다양한 백본에서 기존 PTQ 방법 (RepQ-ViT, FIMA-Q, APHQ-ViT) 보다 우수한 성능을 기록했습니다.
  • W1.58A8 (초저비트): 기존 방법들은 성능이 급격히 떨어지거나 실패하는 반면, 제안된 방법은 ViT-S 에서 68.45% (실제 데이터), 63.71% (합성 데이터) 의 높은 정확도를 유지했습니다.
  • 데이터 프리 성능: 학습된 프롬프트로 생성된 10 만 개의 합성 이미지를 사용한 경우, 실제 데이터 (1 만 개) 를 사용한 경우와 비교해 정확도 차이가 1~2% 이내로 매우 작았습니다.
• 확장성: 칼리브레이션 데이터 크기가 10,000 개까지 증가함에 따라 성능이 지속적으로 향상되었으며, 블록 단위 재구성 방식 (FIMA-Q) 보다 더 큰 데이터셋에서도 효과적이었습니다.
• 다양성 분석: t-SNE 시각화 결과, 학습된 프롬프트로 생성된 이미지는 실제 데이터의 특징 분포 (Manifold) 에 더 가깝고 다양하게 분포하는 반면, 단순 텍스트 프롬프트는 편향된 클러스터를 형성하는 것을 확인했습니다.

5. 의의 및 결론 (Significance)

이 논문은 비전 트랜스포머의 엣지 배포를 위한 핵심 장벽인 고정밀 저비트 양자화와 데이터 의존성 문제를 동시에 해결했습니다.

• 기술적 혁신: 블록 간 의존성을 고려한 전역 최적화와 학습된 생성적 프롬프트를 결합하여, 기존 PTQ 의 한계를 극복했습니다.
• 실용성: 라벨 데이터가 없거나 프라이버시 제약이 있는 환경에서도 고품질의 양자화 모델을 빠르게 (단일 GPU, 1~2 시간) 생성할 수 있음을 입증했습니다.
• 미래 영향: W1.58A8 과 같은 극단적인 저비트 설정에서의 성공은 ViT 기반 모델의 경량화 및 엣지 AI 적용 가능성을 크게 확장시켰으며, 생성형 AI 를 활용한 모델 최적화 (Generative AI for Model Optimization) 의 새로운 방향성을 제시합니다.