Intel과 Hugging Face가 NNCF 기반 Quantization-Aware Training과 Token Merging을 결합해 Stable Diffusion의 CPU 추론 속도를 5.1배 단축하고 모델 크기를 4배 감축

Context

Stable Diffusion은 텍스트-이미지 생성의 핵심 모델이지만, CPU 환경에서 실행 시 매우 오랜 추론 시간이 소요된다. 특히 UNet 모델의 높은 계산 비용이 병목이며, 기존의 post-training 8-bit 양자화 같은 표준 최적화 기법은 픽셀 단위 예측 작업의 낮은 모델 중복성으로 인해 정확도 손실이 크다.

Technical Solution

• Quantization-Aware Training(QAT) 도입: NNCF를 사용해 모델 학습 과정에서 양자화 파라미터를 함께 조정하여 post-training 양자화의 정확도 손실 해결
• Knowledge Distillation 적용: 원본 모델을 teacher로, 최적화 대상 모델을 student로 설정해 보조 손실 함수로 추가하여 최적화 모델의 최종 정확도 향상
• Exponential Moving Average(EMA) 기법 적용: 양자화기를 제외한 모델 파라미터에 EMA 방법을 사용해 학습 과정 안정화
• Token Merging 기법 결합: Vision Transformer 기반 구조에서 중복 토큰을 중요 토큰과 병합하여 self-attention 블록의 계산량을 감소, 8-bit 양자화와 함께 누적 가속 효과 달성
• OpenVINO 런타임 최적화: NNCF로 최적화된 모델을 OpenVINO Toolkit의 고성능 추론 엔진에 배포하여 Intel CPU 기반 실행 최적화

Impact

• 추론 속도 5.1배 가속
• 모델 크기 4배 감축
• 단일 GPU(24GB VRAM)로 4096 이터레이션 최적화 완료(1일 이내)

Key Takeaway

복잡한 픽셀 단위 예측 모델의 최적화는 단순 양자화보다는 학습 과정에 최적화를 통합(QAT)하고 Knowledge Distillation으로 보강하며, 구조적 최적화 기법(Token Merging)과 결합할 때 누적 효과를 얻을 수 있다. 제한된 리소스 환경에서도 체계적인 최적화 워크플로우를 통해 실제 사용 가능한 수준의 성능 개선을 달성할 수 있다.