EfficientDM: Efficient Quantization-Aware Fine-Tuning of Low-Bit Diffusion Models

TL;DR

EfficientDM tackles the high computational cost of diffusion models by delivering a data-free, parameter-efficient fine-tuning framework that achieves QAT-like performance at PTQ-like efficiency. It introduces a quantization-aware low-rank adapter (QALoRA) merged with model weights, plus scale-aware LoRA optimization and temporal activation LSQ (TALSQ) to handle layer-wise scale disparities and step-wise activation distribution shifts. A data-free distillation objective aligns the quantized model’s denoising behavior with the full-precision model, enabling training without the original dataset. Experimental results on CIFAR-10, LSUN, and ImageNet demonstrate state-of-the-art performance for low-bit diffusion models, with substantial speed-ups over QAT-based methods and robust generation quality even at 2-bit weight quantization. The approach offers practical deployment benefits for resource-constrained settings and real-world applications.

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}.

Figures (17)

• Figure 1: An overview of the efficiency-vs-quality tradeoff across various quantization approachs. Data is collected on LDM-8 rombach2021highresolutionLDM with 4-bit weights and activations on LSUN-Churches. The GPU memory consumption is visualized by circle size.
• Figure 2: An overview of the proposed EfficientDM fine-tuning framework. Here, $\mathbf{s}_w$ and $s_x$ represent the learnable quantization scales for weights and activations, respectively. Compared to QLoRA layer, both updated weights and activations in our QALoRA are quantized to enable efficient bitwise operations during inference. Fine-tuning is performed by minimizing the mean squared error between the estimated noises of FP and quantized models.
• Figure 3: The motivation and effect of scale-aware LoRA optimization. Data is collected from the 4-bit LDM-4 model. (a): Due to step-like relationship between $\mathbf{BA}$ and $\hat{\mathbf{W}}$, $\mathbf{BA}$ needs to be large enough to update model weights. Data is collected from the first channel in the $12^{th}$ layer. (b): Significant disparity in weight quantization scales across layers. (c): Mean absolute value of quantized weight updates ($\hat{\mathbf{W}}$-$\hat{\mathbf{W}_0}$) for each layer. Most of them are zero under scale-agnostic training, indicating full-precision LoRA weights are too small to update quantized model weights. The proposed scale-aware LoRA optimization facilitates a more equitable distribution of quantized weight updates across layers.
• Figure A: Ranges of model output across various steps. Results are obtained by LDM-4 model on ImageNet $256\times256$ dataset.
• Figure B: The variance of activations and the learned quantization scales of intermediate layers across different steps. Caption denotes the layer name. Data is collected by W4A4 LDM-4 model on ImageNet $256\times256$.