[논문 리뷰] EfficientDM: Efficient Quantization-Aware Fine-Tuning of Low-Bit Diffusion Models

Yefei He, Jing Liu외 3명|arXiv (Cornell University)|2023. 10. 05.

Advanced Data Compression Techniques인용 수 4

한 줄 요약

EfficientDM는 학습 데이터가 필요 없고, 양자화에 대한 인지 기반이며, 파라미터 효율적인 미세조정 프레임워크를 제안하여 4비트 확산 모델에서 QAT 수준의 성능을 PTQ 수준의 효율성으로 달성한다. QALoRA(양자화에 대한 인지 기반 LoRA 변종), 스케일 인지 최적화, 시간에 따라 학습되는 스텝 크기 양자화를 도입함으로써 어댑터와 모델 가중치를 함께 양자화할 수 있게 하여, ImageNet 256×256에서 4비트 정밀도에서 sFID를 1.56점 감소시키고, QAT에 비해 양자화 시간을 16.2배 단축시켰다.

ABSTRACT

Diffusion models have demonstrated remarkable capabilities in image synthesis and related generative tasks. Nevertheless, their practicality for real-world applications is constrained by substantial computational costs and latency issues. Quantization is a dominant way to compress and accelerate diffusion models, where post-training quantization (PTQ) and quantization-aware training (QAT) are two main approaches, each bearing its own properties. While PTQ exhibits efficiency in terms of both time and data usage, it may lead to diminished performance in low bit-width. On the other hand, QAT can alleviate performance degradation but comes with substantial demands on computational and data resources. In this paper, we introduce a data-free and parameter-efficient fine-tuning framework for low-bit diffusion models, dubbed EfficientDM, to achieve QAT-level performance with PTQ-like efficiency. Specifically, we propose a quantization-aware variant of the low-rank adapter (QALoRA) that can be merged with model weights and jointly quantized to low bit-width. The fine-tuning process distills the denoising capabilities of the full-precision model into its quantized counterpart, eliminating the requirement for training data. We also introduce scale-aware optimization and temporal learned step-size quantization to further enhance performance. Extensive experimental results demonstrate that our method significantly outperforms previous PTQ-based diffusion models while maintaining similar time and data efficiency. Specifically, there is only a 0.05 sFID increase when quantizing both weights and activations of LDM-4 to 4-bit on ImageNet 256x256. Compared to QAT-based methods, our EfficientDM also boasts a 16.2x faster quantization speed with comparable generation quality. Code is available at \href{https://github.com/ThisisBillhe/EfficientDM}{this hrl}.

연구 동기 및 목표

• 학습 데이터가 많고 계산 자원을 많이 요구하는 QAT와 성능이 떨어지는 PTQ 사이의 효율성-품질 트레이드오���을 해결하기 위해.
• 양자화 과정에서 원본 학습 데이터에 의존하지 않도록 하여 자원이 제한된 장치에서의 실용적 구현을 가능하게 하기 위해.
• 기존 PTQ 방법이 소음 제거 능력을 유지하지 못하는 4비트 정밀도에서 가중치와 활성화 모두의 높은 생성 품질을 유지하기 위해.
• QAT 수준의 성능을 달성하면서도 PTQ 수준의 시간과 데이터 효율성을 확보하여, 실세계 응용에 적합한 4비트 확산 모델을 실현하기 위해.

제안 방법

• 모델 가중치와 함께 함께 양자화할 수 있는 저랭크 어댑터의 양자화에 대한 인지 기반 변종인 QALoRA를 도입하여, 효율적이고 4비트 추론을 가능하게 한다.
• 학습 데이터가 필요 없는 디스틸레이션 프레임워크를 적용하여, 정밀도가 높은 모델과 양자화된 모델의 노이즈 예측 간 평균 제곱오차를 최소화한다.
• 층 간에 다른 양자화 스케일을 가진 저비트 환경에서 최적화 효과를 높이기 위해, 층 간에 적응적으로 기울기 스케일을 조정하는 스케일 인지 최적화를 적용한다.
• 소음 제거 단계 동안 활성화의 분포 이동을 다루기 위해 시간에 따라 학습되는 스텝 크기 양자화(TLSQ)를 도입하여 추론 중 안정성과 성능을 향상시킨다.
• QALoRA 가중치를 양자화된 모델에 통합하여, 추가 저장소나 계산 오버헤드 없이 파라미터 효율적인 미세조정을 가능하게 한다.
• 이중 단계 학습 프로세스를 사용한다: 먼저 QALoRA를 통해 정밀도가 높은 모델 지식을 양자화된 모델에 디스틸레이션하고, 이후 TLSQ와 스케일 인지 최적화로 정밀도를 향상시킨다.

실험 결과

연구 질문

• RQ1 학습 데이터에 접근할 수 없더라도, 학습 데이터가 없는 미세조정 방법이 4비트 확산 모델 양자화에서 QAT 수준의 성능을 달성할 수 있는가?
• RQ2 양자화에 대한 인지 기반 저랭크 어댑터는 어떻게 설계되어야 하며, 가중치와 함께 함께 양자화하면서도 4비트 정밀도에서 성능을 유지할 수 있는가?
• RQ3 다양한 양자화 스케일을 가진 네트워크 층 간에서 성능 저하를 완화하기 위해 효과적인 최적화 전략은 무엇인가?
• RQ4 소음 제거 단계에 따라 양자화 파rameter를 시간에 따라 적응적으로 조정하면, 4비트 확산 모델의 모델 안정성과 생성 품질을 향상시킬 수 있는가?
• RQ5 QAT에 비해 시간과 메모리 비용을 크게 줄이고, 4비트 정밀도에서 거의 정밀도 기반 성능을 달성할 수 있는가?

주요 결과

• EfficientDM는 ImageNet 256×256에서 4비트 정밀도(W/A)에서 sFID 7.75를 달성했으며, 전체 정밀도 기반 베이스라인 대비 0.05포인트 뿐만 증가했고, 기존 PTQ 방법이 이 비트 폭에서 소음 제거 능력을 유지하지 못하는 것에 비해 뛰어난 성능을 보였다.
• QAT 기반 접근 방식에 비해 양자화 시간을 16.2배 단축시켜 ImageNet 256×256에서 3.05 GPU 시간을 기록했으며, QAT는 54.5 GPU 시간을 소모했다.
• CIFAR-10에서 EfficientDM는 0.97 GPU 시간에 FID 10.48를 달성했고, QAT 수준의 성능(FID 7.30)을 재현했으며, QAT보다 16.8배 빠르고 PTQ보다 1.1배 더 빠르게 작동했다.
• 스케일 인지 최적화를 통해 QALoRA만 사용했을 때 FID 24.18에서 13.23으로 감소시켜, 다양한 양자화 스케일을 가진 층 간 최적화 효율성을 향상시켰다는 점을 입증했다.
• TLSQ를 추가로 적용함으로써 FID는 13.23에서 7.75로 감소하여, 단계별로 조정되는 양자화가 소음 제거 단계에서 활성화 분포의 이동에 적응함으로써 성능 향상을 이룬다는 것을 보여주었다.
• EfficientDM는 ImageNet 256×256에서 GPU 메모리 19.8GB만을 사용하여, RTX 3090에서 QAT가 OOM을 일으키는 것과는 대조적으로 데이터 의존성을 피하고 실생활 구현에 매우 실용적이었다.