[논문 리뷰] Temporal Dynamic Quantization for Diffusion Models

Junhyuk So, Jungwon Lee외 3명|arXiv (Cornell University)|2023. 06. 04.

Stochastic Gradient Optimization Techniques인용 수 4

한 줄 요약

이 논문은 시간 단계 정보를 기반으로 양자화 간격을 동적으로 조정하는 새로운 양자화 방법인 시간 동적 양자화(TDQ)를 제안한다. 이는 생성 품질을 크게 향상시키며, 추론 오버헤드가 없으며, 후기 양자화(PTQ) 및 양자화 인식 학습(QAT) 모두를 향상시킨다. 4비트 이하의 저비트 폭에서 최신 기술 수준의 성능을 달성한다 (예: CIFAR-10에서 W4A4 FID 4.48).

ABSTRACT

The diffusion model has gained popularity in vision applications due to its remarkable generative performance and versatility. However, high storage and computation demands, resulting from the model size and iterative generation, hinder its use on mobile devices. Existing quantization techniques struggle to maintain performance even in 8-bit precision due to the diffusion model's unique property of temporal variation in activation. We introduce a novel quantization method that dynamically adjusts the quantization interval based on time step information, significantly improving output quality. Unlike conventional dynamic quantization techniques, our approach has no computational overhead during inference and is compatible with both post-training quantization (PTQ) and quantization-aware training (QAT). Our extensive experiments demonstrate substantial improvements in output quality with the quantized diffusion model across various datasets.

연구 동기 및 목표

• 반복적인 노이즈 제거 단계 동안 시간에 따라 변화하는 활성화 분포로 인한 저비트 양자화에서의 성능 저하 문제 해결.
• 정적 또는 기존의 동적 양자화가 활성화 통계의 시간적 이동에 적응하지 못하는 한계를 극복.
• 추론 시 계산 오버헤드 없이 시간에 따라 최적의 양자화 파rameter를 생성하는 양자화 방법 개발.
• 모바일 및 엣지 디바이스에서 넓은 범위의 배포를 위해 후기 양자화(PTQ) 및 양자화 인식 학습(QAT)과의 호환성 확보.
• 완전 정밀 기반 모델에 가까운 성능을 유지하면서도 저정밀도(예: 4비트 가중치 및 활성화)에서 고해상도 이미지 생성 달성.

제안 방법

• 현재 노이즈 제거 단계를 기반으로 각 시간 단계에 대한 양자화 파rameter(스케일 및 제로포인트)를 예측하는 학습 가능한 신경망인 TDQ 모듈을 도입.
• 기존의 PTQ 및 QAT 파이프라인에 손상 없이 원활하게 통합되도록 설계하여 추론 계산 그래프를 수정하지 않음.
• 모든 시간 단계에서 활성화 양자화 오차를 최소화하기 위해 미분 가능한 양자화 방식(예: LSQ)을 사용해 TDQ 모듈을 엔드 투 엔드로 학습.
• 성능와 복잡도의 최적의 트레이드오프를 위해 경험적으로 4층의 다층 퍼셉트론(MLP)을 TDQ의 백본으로 선택.
• 역노이즈 제거 과정 중에 변화하는 활성화 분포에 적응하는 시간에 따라 달라지는 양자화 간격 생성.
• 정방향 전파 중 각 시간 단계당 한 번만 양자화 파rameter를 계산하여 추론 효율성 확보, 런타임 계산 방지.

Figure 1: The forward and reverse processes of diffusion models.

실험 결과

연구 질문

• RQ1 시간 단계별 활성화 분포에 적응하는 동적 양자화가 확산 모델의 저비트 양자화 성능을 크게 향상시킬 수 있는가?
• RQ2 시간에 따라 달라지는 양자화 전략이 정적 및 입력 기반 동적 양자화보다 저비트 폭에서 생성 품질을 유지하는 데 더 우수한가?
• RQ3 제안된 TDQ 모듈이 추론 오버헤드 없이 PTQ 및 QAT 프레임워크에 효과적으로 통합될 수 있는가?
• RQ4 TDQ 모듈의 레이어 수가 최종 모델 성능과 수렴 안정성에 미치는 영향은 어떠한가?
• RQ5 샘플링 스케줄이 변화할 경우(예: 단계 수 감소), TDQ는 어느 정도 성능을 유지하는가?
• RQ6 양자화 간격은 시간 단계에 따라 어떻게 변화하며, 활성화 동역학과 상관관계가 있는가?

주요 결과

• DDIM을 사용한 CIFAR-10에서 W4A4 양자화로 TDQ는 FID 4.48을 달성하여 LSQ(7.30) 및 기타 베이스라인보다 뚜렷이 뛰어나다.
• LDM-Churches 데이터셋에서 TDQ는 W4A4 FID 4.64를 기록하여 정적 양자화(5.06) 및 동적 양자화(5.17)를 모두 초월한다.
• W3A3 양자화로 CIFAR-10에서 TDQ는 FID 6.48을 달성하여 LSQ(7.63) 및 기타 동적 기법보다 뚜렷이 우수하다.
• 제거 분석 결과, 4층 TDQ가 성능와 복잡도의 최적 균형을 이룬 것으로 확인되었으며, FID는 2층일 때 5.18, 8층일 때 4.88에 비해 4.48로 떨어졌다.
• 샘플링 단계 수를 100에서 10으로 줄였을 때도 TDQ는 일관된 성능을 유지하지만, LSQ는 급격히 성능이 저하됨.
• 시각화 결과 TDQ 간격이 시간 단계에 따라 활성화 변동에 적응함을 확인하였으며, 약 70%의 레이어에서 강한 시간적 상관관계를 보였으며, 특히 어텐션 및 잔여 블록에서 두드러졌다.